Initial GPU support

pyproject.toml

@@ -60,6 +60,9 @@ jupyter = [
     'ipytree>=0.2.2',
     'ipywidgets>=8.0.0',
 ]
+gpu = [
+    "cupy>=13.0.0",
+]
 docs = [
     'sphinx',
     'sphinx-autobuild>=2021.3.14',
@@ -220,4 +223,5 @@ filterwarnings = [
     "error:::zarr.*",
     "ignore:PY_SSIZE_T_CLEAN will be required.*:DeprecationWarning",
     "ignore:The loop argument is deprecated since Python 3.8.*:DeprecationWarning",
+    "ignore:Creating a zarr.buffer.gpu.*:UserWarning",
 ]

src/zarr/array.py

@@ -22,7 +22,11 @@
 from zarr.attributes import Attributes
 from zarr.buffer import BufferPrototype, NDArrayLike, NDBuffer, default_buffer_prototype
 from zarr.chunk_grids import RegularChunkGrid
-from zarr.chunk_key_encodings import ChunkKeyEncoding, DefaultChunkKeyEncoding, V2ChunkKeyEncoding
+from zarr.chunk_key_encodings import (
+    ChunkKeyEncoding,
+    DefaultChunkKeyEncoding,
+    V2ChunkKeyEncoding,
+)
 from zarr.codecs import BytesCodec
 from zarr.codecs._v2 import V2Compressor, V2Filters
 from zarr.codecs.pipeline import BatchedCodecPipeline
@@ -76,7 +80,9 @@ def parse_array_metadata(data: Any) -> ArrayV2Metadata | ArrayV3Metadata:
     raise TypeError
 
 
-def create_codec_pipeline(metadata: ArrayV2Metadata | ArrayV3Metadata) -> BatchedCodecPipeline:
+def create_codec_pipeline(
+    metadata: ArrayV2Metadata | ArrayV3Metadata,
+) -> BatchedCodecPipeline:
     if isinstance(metadata, ArrayV3Metadata):
         return BatchedCodecPipeline.from_list(metadata.codecs)
     elif isinstance(metadata, ArrayV2Metadata):
@@ -474,7 +480,10 @@ async def _get_selection(
         return out_buffer.as_ndarray_like()
 
     async def getitem(
-        self, selection: BasicSelection, *, prototype: BufferPrototype = default_buffer_prototype
+        self,
+        selection: BasicSelection,
+        *,
+        prototype: BufferPrototype = default_buffer_prototype,
     ) -> NDArrayLike:
         indexer = BasicIndexer(
             selection,
@@ -502,15 +511,24 @@ async def _set_selection(
 
         # check value shape
         if np.isscalar(value):
-            value = np.asanyarray(value, dtype=self.metadata.dtype)
+            array_like = prototype.buffer.create_zero_length().as_array_like()
+            if isinstance(array_like, np._typing._SupportsArrayFunc):
+                # TODO: need to handle array types that don't support __array_function__
+                # like PyTorch and JAX
+                array_like_ = cast(np._typing._SupportsArrayFunc, array_like)
+            value = np.asanyarray(value, dtype=self.metadata.dtype, like=array_like_)
         else:
             if not hasattr(value, "shape"):
                 value = np.asarray(value, self.metadata.dtype)
             # assert (
             #     value.shape == indexer.shape
             # ), f"shape of value doesn't match indexer shape. Expected {indexer.shape}, got {value.shape}"
             if not hasattr(value, "dtype") or value.dtype.name != self.metadata.dtype.name:
-                value = np.array(value, dtype=self.metadata.dtype, order="A")
+                if hasattr(value, "astype"):
+                    # Handle things that are already NDArrayLike more efficiently
+                    value = value.astype(dtype=self.metadata.dtype, order="A")
+                else:
+                    value = np.array(value, dtype=self.metadata.dtype, order="A")
         value = cast(NDArrayLike, value)
         # We accept any ndarray like object from the user and convert it
         # to a NDBuffer (or subclass). From this point onwards, we only pass

src/zarr/buffer/__init__.py

@@ -0,0 +1,17 @@
+from zarr.buffer.core import (
+    ArrayLike,
+    Buffer,
+    BufferPrototype,
+    NDArrayLike,
+    NDBuffer,
+)
+from zarr.buffer.cpu import default_buffer_prototype
+
+__all__ = [
+    "ArrayLike",
+    "Buffer",
+    "NDArrayLike",
+    "NDBuffer",
+    "BufferPrototype",
+    "default_buffer_prototype",
+]

src/zarr/buffer.py → src/zarr/buffer/core.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 import sys
-from collections.abc import Callable, Iterable, Sequence
+from abc import ABC, abstractmethod
+from collections.abc import Iterable, Sequence
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -106,7 +107,7 @@ def check_item_key_is_1d_contiguous(key: Any) -> None:
         raise ValueError("slice must be contiguous")
 
 
-class Buffer:
+class Buffer(ABC):
     """A flat contiguous memory block
 
     We use Buffer throughout Zarr to represent a contiguous block of memory.
@@ -135,14 +136,15 @@ def __init__(self, array_like: ArrayLike):
         self._data = array_like
 
     @classmethod
+    @abstractmethod
     def create_zero_length(cls) -> Self:
         """Create an empty buffer with length zero
 
         Returns
         -------
             New empty 0-length buffer
         """
-        return cls(np.array([], dtype="b"))
+        ...
 
     @classmethod
     def from_array_like(cls, array_like: ArrayLike) -> Self:
@@ -160,6 +162,7 @@ def from_array_like(cls, array_like: ArrayLike) -> Self:
         return cls(array_like)
 
     @classmethod
+    @abstractmethod
     def from_buffer(cls, buffer: Buffer) -> Self:
         """Create a new buffer of an existing Buffer
 
@@ -179,10 +182,16 @@ def from_buffer(cls, buffer: Buffer) -> Self:
         Returns
         -------
             A new buffer representing the content of the input buffer
+
+        Note
+        ----
+        Subclasses of `Buffer` must override this method to implement
+        more optimal conversions that avoid copies where possible
         """
-        return cls.from_array_like(buffer.as_array_like())
+        ...
 
     @classmethod
+    @abstractmethod
     def from_bytes(cls, bytes_like: BytesLike) -> Self:
         """Create a new buffer of a bytes-like object (host memory)
 
@@ -195,7 +204,7 @@ def from_bytes(cls, bytes_like: BytesLike) -> Self:
         -------
             New buffer representing `bytes_like`
         """
-        return cls.from_array_like(np.frombuffer(bytes_like, dtype="b"))
+        ...
 
     def as_array_like(self) -> ArrayLike:
         """Returns the underlying array (host or device memory) of this buffer
@@ -208,6 +217,7 @@ def as_array_like(self) -> ArrayLike:
         """
         return self._data
 
+    @abstractmethod
     def as_numpy_array(self) -> npt.NDArray[Any]:
         """Returns the buffer as a NumPy array (host memory).
 
@@ -219,7 +229,7 @@ def as_numpy_array(self) -> npt.NDArray[Any]:
         -------
             NumPy array of this buffer (might be a data copy)
         """
-        return np.asanyarray(self._data)
+        ...
 
     def to_bytes(self) -> bytes:
         """Returns the buffer as `bytes` (host memory).
@@ -246,14 +256,10 @@ def __setitem__(self, key: slice, value: Any) -> None:
     def __len__(self) -> int:
         return self._data.size
 
+    @abstractmethod
     def __add__(self, other: Buffer) -> Self:
         """Concatenate two buffers"""
-
-        other_array = other.as_array_like()
-        assert other_array.dtype == np.dtype("b")
-        return self.__class__(
-            np.concatenate((np.asanyarray(self._data), np.asanyarray(other_array)))
-        )
+        ...
 
 
 class NDBuffer:
@@ -287,6 +293,7 @@ def __init__(self, array: NDArrayLike):
         self._data = array
 
     @classmethod
+    @abstractmethod
     def create(
         cls,
         *,
@@ -318,10 +325,7 @@ def create(
         A subclass can overwrite this method to create a ndarray-like object
         other then the default Numpy array.
         """
-        ret = cls(np.empty(shape=tuple(shape), dtype=dtype, order=order))
-        if fill_value is not None:
-            ret.fill(fill_value)
-        return ret
+        ...
 
     @classmethod
     def from_ndarray_like(cls, ndarray_like: NDArrayLike) -> Self:
@@ -339,6 +343,7 @@ def from_ndarray_like(cls, ndarray_like: NDArrayLike) -> Self:
         return cls(ndarray_like)
 
     @classmethod
+    @abstractmethod
     def from_numpy_array(cls, array_like: npt.ArrayLike) -> Self:
         """Create a new buffer of Numpy array-like object
 
@@ -351,7 +356,7 @@ def from_numpy_array(cls, array_like: npt.ArrayLike) -> Self:
         -------
             New buffer representing `array_like`
         """
-        return cls.from_ndarray_like(np.asanyarray(array_like))
+        ...
 
     def as_ndarray_like(self) -> NDArrayLike:
         """Returns the underlying array (host or device memory) of this buffer
@@ -364,6 +369,7 @@ def as_ndarray_like(self) -> NDArrayLike:
         """
         return self._data
 
+    @abstractmethod
     def as_numpy_array(self) -> npt.NDArray[Any]:
         """Returns the buffer as a NumPy array (host memory).
 
@@ -375,7 +381,7 @@ def as_numpy_array(self) -> npt.NDArray[Any]:
         -------
             NumPy array of this buffer (might be a data copy)
         """
-        return np.asanyarray(self._data)
+        ...
 
     @property
     def dtype(self) -> np.dtype[Any]:
@@ -406,13 +412,11 @@ def squeeze(self, axis: tuple[int, ...]) -> Self:
     def astype(self, dtype: npt.DTypeLike, order: Literal["K", "A", "C", "F"] = "K") -> Self:
         return self.__class__(self._data.astype(dtype=dtype, order=order))
 
-    def __getitem__(self, key: Any) -> Self:
-        return self.__class__(np.asanyarray(self._data.__getitem__(key)))
+    @abstractmethod
+    def __getitem__(self, key: Any) -> Self: ...
 
-    def __setitem__(self, key: Any, value: Any) -> None:
-        if isinstance(value, NDBuffer):
-            value = value._data
-        self._data.__setitem__(key, value)
+    @abstractmethod
+    def __setitem__(self, key: Any, value: Any) -> None: ...
 
     def __len__(self) -> int:
         return self._data.__len__()
@@ -433,34 +437,6 @@ def transpose(self, axes: SupportsIndex | Sequence[SupportsIndex] | None) -> Sel
         return self.__class__(self._data.transpose(axes))
 
 
-def as_numpy_array_wrapper(
-    func: Callable[[npt.NDArray[Any]], bytes], buf: Buffer, prototype: BufferPrototype
-) -> Buffer:
-    """Converts the input of `func` to a numpy array and the output back to `Buffer`.
-
-    This function is useful when calling a `func` that only support host memory such
-    as `GZip.decode` and `Blosc.decode`. In this case, use this wrapper to convert
-    the input `buf` to a Numpy array and convert the result back into a `Buffer`.
-
-    Parameters
-    ----------
-    func
-        The callable that will be called with the converted `buf` as input.
-        `func` must return bytes, which will be converted into a `Buffer`
-        before returned.
-    buf
-        The buffer that will be converted to a Numpy array before given as
-        input to `func`.
-    prototype
-        The prototype of the output buffer.
-
-    Returns
-    -------
-        The result of `func` converted to a `Buffer`
-    """
-    return prototype.buffer.from_bytes(func(buf.as_numpy_array()))
-
-
 class BufferPrototype(NamedTuple):
     """Prototype of the Buffer and NDBuffer class
 
@@ -476,7 +452,3 @@ class BufferPrototype(NamedTuple):
 
     buffer: type[Buffer]
     nd_buffer: type[NDBuffer]
-
-
-# The default buffer prototype used throughout the Zarr codebase.
-default_buffer_prototype = BufferPrototype(buffer=Buffer, nd_buffer=NDBuffer)