Skip to content

gpu.rst

Latest commit

 

History

History
34 lines (25 loc) · 1.14 KB

gpu.rst

File metadata and controls

34 lines (25 loc) · 1.14 KB

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.

>>> 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.