Repository: tomwhite/cubed
Pull request #365 (2024-01-29 10:09:46)

Implement `partial_reduce` and `tree_reduce` using generalized blockwise
The new implementation is hidden behind a flag (`use_new_impl`) to make it easier to compare the old and new code.


Patch:
diff --git a/cubed/array_api/statistical_functions.py b/cubed/array_api/statistical_functions.py
index 3d69b6e4..53022cb2 100644
--- a/cubed/array_api/statistical_functions.py
+++ b/cubed/array_api/statistical_functions.py
@@ -25,7 +25,7 @@ def max(x, /, *, axis=None, keepdims=False):
     return reduction(x, nxp.max, axis=axis, dtype=x.dtype, keepdims=keepdims)
 
 
-def mean(x, /, *, axis=None, keepdims=False):
+def mean(x, /, *, axis=None, keepdims=False, use_new_impl=False):
     if x.dtype not in _real_floating_dtypes:
         raise TypeError("Only real floating-point dtypes are allowed in mean")
     # This implementation uses NumPy and Zarr's structured arrays to store a
@@ -46,6 +46,7 @@ def mean(x, /, *, axis=None, keepdims=False):
         intermediate_dtype=intermediate_dtype,
         dtype=dtype,
         keepdims=keepdims,
+        use_new_impl=use_new_impl,
         extra_func_kwargs=extra_func_kwargs,
     )
 
@@ -64,7 +65,7 @@ def _mean_combine(a, **kwargs):
     return {"n": n, "total": total}
 
 
-def _mean_aggregate(a):
+def _mean_aggregate(a, **kwargs):
     return nxp.divide(a["total"], a["n"])
 
 
diff --git a/cubed/core/ops.py b/cubed/core/ops.py
index fcb0dba0..bb71c1b2 100644
--- a/cubed/core/ops.py
+++ b/cubed/core/ops.py
@@ -1,5 +1,8 @@
+import builtins
+import math
 import numbers
 from functools import partial
+from itertools import product
 from numbers import Integral, Number
 from operator import add
 from typing import TYPE_CHECKING, Any, Sequence, Union
@@ -793,9 +796,24 @@ def reduction(
     intermediate_dtype=None,
     dtype=None,
     keepdims=False,
+    use_new_impl=False,
+    split_every=None,
     extra_func_kwargs=None,
 ) -> "Array":
     """Apply a function to reduce an array along one or more axes."""
+    if use_new_impl:
+        return reduction_new(
+            x,
+            func,
+            combine_func,
+            aggegrate_func,
+            axis,
+            intermediate_dtype,
+            dtype,
+            keepdims,
+            split_every,
+            extra_func_kwargs,
+        )
     if combine_func is None:
         combine_func = func
     if axis is None:
@@ -885,6 +903,193 @@ def reduction(
     return result
 
 
+def reduction_new(
+    x: "Array",
+    func,
+    combine_func=None,
+    aggegrate_func=None,
+    axis=None,
+    intermediate_dtype=None,
+    dtype=None,
+    keepdims=False,
+    split_every=None,
+    extra_func_kwargs=None,
+) -> "Array":
+    """Apply a function to reduce an array along one or more axes."""
+    if combine_func is None:
+        combine_func = func
+    if axis is None:
+        axis = tuple(range(x.ndim))
+    if isinstance(axis, Integral):
+        axis = (axis,)
+    axis = validate_axis(axis, x.ndim)
+    if intermediate_dtype is None:
+        intermediate_dtype = dtype
+
+    split_every = _normalize_split_every(split_every, axis)
+
+    result = partial_reduce(
+        x,
+        partial(combine_func, **(extra_func_kwargs or {})),
+        initial_func=partial(
+            func, axis=axis, keepdims=True, **(extra_func_kwargs or {})
+        ),
+        split_every=split_every,
+        dtype=intermediate_dtype,
+    )
+
+    # combine intermediates
+    result = tree_reduce(
+        result,
+        partial(combine_func, **(extra_func_kwargs or {})),
+        axis=axis,
+        dtype=intermediate_dtype,
+        split_every=split_every,
+    )
+
+    # aggregate final chunks
+    if aggegrate_func is not None:
+        result = map_blocks(aggegrate_func, result, dtype=dtype)
+
+    if not keepdims:
+        axis_to_squeeze = tuple(i for i in axis if result.shape[i] == 1)
+        if len(axis_to_squeeze) > 0:
+            result = squeeze(result, axis_to_squeeze)
+
+    from cubed.array_api import astype
+
+    result = astype(result, dtype, copy=False)
+
+    return result
+
+
+def _normalize_split_every(split_every, axis):
+    split_every = split_every or 4
+    if isinstance(split_every, dict):
+        split_every = {k: split_every.get(k, 2) for k in axis}
+    elif isinstance(split_every, Integral):
+        n = builtins.max(int(split_every ** (1 / (len(axis) or 1))), 2)
+        split_every = dict.fromkeys(axis, n)
+    else:
+        raise ValueError("split_every must be a int or a dict")
+    return split_every
+
+
+def tree_reduce(
+    x,
+    func,
+    axis,
+    dtype,
+    split_every=None,
+):
+    """Apply a reduction function repeatedly across multiple axes."""
+    if axis is None:
+        axis = tuple(range(x.ndim))
+    if isinstance(axis, Integral):
+        axis = (axis,)
+    axis = validate_axis(axis, x.ndim)
+
+    split_every = _normalize_split_every(split_every, axis)
+
+    depth = 0
+    for i, n in enumerate(x.numblocks):
+        if i in split_every and split_every[i] != 1:
+            depth = int(builtins.max(depth, math.ceil(math.log(n, split_every[i]))))
+    for _ in range(depth):
+        x = partial_reduce(
+            x,
+            func,
+            split_every=split_every,
+            dtype=dtype,
+        )
+    return x
+
+
+def partial_reduce(x, func, initial_func=None, split_every=None, dtype=None):
+    """Apply a reduction function to multiple blocks across multiple axes.
+
+    Parameters
+    ----------
+    x: Array
+        Array being reduced along one or more axes
+    func: callable
+        Reduction function to apply to each chunk of data, resulting in a chunk
+        with size one in each of the reduction axes.
+    initial_func: callable, optional
+        Function to apply to each chunk of data before reduction.
+    split_every: int >= 2 or dict(axis: int), optional
+        The depth of the recursive aggregation.
+    dtype: DType
+        Output data type.
+    """
+    # map over output chunks
+    chunks = [
+        (1,) * math.ceil(len(c) / split_every[i]) if i in split_every else c
+        for (i, c) in enumerate(x.chunks)
+    ]
+    shape = tuple(map(sum, chunks))
+    axis = tuple(ax for ax in split_every.keys())
+
+    def block_function(out_key):
+        out_coords = out_key[1:]
+
+        # return a tuple with a single item that is an iterator of input keys to be merged
+        in_keys = [
+            list(
+                range(
+                    bi * split_every.get(i, 1),
+                    min((bi + 1) * split_every.get(i, 1), x.numblocks[i]),
+                )
+            )
+            for i, bi in enumerate(out_coords)
+        ]
+        return (iter([(x.name,) + tuple(p) for p in product(*in_keys)]),)
+
+    # Since block_function returns an iterator of input keys, the the array chunks passed to
+    # _partial_reduce are retrieved one at a time. However, we need an extra chunk of memory
+    # to stay within limits (maybe because the iterator doesn't free the previous object
+    # before getting the next).
+    extra_projected_mem = x.chunkmem
+
+    return general_blockwise(
+        _partial_reduce,
+        block_function,
+        x,
+        shape=shape,
+        dtype=dtype,
+        chunks=chunks,
+        extra_projected_mem=extra_projected_mem,
+        reduce_func=func,
+        initial_func=initial_func,
+        axis=axis,
+    )
+
+
+def _partial_reduce(arrays, reduce_func=None, initial_func=None, axis=None):
+    # reduce each array in turn, accumulating in result
+    result = None
+    for array in arrays:
+        if initial_func is not None:
+            array = initial_func(array)
+        reduced_chunk = reduce_func(array, axis=axis, keepdims=True)
+        if result is None:
+            result = reduced_chunk
+        elif isinstance(result, dict):
+            assert result.keys() == reduced_chunk.keys()
+            result = {
+                # only need to concatenate along first axis
+                k: np.concatenate([result[k], reduced_chunk[k]], axis=axis[0])
+                for k in result.keys()
+            }
+            result = reduce_func(result, axis=axis, keepdims=True)
+        else:
+            # only need to concatenate along first axis
+            result = np.concatenate([result, reduced_chunk], axis=axis[0])
+            result = reduce_func(result, axis=axis, keepdims=True)
+
+    return result
+
+
 def arg_reduction(x, /, arg_func, axis=None, *, keepdims=False):
     """A reduction that returns the array indexes, not the values."""
     dtype = np.int64  # index data type


Test patch:
diff --git a/cubed/tests/test_array_api.py b/cubed/tests/test_array_api.py
index 8c2f5f63..7eb51807 100644
--- a/cubed/tests/test_array_api.py
+++ b/cubed/tests/test_array_api.py
@@ -551,11 +551,12 @@ def test_argmin_axis_0(spec):
 # Statistical functions
 
 
-def test_mean_axis_0(spec, executor):
+@pytest.mark.parametrize("use_new_impl", [False, True])
+def test_mean_axis_0(spec, executor, use_new_impl):
     a = xp.asarray(
         [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], chunks=(2, 2), spec=spec
     )
-    b = xp.mean(a, axis=0)
+    b = xp.mean(a, axis=0, use_new_impl=use_new_impl)
     assert_array_equal(
         b.compute(executor=executor),
         np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]).mean(axis=0),
diff --git a/cubed/tests/test_core.py b/cubed/tests/test_core.py
index cc19dd7c..e352886a 100644
--- a/cubed/tests/test_core.py
+++ b/cubed/tests/test_core.py
@@ -10,7 +10,7 @@ import cubed
 import cubed.array_api as xp
 import cubed.random
 from cubed.backend_array_api import namespace as nxp
-from cubed.core.ops import merge_chunks
+from cubed.core.ops import merge_chunks, partial_reduce, tree_reduce
 from cubed.tests.utils import (
     ALL_EXECUTORS,
     MAIN_EXECUTORS,
@@ -352,6 +352,23 @@ def test_reduction_not_enough_memory(tmp_path):
         xp.sum(a, axis=0, dtype=np.uint8)
 
 
+def test_partial_reduce(spec):
+    a = xp.asarray(np.arange(242).reshape((11, 22)), chunks=(3, 4), spec=spec)
+    b = partial_reduce(a, np.sum, split_every={0: 2})
+    c = partial_reduce(b, np.sum, split_every={0: 2})
+    assert_array_equal(
+        c.compute(), np.arange(242).reshape((11, 22)).sum(axis=0, keepdims=True)
+    )
+
+
+def test_tree_reduce(spec):
+    a = xp.asarray(np.arange(242).reshape((11, 22)), chunks=(3, 4), spec=spec)
+    b = tree_reduce(a, np.sum, axis=0, dtype=np.int64, split_every={0: 2})
+    assert_array_equal(
+        b.compute(), np.arange(242).reshape((11, 22)).sum(axis=0, keepdims=True)
+    )
+
+
 @pytest.mark.parametrize(
     "target_chunks, expected_chunksize",
     [
diff --git a/cubed/tests/test_mem_utilization.py b/cubed/tests/test_mem_utilization.py
index bb7013ea..bb04a72e 100644
--- a/cubed/tests/test_mem_utilization.py
+++ b/cubed/tests/test_mem_utilization.py
@@ -3,11 +3,14 @@ import shutil
 
 import pytest
 
+from cubed.core.ops import partial_reduce
+
 pytest.importorskip("lithops")
 
 import cubed
 import cubed.array_api as xp
 import cubed.random
+from cubed.backend_array_api import namespace as nxp
 from cubed.extensions.history import HistoryCallback
 from cubed.runtime.executors.lithops import LithopsDagExecutor
 from cubed.tests.utils import LITHOPS_LOCAL_CONFIG
@@ -205,6 +208,15 @@ def test_mean(tmp_path, spec):
     run_operation(tmp_path, "mean", b)
 
 
+@pytest.mark.slow
+def test_sum_partial_reduce(tmp_path, spec):
+    a = cubed.random.random(
+        (40000, 10000), chunks=(5000, 5000), spec=spec
+    )  # 200MB chunks
+    b = partial_reduce(a, nxp.sum, split_every={0: 8})
+    run_operation(tmp_path, "sum_partial_reduce", b)
+
+
 # Internal functions