Repository: dask/dask

Issue:
Tune dask.array.optimize
In some pathological (yet also real) cases, graph optimization can cost about as much as computation:

I've been benchmarking against this computation

```python
import dask.array as da
n = 12
x = da.random.random((2**n, 2**n), chunks=(2**n, 1))   # 4096 x 4096, 4096 x 1
for i in range(1, n):                                  # 4096 x 4096, 2048 x 2
        x = x.rechunk((2**(n - i), 2**(i + 1)))            # 4096 x 4096, 1024 x 4

y = x.sum()

>>> %time y.compute()
CPU times: user 39 s, sys: 2.45 s, total: 41.5 s
Wall time: 39.1 s

>>> %prun  y._optimize(y.dask, y._keys())
```

```

         29816630 function calls (28569410 primitive calls) in 22.570 seconds

   Ordered by: internal time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
   604153    3.733    0.000    8.800    0.000 core.py:189(get_dependencies)
  6240188    2.504    0.000    2.843    0.000 core.py:27(istask)
  1312754    2.414    0.000    3.091    0.000 core.py:154(_deps)
        3    2.146    0.715    2.146    0.715 core.py:284(<listcomp>)
   299004    1.456    0.000    3.326    0.000 rewrite.py:375(_match)
909302/77823    1.331    0.000    2.017    0.000 optimize.py:211(functions_of)
  2009082    0.625    0.000    0.625    0.000 {built-in method isinstance}
        3    0.618    0.206    2.821    0.940 core.py:275(reverse_dict)
299004/77823    0.553    0.000    5.650    0.000 rewrite.py:363(_bottom_up)
   299004    0.480    0.000    4.206    0.000 rewrite.py:283(iter_matches)
   348156    0.426    0.000    0.426    0.000 rewrite.py:50(__init__)
   874487    0.406    0.000    1.014    0.000 rewrite.py:8(head)
  2303972    0.403    0.000    0.403    0.000 {method 'pop' of 'list' objects}
219136/51200    0.400    0.000    0.990    0.000 core.py:291(subs)
  2250724    0.391    0.000    0.391    0.000 {method 'extend' of 'list' objects}
        2    0.376    0.188    4.363    2.182 optimize.py:117(inline)
    49152    0.341    0.000    0.673    0.000 core.py:307(<listcomp>)
  4671438    0.339    0.000    0.339    0.000 {built-in method callable}
   299004    0.270    0.000    4.477    0.000 rewrite.py:304(_rewrite)
   794616    0.257    0.000    1.211    0.000 rewrite.py:81(current)
        1    0.204    0.204    1.261    1.261 optimize.py:43(fuse)
   745465    0.186    0.000    0.186    0.000 {method 'get' of 'dict' objects}
  1484786    0.168    0.000    0.168    0.000 {method 'append' of 'list' objects}
        2    0.133    0.066    0.752    0.376 core.py:321(_toposort)
        1    0.127    0.127    1.249    1.249 optimize.py:16(cull)
   348156    0.126    0.000    0.126    0.000 {method 'pop' of 'collections.deque' objects}
    79871    0.125    0.000    0.185    0.000 optimization.py:32(is_full_slice)
   745464    0.122    0.000    0.122    0.000 rewrite.py:119(edges)
        1    0.122    0.122    2.202    2.202 optimize.py:200(<listcomp>)
   595958    0.113    0.000    0.113    0.000 {method 'add' of 'set' objects}
        1    0.106    0.106    5.805    5.805 {method 'update' of 'dict' objects}
225278/221182    0.106    0.000    3.578    0.000 rewrite.py:365(<genexpr>)
        1    0.098    0.098    6.363    6.363 optimize.py:237(dealias)
   176126    0.093    0.000    0.175    0.000 rewrite.py:19(args)
    81919    0.086    0.000    0.086    0.000 {built-in method hasattr}
    49152    0.070    0.000    0.135    0.000 rewrite.py:70(next)
    79871    0.069    0.000    0.069    0.000 {built-in method hash}
47103/24575    0.067    0.000    1.377    0.000 rewrite.py:367(<listcomp>)
    79872    0.066    0.000    1.486    0.000 optimize.py:287(<genexpr>)
    79872    0.065    0.000    1.344    0.000 optimize.py:278(<genexpr>)
    79871    0.063    0.000    0.132    0.000 core.py:9(ishashable)
    77824    0.061    0.000    1.311    0.000 optimize.py:197(<genexpr>)
    49152    0.058    0.000    0.084    0.000 rewrite.py:62(copy)
        1    0.058    0.058    6.676    6.676 optimization.py:49(remove_full_slices)
        1    0.054    0.054   22.570   22.570 <string>:1(<module>)
    77823    0.049    0.000    5.699    0.000 rewrite.py:315(rewrite)
    79872    0.044    0.000    0.044    0.000 optimize.py:40(<genexpr>)
    28675    0.042    0.000    0.042    0.000 {method 'values' of 'dict' objects}
        1    0.039    0.039    7.491    7.491 optimize.py:168(inline_functions)
    98304    0.038    0.000    0.038    0.000 optimization.py:46(<genexpr>)
        1    0.037    0.037    0.169    0.169 optimize.py:281(<genexpr>)
        1    0.036    0.036    0.221    0.221 optimization.py:69(<genexpr>)
    79872    0.033    0.000    0.033    0.000 optimization.py:71(<genexpr>)
        1    0.033    0.033   22.515   22.515 optimization.py:14(optimize)
    49152    0.032    0.000    0.780    0.000 core.py:314(<listcomp>)
    79871    0.031    0.000    0.117    0.000 optimize.py:231(unwrap_partial)
    49154    0.026    0.000    0.026    0.000 optimize.py:145(<genexpr>)
    49160    0.018    0.000    0.028    0.000 core.py:246(flatten)
        2    0.011    0.006    0.011    0.006 optimize.py:282(<genexpr>)
    49152    0.011    0.000    0.011    0.000 {method 'extend' of 'collections.deque' objects}
        1    0.009    0.009    0.009    0.009 {method 'copy' of 'dict' objects}
    49152    0.009    0.000    0.037    0.000 {built-in method all}
    77823    0.008    0.000    0.008    0.000 {method 'issubset' of 'set' objects}
    79872    0.008    0.000    0.008    0.000 {built-in method len}
```

I also recommend looking at this with snakeviz

```python
%load_ext snakeviz
%snakeviz y._optimize(y.dask, y._keys())
```


Discussion:
eriknw: Looking now.
eriknw: Great.  I think the tradeoff is worth it for these functions.  Not passing an optional `dependencies` to non-critical functions also seems reasonable as a way to avoid complications.

A downside is that now I don't know the next "obvious" thing to do to optimize these further!  What do profiles look like with these changes?
mrocklin: I recommend pulling out snakeviz and running the following:

```python
In [1]: import dask.array as da
In [2]: n = 12
In [3]: x = da.random.random((2**n, 2**n), chunks=(2**n, 1))   # 4096 x 4096, 4096 x 1
In [4]: for i in range(1, n):                                  # 4096 x 4096, 2048 x 2
            x = x.rechunk((2**(n - i), 2**(i + 1)))            # 4096 x 4096, 1024 x 4
   ...:     
In [5]: y = x.sum()
In [6]: %load_ext snakeviz
In [7]: %snakeviz y._optimize(y.dask, y._keys())
```

Inline_functions takes up about half of the optimization time now.  Really though, I think we're at the point of diminishing returns.  This computation takes around a minute-plus to run on a distributed machine.


mrocklin: If you are looking to do some optimization work we're close to the point where we might want to Cythonize a couple of functions within the distributed scheduler.  If you're interested let me know and I'll produce a benchmark script for you to play with.
eriknw: Okay.  I'm up for Cythonizing some functions.  I'm taking a vacation in a few days, so it'll probably be a few weeks until I can do the work.  Thanks for letting me know this would be beneficial to do.

Patch:
diff --git a/dask/array/optimization.py b/dask/array/optimization.py
index 8d7224fd4..6a201ead5 100644
--- a/dask/array/optimization.py
+++ b/dask/array/optimization.py
@@ -21,10 +21,11 @@ def optimize(dsk, keys, **kwargs):
     keys = list(flatten(keys))
     fast_functions = kwargs.get('fast_functions',
                              set([getarray, np.transpose]))
-    dsk2 = cull(dsk, keys)
-    dsk4 = fuse(dsk2, keys)
+    dsk2, dependencies = cull(dsk, keys)
+    dsk4, dependencies = fuse(dsk2, keys, dependencies)
     dsk5 = optimize_slices(dsk4)
-    dsk6 = inline_functions(dsk5, keys, fast_functions=fast_functions)
+    dsk6 = inline_functions(dsk5, keys, fast_functions=fast_functions,
+            dependencies=dependencies)
     return dsk6
 
 
diff --git a/dask/async.py b/dask/async.py
index 4f9d92b1c..2d8fbb87b 100644
--- a/dask/async.py
+++ b/dask/async.py
@@ -434,7 +434,7 @@ def get_async(apply_async, num_workers, dsk, result, cache=None,
     for f in start_cbs:
         f(dsk)
 
-    dsk = cull(dsk, list(results))
+    dsk, dependencies = cull(dsk, list(results))
 
     keyorder = order(dsk)
 
diff --git a/dask/bag/core.py b/dask/bag/core.py
index 3348be39a..268a17931 100644
--- a/dask/bag/core.py
+++ b/dask/bag/core.py
@@ -64,7 +64,7 @@ def lazify(dsk):
     return valmap(lazify_task, dsk)
 
 
-def inline_singleton_lists(dsk):
+def inline_singleton_lists(dsk, dependencies=None):
     """ Inline lists that are only used once
 
     >>> d = {'b': (list, 'a'),
@@ -74,8 +74,8 @@ def inline_singleton_lists(dsk):
 
     Pairs nicely with lazify afterwards
     """
-
-    dependencies = dict((k, get_dependencies(dsk, k)) for k in dsk)
+    if dependencies is None:
+        dependencies = dict((k, get_dependencies(dsk, k)) for k in dsk)
     dependents = reverse_dict(dependencies)
 
     keys = [k for k, v in dsk.items()
@@ -85,9 +85,9 @@ def inline_singleton_lists(dsk):
 
 def optimize(dsk, keys, **kwargs):
     """ Optimize a dask from a dask.bag """
-    dsk2 = cull(dsk, keys)
-    dsk3 = fuse(dsk2, keys)
-    dsk4 = inline_singleton_lists(dsk3)
+    dsk2, dependencies = cull(dsk, keys)
+    dsk3, dependencies = fuse(dsk2, keys, dependencies)
+    dsk4 = inline_singleton_lists(dsk3, dependencies)
     dsk5 = lazify(dsk4)
     return dsk5
 
diff --git a/dask/core.py b/dask/core.py
index 275748384..ead939fd0 100644
--- a/dask/core.py
+++ b/dask/core.py
@@ -319,7 +319,7 @@ def subs(task, key, val):
     return task[:1] + tuple(newargs)
 
 
-def _toposort(dsk, keys=None, returncycle=False):
+def _toposort(dsk, keys=None, returncycle=False, dependencies=None):
     # Stack-based depth-first search traversal.  This is based on Tarjan's
     # method for topological sorting (see wikipedia for pseudocode)
     if keys is None:
@@ -340,6 +340,9 @@ def _toposort(dsk, keys=None, returncycle=False):
     # that has already been added to `completed`.
     seen = set()
 
+    if dependencies is None:
+        dependencies = dict((k, get_dependencies(dsk, k)) for k in dsk)
+
     for key in keys:
         if key in completed:
             continue
@@ -355,7 +358,7 @@ def _toposort(dsk, keys=None, returncycle=False):
 
             # Add direct descendants of cur to nodes stack
             next_nodes = []
-            for nxt in get_dependencies(dsk, cur):
+            for nxt in dependencies[cur]:
                 if nxt not in completed:
                     if nxt in seen:
                         # Cycle detected!
@@ -385,9 +388,9 @@ def _toposort(dsk, keys=None, returncycle=False):
     return ordered
 
 
-def toposort(dsk):
+def toposort(dsk, dependencies=None):
     """ Return a list of keys of dask sorted in topological order."""
-    return _toposort(dsk)
+    return _toposort(dsk, dependencies=dependencies)
 
 
 def getcycle(d, keys):
diff --git a/dask/dataframe/optimize.py b/dask/dataframe/optimize.py
index 0ae46d54c..d23a7ff2e 100644
--- a/dask/dataframe/optimize.py
+++ b/dask/dataframe/optimize.py
@@ -15,9 +15,9 @@ def fuse_castra_index(dsk):
 
 def optimize(dsk, keys, **kwargs):
     if isinstance(keys, list):
-        dsk2 = cull(dsk, list(core.flatten(keys)))
+        dsk2, dependencies = cull(dsk, list(core.flatten(keys)))
     else:
-        dsk2 = cull(dsk, [keys])
+        dsk2, dependencies = cull(dsk, [keys])
     try:
         from castra import Castra
         dsk3 = fuse_getitem(dsk2, Castra.load_partition, 3)
@@ -25,5 +25,5 @@ def optimize(dsk, keys, **kwargs):
     except ImportError:
         dsk4 = dsk2
     dsk5 = fuse_getitem(dsk4, dataframe_from_ctable, 3)
-    dsk6 = cull(dsk5, keys)
+    dsk6, _ = cull(dsk5, keys)
     return dsk6
diff --git a/dask/dataframe/shuffle.py b/dask/dataframe/shuffle.py
index 0b4c49ab0..fb797a503 100644
--- a/dask/dataframe/shuffle.py
+++ b/dask/dataframe/shuffle.py
@@ -144,7 +144,7 @@ def set_partition(df, index, divisions, compute=False, drop=True, **kwargs):
         dsk.update(index.dask)
 
     if compute:
-        dsk = cull(dsk, list(dsk4.keys()))
+        dsk, _ = cull(dsk, list(dsk4.keys()))
 
     return DataFrame(dsk, name, metadata, divisions)
 
diff --git a/dask/multiprocessing.py b/dask/multiprocessing.py
index dce045998..a2f6b996f 100644
--- a/dask/multiprocessing.py
+++ b/dask/multiprocessing.py
@@ -68,8 +68,9 @@ def get(dsk, keys, optimizations=[], num_workers=None,
                                           func_loads=func_loads)
 
     # Optimize Dask
-    dsk2 = fuse(dsk, keys)
-    dsk3 = pipe(dsk2, partial(cull, keys=keys), *optimizations)
+    dsk2, dependencies = cull(dsk, keys)
+    dsk3, dependencies = fuse(dsk2, keys, dependencies)
+    dsk4 = pipe(dsk3, *optimizations)
 
     try:
         # Run
diff --git a/dask/optimize.py b/dask/optimize.py
index e557ccd73..c34b6e4c7 100644
--- a/dask/optimize.py
+++ b/dask/optimize.py
@@ -22,49 +22,80 @@ def cull(dsk, keys):
     Examples
     --------
     >>> d = {'x': 1, 'y': (inc, 'x'), 'out': (add, 'x', 10)}
-    >>> cull(d, 'out')  # doctest: +SKIP
+    >>> dsk, dependencies = cull(d, 'out')  # doctest: +SKIP
+    >>> dsk  # doctest: +SKIP
     {'x': 1, 'out': (add, 'x', 10)}
+    >>> dependencies  # doctest: +SKIP
+    {'x': set(), 'out': set(['x'])}
+
+    Returns
+    -------
+    dsk: culled dask graph
+    dependencies: Dict mapping {key: [deps]}.  Useful side effect to accelerate
+        other optimizations, notably fuse.
     """
     if not isinstance(keys, (list, set)):
         keys = [keys]
-    nxt = set(flatten(keys))
-    seen = nxt
-    while nxt:
-        cur = nxt
-        nxt = set()
-        for item in cur:
-            for dep in get_dependencies(dsk, item):
-                if dep not in seen:
-                    nxt.add(dep)
-        seen.update(nxt)
-    return dict((k, v) for k, v in dsk.items() if k in seen)
-
-
-def fuse(dsk, keys=None):
-    """ Return new dask with linear sequence of tasks fused together.
+    out = dict()
+    seen = set()
+    dependencies = dict()
+    stack = list(set(flatten(keys)))
+    while stack:
+        key = stack.pop()
+        out[key] = dsk[key]
+        deps = get_dependencies(dsk, key, as_list=True)  # fuse needs lists
+        dependencies[key] = deps
+        unseen = [d for d in deps if d not in seen]
+        stack.extend(unseen)
+        seen.update(unseen)
+    return out, dependencies
+
+
+def fuse(dsk, keys=None, dependencies=None):
+    """ Return new dask graph with linear sequence of tasks fused together.
 
     If specified, the keys in ``keys`` keyword argument are *not* fused.
+    Supply ``dependencies`` from output of ``cull`` if available to avoid
+    recomputing dependencies.
 
-    This may be used as an optimization step.
+    Parameters
+    ----------
+    dsk: dict
+    keys: list
+    dependencies: dict, optional
+        {key: [list-of-keys]}.  Must be a list to provide count of each key
+        This optional input often comes from ``cull``
 
     Examples
     --------
     >>> d = {'a': 1, 'b': (inc, 'a'), 'c': (inc, 'b')}
-    >>> fuse(d)  # doctest: +SKIP
+    >>> dsk, dependencies = fuse(d)
+    >>> dsk # doctest: +SKIP
     {'c': (inc, (inc, 1))}
-    >>> fuse(d, keys=['b'])  # doctest: +SKIP
+    >>> dsk, dependencies = fuse(d, keys=['b'])
+    >>> dsk  # doctest: +SKIP
     {'b': (inc, 1), 'c': (inc, 'b')}
+
+    Returns
+    -------
+    dsk: output graph with keys fused
+    dependencies: dict mapping dependencies after fusion.  Useful side effect
+        to accelerate other downstream optimizations.
     """
     if keys is not None and not isinstance(keys, set):
         if not isinstance(keys, list):
             keys = [keys]
         keys = set(flatten(keys))
 
+    if dependencies is None:
+        dependencies = dict((key, get_dependencies(dsk, key, as_list=True))
+                            for key in dsk)
+
     # locate all members of linear chains
     child2parent = {}
     unfusible = set()
     for parent in dsk:
-        deps = get_dependencies(dsk, parent, as_list=True)
+        deps = dependencies[parent]
         has_many_children = len(deps) > 1
         for child in deps:
             if keys is not None and child in keys:
@@ -94,6 +125,8 @@ def fuse(dsk, keys=None):
             chain.append(child)
         chains.append(chain)
 
+    dependencies = dict((k, set(v)) for k, v in dependencies.items())
+
     # create a new dask with fused chains
     rv = {}
     fused = set()
@@ -102,6 +135,8 @@ def fuse(dsk, keys=None):
         val = dsk[child]
         while chain:
             parent = chain.pop()
+            dependencies[parent].update(dependencies.pop(child))
+            dependencies[parent].remove(child)
             val = subs(dsk[parent], child, val)
             fused.add(child)
             child = parent
@@ -111,10 +146,10 @@ def fuse(dsk, keys=None):
     for key, val in dsk.items():
         if key not in fused:
             rv[key] = val
-    return rv
+    return rv, dependencies
 
 
-def inline(dsk, keys=None, inline_constants=True):
+def inline(dsk, keys=None, inline_constants=True, dependencies=None):
     """ Return new dask with the given keys inlined with their values.
 
     Inlines all constants if ``inline_constants`` keyword is True.
@@ -140,13 +175,17 @@ def inline(dsk, keys=None, inline_constants=True):
     if inline_constants:
         keys.update(k for k, v in dsk.items() if not istask(v))
 
+    if dependencies is None:
+        dependencies = dict((k, get_dependencies(dsk, k)) for k in dsk)
+
     # Keys may depend on other keys, so determine replace order with toposort.
     # The values stored in `keysubs` do not include other keys.
-    replaceorder = toposort(dict((k, dsk[k]) for k in keys if k in dsk))
+    replaceorder = toposort(dict((k, dsk[k]) for k in keys if k in dsk),
+                            dependencies=dependencies)
     keysubs = {}
     for key in replaceorder:
         val = dsk[key]
-        for dep in keys & get_dependencies(dsk, key):
+        for dep in keys & dependencies[key]:
             if dep in keysubs:
                 replace = keysubs[dep]
             else:
@@ -159,13 +198,14 @@ def inline(dsk, keys=None, inline_constants=True):
     for key, val in dsk.items():
         if key in keys:
             continue
-        for item in keys & get_dependencies(dsk, key):
+        for item in keys & dependencies[key]:
             val = subs(val, item, keysubs[item])
         rv[key] = val
     return rv
 
 
-def inline_functions(dsk, output, fast_functions=None, inline_constants=False):
+def inline_functions(dsk, output, fast_functions=None, inline_constants=False,
+        dependencies=None):
     """ Inline cheap functions into larger operations
 
     Examples
@@ -194,7 +234,8 @@ def inline_functions(dsk, output, fast_functions=None, inline_constants=False):
 
     fast_functions = set(fast_functions)
 
-    dependencies = dict((k, get_dependencies(dsk, k)) for k in dsk)
+    if dependencies is None:
+        dependencies = dict((k, get_dependencies(dsk, k)) for k in dsk)
     dependents = reverse_dict(dependencies)
 
     keys = [k for k, v in dsk.items()
@@ -203,7 +244,8 @@ def inline_functions(dsk, output, fast_functions=None, inline_constants=False):
               and dependents[k]
               and k not in output]
     if keys:
-        return inline(dsk, keys, inline_constants=inline_constants)
+        return inline(dsk, keys, inline_constants=inline_constants,
+                dependencies=dependencies)
     else:
         return dsk
 
@@ -515,7 +557,7 @@ def fuse_selections(dsk, head1, head2, merge):
     ...        'y': (getitem, 'x', 'a')}
     >>> merge = lambda t1, t2: (load, t2[1], t2[2], t1[2])
     >>> dsk2 = fuse_selections(dsk, getitem, load, merge)
-    >>> cull(dsk2, 'y')
+    >>> cull(dsk2, 'y')[0]
     {'y': (<function load at ...>, 'store', 'part', 'a')}
     """
     dsk2 = dict()
@@ -548,7 +590,7 @@ def fuse_getitem(dsk, func, place):
     >>> dsk = {'x': (load, 'store', 'part', ['a', 'b']),
     ...        'y': (getitem, 'x', 'a')}
     >>> dsk2 = fuse_getitem(dsk, load, 3)  # columns in arg place 3
-    >>> cull(dsk2, 'y')
+    >>> cull(dsk2, 'y')[0]
     {'y': (<function load at ...>, 'store', 'part', 'a')}
     """
     return fuse_selections(dsk, getitem, func,
diff --git a/dask/order.py b/dask/order.py
index 60fd03787..06f602c5d 100644
--- a/dask/order.py
+++ b/dask/order.py
@@ -58,10 +58,10 @@ from __future__ import absolute_import, division, print_function
 
 from operator import add
 
-from .core import get_deps
+from .core import get_dependencies, reverse_dict, get_deps
 
 
-def order(dsk):
+def order(dsk, dependencies=None):
     """ Order nodes in dask graph
 
     The ordering will be a toposort but will also have other convenient
@@ -74,7 +74,10 @@ def order(dsk):
     >>> order(dsk)
     {'a': 2, 'c': 1, 'b': 3, 'd': 0}
     """
-    dependencies, dependents = get_deps(dsk)
+    if dependencies is None:
+        dependencies = dict((k, get_dependencies(dsk, k)) for k in dsk)
+    dependents = reverse_dict(dependencies)
+
     ndeps = ndependents(dependencies, dependents)
     maxes = child_max(dependencies, dependents, ndeps)
     return dfs(dependencies, dependents, key=maxes.get)


Test patch:
diff --git a/dask/array/tests/test_optimization.py b/dask/array/tests/test_optimization.py
index 37d02d38d..0ff08fb2f 100644
--- a/dask/array/tests/test_optimization.py
+++ b/dask/array/tests/test_optimization.py
@@ -69,14 +69,14 @@ def test_optimize_slicing():
            'e': (getarray, 'd', (slice(None, None, None),))}
 
     expected = {'e': (getarray, (range, 10), (slice(0, 5, None),))}
-    result = optimize_slices(fuse(dsk, []))
+    result = optimize_slices(fuse(dsk, [])[0])
     assert result == expected
 
     # protect output keys
     expected = {'c': (range, 10),
                 'd': (getarray, 'c', (slice(0, 5, None),)),
                 'e': 'd'}
-    result = optimize_slices(fuse(dsk, ['c', 'd', 'e']))
+    result = optimize_slices(fuse(dsk, ['c', 'd', 'e'])[0])
 
     assert result == expected
 
diff --git a/dask/tests/test_optimize.py b/dask/tests/test_optimize.py
index fb713f7f3..749a03c83 100644
--- a/dask/tests/test_optimize.py
+++ b/dask/tests/test_optimize.py
@@ -18,27 +18,27 @@ def double(x):
 def test_cull():
     # 'out' depends on 'x' and 'y', but not 'z'
     d = {'x': 1, 'y': (inc, 'x'), 'z': (inc, 'x'), 'out': (add, 'y', 10)}
-    culled = cull(d, 'out')
+    culled, dependencies = cull(d, 'out')
     assert culled == {'x': 1, 'y': (inc, 'x'), 'out': (add, 'y', 10)}
+    assert dependencies == {'x': [], 'y': ['x'], 'out': ['y']}
+
     assert cull(d, 'out') == cull(d, ['out'])
-    assert cull(d, ['out', 'z']) == d
+    assert cull(d, ['out', 'z'])[0] == d
     assert cull(d, [['out'], ['z']]) == cull(d, ['out', 'z'])
     assert raises(KeyError, lambda: cull(d, 'badkey'))
 
 
 def test_fuse():
-    assert fuse({
-        'w': (inc, 'x'),
-        'x': (inc, 'y'),
-        'y': (inc, 'z'),
-        'z': (add, 'a', 'b'),
-        'a': 1,
-        'b': 2,
-    }) == {
-        'w': (inc, (inc, (inc, (add, 'a', 'b')))),
-        'a': 1,
-        'b': 2,
-    }
+    dsk, dependencies = fuse({'w': (inc, 'x'),
+                              'x': (inc, 'y'),
+                              'y': (inc, 'z'),
+                              'z': (add, 'a', 'b'),
+                              'a': 1,
+                              'b': 2})
+    assert dsk == {'w': (inc, (inc, (inc, (add, 'a', 'b')))),
+                   'a': 1,
+                   'b': 2}
+    assert dependencies == {'a': set(), 'b': set(), 'w': set(['a', 'b'])}
     assert fuse({
         'NEW': (inc, 'y'),
         'w': (inc, 'x'),
@@ -47,13 +47,16 @@ def test_fuse():
         'z': (add, 'a', 'b'),
         'a': 1,
         'b': 2,
-    }) == {
+    }) == ({
         'NEW': (inc, 'y'),
         'w': (inc, (inc, 'y')),
         'y': (inc, (add, 'a', 'b')),
         'a': 1,
         'b': 2,
-    }
+    },
+    {'a': set(), 'b': set(), 'y': set(['a', 'b']),
+     'w': set(['y']), 'NEW': set(['y'])})
+
     assert fuse({
         'v': (inc, 'y'),
         'u': (inc, 'w'),
@@ -65,13 +68,15 @@ def test_fuse():
         'b': (inc, 'd'),
         'c': 1,
         'd': 2,
-    }) == {
+    }) == ({
         'u': (inc, (inc, (inc, 'y'))),
         'v': (inc, 'y'),
         'y': (inc, (add, 'a', 'b')),
         'a': (inc, 1),
         'b': (inc, 2),
-    }
+    },
+    {'a': set(), 'b': set(), 'y': set(['a', 'b']),
+     'v': set(['y']), 'u': set(['y'])})
     assert fuse({
         'a': (inc, 'x'),
         'b': (inc, 'x'),
@@ -79,39 +84,43 @@ def test_fuse():
         'd': (inc, 'c'),
         'x': (inc, 'y'),
         'y': 0,
-    }) == {
+    }) == ({
         'a': (inc, 'x'),
         'b': (inc, 'x'),
         'd': (inc, (inc, 'x')),
         'x': (inc, 0),
-    }
+    },
+    {'x': set(), 'd': set(['x']), 'a': set(['x']), 'b': set(['x'])})
     assert fuse({
         'a': 1,
         'b': (inc, 'a'),
         'c': (add, 'b', 'b')
-    }) == {
+    }) == ({
         'b': (inc, 1),
         'c': (add, 'b', 'b')
-    }
+        }, {'b': set(), 'c': set(['b'])})
 
 
 def test_fuse_keys():
     assert (fuse({'a': 1, 'b': (inc, 'a'), 'c': (inc, 'b')}, keys=['b'])
-            == {'b': (inc, 1), 'c': (inc, 'b')})
-    assert fuse({
+            == ({'b': (inc, 1), 'c': (inc, 'b')},
+                {'b': set(), 'c': set(['b'])}))
+    dsk, dependencies = fuse({
         'w': (inc, 'x'),
         'x': (inc, 'y'),
         'y': (inc, 'z'),
         'z': (add, 'a', 'b'),
         'a': 1,
         'b': 2,
-    }, keys=['x', 'z']) == {
-        'w': (inc, 'x'),
-        'x': (inc, (inc, 'z')),
-        'z': (add, 'a', 'b'),
-        'a': 1,
-        'b': 2,
-    }
+    }, keys=['x', 'z'])
+
+    assert dsk == {'w': (inc, 'x'),
+                   'x': (inc, (inc, 'z')),
+                   'z': (add, 'a', 'b'),
+                   'a': 1,
+                   'b': 2 }
+    assert dependencies == {'a': set(), 'b': set(), 'z': set(['a', 'b']),
+                            'x': set(['z']), 'w': set(['x'])}
 
 
 def test_inline():
@@ -341,7 +350,7 @@ def test_fuse_getitem():
     dsk = {'x': (load, 'store', 'part', ['a', 'b']),
            'y': (getitem, 'x', 'a')}
     dsk2 = fuse_getitem(dsk, load, 3)
-    dsk2 = cull(dsk2, 'y')
+    dsk2, dependencies = cull(dsk2, 'y')
     assert dsk2 == {'y': (load, 'store', 'part', 'a')}
 
 
@@ -352,5 +361,5 @@ def test_fuse_selections():
            'y': (getitem, 'x', 'a')}
     merge = lambda t1, t2: (load, t2[1], t2[2], t1[2])
     dsk2 = fuse_selections(dsk, getitem, load, merge)
-    dsk2 = cull(dsk2, 'y')
+    dsk2, dependencies = cull(dsk2, 'y')
     assert dsk2 == {'y': (load, 'store', 'part', 'a')}
diff --git a/dask/tests/test_order.py b/dask/tests/test_order.py
index d64dee715..3b5f8c3ce 100644
--- a/dask/tests/test_order.py
+++ b/dask/tests/test_order.py
@@ -1,5 +1,6 @@
 from itertools import chain
-from dask.order import dfs, child_max, ndependents, order, inc, get_deps
+from dask.order import dfs, child_max, ndependents, order, inc
+from dask.core import get_deps
 
 
 def issorted(L, reverse=False):