Gpu

zarr.buffer.gpu

__all__ module-attribute

__all__ = ['Buffer', 'NDBuffer', 'buffer_prototype']

buffer_prototype module-attribute

buffer_prototype = BufferPrototype(
    buffer=Buffer, nd_buffer=NDBuffer
)

Buffer

Bases: Buffer

A flat contiguous memory block on the GPU

We use Buffer throughout Zarr to represent a contiguous block of memory.

A Buffer is backed by an underlying array-like instance that represents the memory. The memory type is unspecified; can be regular host memory, CUDA device memory, or something else. The only requirement is that the array-like instance can be copied/converted to a regular Numpy array (host memory).

Notes

This buffer is untyped, so all indexing and sizes are in bytes.

Parameters:

• array_like (ArrayLike) –

  array-like object that must be 1-dim, contiguous, and byte dtype.

Source code in zarr/core/buffer/gpu.py

class Buffer(core.Buffer):
    """A flat contiguous memory block on the GPU

    We use Buffer throughout Zarr to represent a contiguous block of memory.

    A Buffer is backed by an underlying array-like instance that represents
    the memory. The memory type is unspecified; can be regular host memory,
    CUDA device memory, or something else. The only requirement is that the
    array-like instance can be copied/converted to a regular Numpy array
    (host memory).

    Notes
    -----
    This buffer is untyped, so all indexing and sizes are in bytes.

    Parameters
    ----------
    array_like
        array-like object that must be 1-dim, contiguous, and byte dtype.
    """

    def __init__(self, array_like: ArrayLike) -> None:
        if cp is None:
            raise ImportError(
                "Cannot use zarr.buffer.gpu.Buffer without cupy. Please install cupy."
            )

        if array_like.ndim != 1:
            raise ValueError("array_like: only 1-dim allowed")
        if array_like.dtype != np.dtype("B"):
            raise ValueError("array_like: only byte dtype allowed")

        if not hasattr(array_like, "__cuda_array_interface__"):
            # Slow copy based path for arrays that don't support the __cuda_array_interface__
            # TODO: Add a fast zero-copy path for arrays that support the dlpack protocol
            msg = (
                "Creating a zarr.buffer.gpu.Buffer with an array that does not support the "
                "__cuda_array_interface__ for zero-copy transfers, "
                "falling back to slow copy based path"
            )
            warnings.warn(
                msg,
                category=ZarrUserWarning,
                stacklevel=2,
            )
        self._data = cp.asarray(array_like)

    @classmethod
    def create_zero_length(cls) -> Self:
        """Create an empty buffer with length zero

        Returns
        -------
            New empty 0-length buffer
        """
        return cls(cp.array([], dtype="B"))

    @classmethod
    def from_buffer(cls, buffer: core.Buffer) -> Self:
        """Create a GPU Buffer given an arbitrary Buffer
        This will try to be zero-copy if `buffer` is already on the
        GPU and will trigger a copy if not.

        Returns
        -------
            New GPU Buffer constructed from `buffer`
        """
        return cls(buffer.as_array_like())

    @classmethod
    def from_bytes(cls, bytes_like: BytesLike) -> Self:
        return cls.from_array_like(cp.frombuffer(bytes_like, dtype="B"))

    def as_numpy_array(self) -> npt.NDArray[Any]:
        return cast("npt.NDArray[Any]", cp.asnumpy(self._data))

    def combine(self, others: Iterable[core.Buffer]) -> Self:
        data = [cp.asanyarray(self._data)]
        for other in others:
            other_array = other.as_array_like()
            assert other_array.dtype == np.dtype("B")
            gpu_other = Buffer(other_array)
            gpu_other_array = gpu_other.as_array_like()
            data.append(cp.asanyarray(gpu_other_array))
        return self.__class__(cp.concatenate(data))

__add__

__add__(other: Buffer) -> Self

Concatenate two buffers

Source code in zarr/core/buffer/core.py

def __add__(self, other: Buffer) -> Self:
    """Concatenate two buffers"""
    return self.combine([other])

__eq__

__eq__(other: object) -> bool

Source code in zarr/core/buffer/core.py

def __eq__(self, other: object) -> bool:
    # Another Buffer class can override this to choose a more efficient path
    return isinstance(other, Buffer) and np.array_equal(
        self.as_numpy_array(), other.as_numpy_array()
    )

__getitem__

__getitem__(key: slice) -> Self

Source code in zarr/core/buffer/core.py

def __getitem__(self, key: slice) -> Self:
    check_item_key_is_1d_contiguous(key)
    return self.__class__(self._data.__getitem__(key))

__init__

__init__(array_like: ArrayLike) -> None

Source code in zarr/core/buffer/gpu.py

def __init__(self, array_like: ArrayLike) -> None:
    if cp is None:
        raise ImportError(
            "Cannot use zarr.buffer.gpu.Buffer without cupy. Please install cupy."
        )

    if array_like.ndim != 1:
        raise ValueError("array_like: only 1-dim allowed")
    if array_like.dtype != np.dtype("B"):
        raise ValueError("array_like: only byte dtype allowed")

    if not hasattr(array_like, "__cuda_array_interface__"):
        # Slow copy based path for arrays that don't support the __cuda_array_interface__
        # TODO: Add a fast zero-copy path for arrays that support the dlpack protocol
        msg = (
            "Creating a zarr.buffer.gpu.Buffer with an array that does not support the "
            "__cuda_array_interface__ for zero-copy transfers, "
            "falling back to slow copy based path"
        )
        warnings.warn(
            msg,
            category=ZarrUserWarning,
            stacklevel=2,
        )
    self._data = cp.asarray(array_like)

__len__

__len__() -> int

Source code in zarr/core/buffer/core.py

def __len__(self) -> int:
    return self._data.size

__setitem__

__setitem__(key: slice, value: Any) -> None

Source code in zarr/core/buffer/core.py

def __setitem__(self, key: slice, value: Any) -> None:
    check_item_key_is_1d_contiguous(key)
    self._data.__setitem__(key, value)

as_array_like

as_array_like() -> ArrayLike

Returns the underlying array (host or device memory) of this buffer

This will never copy data.

Returns:

• The underlying 1d array such as a NumPy or CuPy array. –

Source code in zarr/core/buffer/core.py

def as_array_like(self) -> ArrayLike:
    """Returns the underlying array (host or device memory) of this buffer

    This will never copy data.

    Returns
    -------
        The underlying 1d array such as a NumPy or CuPy array.
    """
    return self._data

as_buffer_like

as_buffer_like() -> BytesLike

Returns the buffer as an object that implements the Python buffer protocol.

Notes

Might have to copy data, since the implementation uses .as_numpy_array().

Returns:

• An object that implements the Python buffer protocol –

Source code in zarr/core/buffer/core.py

def as_buffer_like(self) -> BytesLike:
    """Returns the buffer as an object that implements the Python buffer protocol.

    Notes
    -----
    Might have to copy data, since the implementation uses `.as_numpy_array()`.

    Returns
    -------
        An object that implements the Python buffer protocol
    """
    return memoryview(self.as_numpy_array())  # type: ignore[arg-type]

as_numpy_array

as_numpy_array() -> NDArray[Any]

Returns the buffer as a NumPy array (host memory).

Notes

Might have to copy data, consider using .as_array_like() instead.

Returns:

• NumPy array of this buffer (might be a data copy) –

Source code in zarr/core/buffer/gpu.py

def as_numpy_array(self) -> npt.NDArray[Any]:
    return cast("npt.NDArray[Any]", cp.asnumpy(self._data))

combine

combine(others: Iterable[Buffer]) -> Self

Concatenate many buffers

Source code in zarr/core/buffer/gpu.py

def combine(self, others: Iterable[core.Buffer]) -> Self:
    data = [cp.asanyarray(self._data)]
    for other in others:
        other_array = other.as_array_like()
        assert other_array.dtype == np.dtype("B")
        gpu_other = Buffer(other_array)
        gpu_other_array = gpu_other.as_array_like()
        data.append(cp.asanyarray(gpu_other_array))
    return self.__class__(cp.concatenate(data))

create_zero_length classmethod

create_zero_length() -> Self

Create an empty buffer with length zero

Returns:

• New empty 0-length buffer –

Source code in zarr/core/buffer/gpu.py

@classmethod
def create_zero_length(cls) -> Self:
    """Create an empty buffer with length zero

    Returns
    -------
        New empty 0-length buffer
    """
    return cls(cp.array([], dtype="B"))

from_array_like classmethod

from_array_like(array_like: ArrayLike) -> Self

Create a new buffer of an array-like object

Parameters:

• array_like (ArrayLike) –

  array-like object that must be 1-dim, contiguous, and byte dtype.

Returns:

• New buffer representing `array_like` –

Source code in zarr/core/buffer/core.py

@classmethod
def from_array_like(cls, array_like: ArrayLike) -> Self:
    """Create a new buffer of an array-like object

    Parameters
    ----------
    array_like
        array-like object that must be 1-dim, contiguous, and byte dtype.

    Returns
    -------
        New buffer representing `array_like`
    """
    return cls(array_like)

from_buffer classmethod

from_buffer(buffer: Buffer) -> Self

Create a GPU Buffer given an arbitrary Buffer This will try to be zero-copy if buffer is already on the GPU and will trigger a copy if not.

Returns:

• New GPU Buffer constructed from `buffer` –

Source code in zarr/core/buffer/gpu.py

@classmethod
def from_buffer(cls, buffer: core.Buffer) -> Self:
    """Create a GPU Buffer given an arbitrary Buffer
    This will try to be zero-copy if `buffer` is already on the
    GPU and will trigger a copy if not.

    Returns
    -------
        New GPU Buffer constructed from `buffer`
    """
    return cls(buffer.as_array_like())

from_bytes classmethod

from_bytes(bytes_like: BytesLike) -> Self

Create a new buffer of a bytes-like object (host memory)

Parameters:

• bytes_like (BytesLike) –

  bytes-like object

Returns:

• New buffer representing `bytes_like` –

Source code in zarr/core/buffer/gpu.py

@classmethod
def from_bytes(cls, bytes_like: BytesLike) -> Self:
    return cls.from_array_like(cp.frombuffer(bytes_like, dtype="B"))

to_bytes

to_bytes() -> bytes

Returns the buffer as bytes (host memory).

Warnings

Will always copy data, only use this method for small buffers such as metadata buffers. If possible, use .as_numpy_array() or .as_array_like() instead.

Returns:

• `bytes` of this buffer (data copy) –

Source code in zarr/core/buffer/core.py

def to_bytes(self) -> bytes:
    """Returns the buffer as `bytes` (host memory).

    Warnings
    --------
    Will always copy data, only use this method for small buffers such as metadata
    buffers. If possible, use `.as_numpy_array()` or `.as_array_like()` instead.

    Returns
    -------
        `bytes` of this buffer (data copy)
    """
    return bytes(self.as_numpy_array())

NDBuffer

Bases: NDBuffer

A n-dimensional memory block on the GPU

We use NDBuffer throughout Zarr to represent a n-dimensional memory block.

An NDBuffer is backed by an underlying ndarray-like instance that represents the memory. The memory type is unspecified; can be regular host memory, CUDA device memory, or something else. The only requirement is that the ndarray-like instance can be copied/converted to a regular Numpy array (host memory).

Notes

The two buffer classes Buffer and NDBuffer are very similar. In fact, Buffer is a special case of NDBuffer where dim=1, stride=1, and dtype="B". However, in order to use Python's type system to differentiate between the contiguous Buffer and the n-dim (non-contiguous) NDBuffer, we keep the definition of the two classes separate.

Parameters:

• array (NDArrayLike) –

  ndarray-like object that is convertible to a regular Numpy array.

Source code in zarr/core/buffer/gpu.py

class NDBuffer(core.NDBuffer):
    """A n-dimensional memory block on the GPU

    We use NDBuffer throughout Zarr to represent a n-dimensional memory block.

    An NDBuffer is backed by an underlying ndarray-like instance that represents
    the memory. The memory type is unspecified; can be regular host memory,
    CUDA device memory, or something else. The only requirement is that the
    ndarray-like instance can be copied/converted to a regular Numpy array
    (host memory).

    Notes
    -----
    The two buffer classes Buffer and NDBuffer are very similar. In fact, Buffer
    is a special case of NDBuffer where dim=1, stride=1, and dtype="B". However,
    in order to use Python's type system to differentiate between the contiguous
    Buffer and the n-dim (non-contiguous) NDBuffer, we keep the definition of the
    two classes separate.

    Parameters
    ----------
    array
        ndarray-like object that is convertible to a regular Numpy array.
    """

    def __init__(self, array: NDArrayLike) -> None:
        if cp is None:
            raise ImportError(
                "Cannot use zarr.buffer.gpu.NDBuffer without cupy. Please install cupy."
            )

        # assert array.ndim > 0
        assert array.dtype != object
        self._data = array

        if not hasattr(array, "__cuda_array_interface__"):
            # Slow copy based path for arrays that don't support the __cuda_array_interface__
            # TODO: Add a fast zero-copy path for arrays that support the dlpack protocol
            msg = (
                "Creating a zarr.buffer.gpu.NDBuffer with an array that does not support the "
                "__cuda_array_interface__ for zero-copy transfers, "
                "falling back to slow copy based path"
            )
            warnings.warn(
                msg,
                stacklevel=2,
            )
        self._data = cp.asarray(array)

    @classmethod
    def create(
        cls,
        *,
        shape: Iterable[int],
        dtype: npt.DTypeLike,
        order: Literal["C", "F"] = "C",
        fill_value: Any | None = None,
    ) -> Self:
        ret = cls(cp.empty(shape=tuple(shape), dtype=dtype, order=order))
        if fill_value is not None:
            ret.fill(fill_value)
        return ret

    @classmethod
    def empty(
        cls, shape: tuple[int, ...], dtype: npt.DTypeLike, order: Literal["C", "F"] = "C"
    ) -> Self:
        return cls(cp.empty(shape=shape, dtype=dtype, order=order))

    @classmethod
    def from_numpy_array(cls, array_like: npt.ArrayLike) -> Self:
        """Create a new buffer of Numpy array-like object

        Parameters
        ----------
        array_like
            Object that can be coerced into a Numpy array

        Returns
        -------
            New buffer representing `array_like`
        """
        return cls(cp.asarray(array_like))

    def as_numpy_array(self) -> npt.NDArray[Any]:
        """Returns the buffer as a NumPy array (host memory).

        Warnings
        --------
        Might have to copy data, consider using `.as_ndarray_like()` instead.

        Returns
        -------
            NumPy array of this buffer (might be a data copy)
        """
        return cast("npt.NDArray[Any]", cp.asnumpy(self._data))

    def __getitem__(self, key: Any) -> Self:
        return self.__class__(self._data.__getitem__(key))

    def __setitem__(self, key: Any, value: Any) -> None:
        if isinstance(value, NDBuffer):
            value = value._data
        elif isinstance(value, core.NDBuffer):
            gpu_value = NDBuffer(value.as_ndarray_like())
            value = gpu_value._data
        self._data.__setitem__(key, value)

byteorder property

byteorder: Endian

dtype property

dtype: dtype[Any]

shape property

shape: tuple[int, ...]

__getitem__

__getitem__(key: Any) -> Self

Source code in zarr/core/buffer/gpu.py

def __getitem__(self, key: Any) -> Self:
    return self.__class__(self._data.__getitem__(key))

__init__

__init__(array: NDArrayLike) -> None

Source code in zarr/core/buffer/gpu.py

def __init__(self, array: NDArrayLike) -> None:
    if cp is None:
        raise ImportError(
            "Cannot use zarr.buffer.gpu.NDBuffer without cupy. Please install cupy."
        )

    # assert array.ndim > 0
    assert array.dtype != object
    self._data = array

    if not hasattr(array, "__cuda_array_interface__"):
        # Slow copy based path for arrays that don't support the __cuda_array_interface__
        # TODO: Add a fast zero-copy path for arrays that support the dlpack protocol
        msg = (
            "Creating a zarr.buffer.gpu.NDBuffer with an array that does not support the "
            "__cuda_array_interface__ for zero-copy transfers, "
            "falling back to slow copy based path"
        )
        warnings.warn(
            msg,
            stacklevel=2,
        )
    self._data = cp.asarray(array)

__len__

__len__() -> int

Source code in zarr/core/buffer/core.py

def __len__(self) -> int:
    return self._data.__len__()

__repr__

__repr__() -> str

Source code in zarr/core/buffer/core.py

def __repr__(self) -> str:
    return f"<NDBuffer shape={self.shape} dtype={self.dtype} {self._data!r}>"

__setitem__

__setitem__(key: Any, value: Any) -> None

Source code in zarr/core/buffer/gpu.py

def __setitem__(self, key: Any, value: Any) -> None:
    if isinstance(value, NDBuffer):
        value = value._data
    elif isinstance(value, core.NDBuffer):
        gpu_value = NDBuffer(value.as_ndarray_like())
        value = gpu_value._data
    self._data.__setitem__(key, value)

all_equal

all_equal(other: Any, equal_nan: bool = True) -> bool

Compare to other using np.array_equal.

Source code in zarr/core/buffer/core.py

def all_equal(self, other: Any, equal_nan: bool = True) -> bool:
    """Compare to `other` using np.array_equal."""
    if other is None:
        # Handle None fill_value for Zarr V2
        return False
    # Handle positive and negative zero by comparing bit patterns:
    if (
        np.asarray(other).dtype.kind == "f"
        and other == 0.0
        and self._data.dtype.kind not in ("U", "S", "T", "O", "V")
    ):
        _data, other = np.broadcast_arrays(self._data, np.asarray(other, self._data.dtype))
        void_dtype = "V" + str(_data.dtype.itemsize)
        return np.array_equal(_data.view(void_dtype), other.view(void_dtype))
    # use array_equal to obtain equal_nan=True functionality
    # Since fill-value is a scalar, isn't there a faster path than allocating a new array for fill value
    # every single time we have to write data?
    _data, other = np.broadcast_arrays(self._data, other)
    return np.array_equal(
        _data,
        other,
        equal_nan=equal_nan
        if self._data.dtype.kind not in ("U", "S", "T", "O", "V")
        else False,
    )

as_ndarray_like

as_ndarray_like() -> NDArrayLike

Returns the underlying array (host or device memory) of this buffer

This will never copy data.

Returns:

• The underlying array such as a NumPy or CuPy array. –

Source code in zarr/core/buffer/core.py

def as_ndarray_like(self) -> NDArrayLike:
    """Returns the underlying array (host or device memory) of this buffer

    This will never copy data.

    Returns
    -------
        The underlying array such as a NumPy or CuPy array.
    """
    return self._data

as_numpy_array

as_numpy_array() -> NDArray[Any]

Returns the buffer as a NumPy array (host memory).

Warnings

Might have to copy data, consider using .as_ndarray_like() instead.

Returns:

• NumPy array of this buffer (might be a data copy) –

Source code in zarr/core/buffer/gpu.py

def as_numpy_array(self) -> npt.NDArray[Any]:
    """Returns the buffer as a NumPy array (host memory).

    Warnings
    --------
    Might have to copy data, consider using `.as_ndarray_like()` instead.

    Returns
    -------
        NumPy array of this buffer (might be a data copy)
    """
    return cast("npt.NDArray[Any]", cp.asnumpy(self._data))

as_scalar

as_scalar() -> ScalarType

Returns the buffer as a scalar value

Source code in zarr/core/buffer/core.py

def as_scalar(self) -> ScalarType:
    """Returns the buffer as a scalar value"""
    if self._data.size != 1:
        raise ValueError("Buffer does not contain a single scalar value")
    return cast("ScalarType", self.as_numpy_array()[()])

astype

astype(
    dtype: DTypeLike,
    order: Literal["K", "A", "C", "F"] = "K",
) -> Self

Source code in zarr/core/buffer/core.py

def astype(self, dtype: npt.DTypeLike, order: Literal["K", "A", "C", "F"] = "K") -> Self:
    return self.__class__(self._data.astype(dtype=dtype, order=order))

copy

copy() -> Self

Source code in zarr/core/buffer/core.py

def copy(self) -> Self:
    return self.__class__(self._data.copy())

create classmethod

create(
    *,
    shape: Iterable[int],
    dtype: DTypeLike,
    order: Literal["C", "F"] = "C",
    fill_value: Any | None = None,
) -> Self

Create a new buffer and its underlying ndarray-like object

Parameters:

• shape (Iterable[int]) –

  The shape of the buffer and its underlying ndarray-like object

• dtype (DTypeLike) –

  The datatype of the buffer and its underlying ndarray-like object

• order (Literal['C', 'F'], default: 'C' ) –

  Whether to store multi-dimensional data in row-major (C-style) or column-major (Fortran-style) order in memory.

• fill_value (Any | None, default: None ) –

  If not None, fill the new buffer with a scalar value.

Returns:

• New buffer representing a new ndarray_like object –

Notes

A subclass can overwrite this method to create an ndarray-like object other then the default Numpy array.

Source code in zarr/core/buffer/gpu.py

@classmethod
def create(
    cls,
    *,
    shape: Iterable[int],
    dtype: npt.DTypeLike,
    order: Literal["C", "F"] = "C",
    fill_value: Any | None = None,
) -> Self:
    ret = cls(cp.empty(shape=tuple(shape), dtype=dtype, order=order))
    if fill_value is not None:
        ret.fill(fill_value)
    return ret

empty classmethod

empty(
    shape: tuple[int, ...],
    dtype: DTypeLike,
    order: Literal["C", "F"] = "C",
) -> Self

Create an empty buffer with the given shape, dtype, and order.

This method can be faster than NDBuffer.create because it doesn't have to initialize the memory used by the underlying ndarray-like object.

Parameters:

• shape (tuple[int, ...]) –

  The shape of the buffer and its underlying ndarray-like object

• dtype (DTypeLike) –

  The datatype of the buffer and its underlying ndarray-like object

• order (Literal['C', 'F'], default: 'C' ) –

  Whether to store multi-dimensional data in row-major (C-style) or column-major (Fortran-style) order in memory.

Returns:

• buffer –

  New buffer representing a new ndarray_like object with empty data.

See Also

NDBuffer.create Create a new buffer with some initial fill value.

Source code in zarr/core/buffer/gpu.py

@classmethod
def empty(
    cls, shape: tuple[int, ...], dtype: npt.DTypeLike, order: Literal["C", "F"] = "C"
) -> Self:
    return cls(cp.empty(shape=shape, dtype=dtype, order=order))

fill

fill(value: Any) -> None

Source code in zarr/core/buffer/core.py

def fill(self, value: Any) -> None:
    self._data.fill(value)

from_ndarray_like classmethod

from_ndarray_like(ndarray_like: NDArrayLike) -> Self

Create a new buffer of an ndarray-like object

Parameters:

• ndarray_like (NDArrayLike) –

  ndarray-like object

Returns:

• New buffer representing `ndarray_like` –

Source code in zarr/core/buffer/core.py

@classmethod
def from_ndarray_like(cls, ndarray_like: NDArrayLike) -> Self:
    """Create a new buffer of an ndarray-like object

    Parameters
    ----------
    ndarray_like
        ndarray-like object

    Returns
    -------
        New buffer representing `ndarray_like`
    """
    return cls(ndarray_like)

from_numpy_array classmethod

from_numpy_array(array_like: ArrayLike) -> Self

Create a new buffer of Numpy array-like object

Parameters:

• array_like (ArrayLike) –

  Object that can be coerced into a Numpy array

Returns:

• New buffer representing `array_like` –

Source code in zarr/core/buffer/gpu.py

@classmethod
def from_numpy_array(cls, array_like: npt.ArrayLike) -> Self:
    """Create a new buffer of Numpy array-like object

    Parameters
    ----------
    array_like
        Object that can be coerced into a Numpy array

    Returns
    -------
        New buffer representing `array_like`
    """
    return cls(cp.asarray(array_like))

reshape

reshape(newshape: tuple[int, ...] | Literal[-1]) -> Self

Source code in zarr/core/buffer/core.py

def reshape(self, newshape: tuple[int, ...] | Literal[-1]) -> Self:
    return self.__class__(self._data.reshape(newshape))

squeeze

squeeze(axis: tuple[int, ...]) -> Self

Source code in zarr/core/buffer/core.py

def squeeze(self, axis: tuple[int, ...]) -> Self:
    newshape = tuple(a for i, a in enumerate(self.shape) if i not in axis)
    return self.__class__(self._data.reshape(newshape))

transpose

transpose(
    axes: SupportsIndex | Sequence[SupportsIndex] | None,
) -> Self

Source code in zarr/core/buffer/core.py

def transpose(self, axes: SupportsIndex | Sequence[SupportsIndex] | None) -> Self:
    return self.__class__(self._data.transpose(axes))