nvmath-python 1.0.0__cp314-cp314t-win_amd64.whl

nvmath/internal/tensor_ifc_cupy.py

@@ -0,0 +1,234 @@
+# 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 cuda.core import Buffer
+
+from . import utils
+from .ndbuffer import NDBuffer
+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"
+    host_tensor_class: type["HostTensor"]  # set at the end of the file
+    device_tensor_class: type["CupyTensor"]  # set at the end of the file
+
+    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="cpu")
+        dst_nd.copy_(src_nd, stream=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)
+        return Buffer.from_handle(alloc.ptr, size, owner=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: type["HostTensor"]  # set at the end of the file
+    device_tensor_class: type["CupyTensor"]  # set at the end of the file
+
+    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, memory_resource=_cupy_allocator)
+            dst_nd.copy_(src_nd, stream=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 NDBuffer.from_cupy(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):
+            self.asndbuffer().copy_(src.asndbuffer(), stream=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)
+
+
+CupyTensor.host_tensor_class = HostTensor
+HostTensor.host_tensor_class = HostTensor
+CupyTensor.device_tensor_class = CupyTensor
+HostTensor.device_tensor_class = CupyTensor

nvmath/internal/tensor_ifc_ndbuffer.py

@@ -0,0 +1,147 @@
+# 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 ndbuffer, typemaps, utils
+from .ndbuffer import NDBuffer, StridedLayout
+from .package_ifc import StreamHolder
+from .tensor_ifc import TensorHolder
+
+
+class NDBufferTensor(TensorHolder[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.
+        """
+        layout = StridedLayout(shape, strides, typemaps.NAME_TO_ITEM_SIZE[dtype])
+        if device_id == "cpu":
+            tensor = NDBuffer.empty(layout, dtype, "cpu", 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(layout, dtype, device_id, 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="cpu", stream=stream_holder)
+                tensor.copy_(self.tensor, stream=stream_holder)
+        else:
+            with utils.device_ctx(device_id):
+                tensor = NDBuffer.empty_like(self.tensor, device_id=device_id, stream=stream_holder)
+                tensor.copy_(self.tensor, stream=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.asndbuffer()
+        if device_id == "cpu":
+            device_id = src_nd.device_id
+        if device_id == "cpu":
+            self.tensor.copy_(src_nd, stream=stream_holder)
+        else:
+            with utils.device_ctx(device_id):
+                self.tensor.copy_(src_nd, stream=stream_holder)
+
+    def istensor(self):
+        """
+        Check if the object is ndarray-like.
+        """
+        return isinstance(self.tensor, 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__(self.tensor.reshape(shape))
+
+    def memory_buffer(self):
+        """Creates a view of the memory buffer as a 1D tensor."""
+        layout = StridedLayout.dense(self.size, itemsize=self.itemsize)
+        return self.__class__(self.tensor.as_strided(layout))
+
+    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."
+        layout = StridedLayout(shape, strides, itemsize=self.itemsize)
+        return self.__class__(self.tensor.as_strided(layout))

nvmath/internal/tensor_ifc_numpy.py

@@ -0,0 +1,184 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Interface to seamlessly use Numpy ndarray objects.
+"""
+
+__all__ = ["NumpyTensor", "CudaTensor"]
+
+from collections.abc import Sequence
+
+import numpy
+import numpy.typing as npt
+
+from nvmath.internal.tensor_ifc_ndbuffer import NDBufferTensor
+
+from . import utils
+from .ndbuffer import NDBuffer
+from .package_ifc import StreamHolder
+from .tensor_ifc import TensorHolder
+
+
+class CudaTensor(NDBufferTensor):
+    """
+    Wraps ndbuffer with data residing on the GPU.
+    It serves as a CUDA counterpart for NumpyTensor.
+    """
+
+    name = "cuda"
+    host_tensor_class: type["NumpyTensor"]  # set once NumpyTensor is defined
+    device_tensor_class: type["CudaTensor"]  # set once CudaTensor is defined
+
+    def __init__(self, tensor):
+        super().__init__(tensor)
+
+    @classmethod
+    def create_from_host(cls, tensor: TensorHolder, device_id: int, stream_holder: StreamHolder):
+        with utils.device_ctx(device_id):
+            src_ndbuffer = tensor.asndbuffer()
+            dst_ndbuffer = NDBuffer.empty_like(src_ndbuffer, device_id=device_id, stream=stream_holder)
+            dst_ndbuffer.copy_(src_ndbuffer, stream=stream_holder)
+            return cls(dst_ndbuffer)
+
+    def to(self, device_id, stream_holder):
+        if device_id == "cpu":
+            with utils.device_ctx(self.device_id):
+                dst = self.host_tensor_class.create_host_from(self, stream_holder)
+                return dst
+        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}.")
+        else:
+            raise ValueError(f"The device must be specified as an integer or 'cpu', not '{device_id}'.")
+
+
+class NumpyTensor(TensorHolder[npt.NDArray]):
+    """
+    TensorHolder for numpy ndarrays.
+    """
+
+    name = "numpy"
+    module = numpy
+    name_to_dtype = TensorHolder.create_name_dtype_map(
+        conversion_function=lambda name: numpy.dtype(name), exception_type=TypeError
+    )
+
+    host_tensor_class: type["NumpyTensor"]  # set once NumpyTensor is defined
+    device_tensor_class: type["CudaTensor"]  # set once CudaTensor is defined
+
+    def __init__(self, tensor):
+        super().__init__(tensor)
+
+    @property
+    def data_ptr(self):
+        return self.tensor.ctypes.data
+
+    @property
+    def device(self):
+        return "cpu"
+
+    @property
+    def device_id(self):
+        return "cpu"
+
+    @property
+    def dtype(self):
+        """Name of the data type"""
+        return self.tensor.dtype.name
+
+    @property
+    def itemsize(self):
+        return self.tensor.itemsize
+
+    @classmethod
+    def empty(cls, shape, device_id="cpu", *, dtype="float32", strides=None, **context):
+        """
+        Create an empty tensor of the specified shape and data type.
+        """
+        assert device_id == "cpu", "Internal Error: Numpy tensors must be allocated with device_id='cpu'"
+        dtype = NumpyTensor.name_to_dtype[dtype]
+        # when strides is not None, it should be of unit counts not bytes
+        return cls(
+            cls.module.ndarray(shape, dtype=dtype, strides=(tuple(s * dtype.itemsize for s in strides) if strides else None))
+        )
+
+    @classmethod
+    def create_host_from(cls, tensor: TensorHolder, stream_holder: StreamHolder):
+        src_nd = tensor.asndbuffer()
+        dst_nd = NDBuffer.empty_like(src_nd, device_id="cpu", stream=stream_holder)
+        dst_nd.copy_(src_nd, stream=stream_holder)
+        return cls(dst_nd.as_numpy())
+
+    @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)
+
+    def asndbuffer(self):
+        return NDBuffer.from_numpy(self.tensor)
+
+    def to(self, device_id, stream_holder):
+        if device_id == "cpu":
+            return self
+        elif isinstance(device_id, int):
+            dst = self.device_tensor_class.create_from_host(self, device_id, stream_holder)
+            return dst
+        else:
+            raise ValueError(f"The device must be specified as an integer or 'cpu', not '{device_id}'.")
+
+    def copy_(self, src, stream_holder):
+        match src.name:
+            case "numpy":
+                numpy.copyto(self.tensor, src.tensor)
+            case _:
+                with utils.device_ctx(src.device_id):
+                    self.asndbuffer().copy_(src.asndbuffer(), stream=stream_holder)
+
+    def istensor(self):
+        """
+        Check if the object is ndarray-like.
+        """
+        return isinstance(self.tensor, numpy.ndarray)
+
+    def reshape(self, shape: Sequence[int], *, copy: bool | None = None):
+        if int(numpy.__version__.split(".")[0]) < 2:
+            return self.__class__(numpy.reshape(self.tensor, shape))
+        else:
+            return self.__class__(numpy.reshape(self.tensor, shape, copy=copy))
+
+    def memory_buffer(self):
+        """Creates a view of the memory buffer as a 1D tensor."""
+        t = self.tensor
+        # TODO: ensure tensor is dense for now, and later support linear
+        # memory with constant stride.
+        v = NumpyTensor(t.ravel(order="K"))
+        assert v.data_ptr == self.data_ptr, "Internal error."
+        return v
+
+    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."
+        t = self.tensor
+        strides = [s * self.itemsize for s in strides]
+        v = NumpyTensor(numpy.lib.stride_tricks.as_strided(t, shape=shape, strides=strides))
+        assert v.data_ptr == self.data_ptr, "Internal error."
+        return v
+
+
+NumpyTensor.host_tensor_class = NumpyTensor
+CudaTensor.host_tensor_class = NumpyTensor
+NumpyTensor.device_tensor_class = CudaTensor
+CudaTensor.device_tensor_class = CudaTensor