Zstd Codec on the GPU

changes/2863.feature.md

@@ -0,0 +1,11 @@
+Added GPU-accelerated Zstd Codec
+
+This adds support for decoding with the Zstd Codec on NVIDIA GPUs using the
+nvidia-nvcomp library.
+
+With `zarr.config.enable_gpu()`, buffers will be decoded using the GPU
+and the output will reside in device memory.
+
+The optional-dependency extra for GPU support is now named `cuda` (``pip install "zarr[cuda]"``)
+instead of `gpu`, since it installs CUDA-specific packages (CuPy and nvCOMP). The `gpu` extra
+has been removed.

docs/user-guide/config.md

@@ -1,4 +1,4 @@
-# Runtime configuration
+# Runtime configuration {#config-runtime-configuration}
 
 [`zarr.config`][] is responsible for managing the configuration of zarr and
 is based on the [donfig](https://github.com/pytroll/donfig) Python library.
@@ -50,6 +50,24 @@ first register the implementations in the registry and then select them in the c
 For example, an implementation of the bytes codec in a class `'custompackage.NewBytesCodec'`,
 requires the value of `codecs.bytes` to be `'custompackage.NewBytesCodec'`.
 
+## Codecs
+
+Zarr and zarr-python split the logical codec definition from the implementation.
+The Zarr metadata serialized in the store specifies just the codec name and
+configuration. To resolve the specific implementation, a Python class, that's
+used at runtime to encode or decode data, zarr-python looks up the codec name
+in the codec registry.
+
+For example, after calling `zarr.config.enable_gpu()`, an nvcomp-based
+codec will be used:
+
+```python exec="true" session="config" source="above" result="ansi"
+with zarr.config.enable_gpu():
+    print(zarr.config.get('codecs.zstd'))
+```
+
+## Default Configuration
+
 This is the current default configuration:
 
 ```python exec="true" session="config" source="above" result="ansi"

docs/user-guide/gpu.md

@@ -2,26 +2,17 @@
 
 Zarr can use GPUs to accelerate your workload by running `zarr.config.enable_gpu()`.
 
-!!! note
-    `zarr-python` currently supports reading the ndarray data into device (GPU)
-    memory as the final stage of the codec pipeline. Data will still be read into
-    or copied to host (CPU) memory for encoding and decoding.
-
-    In the future, codecs will be available for compressing and decompressing data on
-    the GPU, avoiding the need to move data between the host and device for
-    compression and decompression.
-
 ## Installation
 
 Zarr's GPU support requires [CuPy](https://cupy.dev), which in turn requires a
-CUDA-compatible NVIDIA GPU. CuPy can be installed alongside Zarr with the `gpu`
+CUDA-compatible NVIDIA GPU. CuPy and nvCOMP can be installed alongside Zarr with the `cuda`
 extra (see [Installation](installation.md) for the other optional dependency groups):
 
 ```console
-pip install "zarr[gpu]"
+pip install "zarr[cuda]"
 ```
 
-This installs the `cupy-cuda12x` package. If you need a CuPy build for a different
+This installs the `cupy-cuda12x` and `nvidia-nvcomp-cu12` packages. If you need a CuPy build for a different
 CUDA version, see the [CuPy installation guide](https://docs.cupy.dev/en/stable/install.html)
 and install the appropriate package yourself.
 
@@ -44,6 +35,9 @@ assert isinstance(z[:10, :10], cp.ndarray)
 Note that the arrays returned by reads are of type `cupy.ndarray` rather than
 NumPy arrays.
 
+For supported codecs, data will be decoded using the GPU via the [nvcomp] library.
+See [runtime-configuration][config-runtime-configuration] for more. Issues and feature requests for NVIDIA nvCOMP can be reported on the nvcomp [issue tracker].
+
 `zarr.config.enable_gpu()` returns a [donfig](https://donfig.readthedocs.io/en/latest/)
 `ConfigSet`, which can be used as a context manager to enable GPU support for a
 limited scope:
@@ -58,3 +52,6 @@ Under the hood, `enable_gpu()` selects the GPU-backed buffer classes
 `zarr.buffer.gpu.Buffer` and `zarr.buffer.gpu.NDBuffer` via the `buffer` and
 `ndbuffer` configuration keys. See [Custom array buffers](extending.md#custom-array-buffers)
 for more on Zarr's buffer classes, including how to implement your own.
+
+[nvcomp]: https://docs.nvidia.com/cuda/nvcomp/samples/python_samples.html
+[issue tracker]: https://github.com/NVIDIA/CUDALibrarySamples/issues

docs/user-guide/gpu.rst

@@ -0,0 +1,34 @@
+.. _user-guide-gpu:
+
+Using GPUs with Zarr
+====================
+
+Zarr can use GPUs to accelerate your workload by running
+:meth:`zarr.config.enable_gpu`.
+
+Reading data into device memory
+-------------------------------
+
+:meth:`zarr.config.enable_gpu` configures Zarr to use GPU memory for the data
+buffers used internally by Zarr.
+
+.. code-block:: python
+
+   >>> import zarr
+   >>> import cupy as cp  # doctest: +SKIP
+   >>> zarr.config.enable_gpu()  # doctest: +SKIP
+   >>> store = zarr.storage.MemoryStore()  # doctest: +SKIP
+   >>> z = zarr.create_array(  # doctest: +SKIP
+   ...     store=store, shape=(100, 100), chunks=(10, 10), dtype="float32",
+   ... )
+   >>> type(z[:10, :10])  # doctest: +SKIP
+   cupy.ndarray
+
+Note that the output type is a ``cupy.ndarray`` rather than a NumPy array.
+
+For supported codecs, data will be decoded using the GPU via the `nvcomp`_
+library. See :ref:`user-guide-config` for more. Isseus and feature requests
+for NVIDIA nvCOMP can be reported in the `nvcomp issue tracker`_.
+
+.. _nvcomp: https://docs.nvidia.com/cuda/nvcomp/samples/python_samples.html
+.. _nvcomp issue tracker: https://github.com/NVIDIA/CUDALibrarySamples/issues

docs/user-guide/installation.md

@@ -21,10 +21,10 @@ pip install zarr
 ```
 
 There are a number of optional dependency groups you can install for extra functionality.
-These can be installed using `pip install "zarr[<extra>]"`, e.g. `pip install "zarr[gpu]"`
+These can be installed using `pip install "zarr[<extra>]"`, e.g. `pip install "zarr[cuda]"`
 
 - `remote`: support for reading/writing to remote data stores (fsspec, obstore)
-- `gpu`: support for GPUs (cupy)
+- `cuda`: support for NVIDIA CUDA GPUs (cupy, nvidia-nvcomp)
 - `cli`: support for the `zarr` [command-line interface](cli.md) (typer)
 - `optional`: support for path-like access to local and remote stores (universal-pathlib)
 - `cast-value-rs`: support for the `cast_value` codec (cast-value-rs)

pyproject.toml

@@ -80,8 +80,9 @@ remote = [
     "fsspec>=2023.10.0",
     "obstore>=0.5.1",
 ]
-gpu = [
+cuda = [
     "cupy-cuda12x; sys_platform != 'darwin'",
+    "nvidia-nvcomp-cu12; sys_platform != 'darwin'",
 ]
 cast-value-rs = ["cast-value-rs>=0.4.2"]
 cli = ["typer"]
@@ -216,7 +217,7 @@ extra-dependencies = [
     # docs example is never executed on GPU hardware.
     "pytest-examples",
 ]
-features = ["gpu"]
+features = ["cuda"]
 
 [[tool.hatch.envs.gputest.matrix]]
 python = ["3.12", "3.13"]

src/zarr/codecs/__init__.py

@@ -4,6 +4,7 @@
 from zarr.codecs.bytes import BytesCodec, Endian
 from zarr.codecs.cast_value import CastValue
 from zarr.codecs.crc32c_ import Crc32cCodec
+from zarr.codecs.cuda import NvcompZstdCodec
 from zarr.codecs.gzip import GzipCodec
 from zarr.codecs.numcodecs import (
     BZ2,
@@ -44,6 +45,7 @@
     "Crc32cCodec",
     "Endian",
     "GzipCodec",
+    "NvcompZstdCodec",
     "ScaleOffset",
     "ShardingCodec",
     "ShardingCodecIndexLocation",

src/zarr/codecs/cuda.py

@@ -0,0 +1,185 @@
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+from functools import cached_property
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from zarr.abc.codec import BytesBytesCodec
+from zarr.core.common import JSON, parse_named_configuration
+from zarr.registry import register_codec
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+    from typing import Self
+
+    from zarr.core.array_spec import ArraySpec
+    from zarr.core.buffer import Buffer
+
+try:
+    import cupy as cp
+except ImportError:  # pragma: no cover
+    cp = None
+
+try:
+    from nvidia import nvcomp
+except ImportError:  # pragma: no cover
+    nvcomp = None
+
+
+def _parse_zstd_level(data: JSON) -> int:
+    if isinstance(data, int):
+        if data >= 23:
+            raise ValueError(f"Value must be less than or equal to 22. Got {data} instead.")
+        return data
+    raise TypeError(f"Got value with type {type(data)}, but expected an int.")
+
+
+def _parse_checksum(data: JSON) -> bool:
+    if isinstance(data, bool):
+        return data
+    raise TypeError(f"Expected bool. Got {type(data)}.")
+
+
+@dataclass(frozen=True)
+class NvcompZstdCodec(BytesBytesCodec):
+    is_fixed_size = True
+
+    level: int = 0
+    checksum: bool = False
+
+    def __init__(self, *, level: int = 0, checksum: bool = False) -> None:
+        # TODO: Set CUDA device appropriately here and also set CUDA stream
+
+        level_parsed = _parse_zstd_level(level)
+        checksum_parsed = _parse_checksum(checksum)
+
+        object.__setattr__(self, "level", level_parsed)
+        object.__setattr__(self, "checksum", checksum_parsed)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, JSON]) -> Self:
+        _, configuration_parsed = parse_named_configuration(data, "zstd")
+        return cls(**configuration_parsed)  # type: ignore[arg-type]
+
+    def to_dict(self) -> dict[str, JSON]:
+        return {
+            "name": "zstd",
+            "configuration": {"level": self.level, "checksum": self.checksum},
+        }
+
+    @cached_property
+    def _zstd_codec(self) -> nvcomp.Codec:
+        device = cp.cuda.Device()  # Select the current default device
+        stream = cp.cuda.get_current_stream()  # Use the current default stream
+        # Note: this returns an array with dtype=np.dtype("int8")
+        return nvcomp.Codec(
+            algorithm="Zstd",
+            bitstream_kind=nvcomp.BitstreamKind.RAW,
+            device_id=device.id,
+            cuda_stream=stream.ptr,
+        )
+
+    def _convert_to_nvcomp_arrays(
+        self,
+        chunks_and_specs: Iterable[tuple[Buffer | None, ArraySpec]],
+    ) -> tuple[list[nvcomp.Array], list[int]]:
+        none_indices = [i for i, (b, _) in enumerate(chunks_and_specs) if b is None]
+        filtered_inputs = [b.as_array_like() for b, _ in chunks_and_specs if b is not None]
+        # TODO: add CUDA stream here
+        return nvcomp.as_arrays(filtered_inputs), none_indices
+
+    def _convert_from_nvcomp_arrays(
+        self,
+        arrays: Iterable[nvcomp.Array],
+        chunks_and_specs: Iterable[tuple[Buffer | None, ArraySpec]],
+    ) -> Iterable[Buffer | None]:
+        result: list[Buffer | None] = []
+
+        for a, (_, spec) in zip(arrays, chunks_and_specs, strict=True):
+            if a is None:
+                result.append(None)
+            else:
+                a2 = cp.array(a, dtype=a.dtype, copy=False)
+                if a2.dtype != np.dtype("B"):
+                    a2 = a2.view(dtype=np.dtype("B"))
+                result.append(spec.prototype.buffer.from_array_like(a2))
+
+        return result
+
+    async def decode(
+        self,
+        chunks_and_specs: Iterable[tuple[Buffer | None, ArraySpec]],
+    ) -> Iterable[Buffer | None]:
+        """Decodes a batch of chunks.
+        Chunks can be None in which case they are ignored by the codec.
+
+        Parameters
+        ----------
+        chunks_and_specs : Iterable[tuple[Buffer | None, ArraySpec]]
+            Ordered set of encoded chunks with their accompanying chunk spec.
+
+        Returns
+        -------
+        Iterable[Buffer | None]
+        """
+        chunks_and_specs = list(chunks_and_specs)
+
+        # Convert to nvcomp arrays
+        filtered_inputs, none_indices = self._convert_to_nvcomp_arrays(chunks_and_specs)
+
+        outputs = self._zstd_codec.decode(filtered_inputs) if len(filtered_inputs) > 0 else []
+
+        # Record event for synchronization
+        event = cp.cuda.Event()
+        event.record(cp.cuda.get_current_stream())
+        # Wait for decode to complete in a separate async thread
+        await asyncio.to_thread(event.synchronize)
+
+        for index in none_indices:
+            outputs.insert(index, None)
+
+        return self._convert_from_nvcomp_arrays(outputs, chunks_and_specs)
+
+    async def encode(
+        self,
+        chunks_and_specs: Iterable[tuple[Buffer | None, ArraySpec]],
+    ) -> Iterable[Buffer | None]:
+        """Encodes a batch of chunks.
+        Chunks can be None in which case they are ignored by the codec.
+
+        Parameters
+        ----------
+        chunks_and_specs : Iterable[tuple[Buffer | None, ArraySpec]]
+            Ordered set of to-be-encoded chunks with their accompanying chunk spec.
+
+        Returns
+        -------
+        Iterable[Buffer | None]
+        """
+        # TODO: Make this actually async
+        chunks_and_specs = list(chunks_and_specs)
+
+        # Convert to nvcomp arrays
+        filtered_inputs, none_indices = self._convert_to_nvcomp_arrays(chunks_and_specs)
+
+        outputs = self._zstd_codec.encode(filtered_inputs) if len(filtered_inputs) > 0 else []
+
+        # Record event for synchronization
+        event = cp.cuda.Event()
+        event.record(cp.cuda.get_current_stream())
+        # Wait for encode to complete in a separate async thread
+        await asyncio.to_thread(event.synchronize)
+
+        for index in none_indices:
+            outputs.insert(index, None)
+
+        return self._convert_from_nvcomp_arrays(outputs, chunks_and_specs)
+
+    def compute_encoded_size(self, _input_byte_length: int, _chunk_spec: ArraySpec) -> int:
+        raise NotImplementedError
+
+
+register_codec("zstd", NvcompZstdCodec)

src/zarr/core/array.py

@@ -28,7 +28,6 @@
 from zarr.codecs._v2 import V2Codec
 from zarr.codecs.bytes import BytesCodec
 from zarr.codecs.vlen_utf8 import VLenBytesCodec, VLenUTF8Codec
-from zarr.codecs.zstd import ZstdCodec
 from zarr.core._info import ArrayInfo
 from zarr.core._json import buffer_to_json_object
 from zarr.core.array_spec import ArrayConfig, ArrayConfigLike, ArraySpec, parse_array_config
@@ -154,6 +153,7 @@
     _parse_array_array_codec,
     _parse_array_bytes_codec,
     _parse_bytes_bytes_codec,
+    get_codec_class,
     get_pipeline_class,
 )
 from zarr.storage._common import StorePath, make_store_path
@@ -5113,9 +5113,15 @@ def default_compressors_v3(dtype: ZDType[Any, Any]) -> tuple[BytesBytesCodec, ..
     """
     Given a data type, return the default compressors for that data type.
 
-    This is just a tuple containing `ZstdCodec`
+    This is just a tuple containing an instance of the default "zstd" codec class.
+
+    The codec is resolved through the registry (rather than hard-coding ``ZstdCodec``)
+    so that configuration overrides -- e.g. ``zarr.config.enable_gpu()`` mapping
+    "zstd" to ``NvcompZstdCodec`` -- take effect. ``get_codec_class`` returns
+    ``type[Codec]``, so the instance is narrowed to ``BytesBytesCodec`` with a
+    runtime check.
     """
-    return (ZstdCodec(),)
+    return (_parse_bytes_bytes_codec(get_codec_class("zstd")()),)
 
 
 def default_serializer_v3(dtype: ZDType[Any, Any]) -> ArrayBytesCodec: