Basic support for device arrays (#772)

* cupy compatibility

Fix empty allocation

Register cubed.Array in to_numpy

Some tests of Array implementing the array API need to convert
from the active backend to NumPy.

* skip processes for cupy

* array api passing

* fixup

* review

* ignore cupy type annotations

* try to handle jax setitem

* more jax fixes

Author: TomAugspurger

cubed/_testing.py

@@ -0,0 +1,42 @@
+import functools
+import importlib.util
+
+import numpy as np
+import numpy.testing as npt
+
+import cubed
+
+
+@functools.cache
+def has_cupy() -> bool:
+    return importlib.util.find_spec("cupy") is not None
+
+
+@functools.singledispatch
+def to_numpy(a):
+    return np.asarray(a)
+
+
+@to_numpy.register(cubed.Array)
+def _(a: cubed.Array) -> np.ndarray:
+    return to_numpy(a.compute())
+
+
+if has_cupy():
+    import cupy
+
+    @to_numpy.register(cupy.ndarray)
+    def _(a):
+        return a.get()
+
+
+def assert_array_equal(a, b, **kwargs):
+    a = to_numpy(a)
+    b = to_numpy(b)
+    npt.assert_array_equal(a, b, **kwargs)
+
+
+def assert_allclose(a, b, **kwargs):
+    a = to_numpy(a)
+    b = to_numpy(b)
+    npt.assert_allclose(a, b, **kwargs)

cubed/array_api/manipulation_functions.py

@@ -6,6 +6,7 @@
 from toolz import reduce
 
 from cubed.array_api.creation_functions import empty
+from cubed.backend_array_api import IS_IMMUTABLE_ARRAY
 from cubed.backend_array_api import namespace as nxp
 from cubed.core import (
     blockwise,
@@ -213,14 +214,17 @@ def _read_concat_chunk(
     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)
+    out = nxp.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]
+        if IS_IMMUTABLE_ARRAY:
+            out = out.at[lout_selection].set(array[lchunk_selection])
+        else:
+            out[lout_selection] = array[lchunk_selection]
     return out
 
 

cubed/backend_array_api.py

@@ -32,19 +32,32 @@
     import array_api_compat.numpy
 
     namespace = array_api_compat.numpy
+    xp_name = "numpy"
 
 
 # These functions to convert to/from backend arrays
 # assume that no extra memory is allocated, by using the
 # Python buffer protocol.
 # See https://data-apis.org/array-api/latest/API_specification/generated/array_api.asarray.html
+if "cupy" in namespace.__name__:
+    # zarr-python 3.x natively supports some device buffers (currently just cupy,
+    # but https://github.com/zarr-developers/zarr-python/issues/2658 is expanding the
+    # set). For these backends, we *don't* want to copy to the host.
 
+    def backend_array_to_numpy_array(arr):
+        return arr
 
-def backend_array_to_numpy_array(arr):
-    return np.asarray(arr)
+else:
+
+    def backend_array_to_numpy_array(arr):
+        return np.asarray(arr)
 
 
 def numpy_array_to_backend_array(arr, *, dtype=None):
     if isinstance(arr, dict):
         return {k: namespace.asarray(v, dtype=dtype) for k, v in arr.items()}
     return namespace.asarray(arr, dtype=dtype)
+
+
+# jax doesn't support in-place assignment, so we use .at[].set() instead.
+IS_IMMUTABLE_ARRAY = "jax" in xp_name

cubed/core/ops.py

@@ -14,8 +14,8 @@
 from toolz import map
 
 from cubed import config
+from cubed.backend_array_api import IS_IMMUTABLE_ARRAY, numpy_array_to_backend_array
 from cubed.backend_array_api import namespace as nxp
-from cubed.backend_array_api import numpy_array_to_backend_array
 from cubed.core.array import CoreArray, check_array_specs, compute, gensym
 from cubed.core.plan import Plan, new_temp_path
 from cubed.primitive.blockwise import blockwise as primitive_blockwise
@@ -559,14 +559,17 @@ def _assemble_index_chunk(
     indexer = _create_zarr_indexer(in_sel, in_shape, in_chunksize)
 
     shape = indexer.shape
-    out = np.empty(shape, dtype=dtype)
+    out = nxp.empty(shape, dtype=dtype)
 
     if array_size(shape) > 0:
         _, lchunk_selection, lout_selection, *_ = zip(*indexer)
         for ai, chunk_select, out_select in zip(
             arrays, lchunk_selection, lout_selection
         ):
-            out[out_select] = ai[chunk_select]
+            if IS_IMMUTABLE_ARRAY:
+                out = out.at[out_select].set(ai[chunk_select])
+            else:
+                out[out_select] = ai[chunk_select]
 
     if func is not None:
         if has_keyword(func, "block_id"):

cubed/tests/test_array_api.py

@@ -1,12 +1,18 @@
 import fsspec
 import numpy as np
 import pytest
-from numpy.testing import assert_allclose, assert_array_equal
 
 import cubed
 import cubed.array_api as xp
+from cubed._testing import assert_allclose, assert_array_equal
 from cubed.array_api.manipulation_functions import reshape_chunks
-from cubed.tests.utils import ALL_EXECUTORS, MAIN_EXECUTORS, MODAL_EXECUTORS
+from cubed.backend_array_api import namespace as nxp
+from cubed.tests.utils import (
+    ALL_EXECUTORS,
+    MAIN_EXECUTORS,
+    MODAL_EXECUTORS,
+    skip_if_cupy,
+)
 
 
 @pytest.fixture
@@ -20,6 +26,8 @@ def spec(tmp_path):
     ids=[executor.name for executor in MAIN_EXECUTORS],
 )
 def executor(request):
+    if request.param.name == "processes" and "cupy" in nxp.__name__:
+        pytest.skip(reason="CuPy is not supported with 'processes' executor")
     return request.param
 
 
@@ -29,6 +37,8 @@ def executor(request):
     ids=[executor.name for executor in ALL_EXECUTORS],
 )
 def any_executor(request):
+    if request.param.name == "processes" and "cupy" in nxp.__name__:
+        pytest.skip(reason="CuPy is not supported with 'processes' executor")
     return request.param
 
 
@@ -384,6 +394,7 @@ def test_index_slice_unsupported_step(spec):
 
 
 @pytest.mark.parametrize("axis", [0, 1])
+@skip_if_cupy  # ndindex with a cupy.ndarray
 def test_take(spec, axis):
     a = xp.asarray(
         [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]],

cubed/tests/utils.py

@@ -5,6 +5,7 @@
 import numpy as np
 
 from cubed import config
+from cubed.backend_array_api import namespace as nxp
 from cubed.runtime.create import create_executor
 from cubed.runtime.types import Callback
 from cubed.storage.backend import open_backend_array
@@ -125,3 +126,14 @@ def execute_pipeline(pipeline, executor):
     dag = nx.MultiDiGraph()
     dag.add_node("node", pipeline=pipeline)
     executor.execute_dag(dag)
+
+
+try:
+    import pytest
+except ImportError:
+    pass
+else:
+    skip_if_cupy = pytest.mark.skipif(
+        "cupy" in nxp.__name__,
+        reason="CuPy is not supported",
+    )

docs/user-guide/gpus.md

@@ -0,0 +1,24 @@
+# GPU Support
+
+Cubed has experimental support for using GPU-backed ndarrays. With zarr-python's
+[native GPU support], you can load data into GPU memory, perform some cubed
+computation on the GPU, and write the result, while minimizing the number of host
+to device transfers.
+
+```{note}
+Currently only NVIDIA GPUs and [CuPy] arrays are supported.
+```
+
+Set the following environment variables to control whether host or device arrays
+are in Cubed and Zarr.
+
+```shell
+# syntax may differ in your shell
+export CUBED_BACKEND_ARRAY_API_MODULE="array_api_compat.cupy"
+export ZARR_BUFFER="zarr.buffer.gpu.Buffer"
+export ZARR_NDBUFFER="zarr.buffer.gpu.NDBuffer"
+```
+
+
+[native GPU support]: https://zarr.readthedocs.io/en/stable/user-guide/gpu.html
+[CuPy]: https://cupy.dev/

docs/user-guide/index.md

@@ -13,4 +13,5 @@ reliability
 optimization
 scaling
 diagnostics
+gpus
 ```

setup.cfg

@@ -19,6 +19,8 @@ ignore_missing_imports = True
 ignore_missing_imports = True
 [mypy-coiled.*]
 ignore_missing_imports = True
+[mypy-cupy.*]
+ignore_missing_imports = True
 [mypy-dask.*]
 ignore_missing_imports = True
 [mypy-donfig.*]