nvmath-python 0.9.0__cp314-cp314-win_amd64.whl

nvmath/internal/tensor_ifc.py

@@ -0,0 +1,194 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Interface to seamlessly use tensors (or ndarray-like objects) from different libraries.
+"""
+
+from __future__ import annotations  # allows typehint of class methods to return the self class
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Sequence
+from types import ModuleType
+from typing import Any, Generic, Literal, Protocol, TypeVar
+
+from . import typemaps
+from .ndbuffer import ndbuffer
+from .package_ifc import StreamHolder
+
+
+class AnyTensor(Protocol):
+    """Any supported external tensor object such as NumPy, CuPy, and PyTorch arrays."""
+
+    # This empty protocol class is just a placeholder to make type checking more helpful.
+    # Since the class is empty, it doesn't assume anything is true about the implementation
+    # classes. Used to type hint any object that TensorHolder knows how to wrap.
+    pass
+
+
+"""
+A generic type for the third-party Tensor which a given TensorHolder implementation wraps
+around.
+"""
+Tensor = TypeVar("Tensor")
+
+
+class TensorHolder(ABC, Generic[Tensor]):
+    """
+    A simple wrapper type for tensors to make the API package-agnostic.
+
+    Methods of a TensorHolder should always return a TensorHolder instead of a Tensor (one
+    of the wrapped classes) in order to prevent implementation details from the various
+    Tensor implementations from leaking into nvmath-python.
+
+    Tensors from the user should be immediately wrapped with a TensorHolder and should
+    remain wrapped until just before returning to the user.
+    """
+
+    name: str
+    module: ModuleType
+    name_to_dtype: dict[str, Any]
+
+    def __init__(self, tensor: Tensor):
+        self.tensor: Tensor = tensor
+
+    @property
+    @abstractmethod
+    def data_ptr(self) -> int:
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def device(self) -> Literal["cuda"] | Literal["cpu"]:
+        """The type of the device which stores the tensor."""
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def device_id(self) -> int | Literal["cpu"]:
+        """The device ordinal of the device storing the tensor."""
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def dtype(self) -> str:
+        """Name of the data type"""
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def itemsize(self) -> int:
+        """The size of the data type in bytes."""
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def empty(cls, shape: Sequence[int], device_id: int | Literal["cpu"], **context: Any) -> TensorHolder[Tensor]:
+        """Create an empty TensorHolder of the specified shape and data type."""
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def shape(self) -> Sequence[int]:
+        """The extent of each dimension in number of elements."""
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def size(self) -> int:
+        """The total number of elements in the tensor."""
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def strides(self) -> Sequence[int]:
+        """The stride of each dimension in number of elements."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def to(self, device_id: int | Literal["cpu"], stream_holder: StreamHolder | None) -> TensorHolder:
+        """Copy the TensorHolder to a different device.
+
+        No copy is performed if the TensorHolder is already on the requested device.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def asndbuffer(self) -> ndbuffer.NDBuffer:
+        """Wraps the package tensor as a ndbuffer.NDBuffer object."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def copy_(self, src: TensorHolder, stream_holder: StreamHolder | None) -> None:
+        """Overwrite self.tensor (in-place) with a copy of src."""
+        raise NotImplementedError
+
+    @staticmethod
+    def create_name_dtype_map(conversion_function: Callable[[str], Any], exception_type: type[Exception]) -> dict[str, Any]:
+        """
+        Create a map between CUDA data type names and the corresponding package dtypes for
+        supported data types.
+        """
+        names = typemaps.NAME_TO_DATA_TYPE.keys()
+        name_to_dtype = {}
+        for name in names:
+            try:
+                name_to_dtype[name] = conversion_function(name)
+            except exception_type:
+                pass
+        return name_to_dtype
+
+    @abstractmethod
+    def istensor(self) -> bool:
+        """Return whether self.tensor is the expected type."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def reshape(self, shape: Sequence[int], *, copy: bool | None = None) -> TensorHolder[Tensor]:
+        """Reshapes tensor without changing its data.
+
+        Args:
+            shape: a new shape compatible with the original shape.
+
+            copy (Optional[bool]): whether or not to copy the input tensor. If True, the
+                function must always copy. If False, the function must never copy. If None,
+                the function must avoid copying, if possible, and may copy otherwise.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def memory_buffer(self) -> TensorHolder[Tensor]:
+        """Creates a view of the memory buffer as a 1D tensor."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def memory_buffer_to_tensor(self, shape, strides) -> TensorHolder[Tensor]:
+        """
+        Creates a N-D tensor view of the memory buffer according to the specified
+        shape and strides.
+        """
+        raise NotImplementedError
+
+    def _broadcast_to(self, shape: Sequence[int]) -> TensorHolder[Tensor]:
+        """Returns a view of the tensor broadcast to the desired shape.
+
+        .. deprecated:: 0.8.0
+            This method is private, and should not be used. It will be removed in the
+            future once NDBuffer supports broadcasting-aware copy.
+
+        Like reshape, but raises an error if the tensor is not broadcastable to the desired
+        shape.
+
+        Args:
+            shape: a new shape compatible with the original shape.
+        """
+        raise NotImplementedError
+
+    @property
+    def is_conjugate(self) -> bool:
+        """Return True when self.tensor has a no-op conjugation flag enabled.
+
+        The default implementation always returns False.
+        """
+        return False

nvmath/internal/tensor_ifc_cupy.py

@@ -0,0 +1,241 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Interface to seamlessly use Cupy ndarray objects.
+"""
+
+__all__ = ["CupyTensor", "HostTensor"]
+
+from collections.abc import Sequence
+
+import cupy
+import numpy as np
+
+from nvmath.internal.memory import _MemoryPointer
+
+from . import utils
+from .ndbuffer import ndbuffer, package_utils
+from .package_ifc import StreamHolder
+from .tensor_ifc import TensorHolder
+from .tensor_ifc_ndbuffer import NDBufferTensor
+
+
+class HostTensor(NDBufferTensor):
+    """
+    Wraps ndbuffer with data residing on the host.
+    It serves as a host counterpart for CupyTensor.
+    """
+
+    name = "cupy_host"
+    device_tensor_class: type["CupyTensor"]  # set once CupyTensor is defined
+
+    def __init__(self, tensor):
+        super().__init__(tensor)
+
+    @classmethod
+    def create_host_from(cls, tensor: TensorHolder, stream_holder: StreamHolder):
+        src_nd = tensor.asndbuffer()
+        # empty_like (and not empty_numpy_like) is used as we don't need
+        # full-fledged numpy array (with proper layout)
+        dst_nd = ndbuffer.empty_like(src_nd, device_id=ndbuffer.CPU_DEVICE_ID)
+        ndbuffer.copy_into(dst_nd, src_nd, stream_holder)
+        return cls(dst_nd)
+
+    def to(self, device_id, stream_holder):
+        if device_id == "cpu":
+            return self
+        elif isinstance(device_id, int):
+            return self.device_tensor_class.create_from_host(self, device_id, stream_holder)
+        else:
+            raise ValueError(f"The device must be specified as an integer or 'cpu', not '{device_id}'.")
+
+
+class _CupyAllocatorAdapter:
+    def allocate(self, size, stream):
+        # we accept the stream, because the ndbuffer.empty_like
+        # passes it to the allocator, but we don't use it:
+        # cupy.cuda.alloc does not accept the stream, we make sure to set
+        # the correct current stream when calling ndbuffer.empty_like
+        alloc = cupy.cuda.alloc(size)
+        # Once we require cuda.core >= 0.5.0,
+        # use Buffer.from_handle(alloc.ptr, owner=alloc)
+        return _MemoryPointer.from_handle(alloc.ptr, alloc)
+
+
+_cupy_allocator = _CupyAllocatorAdapter()
+
+
+class CupyTensor(TensorHolder[cupy.ndarray]):
+    """
+    TensorHolder for cupy ndarrays.
+    """
+
+    name = "cupy"
+    module = cupy
+    name_to_dtype = TensorHolder.create_name_dtype_map(
+        conversion_function=lambda name: np.dtype(name), exception_type=TypeError
+    )
+    host_tensor_class = HostTensor
+
+    def __init__(self, tensor):
+        super().__init__(tensor)
+
+    @property
+    def data_ptr(self):
+        return self.tensor.data.ptr
+
+    @property
+    def device(self):
+        return "cuda"
+
+    @property
+    def device_id(self):
+        return self.tensor.device.id
+
+    @property
+    def dtype(self):
+        """Name of the data type"""
+        return self.tensor.dtype.name
+
+    @property
+    def itemsize(self):
+        return self.tensor.itemsize
+
+    @property
+    def shape(self):
+        return tuple(self.tensor.shape)
+
+    @property
+    def size(self):
+        return self.tensor.size
+
+    @property
+    def strides(self):
+        return tuple(stride_in_bytes // self.tensor.itemsize for stride_in_bytes in self.tensor.strides)
+
+    @classmethod
+    def empty(
+        cls, shape, device_id="cpu", *, dtype="float32", strides=None, stream_holder: StreamHolder | None = None, **context
+    ):
+        """
+        Create an empty tensor of the specified shape and data type.
+
+        Note, that the strides, if specified, MUST correspond to a dense (possibly permuted)
+        tensor, otherwise the created tensor may be corrupted.
+        """
+        assert isinstance(device_id, int), "Internal Error: Cupy tensors must be allocated with an integer device_id."
+        dtype = CupyTensor.name_to_dtype[dtype]
+
+        # When using the strides, we need an explicit allocation (see below).
+        # If the strides are simple enough, we can still avoid this overhead.
+        order = "C"
+        if strides is not None:
+            if len(strides) == 1 and strides[0] == 1:
+                strides = None
+            elif len(strides) == 2:
+                if strides[0] == 1 and strides[1] == shape[0]:
+                    strides = None
+                    order = "F"
+                elif strides[0] == shape[1] and strides[1] == 1:
+                    strides = None
+
+        assert isinstance(stream_holder, StreamHolder), "Internal Error: CupyTensors must be allocated on a stream."
+        with utils.device_ctx(device_id), stream_holder.ctx:
+            if strides:
+                # need an explicit allocation due to cupy/cupy#7818
+                size = dtype.itemsize
+                for s in shape:
+                    size = size * s
+                ptr = cupy.cuda.alloc(size)
+                # when strides is not None, it should be of unit counts not bytes
+                strides = tuple(s * dtype.itemsize for s in strides)
+                tensor = cupy.ndarray(shape, dtype=dtype, strides=strides, memptr=ptr)
+            else:
+                tensor = cupy.ndarray(shape, dtype=dtype, order=order)
+
+        return cls(tensor)
+
+    @classmethod
+    def create_from_host(cls, tensor: TensorHolder, device_id: int, stream_holder: StreamHolder):
+        with utils.device_ctx(device_id), stream_holder.ctx:
+            src_nd = tensor.asndbuffer()
+            dst_nd = ndbuffer.empty_like(
+                src_nd,
+                device_id=device_id,
+                stream=stream_holder,
+                device_memory_pool=_cupy_allocator,
+            )
+            ndbuffer.copy_into(dst_nd, src_nd, stream_holder)
+            dst = cupy.ndarray(dst_nd.shape, dtype=dst_nd.dtype_name, strides=dst_nd.strides_in_bytes, memptr=dst_nd.data.owner)
+            return cls(dst)
+
+    def asndbuffer(self):
+        return package_utils.wrap_cupy_array(self.tensor)
+
+    def to(self, device_id, stream_holder):
+        if device_id == "cpu":
+            with utils.device_ctx(self.device_id):
+                return self.host_tensor_class.create_host_from(self, stream_holder)
+        elif device_id == self.device_id:
+            return self
+        elif isinstance(device_id, int):
+            raise ValueError(f"Unsupported copy between different devices {self.device_id} and {device_id}.")
+        raise ValueError(f"The device must be specified as an integer or 'cpu', not '{device_id}'.")
+
+    def copy_(self, src, stream_holder):
+        """
+        Inplace copy of src (copy the data from src into self).
+        """
+        with utils.device_ctx(self.device_id):
+            ndbuffer.copy_into(self.asndbuffer(), src._broadcast_to(self.shape).asndbuffer(), stream_holder)
+
+    def istensor(self):
+        """
+        Check if the object is ndarray-like.
+        """
+        return isinstance(self.tensor, cupy.ndarray)
+
+    def reshape(self, shape: Sequence[int], *, copy: bool | None = None):
+        if copy:
+            raise NotImplementedError("reshape with copy=True not implemented")
+        if copy is False:
+            try:
+                reshaped_tensor = self.tensor.view()
+                reshaped_tensor.shape = shape
+            except AttributeError as e:
+                raise ValueError(
+                    f"Could not reshape cupy array without copy: current shape={self.shape}, new shape={shape}"
+                ) from e
+        else:
+            reshaped_tensor = self.tensor.reshape(shape)
+        return self.__class__(reshaped_tensor)
+
+    def memory_buffer(self):
+        """Creates a view of the memory buffer as a 1D tensor."""
+        m = cupy.cuda.UnownedMemory(self.data_ptr, self.size * self.itemsize, owner=self.tensor, device_id=self.device_id)
+        p = cupy.cuda.MemoryPointer(m, 0)
+        # TODO: ensure tensor is dense for now, and later support linear
+        # memory with constant stride.
+        t = cupy.ndarray(shape=(self.size,), dtype=self.dtype, memptr=p)
+        return CupyTensor(t)
+
+    def memory_buffer_to_tensor(self, shape, strides):
+        """
+        Creates a N-D tensor view of the memory buffer according to the specified
+        shape and strides.
+        """
+        assert len(self.shape) == 1, "Internal error."
+        m = cupy.cuda.UnownedMemory(self.data_ptr, self.size * self.itemsize, owner=self.tensor)
+        p = cupy.cuda.MemoryPointer(m, 0)
+        strides = [s * self.itemsize for s in strides]
+        t = cupy.ndarray(shape=shape, dtype=self.dtype, memptr=p, strides=strides)
+        return CupyTensor(t)
+
+    def _broadcast_to(self, shape):
+        reshaped_tensor = cupy.broadcast_to(self.tensor, shape)
+        return self.__class__(reshaped_tensor)
+
+
+HostTensor.device_tensor_class = CupyTensor

nvmath/internal/tensor_ifc_ndbuffer.py

@@ -0,0 +1,153 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+TensorHolder implementation that adapts ndbuffer for use with
+the tensor_ifc interface. The class is meant for internal use only
+(e.g. for internal operands of host APIs), but the NDBufferTensor
+or underlying ndbuffer should never be returned to the user, because
+1. the ndbuffer is opaque, i.e. it missies most of typical tensor functionality, and
+2. we make no guarantees about the ndbuffer's API stability.
+"""
+
+__all__ = ["NDBufferTensor"]
+
+from collections.abc import Sequence
+
+from . import typemaps, utils
+from .ndbuffer import ndbuffer
+from .package_ifc import StreamHolder
+from .tensor_ifc import TensorHolder
+
+
+class NDBufferTensor(TensorHolder[ndbuffer.NDBuffer]):
+    """
+    TensorHolder for ndbuffer ndarrays.
+    """
+
+    name = "ndbuffer"
+    module = ndbuffer
+    name_to_dtype = TensorHolder.create_name_dtype_map(conversion_function=lambda name: name, exception_type=TypeError)
+
+    def __init__(self, tensor):
+        super().__init__(tensor)
+
+    @property
+    def data_ptr(self):
+        return self.tensor.data_ptr
+
+    @property
+    def device(self):
+        return self.tensor.device
+
+    @property
+    def device_id(self):
+        return self.tensor.device_id
+
+    @property
+    def dtype(self):
+        """Name of the data type"""
+        return self.tensor.dtype_name
+
+    @property
+    def itemsize(self):
+        return self.tensor.itemsize
+
+    @property
+    def shape(self):
+        return self.tensor.shape
+
+    @property
+    def size(self):
+        return self.tensor.size
+
+    @property
+    def strides(self):
+        return self.tensor.strides
+
+    @classmethod
+    def empty(
+        cls, shape, device_id="cpu", *, dtype="float32", strides=None, stream_holder: StreamHolder | None = None, **context
+    ):
+        """
+        Create an empty tensor of the specified shape and data type.
+
+        Note, that the strides, if specified, MUST correspond to a dense (possibly permuted)
+        tensor, otherwise the created tensor may be corrupted.
+        """
+        itemsize = typemaps.NAME_TO_ITEM_SIZE[dtype]
+        if device_id == "cpu":
+            tensor = ndbuffer.empty(
+                shape, ndbuffer.CPU_DEVICE_ID, dtype, itemsize, strides=strides, stream=stream_holder, **context
+            )
+        else:
+            assert isinstance(device_id, int), "Internal Error: Cuda tensors must be allocated with an integer device_id."
+            with utils.device_ctx(device_id):
+                tensor = ndbuffer.empty(shape, device_id, dtype, itemsize, strides=strides, stream=stream_holder, **context)
+        return cls(tensor)
+
+    def asndbuffer(self):
+        return self.tensor
+
+    def to(self, device_id, stream_holder):
+        src_device_id = self.tensor.device_id
+        if src_device_id == device_id:
+            return self
+
+        if stream_holder is None:
+            raise ValueError("Stream holder is required for h2d/d2h transfers.")
+
+        if device_id == "cpu":
+            with utils.device_ctx(src_device_id):
+                tensor = ndbuffer.empty_like(self.tensor, device_id=ndbuffer.CPU_DEVICE_ID, stream=stream_holder)
+                ndbuffer.copy_into(tensor, self.tensor, stream_holder)
+        else:
+            with utils.device_ctx(device_id):
+                tensor = ndbuffer.empty_like(self.tensor, device_id=device_id, stream=stream_holder)
+                ndbuffer.copy_into(tensor, self.tensor, stream_holder)
+
+        return NDBufferTensor(tensor)
+
+    def copy_(self, src, stream_holder):
+        """
+        Inplace copy of src (copy the data from src into self).
+        """
+        device_id = self.tensor.device_id
+        src_nd = src._broadcast_to(self.shape).asndbuffer()
+        if device_id == "cpu":
+            device_id = src_nd.device_id
+        if device_id == "cpu":
+            ndbuffer.copy_into(self.tensor, src_nd, stream_holder)
+        else:
+            with utils.device_ctx(device_id):
+                ndbuffer.copy_into(self.tensor, src_nd, stream_holder)
+
+    def istensor(self):
+        """
+        Check if the object is ndarray-like.
+        """
+        return isinstance(self.tensor, ndbuffer.NDBuffer)
+
+    def reshape(self, shape: Sequence[int], *, copy: bool | None = None):
+        if copy:
+            raise NotImplementedError("Reshape with copy is not supported for ndbuffer")
+        return self.__class__(ndbuffer.reshaped_view(self.tensor, shape))
+
+    def memory_buffer(self):
+        """Creates a view of the memory buffer as a 1D tensor."""
+        raise NotImplementedError
+
+    def memory_buffer_to_tensor(self, shape, strides):
+        """
+        Creates a N-D tensor view of the memory buffer according to the specified
+        shape and strides.
+        """
+        raise NotImplementedError
+
+    def _broadcast_to(self, shape):
+        if self.tensor.shape == shape:
+            return self
+        raise NotImplementedError(
+            f"_broadcast_to is not supported for ndbuffer. Current shape: {self.tensor.shape}, requested shape: {shape}."
+        )