Repository: dask/dask
Pull request #11549 (2024-11-20 22:19:00)

Fuse data nodes in linear fusion too
- [ ] Closes #xxxx
- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`


github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/dask/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

     13 files  ± 0       13 suites  ±0   3h 26m 42s ⏱️ + 9m 4s
 13 256 tests + 1   12 194 ✅ + 1   1 062 💤 ±0  0 ❌ ±0 
138 328 runs  +13  119 332 ✅ +13  18 996 💤 ±0  0 ❌ ±0 

Results for commit c1151c1e. ± Comparison against base commit b63ca061.

[test-results]:data:application/gzip;base64,H4sIAMRpPmcC/02MSw7DIAwFrxKx7gJDCHEvEwWKJdR8KgKrqncv+ZEsZ549X0Z+cAt7ViAfFVuSjxe9Uuijn6eVRc1FVnmN+y5Uc3K3JGu3I8D6km//WSVvymNHvR+y40W4EOZwmJCmPd1K0R5c0oBSiiKPdIvYnO6W3vhetvM4+piBWQAFFpwCTk6jUdAr7SQZ5FZL7UAhGSLDfn/CoVNuGgEAAA==


Patch:
diff --git a/dask/_task_spec.py b/dask/_task_spec.py
index 44562759b..82b896b61 100644
--- a/dask/_task_spec.py
+++ b/dask/_task_spec.py
@@ -849,13 +849,6 @@ def fuse_linear_task_spec(dsk, keys):
     seen = set()
     result = {}
 
-    def _check_non_fusable_node(tsk):
-        return (
-            tsk.block_fusion
-            or isinstance(tsk, DataNode)
-            and not isinstance(tsk.value, str)
-        )
-
     for key in dsk:
         if key in seen:
             continue
@@ -865,20 +858,13 @@ def fuse_linear_task_spec(dsk, keys):
         deps = dependencies[key]
         dependents_key = dependents[key]
 
-        if (
-            len(deps) != 1
-            and len(dependents_key) != 1
-            or _check_non_fusable_node(dsk[key])
-        ):
+        if len(deps) != 1 and len(dependents_key) != 1 or dsk[key].block_fusion:
             result[key] = dsk[key]
             continue
 
         linear_chain = [dsk[key]]
         top_key = key
 
-        # TODO: The old implementation doesn't fuse data nodes
-        # It claims that a data node is easier to serialize as a raw value
-
         # Walk towards the leafs as long as the nodes have a single dependency
         # and a single dependent, we can't fuse two nodes of an intermediate node
         # is the source for 2 dependents
@@ -894,7 +880,7 @@ def fuse_linear_task_spec(dsk, keys):
                 break
             if (
                 len(dependents[new_key]) != 1
-                or _check_non_fusable_node(dsk[new_key])
+                or dsk[new_key].block_fusion
                 or new_key in keys
             ):
                 result[new_key] = dsk[new_key]
@@ -911,7 +897,7 @@ def fuse_linear_task_spec(dsk, keys):
             if new_key in seen:
                 break
             seen.add(new_key)
-            if len(dependencies[new_key]) != 1 or _check_non_fusable_node(dsk[new_key]):
+            if len(dependencies[new_key]) != 1 or dsk[new_key].block_fusion:
                 # Exit if the dependent has multiple dependencies, triangle
                 result[new_key] = dsk[new_key]
                 break
diff --git a/dask/array/_shuffle.py b/dask/array/_shuffle.py
index 7ae19568d..2c926d317 100644
--- a/dask/array/_shuffle.py
+++ b/dask/array/_shuffle.py
@@ -10,6 +10,7 @@ from typing import Literal
 import numpy as np
 
 from dask import config
+from dask._task_spec import DataNode, List, Task, TaskRef
 from dask.array.chunk import getitem
 from dask.array.core import Array, unknown_chunk_message
 from dask.array.dispatch import concatenate_lookup, take_lookup
@@ -226,16 +227,17 @@ def _shuffle(chunks, indexer, axis, in_name, out_name, token):
     sorter_name = "shuffle-sorter-"
     taker_name = "shuffle-taker-"
 
-    old_blocks = np.empty([len(c) for c in chunks], dtype="O")
-    for old_index in np.ndindex(old_blocks.shape):
-        old_blocks[old_index] = (in_name,) + old_index
+    old_blocks = {
+        old_index: (in_name,) + old_index
+        for old_index in np.ndindex(tuple([len(c) for c in chunks]))
+    }
 
     for new_chunk_idx, new_chunk_taker in enumerate(new_chunks):
         new_chunk_taker = np.array(new_chunk_taker)
         sorter = np.argsort(new_chunk_taker).astype(dtype)
         sorter_key = sorter_name + tokenize(sorter)
         # low level fusion can't deal with arrays on first position
-        merges[sorter_key] = (1, sorter)
+        merges[sorter_key] = DataNode(sorter_key, (1, sorter))
 
         sorted_array = new_chunk_taker[sorter]
         source_chunk_nr, taker_boundary = np.unique(
@@ -271,22 +273,30 @@ def _shuffle(chunks, indexer, axis, in_name, out_name, token):
 
                     taker_key = taker_name + tokenize(this_slice)
                     # low level fusion can't deal with arrays on first position
-                    intermediates[taker_key] = (1, tuple(this_slice))
+                    intermediates[taker_key] = DataNode(
+                        taker_key, (1, tuple(this_slice))
+                    )
                     taker_cache[c] = taker_key
 
-                intermediates[name] = _getitem, old_blocks[chunk_key], taker_key
+                intermediates[name] = Task(
+                    name, _getitem, TaskRef(old_blocks[chunk_key]), TaskRef(taker_key)
+                )
                 merge_keys.append(name)
 
             merge_suffix = convert_key(chunk_tuple, new_chunk_idx, axis)
+            out_name_merge = (out_name,) + merge_suffix
             if len(merge_keys) > 1:
-                merges[(out_name,) + merge_suffix] = (
+                merges[out_name_merge] = Task(
+                    out_name_merge,
                     concatenate_arrays,
-                    merge_keys,
-                    sorter_key,
+                    List(*(TaskRef(m) for m in merge_keys)),
+                    TaskRef(sorter_key),
                     axis,
                 )
             elif len(merge_keys) == 1:
-                merges[(out_name,) + merge_suffix] = intermediates.pop(merge_keys[0])
+                t = intermediates.pop(merge_keys[0])
+                t.key = out_name_merge
+                merges[out_name_merge] = t
             else:
                 raise NotImplementedError
 
diff --git a/dask/array/slicing.py b/dask/array/slicing.py
index ec12479cf..d3bd3a9d9 100644
--- a/dask/array/slicing.py
+++ b/dask/array/slicing.py
@@ -12,6 +12,7 @@ import numpy as np
 from tlz import concat, memoize, merge, pluck
 
 from dask import core
+from dask._task_spec import Alias, Task, TaskRef
 from dask.array.chunk import getitem
 from dask.base import is_dask_collection, tokenize
 from dask.highlevelgraph import HighLevelGraph
@@ -141,10 +142,8 @@ def slice_array(out_name, in_name, blockdims, index, itemsize):
     >>> dsk, blockdims = slice_array('y', 'x', [(20, 20, 20, 20, 20)],
     ...                              (slice(10, 35),), 8)
     >>> pprint(dsk)  # doctest: +ELLIPSIS +NORMALIZE_WHITESPACE
-    {('y', 0): (<function getitem at ...>,
-                ('x', 0),
-                (slice(10, 20, 1),)),
-     ('y', 1): (<function getitem at ...>, ('x', 1), (slice(0, 15, 1),))}
+    {('y', 0): <Task ('y', 0) getitem(Alias(('x', 0)->('x', 0)), (slice(10, 20, 1),))>,
+     ('y', 1): <Task ('y', 1) getitem(Alias(('x', 1)->('x', 1)), (slice(0, 15, 1),))>}
     >>> blockdims
     ((10, 15),)
 
@@ -164,7 +163,9 @@ def slice_array(out_name, in_name, blockdims, index, itemsize):
         isinstance(index, slice) and index == slice(None, None, None) for index in index
     ):
         suffixes = product(*[range(len(bd)) for bd in blockdims])
-        dsk = {(out_name,) + s: (in_name,) + s for s in suffixes}
+        dsk = {
+            (out_name,) + s: Alias((out_name,) + s, (in_name,) + s) for s in suffixes
+        }
         return dsk, blockdims
 
     # Add in missing colons at the end as needed.  x[5] -> x[5, :, :]
@@ -204,19 +205,17 @@ def slice_with_newaxes(out_name, in_name, blockdims, index):
         expand_orig = expander(where_none_orig)
 
         # Insert ",0" into the key:  ('x', 2, 3) -> ('x', 0, 2, 0, 3)
-        dsk2 = {
-            (out_name,) + expand(k[1:], 0): (v[:2] + (expand_orig(v[2], None),))
-            for k, v in dsk.items()
-            if k[0] == out_name
-        }
-
-        # Add back intermediate parts of the dask that weren't the output
-        dsk3 = merge(dsk2, {k: v for k, v in dsk.items() if k[0] != out_name})
+        dsk2 = {}
+        for k, v in dsk.items():
+            if k[0] == out_name:
+                k2 = (out_name,) + expand(k[1:], 0)
+                dsk2[k2] = Task(k2, v.func, v.args[0], expand_orig(v.args[1], None))
+            else:
+                dsk2[k] = v
 
         # Insert (1,) into blockdims:  ((2, 2), (3, 3)) -> ((2, 2), (1,), (3, 3))
         blockdims3 = expand(blockdims2, (1,))
-
-        return dsk3, blockdims3
+        return dsk2, blockdims3
 
     else:
         return dsk, blockdims2
@@ -340,10 +339,10 @@ def slice_slices_and_integers(
 
     dsk_out = {
         out_name: (
-            (getitem, in_name, slices)
+            Task(out_name, getitem, TaskRef(in_name), slices)
             if not allow_getitem_optimization
             or not all(sl == slice(None, None, None) for sl in slices)
-            else in_name
+            else Alias(out_name, in_name)
         )
         for out_name, in_name, slices in zip(out_names, in_names, all_slices)
     }
@@ -587,7 +586,10 @@ def take(outname, inname, chunks, index, axis=0):
             # TODO: This should be a real no-op, but the call stack is
             # too deep to do this efficiently for now
             chunk_tuples = list(product(*(range(len(c)) for i, c in enumerate(chunks))))
-            graph = {(outname,) + c: (inname,) + c for c in chunk_tuples}
+            graph = {
+                (outname,) + c: Alias((outname,) + c, (inname,) + c)
+                for c in chunk_tuples
+            }
             return tuple(chunks), graph
 
         average_chunk_size = int(sum(chunks[axis]) / len(chunks[axis]))
@@ -610,7 +612,9 @@ def take(outname, inname, chunks, index, axis=0):
         slices = tuple(slices)
         chunk_tuples = list(product(*(range(len(c)) for i, c in enumerate(chunks))))
         dsk = {
-            (outname,) + ct: (getitem, (inname,) + ct, slices) for ct in chunk_tuples
+            (outname,)
+            + ct: Task((outname,) + ct, getitem, TaskRef((inname,) + ct), slices)
+            for ct in chunk_tuples
         }
         return chunks, dsk
     else:
diff --git a/dask/layers.py b/dask/layers.py
index c0b1384f3..5fe1bb482 100644
--- a/dask/layers.py
+++ b/dask/layers.py
@@ -110,7 +110,7 @@ class ArrayValuesDep(ArrayBlockwiseDep):
 
 @normalize_token.register(ArraySliceDep)
 def normalize_array_slice_dep(dep):
-    return "ArraySliceDep", dep.chunks, dep.starts
+    return "ArraySliceDep", dep.chunks
 
 
 @normalize_token.register(ArrayBlockIdDep)


Test patch:
diff --git a/dask/array/tests/test_optimization.py b/dask/array/tests/test_optimization.py
index 66d25551d..092c51cc0 100644
--- a/dask/array/tests/test_optimization.py
+++ b/dask/array/tests/test_optimization.py
@@ -194,7 +194,7 @@ def test_fuse_slices_with_alias(getter, getitem):
     }
     keys = [("dx2", 0)]
     dsk2 = optimize(dsk, keys)
-    assert len(dsk2) == 3
+    assert len(dsk2) == 2
 
 
 @pytest.mark.parametrize("chunks", [10, 5, 3])
@@ -203,10 +203,13 @@ def test_fuse_getter_with_asarray(chunks):
     y = da.ones(10, chunks=chunks)
     z = x + y
     dsk = z.__dask_optimize__(z.dask, z.__dask_keys__())
-    assert any(
-        isinstance(v, DataNode) and isinstance(v.value, np.ndarray)
-        for v in dsk.values()
-    )
+    if chunks == 10:
+        assert len(dsk) == 2 and any(isinstance(t, Alias) for t in dsk.values())
+    else:
+        assert any(
+            isinstance(v, DataNode) and isinstance(v.value, np.ndarray)
+            for v in dsk.values()
+        )
     assert_eq(z, x + 1)
 
 
diff --git a/dask/array/tests/test_shuffle.py b/dask/array/tests/test_shuffle.py
index 0b4aebbc4..fda57f41f 100644
--- a/dask/array/tests/test_shuffle.py
+++ b/dask/array/tests/test_shuffle.py
@@ -148,7 +148,7 @@ def test_dtype_taker(arr, darr):
     expected = arr[:, [1, 5, 6, 0, 3, 4, 2, 7]]
     assert_eq(result, expected)
     assert all(
-        v[1][1].dtype == np.uint8
+        v.value[1][1].dtype == np.uint8
         for k, v in dict(result.dask).items()
         if "shuffle-taker" in k
     )
diff --git a/dask/array/tests/test_slicing.py b/dask/array/tests/test_slicing.py
index 2d1ab3006..509c2e3aa 100644
--- a/dask/array/tests/test_slicing.py
+++ b/dask/array/tests/test_slicing.py
@@ -5,6 +5,8 @@ import warnings
 
 import pytest
 
+from dask._task_spec import Alias, Task, TaskRef
+
 np = pytest.importorskip("numpy")
 
 import dask
@@ -177,10 +179,10 @@ def test_slice_singleton_value_on_boundary():
 def test_slice_array_1d():
     # x[24::2]
     expected = {
-        ("y", 0): (getitem, ("x", 0), (slice(24, 25, 2),)),
-        ("y", 1): (getitem, ("x", 1), (slice(1, 25, 2),)),
-        ("y", 2): (getitem, ("x", 2), (slice(0, 25, 2),)),
-        ("y", 3): (getitem, ("x", 3), (slice(1, 25, 2),)),
+        ("y", 0): Task(("y", 0), getitem, TaskRef(("x", 0)), (slice(24, 25, 2),)),
+        ("y", 1): Task(("y", 1), getitem, TaskRef(("x", 1)), (slice(1, 25, 2),)),
+        ("y", 2): Task(("y", 2), getitem, TaskRef(("x", 2)), (slice(0, 25, 2),)),
+        ("y", 3): Task(("y", 3), getitem, TaskRef(("x", 3)), (slice(1, 25, 2),)),
     }
     result, chunks = slice_array("y", "x", [[25] * 4], [slice(24, None, 2)], 8)
 
@@ -188,9 +190,9 @@ def test_slice_array_1d():
 
     # x[26::2]
     expected = {
-        ("y", 0): (getitem, ("x", 1), (slice(1, 25, 2),)),
-        ("y", 1): (getitem, ("x", 2), (slice(0, 25, 2),)),
-        ("y", 2): (getitem, ("x", 3), (slice(1, 25, 2),)),
+        ("y", 0): Task(("y", 0), getitem, TaskRef(("x", 1)), (slice(1, 25, 2),)),
+        ("y", 1): Task(("y", 1), getitem, TaskRef(("x", 2)), (slice(0, 25, 2),)),
+        ("y", 2): Task(("y", 2), getitem, TaskRef(("x", 3)), (slice(1, 25, 2),)),
     }
 
     result, chunks = slice_array("y", "x", [[25] * 4], [slice(26, None, 2)], 8)
@@ -198,10 +200,10 @@ def test_slice_array_1d():
 
     # x[24::2]
     expected = {
-        ("y", 0): (getitem, ("x", 0), (slice(24, 25, 2),)),
-        ("y", 1): (getitem, ("x", 1), (slice(1, 25, 2),)),
-        ("y", 2): (getitem, ("x", 2), (slice(0, 25, 2),)),
-        ("y", 3): (getitem, ("x", 3), (slice(1, 25, 2),)),
+        ("y", 0): Task(("y", 0), getitem, TaskRef(("x", 0)), (slice(24, 25, 2),)),
+        ("y", 1): Task(("y", 1), getitem, TaskRef(("x", 1)), (slice(1, 25, 2),)),
+        ("y", 2): Task(("y", 2), getitem, TaskRef(("x", 2)), (slice(0, 25, 2),)),
+        ("y", 3): Task(("y", 3), getitem, TaskRef(("x", 3)), (slice(1, 25, 2),)),
     }
     result, chunks = slice_array("y", "x", [(25,) * 4], (slice(24, None, 2),), 8)
 
@@ -209,9 +211,9 @@ def test_slice_array_1d():
 
     # x[26::2]
     expected = {
-        ("y", 0): (getitem, ("x", 1), (slice(1, 25, 2),)),
-        ("y", 1): (getitem, ("x", 2), (slice(0, 25, 2),)),
-        ("y", 2): (getitem, ("x", 3), (slice(1, 25, 2),)),
+        ("y", 0): Task(("y", 0), getitem, TaskRef(("x", 1)), (slice(1, 25, 2),)),
+        ("y", 1): Task(("y", 1), getitem, TaskRef(("x", 2)), (slice(0, 25, 2),)),
+        ("y", 2): Task(("y", 2), getitem, TaskRef(("x", 3)), (slice(1, 25, 2),)),
     }
 
     result, chunks = slice_array("y", "x", [(25,) * 4], (slice(26, None, 2),), 8)
@@ -221,16 +223,17 @@ def test_slice_array_1d():
 def test_slice_array_2d():
     # 2d slices: x[13::2,10::1]
     expected = {
-        ("y", 0, 0): (getitem, ("x", 0, 0), (slice(13, 20, 2), slice(10, 20, 1))),
-        ("y", 0, 1): (
+        ("y", 0, 0): Task(
+            ("y", 0, 0),
             getitem,
-            ("x", 0, 1),
-            (slice(13, 20, 2), slice(None, None, None)),
+            TaskRef(("x", 0, 0)),
+            (slice(13, 20, 2), slice(10, 20, 1)),
         ),
-        ("y", 0, 2): (
-            getitem,
-            ("x", 0, 2),
-            (slice(13, 20, 2), slice(None, None, None)),
+        ("y", 0, 1): Task(
+            ("y", 0, 1), getitem, TaskRef(("x", 0, 1)), (slice(13, 20, 2), slice(None))
+        ),
+        ("y", 0, 2): Task(
+            ("y", 0, 2), getitem, TaskRef(("x", 0, 2)), (slice(13, 20, 2), slice(None))
         ),
     }
 
@@ -246,9 +249,9 @@ def test_slice_array_2d():
 
     # 2d slices with one dimension: x[5,10::1]
     expected = {
-        ("y", 0): (getitem, ("x", 0, 0), (5, slice(10, 20, 1))),
-        ("y", 1): (getitem, ("x", 0, 1), (5, slice(None, None, None))),
-        ("y", 2): (getitem, ("x", 0, 2), (5, slice(None, None, None))),
+        ("y", 0): Task(("y", 0), getitem, TaskRef(("x", 0, 0)), (5, slice(10, 20, 1))),
+        ("y", 1): Task(("y", 1), getitem, TaskRef(("x", 0, 1)), (5, slice(None))),
+        ("y", 2): Task(("y", 2), getitem, TaskRef(("x", 0, 2)), (5, slice(None))),
     }
 
     result, chunks = slice_array(
@@ -260,12 +263,16 @@ def test_slice_array_2d():
 
 def test_slice_optimizations():
     # bar[:]
-    expected = {("foo", 0): ("bar", 0)}
+    expected = {("foo", 0): Alias(("foo", 0), ("bar", 0))}
     result, chunks = slice_array("foo", "bar", [[100]], (slice(None, None, None),), 8)
     assert expected == result
 
     # bar[:,:,:]
-    expected = {("foo", 0): ("bar", 0), ("foo", 1): ("bar", 1), ("foo", 2): ("bar", 2)}
+    expected = {
+        ("foo", 0): Alias(("foo", 0), ("bar", 0)),
+        ("foo", 1): Alias(("foo", 1), ("bar", 1)),
+        ("foo", 2): Alias(("foo", 2), ("bar", 2)),
+    }
     result, chunks = slice_array(
         "foo",
         "bar",
@@ -281,7 +288,11 @@ def test_slicing_with_singleton_indices():
         "y", "x", ([5, 5], [5, 5]), (slice(0, 5), 8), itemsize=8
     )
 
-    expected = {("y", 0): (getitem, ("x", 0, 1), (slice(None, None, None), 3))}
+    expected = {
+        ("y", 0): Task(
+            ("y", 0), getitem, TaskRef(("x", 0, 1)), (slice(None, None, None), 3)
+        )
+    }
 
     assert expected == result
 
@@ -296,14 +307,16 @@ def test_slicing_with_newaxis():
     )
 
     expected = {
-        ("y", 0, 0, 0): (
+        ("y", 0, 0, 0): Task(
+            ("y", 0, 0, 0),
             getitem,
-            ("x", 0, 0),
+            TaskRef(("x", 0, 0)),
             (slice(0, 3, 1), None, slice(None, None, None)),
         ),
-        ("y", 0, 0, 1): (
+        ("y", 0, 0, 1): Task(
+            ("y", 0, 0, 1),
             getitem,
-            ("x", 0, 1),
+            TaskRef(("x", 0, 1)),
             (slice(0, 3, 1), None, slice(None, None, None)),
         ),
     }
@@ -450,15 +463,19 @@ def test_slicing_exhaustively():
 def test_slicing_with_negative_step_flops_keys():
     x = da.arange(10, chunks=5)
     y = x[:1:-1]
-    assert (x.name, 1) in y.dask[(y.name, 0)]
-    assert (x.name, 0) in y.dask[(y.name, 1)]
+    assert (x.name, 1) in y.dask[(y.name, 0)].dependencies
+    assert (x.name, 0) in y.dask[(y.name, 1)].dependencies
 
     assert_eq(y, np.arange(10)[:1:-1])
 
     assert y.chunks == ((5, 3),)
 
-    assert y.dask[(y.name, 0)] == (getitem, (x.name, 1), (slice(-1, -6, -1),))
-    assert y.dask[(y.name, 1)] == (getitem, (x.name, 0), (slice(-1, -4, -1),))
+    assert y.dask[(y.name, 0)] == Task(
+        (y.name, 0), getitem, TaskRef((x.name, 1)), (slice(-1, -6, -1),)
+    )
+    assert y.dask[(y.name, 1)] == Task(
+        (y.name, 1), getitem, TaskRef((x.name, 0)), (slice(-1, -4, -1),)
+    )
 
 
 def test_empty_slice():
@@ -559,8 +576,18 @@ def test_uneven_blockdims():
     index = (slice(240, 270), slice(None))
     dsk_out, bd_out = slice_array("in", "out", blockdims, index, itemsize=8)
     sol = {
-        ("in", 0, 0): (getitem, ("out", 7, 0), (slice(28, 31, 1), slice(None))),
-        ("in", 1, 0): (getitem, ("out", 8, 0), (slice(0, 27, 1), slice(None))),
+        ("in", 0, 0): Task(
+            ("in", 0, 0),
+            getitem,
+            TaskRef(("out", 7, 0)),
+            (slice(28, 31, 1), slice(None)),
+        ),
+        ("in", 1, 0): Task(
+            ("in", 1, 0),
+            getitem,
+            TaskRef(("out", 8, 0)),
+            (slice(0, 27, 1), slice(None)),
+        ),
     }
     assert dsk_out == sol
     assert bd_out == ((3, 27), (100,))
@@ -569,12 +596,42 @@ def test_uneven_blockdims():
     index = (slice(240, 270), slice(180, 230))
     dsk_out, bd_out = slice_array("in", "out", blockdims, index, itemsize=8)
     sol = {
-        ("in", 0, 0): (getitem, ("out", 7, 5), (slice(28, 31, 1), slice(29, 30, 1))),
-        ("in", 0, 1): (getitem, ("out", 7, 6), (slice(28, 31, 1), slice(None))),
-        ("in", 0, 2): (getitem, ("out", 7, 7), (slice(28, 31, 1), slice(0, 18, 1))),
-        ("in", 1, 0): (getitem, ("out", 8, 5), (slice(0, 27, 1), slice(29, 30, 1))),
-        ("in", 1, 1): (getitem, ("out", 8, 6), (slice(0, 27, 1), slice(None))),
-        ("in", 1, 2): (getitem, ("out", 8, 7), (slice(0, 27, 1), slice(0, 18, 1))),
+        ("in", 0, 0): Task(
+            ("in", 0, 0),
+            getitem,
+            TaskRef(("out", 7, 5)),
+            (slice(28, 31, 1), slice(29, 30, 1)),
+        ),
+        ("in", 0, 1): Task(
+            ("in", 0, 0),
+            getitem,
+            TaskRef(("out", 7, 6)),
+            (slice(28, 31, 1), slice(None)),
+        ),
+        ("in", 0, 2): Task(
+            ("in", 0, 0),
+            getitem,
+            TaskRef(("out", 7, 7)),
+            (slice(28, 31, 1), slice(0, 18, 1)),
+        ),
+        ("in", 1, 0): Task(
+            ("in", 0, 0),
+            getitem,
+            TaskRef(("out", 8, 5)),
+            (slice(0, 27, 1), slice(29, 30, 1)),
+        ),
+        ("in", 1, 1): Task(
+            ("in", 0, 0),
+            getitem,
+            TaskRef(("out", 8, 6)),
+            (slice(0, 27, 1), slice(None)),
+        ),
+        ("in", 1, 2): Task(
+            ("in", 0, 0),
+            getitem,
+            TaskRef(("out", 8, 7)),
+            (slice(0, 27, 1), slice(0, 18, 1)),
+        ),
     }
     assert dsk_out == sol
     assert bd_out == ((3, 27), (1, 31, 18))
@@ -1029,7 +1086,7 @@ def test_take_sorted_indexer():
     assert {
         **dict(arr.dask),
         **{
-            k: k2
+            k: Alias(k, k2)
             for k, k2 in zip(
                 [k for k in dict(result.dask) if "getitem" in k[0]],
                 dict(arr.dask).keys(),
@@ -1044,5 +1101,5 @@ def test_all_none_slices_just_mappings():
     dsk = dict(result.dask)
     assert len([k for k in dsk if "getitem" in k[0]]) == 12
     # check that we are just mapping the keys
-    assert all(v in dsk for k, v in dsk.items() if "getitem" in k[0])
+    assert all(isinstance(v, Alias) for k, v in dsk.items() if "getitem" in k[0])
     assert_eq(result, np.ones((6, 10)))
diff --git a/dask/tests/test_task_spec.py b/dask/tests/test_task_spec.py
index 9ee6051e3..0b41e673b 100644
--- a/dask/tests/test_task_spec.py
+++ b/dask/tests/test_task_spec.py
@@ -898,20 +898,20 @@ def test_linear_fusion():
     ]
     dsk = {t.key: t for t in tasks}
     result = fuse_linear_task_spec(dsk, {"fourth"})
-    assert len(result) == 3
+    assert len(result) == 2
     assert isinstance(result["fourth"], Alias)
     assert (
-        isinstance(result["second-third-fourth"], Task)
-        and funcname(result["second-third-fourth"].func) == "_execute_subgraph"
+        isinstance(result["foo-second-third-fourth"], Task)
+        and funcname(result["foo-second-third-fourth"].func) == "_execute_subgraph"
     )
-    assert isinstance(result["foo"], DataNode)
+    assert "foo" not in result
 
     # Branch, so no fusion
     tasks.append(Task("branch", func, third.ref()))
     dsk = {t.key: t for t in tasks}
     result = fuse_linear_task_spec(dsk, {"fourth"})
-    assert len(result) == 5
-    assert "second-third" in result
+    assert len(result) == 4
+    assert "foo-second-third" in result
     assert isinstance(result["third"], Alias)
 
     # Branch, so no fusion at all
@@ -932,7 +932,7 @@ def test_linear_fusion_intermediate_branch():
     ]
     dsk = {t.key: t for t in tasks}
     result = fuse_linear_task_spec(dsk, {"fourth"})
-    assert "second-third" in result
+    assert "foo-second-third" in result
     assert isinstance(result["fourth"], Task)
 
 
@@ -1000,6 +1000,27 @@ def test_nested_containers():
     assert t({"b": "b"}) == {"k": "a-b", ("v", 1): "c-d"}
 
 
+def test_block_io_fusion():
+
+    class SubTask(Task):
+
+        @property
+        def block_fusion(self) -> bool:
+            return True
+
+    tasks = [
+        io := DataNode("foo", 1),
+        second := Task("second", func, io.ref()),
+        third := SubTask("third", func, second.ref()),
+        Task("fourth", func, third.ref()),
+    ]
+    dsk = {t.key: t for t in tasks}
+    result = fuse_linear_task_spec(dsk, {"fourth"})
+    assert "foo-second" in result
+    assert isinstance(result["fourth"], Task)
+    assert isinstance(result["third"], SubTask)
+
+
 @pytest.mark.parametrize(
     "task_type, python_type",
     [