nvmath-python 0.5.0__cp313-cp313-win_amd64.whl

nvmath/internal/tensor_ifc_cupy.py

@@ -0,0 +1,185 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Interface to seamlessly use Cupy ndarray objects.
+"""
+
+__all__ = ["CupyTensor"]
+
+from collections.abc import Sequence
+
+import cupy
+import numpy as np
+
+from . import utils
+from .tensor_ifc import TensorHolder
+from .tensor_ifc_numpy import NumpyTensor
+from .package_ifc import StreamHolder
+
+
+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
+    )
+
+    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 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 numpy(self, stream_holder: StreamHolder):
+        stream = stream_holder.external
+        with stream:
+            out = self.tensor.get(stream=stream)
+        # cupy/cupy#7820
+        if stream is not None:
+            stream.synchronize()
+        return NumpyTensor(out)
+
+    @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)
+
+    def to(self, device_id, stream_holder):
+        if not (device_id == "cpu" or isinstance(device_id, int)):
+            raise ValueError(f"The device must be specified as an integer or 'cpu', not '{device_id}'.")
+
+        if device_id == "cpu":
+            return self.numpy(stream_holder=stream_holder)
+
+        with utils.device_ctx(device_id), stream_holder.ctx:
+            return CupyTensor(cupy.asarray(self.tensor))
+
+    def _c2c_copy_(self, src: cupy.ndarray, stream_holder: StreamHolder):
+        """
+        Inplace copy of src (copy the data from src into self).
+        The src must by cupy ndarray
+        """
+        with stream_holder.ctx:
+            cupy.copyto(self.tensor, src)
+
+    def _n2c_copy_(self, src: np.ndarray, stream_holder: StreamHolder):
+        """
+        Inplace copy of src (copy the data from src into self).
+        The src must by numpy ndarray
+        """
+        stream = stream_holder.external
+        try:
+            self.tensor.set(src, stream=stream)
+        except RuntimeError as e:
+            # If self is a strided tensor (neither c nor f layout)
+            # cupy refuses to copy from numpy array
+            if "set to non-contiguous array" not in str(e):
+                raise
+            else:
+                with stream_holder.ctx:
+                    src_gpu = cupy.asarray(src)
+                    cupy.copyto(self.tensor, src_gpu)
+        # cupy/cupy#7820
+        if stream is not None:
+            stream.synchronize()
+
+    def copy_(self, src, stream_holder):
+        """
+        Inplace copy of src (copy the data from src into self).
+        """
+        match src.name:
+            case "cupy":
+                self._c2c_copy_(src.tensor, stream_holder)
+            case "numpy":
+                self._n2c_copy_(src.tensor, stream_holder)
+            case _:
+                msg = f"CupyTensor does not convert from {src.name}."
+                raise NotImplementedError(msg)
+
+    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:
+                raise ValueError(f"Could not reshape cupy array without copy: current shape={self.shape}, new shape={shape}")
+        else:
+            reshaped_tensor = self.tensor.reshape(shape)
+        return self.__class__(reshaped_tensor)

nvmath/internal/tensor_ifc_numpy.py

@@ -0,0 +1,157 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Interface to seamlessly use Numpy ndarray objects.
+"""
+
+__all__ = ["NumpyTensor"]
+
+try:
+    import cupy  # type: ignore
+except ImportError:
+
+    class cupy:  # type: ignore
+        """A placeholder for the cupy module when it is unavailable."""
+
+        @classmethod
+        def asnumpy(cls, *args, **kwargs):
+            raise ImportError("Cannot convert cupy to numpy array when cupy is not installed!")
+
+        @classmethod
+        def asarray(cls, *args, **kwargs):
+            raise ImportError("Cannot convert numpy to cupy array when cupy is not installed!")
+
+
+from collections.abc import Sequence
+
+import numpy
+import numpy.typing as npt
+
+from . import utils
+from .package_ifc import StreamHolder
+from .tensor_ifc import TensorHolder
+
+
+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
+    )
+
+    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
+
+    @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))
+        )
+
+    @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 to(self, device_id, stream_holder):
+        if not (device_id == "cpu" or isinstance(device_id, int)):
+            raise ValueError(f"The device must be specified as an integer or 'cpu', not '{device_id}'.")
+
+        if device_id == "cpu":
+            return NumpyTensor(self.tensor)
+
+        # FIXME: Replace with native tensor implementation to avoid required dep on CuPy
+        from .tensor_ifc_cupy import CupyTensor
+
+        with utils.device_ctx(device_id), stream_holder.ctx:
+            return CupyTensor(cupy.asarray(self.tensor))
+
+    def _n2n_copy_(self, src: npt.NDArray) -> None:
+        """
+        Inplace copy of src (copy the data from src into self).
+        The src must by numpy ndarray
+        """
+        numpy.copyto(self.tensor, src)
+
+    def _c2n_copy_(self, src, stream_holder: StreamHolder) -> None:
+        """
+        Inplace copy of src (copy the data from src into self).
+        The src must by cupy ndarray
+        """
+        stream = stream_holder.external
+        try:
+            with stream:
+                src.get(stream=stream, out=self.tensor)
+        except RuntimeError as e:
+            # If self is a strided tensor (neither c nor f layout)
+            # cupy refuses to copy to numpy array
+            if "copying to non-contiguous ndarray" not in str(e):
+                raise
+            else:
+                # we cannot simply use blocking=True, as it is
+                # not supported by older cupy releases (<13)
+                src_cpu = cupy.asnumpy(src, stream=stream)
+                self._n2n_copy_(src_cpu)
+        # cupy/cupy#7820
+        if stream is not None:
+            stream.synchronize()
+
+    def copy_(self, src, stream_holder):
+        # Handle NumPy <=> CuPy CPU-GPU ndarray asymmetry.
+        match src.name:
+            case "cupy":
+                assert stream_holder is not None
+                self._c2n_copy_(src.tensor, stream_holder)
+            case "numpy":
+                self._n2n_copy_(src.tensor)
+            case _:
+                msg = f"NumpyTensor does not convert from {src.name}."
+                raise NotImplementedError(msg)
+
+    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))

nvmath/internal/tensor_ifc_torch.py

@@ -0,0 +1,115 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Interface to seamlessly use Torch tensor objects.
+"""
+
+__all__ = ["TorchTensor"]
+
+import contextlib
+from collections.abc import Sequence
+
+import torch
+
+from .tensor_ifc import TensorHolder
+from .package_ifc import StreamHolder
+
+
+class TorchTensor(TensorHolder[torch.Tensor]):
+    """
+    TensorHolder for Torch Tensors.
+    """
+
+    name = "torch"
+    module = torch
+    name_to_dtype = TensorHolder.create_name_dtype_map(
+        conversion_function=lambda name: getattr(torch, name), exception_type=AttributeError
+    )
+
+    def __init__(self, tensor):
+        super().__init__(tensor)
+
+    @property
+    def data_ptr(self):
+        return self.tensor.data_ptr()
+
+    @property
+    def device(self):
+        return str(self.tensor.device).split(":")[0]
+
+    @property
+    def device_id(self):
+        index = self.tensor.device.index
+        return "cpu" if index is None else index
+
+    @property
+    def dtype(self):
+        """Name of the data type"""
+        return str(self.tensor.dtype).split(".")[-1]
+
+    @property
+    def shape(self):
+        return tuple(self.tensor.shape)
+
+    @property
+    def size(self):
+        return self.tensor.nelement()
+
+    @property
+    def strides(self):
+        return self.tensor.stride()
+
+    @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 on the specified device
+        (None, 'cpu', or device id).
+
+        Note, that the strides, if specified, should correspond to a dense
+        (possibly permuted) tensor and MUST NOT overlap.
+        Otherwise, the behaviour is not defined.
+        """
+        dtype = TorchTensor.name_to_dtype[dtype]
+        stream_ctx = contextlib.nullcontext() if stream_holder is None else stream_holder.ctx
+        with stream_ctx:
+            if strides:
+                # note: torch strides is not scaled by bytes
+                tensor = torch.empty_strided(shape, strides, dtype=dtype, device=device_id)
+            else:
+                tensor = torch.empty(shape, dtype=dtype, device=device_id)
+
+        return cls(tensor)
+
+    def to(self, device_id, stream_holder):
+        if not (device_id == "cpu" or isinstance(device_id, int)):
+            raise ValueError(f"The device must be specified as an integer or 'cpu', not '{device_id}'.")
+
+        assert stream_holder is not None, "Internal Error: moving TorchTensor requires a stream."
+        with stream_holder.ctx:
+            tensor_device = self.tensor.to(device=device_id, non_blocking=(device_id != "cpu"))
+
+        return TorchTensor(tensor_device)
+
+    def copy_(self, src, stream_holder):
+        """
+        Inplace copy of src (copy the data from src into self).
+        """
+        with (stream_holder and stream_holder.ctx) or contextlib.nullcontext():
+            self.tensor.copy_(src.tensor)
+
+    def istensor(self):
+        """
+        Check if the object is ndarray-like.
+        """
+        return isinstance(self.tensor, torch.Tensor)
+
+    def reshape(self, shape: Sequence[int], *, copy: bool | None = None):
+        if copy:
+            raise NotImplementedError("reshape with copy=True not implemented")
+        if copy is False:
+            return self.__class__(self.tensor.view(shape))
+        return self.__class__(self.tensor.reshape(shape))

nvmath/internal/tensor_wrapper.py

@@ -0,0 +1,160 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Entry point to using tensors from different libraries seamlessly.
+"""
+
+__all__ = ["infer_tensor_package", "wrap_operand", "wrap_operands", "to", "copy_", "TensorHolder", "AnyTensor"]
+
+from collections.abc import Sequence
+import warnings
+
+import numpy as np
+
+from . import formatters
+from .tensor_ifc import TensorHolder, Tensor, AnyTensor
+from .tensor_ifc_numpy import NumpyTensor
+
+
+_TENSOR_TYPES: dict[str, type[TensorHolder]] = {"numpy": NumpyTensor}
+
+_SUPPORTED_PACKAGES = tuple(_TENSOR_TYPES.keys())
+
+
+def infer_tensor_package(tensor):
+    """
+    Infer the package that defines this tensor.
+    """
+    if issubclass(tensor.__class__, np.ndarray):
+        return "numpy"
+    module = tensor.__class__.__module__
+    package = module.split(".")[0]
+    maybe_register_package(package)
+    return package
+
+
+def maybe_register_package(package):
+    global _SUPPORTED_PACKAGES
+    if package not in _SUPPORTED_PACKAGES:
+        global _TENSOR_TYPES
+        from . import package_wrapper
+        from .. import memory
+
+        if package == "torch":
+            from .tensor_ifc_torch import TorchTensor
+            from .package_ifc_torch import TorchPackage
+
+            _TENSOR_TYPES[package] = TorchTensor
+            package_wrapper.PACKAGE[package] = TorchPackage
+            memory.lazy_load_torch()
+        elif package == "cupy":
+            from .tensor_ifc_cupy import CupyTensor
+            from .package_ifc_cupy import CupyPackage
+
+            _TENSOR_TYPES[package] = CupyTensor
+            package_wrapper.PACKAGE[package] = CupyPackage
+            memory.lazy_load_cupy()
+        else:
+            message = f"""{package} not supported yet. Currently must be one of ['numpy', 'cupy', 'torch']"""
+            raise ValueError(message)
+        _SUPPORTED_PACKAGES = tuple(_TENSOR_TYPES.keys())
+
+
+def wrap_operand(native_operand: Tensor) -> TensorHolder[Tensor]:
+    """
+    Wrap one "native" operand so that package-agnostic API can be used.
+    """
+    if isinstance(native_operand, TensorHolder):
+        msg = (
+            "wrap_operand() is being called unnecessarily because the input is already a TensorHolder. "
+            "Only public facing APIs should call wrap_operand(). "
+            "Internal APIs should assume the operands are TensorHolder already."
+            "Trying to wrap a TensorHolder will become an error in the future."
+        )
+        warnings.warn(msg, DeprecationWarning)
+        return native_operand
+    wrapped_operand = _TENSOR_TYPES[infer_tensor_package(native_operand)](native_operand)
+    return wrapped_operand
+
+
+def check_valid_package(native_operands):
+    """
+    Check if the operands belong to one of the supported packages.
+    """
+    operands_pkg = [infer_tensor_package(o) for o in native_operands]
+    checks = [p in _SUPPORTED_PACKAGES for p in operands_pkg]
+    if not all(checks):
+        unknown = [f"{location}: {operands_pkg[location]}" for location, predicate in enumerate(checks) if predicate is False]
+        unknown = formatters.array2string(unknown)
+        message = f"""The operands should be ndarray-like objects from one of {_SUPPORTED_PACKAGES} packages.
+The unsupported operands as a sequence of "position: package" is: \n{unknown}"""
+        raise ValueError(message)
+
+    return operands_pkg
+
+
+def check_valid_operand_type(wrapped_operands):
+    """
+    Check if the wrapped operands are ndarray-like.
+    """
+    istensor = [o.istensor() for o in wrapped_operands]
+    if not all(istensor):
+        unknown = [
+            f"{location}: {type(wrapped_operands[location].tensor)}"
+            for location, predicate in enumerate(istensor)
+            if predicate is False
+        ]
+        unknown = formatters.array2string(unknown)
+        message = f"""The operands should be ndarray-like objects from one of {_SUPPORTED_PACKAGES} packages.
+The unsupported operands as a sequence of "position: type" is: \n{unknown}"""
+        raise ValueError(message)
+
+
+def _wrapper_helper(operand: Tensor | TensorHolder, index: int, packages: list[str]) -> TensorHolder[Tensor]:
+    if isinstance(operand, TensorHolder):
+        msg = (
+            "wrap_operands() is being called unnecessarily because at least one input is already a TensorHolder. "
+            "Only public facing APIs should call wrap_operands(). "
+            "Internal APIs should assume the operands are TensorHolder already."
+            "Trying to wrap a TensorHolder will become an error in the future."
+        )
+        warnings.warn(msg, DeprecationWarning)
+        return operand
+    return _TENSOR_TYPES[packages[index]](operand)
+
+
+def wrap_operands(native_operands: Sequence[Tensor]) -> Sequence[TensorHolder[Tensor]]:
+    """
+    Wrap the "native" operands so that package-agnostic API can be used.
+    """
+
+    operands_pkg = check_valid_package(native_operands)
+
+    wrapped_operands = tuple(_wrapper_helper(o, i, operands_pkg) for i, o in enumerate(native_operands))
+
+    check_valid_operand_type(wrapped_operands)
+
+    return wrapped_operands
+
+
+def to(operands: Sequence[TensorHolder], device_id, stream_holder) -> Sequence[TensorHolder]:
+    """
+    Copy the wrapped operands to the specified device_id (None or int) and return the
+    wrapped operands on the device.
+    """
+    assert isinstance(operands, Sequence) and isinstance(operands[0], TensorHolder)
+    return tuple(o.to(device_id, stream_holder) for o in operands)
+
+
+def copy_(src: Sequence[TensorHolder], dest: Sequence[TensorHolder], stream_holder):
+    """
+    Copy the wrapped operands in dest to the corresponding wrapped operands in src.
+    """
+    for s, d in zip(src, dest, strict=True):
+        if s.device_id == "cpu":
+            # FIXME: This is probably an extra step because it's supported to copy directly
+            # from cpu to device?
+            s = s.to(d.device_id, stream_holder=stream_holder)
+        d.copy_(s, stream_holder=stream_holder)