Initial GPU support

.github/workflows/gpu_test.yml

@@ -0,0 +1,48 @@
+# This workflow will install Python dependencies, run tests and lint with a variety of Python versions
+# For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions
+
+name: GPU Test V3
+
+on:
+  push:
+    branches: [ v3 ]
+  pull_request:
+    branches: [ v3 ]
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  test:
+    name: py=${{ matrix.python-version }}, np=${{ matrix.numpy-version }}, deps=${{ matrix.dependency-set }}
+
+    runs-on: gpu-runner
+    strategy:
+      matrix:
+        python-version: ['3.11']
+        numpy-version: ['2.0']
+        dependency-set: ["gpu"]
+
+    steps:
+    - uses: actions/checkout@v4
+    - name: GPU check
+      run: |
+        nvidia-smi
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+        cache: 'pip'
+    - name: Install Hatch and CuPy
+      run: |
+        python -m pip install --upgrade pip 
+        pip install hatch cupy-cuda12x
+    - name: Set Up Hatch Env
+      run: |
+        hatch env create test.py${{ matrix.python-version }}-${{ matrix.numpy-version }}-${{ matrix.dependency-set }}
+        hatch env run -e test.py${{ matrix.python-version }}-${{ matrix.numpy-version }}-${{ matrix.dependency-set }} list-env
+    - name: Run Tests
+      run: |
+        hatch env run --env test.py${{ matrix.python-version }}-${{ matrix.numpy-version }}-${{ matrix.dependency-set }} run-coverage-gpu

pyproject.toml

@@ -74,6 +74,9 @@ jupyter = [
     'ipytree>=0.2.2',
     'ipywidgets>=8.0.0',
 ]
+gpu = [
+    "cupy-cuda12x",
+]
 docs = [
     'sphinx',
     'sphinx-autobuild>=2021.3.14',
@@ -120,9 +123,9 @@ build.hooks.vcs.version-file = "src/zarr/_version.py"
 [tool.hatch.envs.test]
 dependencies = [
     "numpy~={matrix:numpy}",
-    "universal_pathlib"
+    "universal_pathlib",
 ]
-features = ["test", "extra"]
+features = ["test", "extra", "gpu"]
 
 [[tool.hatch.envs.test.matrix]]
 python = ["3.10", "3.11", "3.12"]
@@ -134,8 +137,14 @@ python = ["3.10", "3.11", "3.12"]
 numpy = ["1.24", "1.26", "2.0"]
 features = ["optional"]
 
+[[tool.hatch.envs.test.matrix]]
+python = ["3.10", "3.11", "3.12"]
+numpy = ["1.24", "1.26", "2.0"]
+features = ["gpu"]
+
 [tool.hatch.envs.test.scripts]
 run-coverage = "pytest --cov-config=pyproject.toml --cov=pkg --cov=tests"
+run-coverage-gpu = "pytest -m gpu --cov-config=pyproject.toml --cov=pkg --cov=tests"
 run = "run-coverage --no-cov"
 run-verbose = "run-coverage --verbose"
 run-mypy = "mypy src"
@@ -223,4 +232,8 @@ 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",
+]
+markers = [
+    "gpu: mark a test as requiring CuPy and GPU"
 ]

src/zarr/codecs/blosc.py

@@ -10,7 +10,8 @@
 
 from zarr.abc.codec import BytesBytesCodec
 from zarr.core.array_spec import ArraySpec
-from zarr.core.buffer import Buffer, as_numpy_array_wrapper
+from zarr.core.buffer import Buffer
+from zarr.core.buffer.cpu import as_numpy_array_wrapper
 from zarr.core.common import JSON, parse_enum, parse_named_configuration, to_thread
 from zarr.registry import register_codec
 

src/zarr/codecs/gzip.py

@@ -7,7 +7,8 @@
 
 from zarr.abc.codec import BytesBytesCodec
 from zarr.core.array_spec import ArraySpec
-from zarr.core.buffer import Buffer, as_numpy_array_wrapper
+from zarr.core.buffer import Buffer
+from zarr.core.buffer.cpu import as_numpy_array_wrapper
 from zarr.core.common import JSON, parse_named_configuration, to_thread
 from zarr.registry import register_codec
 

src/zarr/codecs/zstd.py

@@ -9,7 +9,8 @@
 
 from zarr.abc.codec import BytesBytesCodec
 from zarr.core.array_spec import ArraySpec
-from zarr.core.buffer import Buffer, as_numpy_array_wrapper
+from zarr.core.buffer import Buffer
+from zarr.core.buffer.cpu import as_numpy_array_wrapper
 from zarr.core.common import JSON, parse_named_configuration, to_thread
 from zarr.registry import register_codec
 

src/zarr/core/array.py

@@ -512,15 +512,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/core/buffer/__init__.py

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