GPU blosc zstd decode

docs/user-guide/gpu.md

@@ -1,15 +1,18 @@
 # Using GPUs with Zarr
 
-Zarr can use GPUs to accelerate your workload by running `zarr.Config.enable_gpu`.
+Zarr can use GPUs to accelerate your workload by enabling GPU buffers and codecs
+with `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.
+    With `enable_gpu()`:
 
-    In the future, codecs will be available compressing and decompressing data on
-    the GPU, avoiding the need to move data between the host and device for
-    compression and decompression.
+    - array reads return `cupy.ndarray` values
+    - the default Zarr v3 `zstd` codec is replaced with a GPU-backed nvCOMP
+      implementation
+    - Blosc chunks can be decoded on the GPU when `cname="zstd"` and `shuffle`
+      is `bitshuffle` or `noshuffle`
+
+    Blosc GPU acceleration currently applies to decode only.
 
 ## Reading data into device memory
 

pyproject.toml

@@ -68,6 +68,7 @@ remote = [
 ]
 gpu = [
     "cupy-cuda12x",
+    "nvidia-nvcomp-cu12",
 ]
 cli = ["typer"]
 # Testing extras
@@ -154,6 +155,9 @@ omit = [
 [tool.hatch]
 version.source = "vcs"
 
+[tool.hatch.version.raw-options]
+fallback_version = "3.1.0"
+
 [tool.hatch.build]
 hooks.vcs.version-file = "src/zarr/_version.py"
 

src/zarr/codecs/__init__.py

@@ -3,6 +3,7 @@
 from zarr.codecs.blosc import BloscCname, BloscCodec, BloscShuffle
 from zarr.codecs.bytes import BytesCodec, Endian
 from zarr.codecs.crc32c_ import Crc32cCodec
+from zarr.codecs.gpu import NvcompBloscCodec, NvcompZstdCodec
 from zarr.codecs.gzip import GzipCodec
 from zarr.codecs.numcodecs import (
     BZ2,
@@ -41,6 +42,8 @@
     "Crc32cCodec",
     "Endian",
     "GzipCodec",
+    "NvcompBloscCodec",
+    "NvcompZstdCodec",
     "ShardingCodec",
     "ShardingCodecIndexLocation",
     "TransposeCodec",