Repository: cubed-dev/cubed
Pull request #607 (2024-11-05 11:43:16)

Implement concat using general_blockwise
Part of #464


Patch:
diff --git a/cubed/array_api/manipulation_functions.py b/cubed/array_api/manipulation_functions.py
index f30bb6da..35eb76bc 100644
--- a/cubed/array_api/manipulation_functions.py
+++ b/cubed/array_api/manipulation_functions.py
@@ -10,10 +10,10 @@ from cubed.backend_array_api import namespace as nxp
 from cubed.core import squeeze  # noqa: F401
 from cubed.core import blockwise, rechunk, unify_chunks
 from cubed.core.ops import (
+    _create_zarr_indexer,
     elemwise,
     general_blockwise,
     map_blocks,
-    map_direct,
     map_selection,
 )
 from cubed.utils import block_id_to_offset, get_item, offset_to_block_id, to_chunksize
@@ -95,6 +95,7 @@ def concat(arrays, /, *, axis=0, chunks=None):
 
     # offsets along axis for the start of each array
     offsets = [0] + list(tlz.accumulate(add, [a.shape[axis] for a in arrays]))
+    in_shapes = tuple(array.shape for array in arrays)
 
     axis = validate_axis(axis, a.ndim)
     shape = a.shape[:axis] + (offsets[-1],) + a.shape[axis + 1 :]
@@ -104,47 +105,93 @@ def concat(arrays, /, *, axis=0, chunks=None):
     else:
         chunks = normalize_chunks(chunks, shape=shape, dtype=dtype)
 
-    # memory allocated by reading one chunk from input array
-    # note that although the output chunk will overlap multiple input chunks,
-    # the chunks are read in series, reusing memory
-    extra_projected_mem = a.chunkmem
+    def key_function(out_key):
+        out_coords = out_key[1:]
+        block_id = out_coords
+
+        # determine the start and stop indexes for this block along the axis dimension
+        chunksize = to_chunksize(chunks)
+        start = block_id[axis] * chunksize[axis]
+        stop = start + chunksize[axis]
+        stop = min(stop, shape[axis])
+
+        # produce a key that has slices (except for axis dimension, which is replaced below)
+        idx = tuple(0 if i == axis else v for i, v in enumerate(block_id))
+        key = get_item(chunks, idx)
+
+        # find slices of the arrays
+        in_keys = []
+        for ai, sl in _array_slices(offsets, start, stop):
+            key = tuple(sl if i == axis else k for i, k in enumerate(key))
+
+            # use a Zarr BasicIndexer to convert this to input coordinates
+            a = arrays[ai]
+            indexer = _create_zarr_indexer(key, a.shape, a.chunksize)
+
+            in_keys.extend(
+                [(a.name,) + chunk_coords for (chunk_coords, _, _) in indexer]
+            )
+
+        return (iter(tuple(in_key for in_key in in_keys)),)
 
-    return map_direct(
+    num_input_blocks = (1,) * len(arrays)
+    iterable_input_blocks = (True,) * len(arrays)
+
+    # We have to mark this as fusable=False since the number of input args to
+    # the _read_concat_chunk function is *not* the same as the number of
+    # predecessor nodes in the DAG, and the fusion functions in blockwise
+    # assume they are the same. See https://github.com/cubed-dev/cubed/issues/414
+    # This also affects stack.
+    return general_blockwise(
         _read_concat_chunk,
+        key_function,
         *arrays,
-        shape=shape,
-        dtype=dtype,
-        chunks=chunks,
-        extra_projected_mem=extra_projected_mem,
+        shapes=[shape],
+        dtypes=[dtype],
+        chunkss=[chunks],
+        num_input_blocks=num_input_blocks,
+        iterable_input_blocks=iterable_input_blocks,
+        extra_func_kwargs=dict(dtype=dtype),
+        target_shape=shape,
         target_chunks=chunks,
         axis=axis,
         offsets=offsets,
+        in_shapes=in_shapes,
+        function_nargs=1,
+        fusable=False,
     )
 
 
 def _read_concat_chunk(
-    x, *arrays, target_chunks=None, axis=None, offsets=None, block_id=None
+    arrays,
+    dtype=None,
+    target_shape=None,
+    target_chunks=None,
+    axis=None,
+    offsets=None,
+    in_shapes=None,
+    block_id=None,
 ):
     # determine the start and stop indexes for this block along the axis dimension
-    chunks = target_chunks
-    chunksize = to_chunksize(chunks)
+    chunksize = to_chunksize(target_chunks)
     start = block_id[axis] * chunksize[axis]
-    stop = start + x.shape[axis]
-
-    # produce a key that has slices (except for axis dimension, which is replaced below)
-    idx = tuple(0 if i == axis else v for i, v in enumerate(block_id))
-    key = get_item(chunks, idx)
-
-    # concatenate slices of the arrays
-    parts = []
-    for ai, sl in _array_slices(offsets, start, stop):
-        key = tuple(sl if i == axis else k for i, k in enumerate(key))
-        parts.append(arrays[ai].zarray[key])
-    return nxp.concat(parts, axis=axis)
+    stop = start + chunksize[axis]
+    stop = min(stop, target_shape[axis])
+
+    chunk_shape = tuple(ch[bi] for ch, bi in zip(target_chunks, block_id))
+    out = np.empty(chunk_shape, dtype=dtype)
+    for array, (lchunk_selection, lout_selection) in zip(
+        arrays,
+        _chunk_slices(
+            offsets, start, stop, target_chunks, chunksize, in_shapes, axis, block_id
+        ),
+    ):
+        out[lout_selection] = array[lchunk_selection]
+    return out
 
 
 def _array_slices(offsets, start, stop):
-    """Return pairs of array index and slice to slice from start to stop in the concatenated array."""
+    """Return pairs of array index and array slice to slice from start to stop in the concatenated array."""
     slice_start = start
     while slice_start < stop:
         # find array that slice_start falls in
@@ -154,6 +201,33 @@ def _array_slices(offsets, start, stop):
         slice_start = slice_stop
 
 
+def _chunk_slices(
+    offsets, start, stop, target_chunks, chunksize, in_shapes, axis, block_id
+):
+    """Return pairs of chunk slices to slice input array chunks and output concatenated chunk."""
+
+    # an output chunk may have selections from more than one array, so we need an offset per array
+    arr_sel_offset = 0  # offset along axis
+
+    # produce a key that has slices (except for axis dimension, which is replaced below)
+    idx = tuple(0 if i == axis else v for i, v in enumerate(block_id))
+    key = get_item(target_chunks, idx)
+
+    for ai, sl in _array_slices(offsets, start, stop):
+        key = tuple(sl if i == axis else k for i, k in enumerate(key))
+        indexer = _create_zarr_indexer(key, in_shapes[ai], chunksize)
+        for _, lchunk_selection, lout_selection in indexer:
+            lout_selection_with_offset = tuple(
+                sl
+                if ax != axis
+                else slice(sl.start + arr_sel_offset, sl.stop + arr_sel_offset)
+                for ax, sl in enumerate(lout_selection)
+            )
+            yield lchunk_selection, lout_selection_with_offset
+
+        arr_sel_offset += lout_selection[axis].stop
+
+
 def expand_dims(x, /, *, axis):
     if not isinstance(axis, tuple):
         axis = (axis,)
@@ -420,6 +494,7 @@ def stack(arrays, /, *, axis=0):
         dtypes=[dtype],
         chunkss=[chunks],
         axis=axis,
+        function_nargs=1,
         fusable=False,
     )
 
diff --git a/cubed/core/groupby.py b/cubed/core/groupby.py
index 8ec97045..84a491a6 100644
--- a/cubed/core/groupby.py
+++ b/cubed/core/groupby.py
@@ -127,7 +127,7 @@ def groupby_blockwise(
     axis=None,
     dtype=None,
     num_groups=None,
-    extra_func_kwargs=None,
+    **kwargs,
 ):
     """A blockwise operation that performs groupby aggregations.
 
@@ -148,8 +148,6 @@ def groupby_blockwise(
         Data type of output.
     num_groups: int
         The number of groups in the grouping array ``by``.
-    extra_func_kwargs: dict, optional
-        Extra keyword arguments to pass to ``func``.
     """
 
     if by.ndim != 1:
@@ -203,6 +201,7 @@ def groupby_blockwise(
         by_read_chunks=by_read_chunks,
         target_chunks=target_chunks,
         groups_per_chunk=groups_per_chunk,
+        **kwargs,
     )
 
 
diff --git a/cubed/core/ops.py b/cubed/core/ops.py
index 5896aad6..23ea169b 100644
--- a/cubed/core/ops.py
+++ b/cubed/core/ops.py
@@ -367,6 +367,9 @@ def general_blockwise(
 
             return wrap
 
+        function_nargs = kwargs.pop("function_nargs", None)
+        if function_nargs is not None:
+            function_nargs = function_nargs + 1  # for offsets array
         num_input_blocks = kwargs.pop("num_input_blocks", None)
         if num_input_blocks is not None:
             num_input_blocks = num_input_blocks + (1,)  # for offsets array
@@ -386,6 +389,7 @@ def general_blockwise(
             target_stores=target_stores,
             target_paths=target_paths,
             extra_func_kwargs=extra_func_kwargs,
+            function_nargs=function_nargs,
             num_input_blocks=num_input_blocks,
             iterable_input_blocks=iterable_input_blocks,
             **kwargs,
diff --git a/cubed/core/optimization.py b/cubed/core/optimization.py
index 86f3b4b5..346d8e0c 100644
--- a/cubed/core/optimization.py
+++ b/cubed/core/optimization.py
@@ -164,7 +164,7 @@ def can_fuse_predecessors(
     # if node itself can't be fused then there is nothing to fuse
     if not is_fusable(nodes[name]):
         logger.debug(
-            "can't fuse %s since it is not a primitive operation, or it uses map_direct",
+            "can't fuse %s since it is not a primitive operation, or it uses an operation that can't be fused (concat or stack)",
             name,
         )
         return False
@@ -224,10 +224,10 @@ def can_fuse_predecessors(
             )
             return False
 
+    # if a predecessor has no primitive op then just use None
     predecessor_primitive_ops = [
-        nodes[pre]["primitive_op"]
+        nodes[pre]["primitive_op"] if can_fuse else None
         for pre, _, can_fuse in predecessor_ops_and_arrays(dag, name)
-        if can_fuse
     ]
     return can_fuse_multiple_primitive_ops(
         name,
diff --git a/cubed/primitive/blockwise.py b/cubed/primitive/blockwise.py
index 53817e10..97296228 100644
--- a/cubed/primitive/blockwise.py
+++ b/cubed/primitive/blockwise.py
@@ -52,7 +52,9 @@ class BlockwiseSpec:
     function : Callable
         A function that maps input chunks to an output chunk.
     function_nargs: int
-        The number of array arguments that ``function`` takes.
+        The number of array arguments that ``function`` takes. Note that for some
+        functions (``concat``, ``stack``) this may be different to the number of
+        arrays that the operation depends on.
     num_input_blocks: Tuple[int, ...]
         The number of input blocks read from each input array.
     iterable_input_blocks: Tuple[int, ...]
@@ -206,14 +208,14 @@ def blockwise(
     inds: Sequence[Union[str, int]] = args[1::2]
 
     numblocks: Dict[str, Tuple[int, ...]] = {}
-    for name, array in zip(array_names, arrays):
+    for name, array in zip(array_names, arrays, strict=True):
         input_chunks = normalize_chunks(
             array.chunks, shape=array.shape, dtype=array.dtype
         )
         numblocks[name] = tuple(map(len, input_chunks))
 
     argindsstr: List[Any] = []
-    for name, ind in zip(array_names, inds):
+    for name, ind in zip(array_names, inds, strict=True):
         argindsstr.extend((name, ind))
 
     key_function = make_blockwise_key_function_flattened(
@@ -265,6 +267,7 @@ def general_blockwise(
     extra_projected_mem: int = 0,
     extra_func_kwargs: Optional[Dict[str, Any]] = None,
     fusable: bool = True,
+    function_nargs: Optional[int] = None,
     num_input_blocks: Optional[Tuple[int, ...]] = None,
     iterable_input_blocks: Optional[Tuple[bool, ...]] = None,
     **kwargs,
@@ -308,10 +311,11 @@ def general_blockwise(
     PrimitiveOperation to run the operation
     """
     array_names = in_names or [f"in_{i}" for i in range(len(arrays))]
-    array_map = {name: array for name, array in zip(array_names, arrays)}
+    array_map = {name: array for name, array in zip(array_names, arrays, strict=True)}
 
     func_kwargs = extra_func_kwargs or {}
     func_with_kwargs = partial(func, **{**kwargs, **func_kwargs})
+    function_nargs = function_nargs or len(arrays)
     num_input_blocks = num_input_blocks or (1,) * len(arrays)
     iterable_input_blocks = iterable_input_blocks or (False,) * len(arrays)
     read_proxies = {
@@ -361,7 +365,7 @@ def general_blockwise(
     spec = BlockwiseSpec(
         key_function,
         func_with_kwargs,
-        len(arrays),
+        function_nargs,
         num_input_blocks,
         iterable_input_blocks,
         read_proxies,
@@ -438,7 +442,7 @@ def can_fuse_multiple_primitive_ops(
     max_total_num_input_blocks: Optional[int] = None,
 ) -> bool:
     if is_fuse_candidate(primitive_op) and all(
-        is_fuse_candidate(p) for p in predecessor_primitive_ops
+        p is None or is_fuse_candidate(p) for p in predecessor_primitive_ops
     ):
         # If the peak projected memory for running all the predecessor ops in
         # order is larger than allowed_mem then we can't fuse.
@@ -455,7 +459,9 @@ def can_fuse_multiple_primitive_ops(
             # If max total input blocks not specified, then only fuse if num
             # tasks of predecessor ops match.
             ret = all(
-                primitive_op.num_tasks == p.num_tasks for p in predecessor_primitive_ops
+                primitive_op.num_tasks == p.num_tasks
+                for p in predecessor_primitive_ops
+                if p is not None
             )
             if ret:
                 logger.debug(
@@ -470,7 +476,9 @@ def can_fuse_multiple_primitive_ops(
         else:
             num_input_blocks = primitive_op.pipeline.config.num_input_blocks
             total_num_input_blocks = 0
-            for ni, p in zip(num_input_blocks, predecessor_primitive_ops):
+            for ni, p in zip(num_input_blocks, predecessor_primitive_ops, strict=True):
+                if p is None:
+                    continue
                 for nj in p.pipeline.config.num_input_blocks:
                     total_num_input_blocks += ni * nj
             ret = total_num_input_blocks <= max_total_num_input_blocks
@@ -500,6 +508,8 @@ def peak_projected_mem(primitive_ops):
     and retaining their return values in memory."""
     memory_modeller = MemoryModeller()
     for p in primitive_ops:
+        if p is None:
+            continue
         memory_modeller.allocate(p.projected_mem)
         chunkmem = chunk_memory(p.target_array)
         memory_modeller.free(p.projected_mem - chunkmem)
@@ -659,7 +669,7 @@ def fuse_blockwise_specs(
         itertools.chain(
             *(
                 tuple(n * m for m in nb)
-                for n, nb in zip(num_input_blocks, predecessor_num_blocks)
+                for n, nb in zip(num_input_blocks, predecessor_num_blocks, strict=True)
             )
         )
     )
@@ -692,12 +702,15 @@ def apply_blockwise_key_func(key_function, arg):
         # more than one input block is being read from arg
         assert isinstance(arg, (list, Iterator))
         if isinstance(arg, list):
-            return tuple(list(item) for item in zip(*(key_function(a) for a in arg)))
+            return tuple(
+                list(item) for item in zip(*(key_function(a) for a in arg), strict=True)
+            )
         else:
             # Return iterators to avoid materializing all array blocks at
             # once.
             return tuple(
-                iter(list(item)) for item in zip(*(key_function(a) for a in arg))
+                iter(list(item))
+                for item in zip(*(key_function(a) for a in arg), strict=True)
             )
 
 
@@ -708,7 +721,7 @@ def apply_blockwise_func(func, is_iterable, *args):
         # More than one input block is being read from this group of args to primitive op.
         # Note that it is important that a list is not returned to avoid materializing all
         # array blocks at once.
-        ret = map(lambda item: func(*item), zip(*args))
+        ret = map(lambda item: func(*item), zip(*args, strict=True))
     return ret
 
 
@@ -720,7 +733,7 @@ def make_fused_key_function(
         # split all args to the fused function into groups, one for each predecessor function
         func_args = tuple(
             item
-            for pkf, a in zip(predecessor_key_functions, args)
+            for pkf, a in zip(predecessor_key_functions, args, strict=True)
             for item in apply_blockwise_key_func(pkf, a)
         )
         return split_into(func_args, predecessor_funcs_nargs)
@@ -733,7 +746,7 @@ def make_fused_function(function, predecessor_functions, iterable_input_blocks):
         # args are grouped appropriately so they can be called by each predecessor function
         func_args = [
             apply_blockwise_func(pf, iterable_input_blocks[i], *a)
-            for i, (pf, a) in enumerate(zip(predecessor_functions, args))
+            for i, (pf, a) in enumerate(zip(predecessor_functions, args, strict=True))
         ]
         return function(*func_args)
 
@@ -742,7 +755,7 @@ def make_fused_function(function, predecessor_functions, iterable_input_blocks):
         # args are grouped appropriately so they can be called by each predecessor function
         func_args = [
             apply_blockwise_func(pf, iterable_input_blocks[i], *a)
-            for i, (pf, a) in enumerate(zip(predecessor_functions, args))
+            for i, (pf, a) in enumerate(zip(predecessor_functions, args, strict=True))
         ]
         yield from function(*func_args)
 
@@ -783,7 +796,7 @@ def make_blockwise_key_function(
         False,
     )
 
-    for axes, (arg, _) in zip(concat_axes, argpairs):
+    for axes, (arg, _) in zip(concat_axes, argpairs, strict=True):
         for ax in axes:
             if numblocks[arg][ax] > 1:
                 raise ValueError(
@@ -799,7 +812,9 @@ def make_blockwise_key_function(
         deps = set()
         coords = out_coords + dummies
         args = []
-        for cmap, axes, (arg, ind) in zip(coord_maps, concat_axes, argpairs):
+        for cmap, axes, (arg, ind) in zip(
+            coord_maps, concat_axes, argpairs, strict=True
+        ):
             if ind is None:
                 args.append(arg)
             else:


Test patch:
diff --git a/.github/workflows/tensorstore-tests.yml b/.github/workflows/tensorstore-tests.yml
index 74bfce99..827ff1df 100644
--- a/.github/workflows/tensorstore-tests.yml
+++ b/.github/workflows/tensorstore-tests.yml
@@ -1,6 +1,7 @@
 name: TensorStore tests
 
 on:
+  pull_request:
   schedule:
     # Every weekday at 03:58 UTC, see https://crontab.guru/
     - cron: "58 3 * * 1-5"
diff --git a/cubed/tests/test_groupby.py b/cubed/tests/test_groupby.py
index ab047893..6ca34c94 100644
--- a/cubed/tests/test_groupby.py
+++ b/cubed/tests/test_groupby.py
@@ -103,6 +103,7 @@ def test_groupby_blockwise_axis0():
         axis=0,
         dtype=nxp.int64,
         num_groups=6,
+        groupby_dtype=nxp.int32,
     )
     assert_array_equal(
         c.compute(),
@@ -129,6 +130,7 @@ def test_groupby_blockwise_axis1():
         axis=1,
         dtype=nxp.int64,
         num_groups=6,
+        groupby_dtype=nxp.int32,
     )
     assert_array_equal(
         c.compute(),
@@ -142,7 +144,9 @@ def test_groupby_blockwise_axis1():
     )
 
 
-def _sum_reduction_func(arr, by, axis, start_group, num_groups):
+def _sum_reduction_func(arr, by, axis, start_group, num_groups, groupby_dtype):
     # change 'by' so it starts from 0 for each chunk
     by = by - start_group
-    return npg.aggregate(by, arr, func="sum", axis=axis, size=num_groups)
+    return npg.aggregate(
+        by, arr, func="sum", dtype=groupby_dtype, axis=axis, size=num_groups
+    )
diff --git a/cubed/tests/test_optimization.py b/cubed/tests/test_optimization.py
index 32101024..6d635f0c 100644
--- a/cubed/tests/test_optimization.py
+++ b/cubed/tests/test_optimization.py
@@ -116,9 +116,9 @@ def test_fusion_transpose(spec, opt_fn):
     )
 
 
-def test_fusion_map_direct(spec):
-    # test that operations after a map_direct operation (indexing) can be fused
-    # with the map_direct operation
+def test_fusion_map_selection(spec):
+    # test that operations after a map_selection operation (indexing) can be fused
+    # with the map_selection operation
     # this is only true for the (default) multiple_inputs_optimize_dag optimize function
     a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2), spec=spec)
     b = a[1:, :]
@@ -1174,5 +1174,16 @@ def test_optimize_stack(spec):
     a = cubed.random.random((10, 10), chunks=(5, 5), spec=spec)
     b = cubed.random.random((10, 10), chunks=(5, 5), spec=spec)
     c = xp.stack((a, b), axis=0)
-    # try to fuse all ops into one
-    c.compute(optimize_function=fuse_multiple_levels(max_total_num_input_blocks=10))
+    d = c + 1
+    # try to fuse all ops into one (d will fuse with c, but c won't fuse with a and b)
+    d.compute(optimize_function=fuse_multiple_levels(max_total_num_input_blocks=10))
+
+
+def test_optimize_concat(spec):
+    # This test fails if concat's general_blockwise call doesn't have fusable=False
+    a = cubed.random.random((10, 10), chunks=(5, 5), spec=spec)
+    b = cubed.random.random((10, 10), chunks=(5, 5), spec=spec)
+    c = xp.concat((a, b), axis=0)
+    d = c + 1
+    # try to fuse all ops into one (d will fuse with c, but c won't fuse with a and b)
+    d.compute(optimize_function=fuse_multiple_levels(max_total_num_input_blocks=10))