Repository: dask/dask

Issue:
Dask array rechunk should do intermediate sized rechunking when appropriate
Consider rechunking from chunks=(1, 100) to chunks=(100, 1) on an array of shape (100, 100).

A direct rechunk operation creates ~10,000 tasks corresponding to intermediate arrays of shape (1,1).

However, if we rechunk first to shape (10,10) (the same array size as all of our chunks) and then to our final chunk shape, there would be only ~2000 tasks, because the smallest intermediate arrays have 10 elements.

I don't know what the general rules are for how to optimize rechunkings but this makes a big difference in computational efficiency -- O(n^2) vs O(n^1.5) for square arrays with a side of length n. For n=10,000, this is a ~50x speed up (assuming that the task overhead dominates for small tasks).

Credit to @freeman-lab for this idea, which apparently works really nicely in Thunder. Jeremy, if you could share how you calculate intermediate chunks sizes that would be helpful!



Discussion:
I've run into this with other ndarray-on-spark systems and have heard that it performs decently.

@shoyer from your description I would propose:
 given final chunk size  `n = final_chunk.size`, and the array has `ndim` number of dimensions. Then the intermediate chunks can be square, where each dimension's length is the _ndim_th root of _n_. Like `intermediate_chunk = (floor(n**(1./ndim)),) * ndim`

But I'm not familiar enough with the rechunking internals to know if this is too naive.

If you do the rechunking repeatedly, each time scaling the longer chunk axis by a factor of two, you can get the number of intermediate chunks created down to O(n log(n)). OTOH, I have no idea how this scales in practice.

@shoyer I'm not sure I understand. Could you give an example?

Again, suppose you have an n x n array completely chunked along one axis with chunks of size (n, 1). In the first chunking operation, create chunks of size (n/2, 2). This involves 2n intermediate arrays. Then repeat the process log(n) times, and you have your final desired chunk size of (1, n). This involves only O(n log(n)) tasks. But you also created O(log(n)) temporary copies of the full dataset.

On Sat, Jul 18, 2015 at 5:34 PM, Blake Griffith notifications@github.com
wrote:

> ## @shoyer I'm not sure I understand. Could you give an example?
> 
> Reply to this email directly or view it on GitHub:
> https://github.com/ContinuumIO/dask/issues/417#issuecomment-122608426

I did a few benchmarks my machine:
http://nbviewer.ipython.org/gist/shoyer/425cb3dc9101c235fe86

For 1024 x 1024 arrays, I get a 15x speedup using one intermediate rechunking, and a 30x speedup using 3 intermediates -- bringing the calculation down from 4 minutes to 7 seconds. So it looks like a pretty clear win.

This is very cool @shoyer ! The size estimation we've done for this kind of chunking has been more about optimizing for axis swaps in large shuffles, and in that case chunks of approximately 100-150mb have empirically proven optimal, in that smaller chunks are slower (due to increased shuffle complexity) and larger chunks run into memory issues.

The optimal sizes here clearly have different constraints, but cool that the concept proves useful! 

@freeman-lab Do you guys use multiple intermediates in axis swaps or just one?

@shoyer just one... at least in Spark I'm fairly certain the time incurred just from processing the additional shuffles would outweigh any other gains, but I don't have hard numbers.

Btw, we've just about finished moving all the custom chunking logic into Bolt (https://github.com/bolt-project/bolt/pull/67), if you want to check out more of the details there.

@shoyer how do these speeds compare to just throwing the array into an empty 1024x1024 numpy array and then pulling an array with different chunks back out?  How would your benchmarks change with an array that didn't fit in memory?  Are the current scheduling routines smart enough to manage a reshuffle like this without filling up the cache and writing back to disk?

> how do these speeds compare to just throwing the array into an empty 1024x1024 numpy array and then pulling an array with different chunks back out?

I'm not quite sure what you're getting at here -- I think we can safely presume that creating the 1024x1024 numpy array directly would be far, far faster.

> How would your benchmarks change with an array that didn't fit in memory?

It takes 53 seconds merely to compute the _graph_ to do the naive rechunking, before any of the values are computed. Naive rechunking will be entirely unworkable for arrays that don't fit in memory.

> Are the current scheduling routines smart enough to manage a reshuffle like this without filling up the cache and writing back to disk?

My sense is that doing a full rechunking will pretty much require spilling to disk, regardless -- or distributing an array in memory across many machines.

> I'm not quite sure what you're getting at here -- I think we can safely presume that creating the 1024x1024 numpy array directly would be far, far faster.
> 
> My sense is that doing a full rechunking will pretty much require spilling to disk, regardless -- or distributing an array in memory across many machines.

In the single machine case if this requires spilling to disk and if creating the temporary explicitly is faster then maybe the optimal approach is to explicitly store the result into a shared array and then pull out from it again.

This issue has become a bit more interesting now with distributed scheduling.  I plan to run a benchmark on a few approaches soon.

Verified that in practice with the distributed scheduler the log-based rechunking seems nicer.  I also found that an explicit `rebalance` operation after rechunking was cheap and useful.  Things tend to get a bit unbalanced during computation.  I have yet to try this on a large cluster.  

It's possible that we'll be scheduler bound.

@freeman-lab roughly what kind of performance does Thunder get on Spark for tranposing the direction of partitioning?  Do you have ballpark numbers for a common dataset shape that you often use?

@mrocklin we're putting some numbers together now with recent versions of everything, will post here in just a bit

@freeman-lab even order-of-magnitude numbers would be really valuable here.  Something along the lines of the following would be really helpful.
-  Our data is usually shaped like 10000 by 10000000
-  A transpose of chunk shape from (10000, 1) to (1, 10000000) usually takes around 10s-2m on 20 m3.2xlarges

At the moment I'm currently curious if I'm anywhere in the same ballpark or not.  I'm trying to figure out how far I can get without building an explicit shuffle operation and just using task scheduling throughout.  My hope is to get a quick "nope, definitely not close" so that I know if I should continue to pursue my current path.

cc @broxtronix in case he has ballpark numbers 

@mrocklin cool gotcha so a typical medium-sized dataset for us might be (10000, 512, 512, 4), that's around ~20 GB as `uint16` and we'll often go from 10000 chunks of size (1, 1000000) to 1000000 chunks of size (10000, 1). The complete operation of loading from disk and doing the rechunking on 10 nodes might take ~2 minutes. Those numbers are from an internal cluster, which based on other benchmarking we've done might be 2-4x faster than the generic `m3`, but the exact factor there depends a lot on whether it's compute / network / io bound.

Is there a sequence of commands I can run in Thunder to do this on random data?

Do you have multiple chunks per spark partition?

yup yup, with v1.0.0 it's as easy as

```
import thunder as td

data = td.images.fromrandom(shape=(10000,512,512,4), npartitions=320, engine=sc)
data.toseries().compute()
```

where `sc` is a valid `SparkContext`

We typically do have multiple chunks per spark partition, assuming we set the partitioning to be roughly 2-3x the number of cores, which generally gives the best performance.

We just confirmed the 2 minutes is correct for that size on 10 nodes, also confirmed that doing the same thing for `(2000, 1000, 2000, 30)` data, going from `(1,60000000)` chunks to `(2000,1)` chunks is about 10 minutes on 20 nodes.

cc @jwittenbach who's been running these tests

Thank you.  These numbers are great.  Will crunch away on dask stuff and see how close I can get.

cool, very curious to hear how it goes!

OK, using a cascading transpose as suggested by @shoyer we can flip around the 20GB array in roughly two minutes.  

https://gist.github.com/mrocklin/01f39076b86fe73c71c353be830e4d0f

Looking at the Dask web UI this _does_ seem to take a fair amount of RAM.  It doesn't appear to be particularly sensitive to how many steps we use in the transpose-cascade.  Using only every second or third step or using significantly more partitions results in only a marginal increase (around 50%).  The scheduler is nowhere near peak when doing these operations, so we should be able to scale quite a bit larger.

The interesting thing we've learned here is that, as long as we keep the scheduler fast and pay attention to batching, it seems to be reasonable to avoid building a separate shuffle operation, at least for now.

@shoyer this might interest you: https://www.youtube.com/watch?v=W_VBpv27soU&feature=youtu.be

@mrocklin Yes, interesting to see this working! I would definitely try to benchmark against Bolt on the same cluster. I would not be surprised at all if EC2 has significantly worse networking than @freeman-lab's cluster, even with enhanced networking.

I can get this down to three minutes by splitting with a finer resolution, `2**5` rather than `2**2`.  This shifts burden from the network to the scheduler (fewer copies, more tasks) but the scheduler seems to be able to handle it just fine.  

I tested network bandwidth by trying `(x + x.T)` on a large square matrix and got up to around 300 MB/s steady output averaged across the network.

@mrocklin Something is wrong with the formatting on your last comment. 

Fixed

By resolution you mean the step size for each rechunk? That makes sense. 

Yeah, this chunks more finely and so accomplishes the movement in fewer stages with more tasks per stage.  We're nowhere near saturating the scheduler though (which is now quite speedy) so all is well.

Changed a couple things in the scheduler gets us down to around 2 minutes.  

Also, I realized that my test cluster didn't have any compression libraries installed.  After conda installing lz4 doing the transpose in more steps wasn't as bad, and came down to about 2.5 minutes.

In the two-step transpose I'm nowhere near saturating the bandwidth of my network.  The scheduler also doesn't appear to be particularly burdened.  Not sure where the current bottleneck is.

Actually, I just measured bandwidth by doing `x + x.T` on a large square matrix and came up with the same bandwidth measurement.  It could be that EC2 is a bit bogged down at the moment.  It would be interesting to run the above computation on a dedicated cluster.


Resolving patch:
diff --git a/dask/array/core.py b/dask/array/core.py
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -1594,10 +1594,10 @@ def cumprod(self, axis, dtype=None):
     def squeeze(self):
         return squeeze(self)
 
-    def rechunk(self, chunks):
+    def rechunk(self, chunks, threshold=None, block_size_limit=None):
         """ See da.rechunk for docstring """
-        from .rechunk import rechunk
-        return rechunk(self, chunks)
+        from . import rechunk   # avoid circular import
+        return rechunk(self, chunks, threshold, block_size_limit)
 
     @property
     def real(self):
diff --git a/dask/array/rechunk.py b/dask/array/rechunk.py
--- a/dask/array/rechunk.py
+++ b/dask/array/rechunk.py
@@ -7,11 +7,13 @@
 """
 from __future__ import absolute_import, division, print_function
 
+import heapq
+
 from itertools import product, chain
-from operator import getitem, add
+from operator import getitem, add, mul, itemgetter
 
 import numpy as np
-from toolz import merge, accumulate
+from toolz import merge, accumulate, reduce
 
 from ..base import tokenize
 from .core import concatenate3, Array, normalize_chunks
@@ -25,7 +27,7 @@ def cumdims_label(chunks, const):
      (('n', 0), ('n', 2), ('n', 4), ('n', 5))]
     """
     return [tuple(zip((const,) * (1 + len(bds)),
-                  list(accumulate(add, (0,) + bds))))
+                      accumulate(add, (0,) + bds)))
             for bds in chunks]
 
 
@@ -40,7 +42,7 @@ def _breakpoints(cumold, cumnew):
     >>> _breakpoints(new[1], old[1])
     (('n', 0), ('o', 0), ('n', 2), ('n', 4), ('n', 5), ('o', 5))
     """
-    return tuple(sorted(tuple(cumold) + tuple(cumnew), key=lambda x:x[1]))
+    return tuple(sorted(cumold + cumnew, key=itemgetter(1)))
 
 
 def _intersect_1d(breaks):
@@ -72,31 +74,28 @@ def _intersect_1d(breaks):
     start = 0
     last_end = 0
     old_idx = 0
-    lastbi = ('n',0)
-    ret = [[]]
+    ret = []
     for idx in range(1, len(breaks)):
-        bi = breaks[idx]
-        lastbi = breaks[idx - 1]
-        if 'n' in lastbi[0] and bi[1]:
+        label, br = breaks[idx]
+        last_label, last_br = breaks[idx - 1]
+        if last_label == 'n':
             ret.append([])
-        if 'o' in lastbi[0]:
+        if last_label == 'o':
             start = 0
         else:
             start = last_end
-        end = bi[1] - lastbi[1] + start
+        end = br - last_br + start
         last_end = end
-        if bi[1] == lastbi[1]:
+        if br == last_br:
             continue
         ret[-1].append((old_idx, slice(start, end)))
-        if bi[0] == 'o':
+        if label == 'o':
             old_idx += 1
             start = 0
     return tuple(map(tuple, filter(None, ret)))
 
 
-def intersect_chunks(old_chunks=None,
-                     new_chunks=None,
-                     shape=None):
+def intersect_chunks(old_chunks, new_chunks):
     """
     Make dask.array slices as intersection of old and new chunks.
 
@@ -114,18 +113,7 @@ def intersect_chunks(old_chunks=None,
         block sizes along each dimension (convert from old_chunks)
     new_chunks: iterable of tuples
         block sizes along each dimension (converts to new_chunks)
-    shape : tuple of ints
-        Shape of the entire array (not needed if using chunks)
-    old_blockshape: size of each old block as tuple
-        (converts from this old_blockshape)
-    new_blockshape: size of each new block as tuple
-        (converts to this old_blockshape)
-
-    Note: shape is only required when using old_blockshape or new_blockshape.
     """
-    old_chunks = normalize_chunks(old_chunks, shape)
-    new_chunks = normalize_chunks(new_chunks, shape)
-
     cmo = cumdims_label(old_chunks, 'o')
     cmn = cumdims_label(new_chunks, 'n')
     sums = [sum(o) for o in old_chunks]
@@ -166,7 +154,12 @@ def blockshape_dict_to_tuple(old_chunks, d):
     return tuple(new_chunks)
 
 
-def rechunk(x, chunks):
+DEFAULT_THRESHOLD = 4
+DEFAULT_BLOCK_SIZE_LIMIT = 1e8
+
+
+def rechunk(x, chunks, threshold=DEFAULT_THRESHOLD,
+            block_size_limit=DEFAULT_BLOCK_SIZE_LIMIT):
     """
     Convert blocks in dask array x for new chunks.
 
@@ -188,8 +181,18 @@ def rechunk(x, chunks):
     ----------
 
     x:   dask array
-    chunks:  the new block dimensions to create
+    chunks:  tuple
+        The new block dimensions to create
+    threshold: int
+        The graph growth factor under which we don't bother
+        introducing an intermediate step
+    block_size_limit: int
+        The maximum block size (in bytes) we want to produce during an
+        intermediate step
     """
+    threshold = threshold or DEFAULT_THRESHOLD
+    block_size_limit = block_size_limit or DEFAULT_BLOCK_SIZE_LIMIT
+
     if isinstance(chunks, dict):
         if not chunks or isinstance(next(iter(chunks.values())), int):
             chunks = blockshape_dict_to_tuple(x.chunks, chunks)
@@ -205,6 +208,269 @@ def rechunk(x, chunks):
     if not len(chunks) == ndim or tuple(map(sum, chunks)) != x.shape:
         raise ValueError("Provided chunks are not consistent with shape")
 
+    steps = plan_rechunk(x.chunks, chunks, x.dtype.itemsize,
+                         threshold, block_size_limit)
+    for c in steps:
+        x = _compute_rechunk(x, c)
+
+    return x
+
+
+def _number_of_blocks(chunks):
+    return reduce(mul, map(len, chunks))
+
+
+def _largest_block_size(chunks):
+    return reduce(mul, map(max, chunks))
+
+
+def estimate_graph_size(old_chunks, new_chunks):
+    """ Estimate the graph size during a rechunk computation.
+    """
+    # Estimate the number of intermediate blocks that will be produced
+    # (we don't use intersect_chunks() which is much more expensive)
+    crossed_size = reduce(mul, (len(oc) + len(nc)
+                                for oc, nc in zip(old_chunks, new_chunks)))
+    return crossed_size
+
+
+def divide_to_width(desired_chunks, max_width):
+    """ Minimally divide the given chunks so as to make the largest chunk
+    width less or equal than *max_width*.
+    """
+    chunks = []
+    for c in desired_chunks:
+        nb_divides = int(np.ceil(c / max_width))
+        for i in range(nb_divides):
+            n = c // (nb_divides - i)
+            chunks.append(n)
+            c -= n
+        assert c == 0
+    return tuple(chunks)
+
+
+def merge_to_number(desired_chunks, max_number):
+    """ Minimally merge the given chunks so as to drop the number of
+    chunks below *max_number*, while minimizing the largest width.
+    """
+    if len(desired_chunks) <= max_number:
+        return desired_chunks
+
+    distinct = set(desired_chunks)
+    if len(distinct) == 1:
+        # Fast path for homogeneous target, also ensuring a regular result
+        w = distinct.pop()
+        n = len(desired_chunks)
+        total = n * w
+
+        desired_width = total // max_number
+        width = w * (desired_width // w)
+        adjust = (total - max_number * width) // w
+
+        return (width + w,) * adjust + (width,) * (max_number - adjust)
+
+    desired_width = sum(desired_chunks) // max_number
+    nmerges = len(desired_chunks) - max_number
+
+    heap = [(desired_chunks[i] + desired_chunks[i + 1], i, i + 1)
+            for i in range(len(desired_chunks) - 1)]
+    heapq.heapify(heap)
+
+    chunks = list(desired_chunks)
+
+    while nmerges > 0:
+        # Find smallest interval to merge
+        width, i, j = heapq.heappop(heap)
+        # If interval was made invalid by another merge, recompute
+        # it, re-insert it and retry.
+        if chunks[j] == 0:
+            j += 1
+            while chunks[j] == 0:
+                j += 1
+            heapq.heappush(heap, (chunks[i] + chunks[j], i, j))
+            continue
+        elif chunks[i] + chunks[j] != width:
+            heapq.heappush(heap, (chunks[i] + chunks[j], i, j))
+            continue
+        # Merge
+        assert chunks[i] != 0
+        chunks[i] = 0   # mark deleted
+        chunks[j] = width
+        nmerges -= 1
+
+    return tuple(filter(None, chunks))
+
+
+def find_merge_rechunk(old_chunks, new_chunks, block_size_limit):
+    """
+    Find an intermediate rechunk that would merge some adjacent blocks
+    together in order to get us nearer the *new_chunks* target, without
+    violating the *block_size_limit* (in number of elements).
+    """
+    ndim = len(old_chunks)
+
+    old_largest_width = [max(c) for c in old_chunks]
+    new_largest_width = [max(c) for c in new_chunks]
+
+    graph_size_effect = {
+        dim: len(nc) / len(oc)
+        for dim, (oc, nc) in enumerate(zip(old_chunks, new_chunks))
+    }
+
+    block_size_effect = {
+        dim: new_largest_width[dim] / old_largest_width[dim]
+        for dim in range(ndim)
+    }
+
+    # Our goal is to reduce the number of nodes in the rechunk graph
+    # by merging some adjacent chunks, so consider dimensions where we can
+    # reduce the # of chunks
+    merge_candidates = [dim for dim in range(ndim)
+                        if graph_size_effect[dim] <= 1.0]
+
+    # Merging along each dimension reduces the graph size by a certain factor
+    # and increases memory largest block size by a certain factor.
+    # We want to optimize the graph size while staying below the given
+    # block_size_limit.  This is in effect a knapsack problem, except with
+    # multiplicative values and weights.  Just use a greedy algorithm
+    # by trying dimensions in decreasing value / weight order.
+    sorted_candidates = sorted(
+        merge_candidates,
+        key=lambda k: np.log(graph_size_effect[k]) / np.log(block_size_effect[k]),
+    )
+
+    largest_block_size = reduce(mul, old_largest_width)
+
+    chunks = list(old_chunks)
+    memory_limit_hit = False
+
+    for dim in sorted_candidates:
+        # Examine this dimension for possible graph reduction
+        new_largest_block_size = (
+            largest_block_size * new_largest_width[dim] // old_largest_width[dim])
+        if new_largest_block_size <= block_size_limit:
+            # Full replacement by new chunks is possible
+            chunks[dim] = new_chunks[dim]
+            largest_block_size = new_largest_block_size
+        else:
+            # Try a partial rechunk, dividing the new chunks into
+            # smaller pieces
+            largest_width = old_largest_width[dim]
+            chunk_limit = int(block_size_limit * largest_width / largest_block_size)
+            c = divide_to_width(new_chunks[dim], chunk_limit)
+            if len(c) <= len(old_chunks[dim]):
+                # We manage to reduce the number of blocks, so do it
+                chunks[dim] = c
+                largest_block_size = largest_block_size * max(c) // largest_width
+
+            memory_limit_hit = True
+
+    assert largest_block_size == _largest_block_size(chunks)
+    assert largest_block_size <= block_size_limit
+    return tuple(chunks), memory_limit_hit
+
+
+def find_split_rechunk(old_chunks, new_chunks, graph_size_limit):
+    """
+    Find an intermediate rechunk that would split some chunks to
+    get us nearer *new_chunks*, without violating the *graph_size_limit*.
+    """
+    ndim = len(old_chunks)
+
+    chunks = list(old_chunks)
+
+    for dim in range(ndim):
+        graph_size = estimate_graph_size(chunks, new_chunks)
+        if graph_size > graph_size_limit:
+            break
+        if len(old_chunks[dim]) > len(new_chunks[dim]):
+            # It's not interesting to split
+            continue
+        # Merge the new chunks so as to stay within the graph size budget
+        max_number = int(len(old_chunks[dim]) * graph_size_limit / graph_size)
+        c = merge_to_number(new_chunks[dim], max_number)
+        assert len(c) <= max_number
+        # Consider the merge successful if its result has a greater length
+        # and smaller max width than the old chunks
+        if len(c) >= len(old_chunks[dim]) and max(c) <= max(old_chunks[dim]):
+            chunks[dim] = c
+
+    return tuple(chunks)
+
+
+def plan_rechunk(old_chunks, new_chunks, itemsize,
+                 threshold=DEFAULT_THRESHOLD,
+                 block_size_limit=DEFAULT_BLOCK_SIZE_LIMIT):
+    """ Plan an iterative rechunking from *old_chunks* to *new_chunks*.
+    The plan aims to minimize the rechunk graph size.
+
+    Parameters
+    ----------
+    itemsize: int
+        The item size of the array
+    threshold: int
+        The graph growth factor under which we don't bother
+        introducing an intermediate step
+    block_size_limit: int
+        The maximum block size (in bytes) we want to produce during an
+        intermediate step
+    """
+    ndim = len(new_chunks)
+    steps = []
+    if ndim <= 1 or not all(new_chunks):
+        # Trivial array => no need for an intermediate
+        return steps + [new_chunks]
+
+    # Make it a number ef elements
+    block_size_limit /= itemsize
+
+    # Fix block_size_limit if too small for either old_chunks or new_chunks
+    largest_old_block = _largest_block_size(old_chunks)
+    largest_new_block = _largest_block_size(new_chunks)
+    block_size_limit = max([block_size_limit,
+                            largest_old_block,
+                            largest_new_block,
+                            ])
+
+    # The graph size above which to optimize
+    graph_size_threshold = threshold * (_number_of_blocks(old_chunks) +
+                                        _number_of_blocks(new_chunks))
+
+    current_chunks = old_chunks
+    first_pass = True
+
+    while True:
+        graph_size = estimate_graph_size(current_chunks, new_chunks)
+        if graph_size < graph_size_threshold:
+            break
+
+        if first_pass:
+            chunks = current_chunks
+        else:
+            # We hit the block_size_limit in a previous merge pass =>
+            # accept a significant increase in graph size in exchange for
+            # 1) getting nearer the goal 2) reducing the largest block size
+            # to make place for the following merge.
+            # To see this pass in action, make the block_size_limit very small.
+            chunks = find_split_rechunk(current_chunks, new_chunks,
+                                        graph_size * threshold)
+        chunks, memory_limit_hit = find_merge_rechunk(chunks, new_chunks,
+                                                      block_size_limit)
+        if chunks == current_chunks or chunks == new_chunks:
+            break
+        steps.append(chunks)
+        current_chunks = chunks
+        if not memory_limit_hit:
+            break
+        first_pass = False
+
+    return steps + [new_chunks]
+
+
+def _compute_rechunk(x, chunks):
+    """ Compute the rechunk of *x* to the given *chunks*.
+    """
+    ndim = x.ndim
     crossed = intersect_chunks(x.chunks, chunks)
     x2 = dict()
     intermediates = dict()
@@ -236,3 +502,76 @@ def rechunk(x, chunks):
         x2[key] = (concatenate3, rec_cat_arg)
     x2 = merge(x.dask, x2, intermediates)
     return Array(x2, temp_name, chunks, dtype=x.dtype)
+
+
+class _PrettyBlocks(object):
+
+    def __init__(self, blocks):
+        self.blocks = blocks
+
+    def __str__(self):
+        runs = []
+        run = []
+        repeats = 0
+        for c in self.blocks:
+            if run and run[-1] == c:
+                if repeats == 0 and len(run) > 1:
+                    runs.append((None, run[:-1]))
+                    run = run[-1:]
+                repeats += 1
+            else:
+                if repeats > 0:
+                    assert len(run) == 1
+                    runs.append((repeats + 1, run[-1]))
+                    run = []
+                    repeats = 0
+                run.append(c)
+        if run:
+            if repeats == 0:
+                runs.append((None, run))
+            else:
+                assert len(run) == 1
+                runs.append((repeats + 1, run[-1]))
+
+        parts = []
+        for repeats, run in runs:
+            if repeats is None:
+                parts.append(str(run))
+            else:
+                parts.append("%d*[%s]" % (repeats, run))
+        return " | ".join(parts)
+
+    __repr__ = __str__
+
+
+def format_blocks(blocks):
+    """
+    Pretty-format *blocks*.
+
+    >>> format_blocks((10, 10, 10))
+    3*[10]
+    >>> format_blocks((2, 3, 4))
+    [2, 3, 4]
+    >>> format_blocks((10, 10, 5, 6, 2, 2, 2, 7))
+    2*[10] | [5, 6] | 3*[2] | [7]
+    """
+    assert (isinstance(blocks, tuple) and
+            all(isinstance(x, int) for x in blocks))
+    return _PrettyBlocks(blocks)
+
+
+def format_chunks(chunks):
+    """
+    >>> format_chunks((10 * (3,), 3 * (10,)))
+    (10*[3], 3*[10])
+    """
+    assert isinstance(chunks, tuple)
+    return tuple(format_blocks(c) for c in chunks)
+
+
+def format_plan(plan):
+    """
+    >>> format_plan([((10, 10, 10), (15, 15)), ((30,), (10, 10, 10))])
+    [(3*[10], 2*[15]), ([30], 3*[10])]
+    """
+    return [format_chunks(c) for c in plan]


Test patch:
diff --git a/dask/array/tests/test_optimization.py b/dask/array/tests/test_optimization.py
--- a/dask/array/tests/test_optimization.py
+++ b/dask/array/tests/test_optimization.py
@@ -1,7 +1,6 @@
 import pytest
 pytest.importorskip('numpy')
 
-import dask.array as da
 from dask.optimize import fuse
 from dask.array.optimization import (getitem, optimize, optimize_slices,
                                      fuse_slice)
@@ -117,13 +116,6 @@ def test_hard_fuse_slice_cases():
     assert optimize_slices(dsk) == {'x': (getarray, 'x', (None, 5))}
 
 
-def test_dont_fuse_different_slices():
-    x = da.random.random(size=(10, 10), chunks=(10, 1))
-    y = x.rechunk((1, 10))
-    dsk = optimize(y.dask, y._keys())
-    assert len(dsk) > 100
-
-
 def test_dont_fuse_fancy_indexing_in_getarray_nofancy():
     dsk = {'a': (getitem, (getarray_nofancy, 'x', (slice(10, 20, None), slice(100, 200, None))),
                  ([1, 3], slice(50, 60, None)))}
diff --git a/dask/array/tests/test_rechunk.py b/dask/array/tests/test_rechunk.py
--- a/dask/array/tests/test_rechunk.py
+++ b/dask/array/tests/test_rechunk.py
@@ -5,6 +5,7 @@
 import numpy as np
 from dask.array.rechunk import intersect_chunks, rechunk, normalize_chunks
 from dask.array.rechunk import cumdims_label, _breakpoints, _intersect_1d
+from dask.array.rechunk import plan_rechunk, divide_to_width, merge_to_number
 import dask.array as da
 
 
@@ -188,3 +189,173 @@ def test_rechunk_intermediates():
     x = da.random.normal(10, 0.1, (10, 10), chunks=(10, 1))
     y = x.rechunk((1, 10))
     assert len(y.dask) > 30
+
+
+def test_divide_to_width():
+    chunks = divide_to_width((8, 9, 10), 10)
+    assert chunks == (8, 9, 10)
+    chunks = divide_to_width((8, 2, 9, 10, 11, 12), 4)
+    # Note how 9 gives (3, 3, 3), not (4, 4, 1) or whatever
+    assert chunks == (4, 4,
+                      2,
+                      3, 3, 3,
+                      3, 3, 4,
+                      3, 4, 4,
+                      4, 4, 4,
+                      )
+
+
+def test_merge_to_number():
+    chunks = merge_to_number((10,) * 4, 5)
+    assert chunks == (10, 10, 10, 10)
+    chunks = merge_to_number((10,) * 4, 4)
+    assert chunks == (10, 10, 10, 10)
+    chunks = merge_to_number((10,) * 4, 3)
+    assert chunks == (20, 10, 10)
+    chunks = merge_to_number((10,) * 4, 2)
+    assert chunks == (20, 20)
+    chunks = merge_to_number((10,) * 4, 1)
+    assert chunks == (40,)
+
+    chunks = merge_to_number((10,) * 10, 2)
+    assert chunks == (50,) * 2
+    chunks = merge_to_number((10,) * 10, 3)
+    assert chunks == (40, 30, 30)
+
+    chunks = merge_to_number((5, 1, 1, 15, 10), 4)
+    assert chunks == (5, 2, 15, 10)
+    chunks = merge_to_number((5, 1, 1, 15, 10), 3)
+    assert chunks == (7, 15, 10)
+    chunks = merge_to_number((5, 1, 1, 15, 10), 2)
+    assert chunks == (22, 10)
+    chunks = merge_to_number((5, 1, 1, 15, 10), 1)
+    assert chunks == (32,)
+
+    chunks = merge_to_number((1, 1, 1, 1, 3, 1, 1), 6)
+    assert chunks == (2, 1, 1, 3, 1, 1)
+    chunks = merge_to_number((1, 1, 1, 1, 3, 1, 1), 5)
+    assert chunks == (2, 2, 3, 1, 1)
+    chunks = merge_to_number((1, 1, 1, 1, 3, 1, 1), 4)
+    assert chunks == (2, 2, 3, 2)
+    chunks = merge_to_number((1, 1, 1, 1, 3, 1, 1), 3)
+    assert chunks == (4, 3, 2)
+    chunks = merge_to_number((1, 1, 1, 1, 3, 1, 1), 2)
+    assert chunks == (4, 5)
+    chunks = merge_to_number((1, 1, 1, 1, 3, 1, 1), 1)
+    assert chunks == (9,)
+
+
+def _plan(old_chunks, new_chunks, itemsize=1, block_size_limit=1e7):
+    return plan_rechunk(old_chunks, new_chunks,
+                        itemsize=itemsize,
+                        block_size_limit=block_size_limit)
+
+
+def _assert_steps(steps, expected):
+    assert len(steps) == len(expected)
+    assert steps == expected
+
+
+def test_plan_rechunk():
+    c = ((20,) * 2)   # coarse
+    f = ((2,) * 20)   # fine
+
+    # Trivial cases
+    steps = _plan((), ())
+    _assert_steps(steps, [()])
+    steps = _plan((c, ()), (f, ()))
+    _assert_steps(steps, [(f, ())])
+
+    # No intermediate required
+    steps = _plan((c,), (f,))
+    _assert_steps(steps, [(f,)])
+    steps = _plan((f,), (c,))
+    _assert_steps(steps, [(c,)])
+    steps = _plan((c, c), (f, f))
+    _assert_steps(steps, [(f, f)])
+    steps = _plan((f, f), (c, c))
+    _assert_steps(steps, [(c, c)])
+    steps = _plan((f, c), (c, c))
+    _assert_steps(steps, [(c, c)])
+
+    # An intermediate is used to reduce graph size
+    steps = _plan((f, c), (c, f))
+    _assert_steps(steps, [(c, c), (c, f)])
+
+    steps = _plan((c + c, c + f), (f + f, c + c))
+    _assert_steps(steps, [(c + c, c + c), (f + f, c + c)])
+
+    # Just at the memory limit => an intermediate is used
+    steps = _plan((f, c), (c, f), block_size_limit=400)
+    _assert_steps(steps, [(c, c), (c, f)])
+
+    # Hitting the memory limit => partial merge
+    m = ((10,) * 4)   # mid
+
+    steps = _plan((f, c), (c, f), block_size_limit=399)
+    _assert_steps(steps, [(m, c), (c, f)])
+
+    steps2 = _plan((f, c), (c, f), block_size_limit=3999, itemsize=10)
+    _assert_steps(steps2, steps)
+
+    # Memory limit too low => no intermediate
+    steps = _plan((f, c), (c, f), block_size_limit=40)
+    _assert_steps(steps, [(c, f)])
+
+    # Larger problem size => more intermediates
+    c = ((1000,) * 2)   # coarse
+    f = ((2,) * 1000)   # fine
+
+    steps = _plan((f, c), (c, f), block_size_limit=99999)
+    assert len(steps) == 3
+    assert steps[-1] == (c, f)
+    for i in range(len(steps) - 1):
+        prev = steps[i]
+        succ = steps[i + 1]
+        # Merging on the first dim, splitting on the second dim
+        assert len(succ[0]) <= len(prev[0]) / 2.0
+        assert len(succ[1]) >= len(prev[1]) * 2.0
+
+
+def test_plan_rechunk_5d():
+    # 5d problem
+    c = ((10,) * 1)    # coarse
+    f = ((1,) * 10)    # fine
+
+    steps = _plan((c, c, c, c, c), (f, f, f, f, f))
+    _assert_steps(steps, [(f, f, f, f, f)])
+    steps = _plan((f, f, f, f, c), (c, c, c, f, f))
+    _assert_steps(steps, [(c, c, c, f, c), (c, c, c, f, f)])
+    # Only 1 dim can be merged at first
+    steps = _plan((c, c, f, f, c), (c, c, c, f, f), block_size_limit=2e4)
+    _assert_steps(steps, [(c, c, c, f, c), (c, c, c, f, f)])
+
+
+def test_plan_rechunk_heterogenous():
+    c = ((10,) * 1)    # coarse
+    f = ((1,) * 10)    # fine
+    cf = c + f
+    cc = c + c
+    ff = f + f
+    fc = f + c
+
+    # No intermediate required
+    steps = _plan((cc, cf), (ff, ff))
+    _assert_steps(steps, [(ff, ff)])
+    steps = _plan((cf, fc), (ff, cf))
+    _assert_steps(steps, [(ff, cf)])
+
+    # An intermediate is used to reduce graph size
+    steps = _plan((cc, cf), (ff, cc))
+    _assert_steps(steps, [(cc, cc), (ff, cc)])
+
+    steps = _plan((cc, cf, cc), (ff, cc, cf))
+    _assert_steps(steps, [(cc, cc, cc), (ff, cc, cf)])
+
+    # Imposing a memory limit => the first intermediate is constrained:
+    #  * cc -> ff would increase the graph size: no
+    #  * ff -> cf would increase the block size too much: no
+    #  * cf -> cc fits the bill (graph size /= 10, block size neutral)
+    #  * cf -> fc also fits the bill (graph size and block size neutral)
+    steps = _plan((cc, ff, cf), (ff, cf, cc), block_size_limit=100)
+    _assert_steps(steps, [(cc, ff, cc), (ff, cf, cc)])