Update GPU handling

This updates how we handle GPU buffers. See the new docs page for a
simple example.

The basic idea, as discussed in ..., is to use host buffers for all
metadata objects and device buffers for data.

Zarr has two types of buffers: plain buffers (used for a stream of
bytes) and ndbuffers (used for bytes that represent ndarrays). To make
it easier for users, I've added a new config option
`zarr.config.enable_gpu()` that can be used to update those both. If
we need additional customizations in the future, we can add them here.

Author: TomAugspurger

docs/user-guide/gpu.rst

@@ -0,0 +1,25 @@
+.. _user-guide-gpu:
+
+Using GPUs with Zarr
+====================
+
+Zarr can be used along with GPUs to accelerate your workload. Currently,
+Zarr supports reading data into GPU memory. In the future, Zarr will
+support GPU-accelerated codecs and file IO.
+
+Reading data into device memory
+-------------------------------
+
+.. code-block:: python
+
+   >>> import zarr
+   >>> import cupy as cp
+   >>> zarr.config.enable_cuda()
+   >>> store = zarr.storage.MemoryStore()
+   >>> type(z[:10, :10])
+   cupy.ndarray
+
+:meth:`zarr.config.enable_cuda` updates the Zarr configuration to use device
+memory for all data buffers used by Zarr. This means that any reads from a Zarr
+store will return a CuPy ndarray rather than a NumPy ndarray. Any buffers used
+for metadata will be on the host.

docs/user-guide/index.rst

@@ -23,6 +23,7 @@ Advanced Topics
     performance
     consolidated_metadata
     extending
+    gpu
 
 
 .. Coming soon

src/zarr/core/array.py

@@ -38,6 +38,7 @@
     NDBuffer,
     default_buffer_prototype,
 )
+from zarr.core.buffer.cpu import buffer_prototype as cpu_buffer_prototype
 from zarr.core.chunk_grids import RegularChunkGrid, _auto_partition, normalize_chunks
 from zarr.core.chunk_key_encodings import (
     ChunkKeyEncoding,
@@ -163,19 +164,20 @@ async def get_array_metadata(
 ) -> dict[str, JSON]:
     if zarr_format == 2:
         zarray_bytes, zattrs_bytes = await gather(
-            (store_path / ZARRAY_JSON).get(), (store_path / ZATTRS_JSON).get()
+            (store_path / ZARRAY_JSON).get(prototype=cpu_buffer_prototype),
+            (store_path / ZATTRS_JSON).get(prototype=cpu_buffer_prototype),
         )
         if zarray_bytes is None:
             raise FileNotFoundError(store_path)
     elif zarr_format == 3:
-        zarr_json_bytes = await (store_path / ZARR_JSON).get()
+        zarr_json_bytes = await (store_path / ZARR_JSON).get(prototype=cpu_buffer_prototype)
         if zarr_json_bytes is None:
             raise FileNotFoundError(store_path)
     elif zarr_format is None:
         zarr_json_bytes, zarray_bytes, zattrs_bytes = await gather(
-            (store_path / ZARR_JSON).get(),
-            (store_path / ZARRAY_JSON).get(),
-            (store_path / ZATTRS_JSON).get(),
+            (store_path / ZARR_JSON).get(prototype=cpu_buffer_prototype),
+            (store_path / ZARRAY_JSON).get(prototype=cpu_buffer_prototype),
+            (store_path / ZATTRS_JSON).get(prototype=cpu_buffer_prototype),
         )
         if zarr_json_bytes is not None and zarray_bytes is not None:
             # warn and favor v3
@@ -1295,7 +1297,7 @@ async def _save_metadata(self, metadata: ArrayMetadata, ensure_parents: bool = F
         """
         Asynchronously save the array metadata.
         """
-        to_save = metadata.to_buffer_dict(default_buffer_prototype())
+        to_save = metadata.to_buffer_dict(cpu_buffer_prototype)
         awaitables = [set_or_delete(self.store_path / key, value) for key, value in to_save.items()]
 
         if ensure_parents:
@@ -1307,7 +1309,7 @@ async def _save_metadata(self, metadata: ArrayMetadata, ensure_parents: bool = F
                     [
                         (parent.store_path / key).set_if_not_exists(value)
                         for key, value in parent.metadata.to_buffer_dict(
-                            default_buffer_prototype()
+                            cpu_buffer_prototype
                         ).items()
                     ]
                 )

src/zarr/core/buffer/gpu.py

@@ -13,6 +13,10 @@
 
 from zarr.core.buffer import core
 from zarr.core.buffer.core import ArrayLike, BufferPrototype, NDArrayLike
+from zarr.registry import (
+    register_buffer,
+    register_ndbuffer,
+)
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
@@ -215,3 +219,6 @@ def __setitem__(self, key: Any, value: Any) -> None:
 
 
 buffer_prototype = BufferPrototype(buffer=Buffer, nd_buffer=NDBuffer)
+
+register_buffer(Buffer)
+register_ndbuffer(NDBuffer)

src/zarr/core/config.py

@@ -29,10 +29,13 @@
 
 from __future__ import annotations
 
-from typing import Any, Literal, cast
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 from donfig import Config as DConfig
 
+if TYPE_CHECKING:
+    from donfig.config_obj import ConfigSet
+
 
 class BadConfigError(ValueError):
     _msg = "bad Config: %r"
@@ -56,6 +59,14 @@ def reset(self) -> None:
         self.clear()
         self.refresh()
 
+    def enable_gpu(self) -> ConfigSet:
+        """
+        Configure Zarr to use GPUs where possible.
+        """
+        return self.set(
+            {"buffer": "zarr.core.buffer.gpu.Buffer", "ndbuffer": "zarr.core.buffer.gpu.NDBuffer"}
+        )
+
 
 # The default configuration for zarr
 config = Config(

src/zarr/testing/utils.py

@@ -38,7 +38,7 @@ def has_cupy() -> bool:
         return False
 
 
-T_Callable = TypeVar("T_Callable", bound=Callable[[], Coroutine[Any, Any, None]])
+T_Callable = TypeVar("T_Callable", bound=Callable[..., Coroutine[Any, Any, None] | None])
 
 
 # Decorator for GPU tests

tests/test_api.py

@@ -27,6 +27,7 @@
 from zarr.errors import MetadataValidationError
 from zarr.storage import MemoryStore
 from zarr.storage._utils import normalize_path
+from zarr.testing.utils import gpu_test
 
 
 def test_create(memory_store: Store) -> None:
@@ -1121,3 +1122,38 @@ def test_open_array_with_mode_r_plus(store: Store) -> None:
     assert isinstance(z2, Array)
     assert (z2[:] == 1).all()
     z2[:] = 3
+
+
+@gpu_test
+@pytest.mark.parametrize(
+    "store",
+    ["local", "memory", "zip"],
+    indirect=True,
+)
+@pytest.mark.parametrize("zarr_format", [None, 2, 3])
+def test_gpu_basic(store: Store, zarr_format: ZarrFormat | None) -> None:
+    import cupy as cp
+
+    if zarr_format == 2:
+        # Without this, the zstd codec attempts to convert the cupy
+        # array to bytes.
+        compressors = None
+    else:
+        compressors = "auto"
+
+    with zarr.config.enable_gpu():
+        src = cp.random.uniform(size=(100, 100))  # allocate on the device
+        z = zarr.create_array(
+            store,
+            name="a",
+            shape=src.shape,
+            chunks=(10, 10),
+            dtype=src.dtype,
+            overwrite=True,
+            zarr_format=zarr_format,
+            compressors=compressors,
+        )
+        z[:10, :10] = src[:10, :10]
+
+        result = z[:10, :10]
+        cp.testing.assert_array_equal(result, src[:10, :10])