Fuse operations with different numbers of tasks

[tomwhite]

Fixes #284

This adds the necessary logic to the fuse code to fuse two operations that have a different number of tasks, but fusion must be controlled manually using the always_fuse argument to the optimization functions. This is fine for experimenting with the optimizer, but we'll probably want to add better heuristics about when to fuse these types of operations automatically in the future.

[TomNicholas]

Is this split_every argument a temporary implementation detail whilst we figure out fusing heuristics? It would be nice to relate the meaning of this argument back to the discussion in #284 (as I currently don't quite understand what it means)

[tomwhite]

The split_every argument is what we referred to as "fan-in" in #284 - the number of chunks read by one task doing the reduction step. It's the same in Dask and can be a dictionary indicating the number of chunks to read in each dimension.

I hope it is something that we can get better heuristics for (or at least good defaults) - possibly by measuring trade offs like they did in the Primula paper, see #331.

[TomNicholas]

This is very exciting!

[tomwhite]

I've generated a few plan visualizations to give an idea of what the code does.

Running this cut-down quadratic means example

https://github.com/tomwhite/cubed/blob/1762d3c685e7a1333c2809d2d9e419ff6508c5ce/cubed/tests/test_core.py#L527-L537

gives the following plan using the old optimization algorithm (for t_length=500):

The reduce rounds (for the mean operations) have 50, 4, and 1 task(s).

If we use the new optimization algorithm then it will fuse the first mean operation that has two inputs into one:

    m = xp.mean(uv, axis=0, split_every=10, use_new_impl=True)

    m.visualize(
        filename=tmp_path / "quad_means",
        optimize_function=multiple_inputs_optimize_dag
    )

Notice that the number of tasks has gone down from 160 to 59, and the amount of intermediate data stored has gone down from 16.8GB to 7.7GB. The reduce rounds have 50, 5, and 1 task(s) - this is controlled by the split_every=10 argument.

If we now fuse the first two operations (purple boxes with rounded corners) then we get make the plan even smaller. Note in the code below we have to explicitly ask for this operation to be fused.

This is using the logic from this PR which is able to fuse operations with a different number of tasks. The first operation (marked __mul__) has 50 tasks, and the second (mean) has 5.

    m = xp.mean(uv, axis=0, split_every=10, use_new_impl=True)

    m.visualize(
        filename=tmp_path / "quad_means",
        optimize_function=partial(multiple_inputs_optimize_dag, always_fuse=["op-008"]),
    )

There are now only two reduce rounds: with 5 and 1 task(s). And a total of only 8 tasks, and 166.8MB of intermediate data - a significant saving.

[tomwhite]

The mypy failures suggested to me that num_input_blocks belongs to BlockwiseSpec, not PrimitiveOperation, since it doesn't apply to a rechunk operation, for example. This change means that num_input_blocks also sits next to function_nargs, which it is related to, so that seems better.

[TomNicholas]

Wow that's a big difference. Thanks for writing this out @tomwhite.

I was about to try and estimate the projected speedup here (assuming IO dominates the execution time), but without being able to see the number of tasks in each stage it's not trivial to reason about. (Half the number of zarr stores to be written, but not necessarily twice as fast because do array-008 and array-007 have the same amount of parallelism that they did before optimization?) Not seeing this information in the visualization also made coiled/feedback#271 (comment) less clear - perhaps this information should be included in the graph visualization itself?

[tomwhite]

It's definitely hard to reason about, even with the number of tasks in each stage. (BTW I have listed the number of tasks in each stage in the comment above.) The total number of stages is probably the biggest factor in how long the computation takes. I plan to run some benchmarks to measure the performance of these workloads.

But I agree that it would be useful to put the number of tasks in the visualization. It does appear in the tooltip for nodes, but for some reason that doesn't appear when the SVG is embedded in another page.