nvmath-python 1.0.0__cp314-cp314t-win_amd64.whl

nvmath/internal/tensor_ifc_torch.py

@@ -0,0 +1,178 @@
+# Copyright (c) 2024-2026, 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 .package_ifc import StreamHolder
+from .tensor_ifc import TensorHolder
+
+
+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
+    )
+    host_tensor_class: type["TorchTensor"]  # set once NumpyTensor is defined
+    device_tensor_class: type["TorchTensor"]  # set once CudaTensor is defined
+
+    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 itemsize(self):
+        return self.tensor.itemsize
+
+    @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 asndbuffer(self):
+        raise RuntimeError("Converting torch tensor to ndbuffer is not supported")
+
+    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}'.")
+
+        # For h2d the data on the CPU is available and we block on `to()`, so a stream
+        # is not needed. It is currently ignored if provided.
+
+        # For d2h and d2d, we require a stream to ensure that the source data is ready
+        # before the copy operation is launched.
+        if self.device_id != "cpu" and stream_holder is None:
+            raise AssertionError("Internal error: a stream holder should be provided for d2h or d2d copies.")
+
+        # Always block for h2d and d2h copies to ensure that the copied data is ready
+        # for consumption when `to()` return.
+
+        # For d2d, the `torch.to()` operation is launched on the source device on the
+        # specified stream, so this stream must be on the source device. Their
+        # implementation ensures correct ordering between the current streams on the
+        # source and target devices.
+
+        blocking = self.device_id == "cpu" or device_id == "cpu"
+
+        stream_ctx = contextlib.nullcontext() if stream_holder is None else stream_holder.ctx
+        with stream_ctx:
+            tensor_device = self.tensor.to(device=device_id, non_blocking=not blocking)
+
+        return TorchTensor(tensor_device)
+
+    def copy_(self, src, stream_holder):
+        """
+        Inplace copy of src (copy the data from src into self).
+        """
+
+        # For d2h and d2d, we require a stream to ensure that the source data is ready
+        # before the copy operation is launched.
+        if src.device_id != "cpu" and stream_holder is None:
+            raise AssertionError("Internal error: a stream holder should be provided for d2h or d2d copies.")
+
+        stream_ctx = contextlib.nullcontext() if stream_holder is None else stream_holder.ctx
+        with stream_ctx:
+            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))
+
+    def memory_buffer(self):
+        """Creates a view of the memory buffer as a 1D tensor."""
+        storage = self.tensor.untyped_storage()
+        size = self.size
+        v = self.tensor.view(self.shape)
+        # TODO: ensure tensor is dense for now, and later support linear
+        # memory with constant stride.
+        v.set_(storage, self.tensor.storage_offset(), (size,), (1,))
+        return TorchTensor(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."
+        storage = self.tensor.untyped_storage()
+        t = self.tensor.view(self.shape)
+        t.set_(storage, self.tensor.storage_offset(), shape, strides)
+        return TorchTensor(t)
+
+    @property
+    def is_conjugate(self) -> bool:
+        return self.tensor.is_conj()
+
+
+TorchTensor.host_tensor_class = TorchTensor
+TorchTensor.device_tensor_class = TorchTensor

nvmath/internal/tensor_wrapper.py

@@ -0,0 +1,160 @@
+# Copyright (c) 2024-2026, 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"]
+
+import warnings
+from collections.abc import Sequence
+
+import numpy as np
+
+from . import formatters
+from .ndbuffer import NDBuffer
+from .tensor_ifc import AnyTensor, Tensor, TensorHolder
+from .tensor_ifc_numpy import CudaTensor, NumpyTensor
+
+_TENSOR_TYPES: dict[str, type[TensorHolder]] = {"numpy": NumpyTensor, "cuda": CudaTensor}
+
+_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"
+    if type(tensor) is NDBuffer:
+        return "cuda" if tensor.device == "cuda" else "cupy_host"
+    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 memory
+        from . import package_wrapper
+
+        if package == "torch":
+            from .package_ifc_torch import TorchPackage
+            from .tensor_ifc_torch import TorchTensor
+
+            _TENSOR_TYPES[package] = TorchTensor
+            package_wrapper.PACKAGE[package] = TorchPackage
+            memory.lazy_load_torch()
+        elif package == "cupy":
+            from .package_ifc_cupy import CupyPackage
+            from .tensor_ifc_cupy import CupyTensor, HostTensor
+
+            _TENSOR_TYPES[package] = CupyTensor
+            _TENSOR_TYPES["cupy_host"] = HostTensor
+            package_wrapper.PACKAGE[package] = CupyPackage
+            package_wrapper.PACKAGE["cupy_host"] = 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 src to the corresponding wrapped operands in dest.
+    """
+    for s, d in zip(src, dest, strict=True):
+        d.copy_(s, stream_holder=stream_holder)

nvmath/internal/typemaps.py

@@ -0,0 +1,113 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Functions to link type names with CUDA data and compute types.
+"""
+
+__all__ = ["COMPUTE_TYPE_TO_NAME", "DATA_TYPE_TO_NAME", "NAME_TO_DATA_TYPE", "NAME_TO_COMPUTE_TYPE", "NAME_TO_DATA_WIDTH"]
+
+import re
+
+from nvmath._utils import ComputeType
+from nvmath._utils import CudaDataType as cudaDataType
+
+
+def create_cuda_data_type_map(cuda_data_type_enum_class):
+    """
+    Map the data type name to the corresponding CUDA data type.
+    """
+    cuda_data_type_pattern = re.compile(r"CUDA_(?P<cr>C|R)_(?P<width>\d+)(?P<type>F|I|U|BF)_?(?P<kind>(E\dM\d)?)")
+
+    type_code_map = {"i": "int", "u": "uint", "f": "float", "bf": "bfloat"}
+    # A map from (width, exponent kind) to qualifiers (finite, unsigned zero, ...) for data
+    # types.
+    type_qualifier_map = {
+        (8, "e4m3"): "fn",
+        # FP4: finite-only, packed 2 per byte, see for example:
+        # - https://github.com/pytorch/pytorch/issues/146414
+        # - https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf
+        (4, "e2m1"): "fn_x2",
+    }
+
+    complex_types = {"float": "complex", "bfloat": "bcomplex"}
+
+    cuda_data_type_map = {}
+    data_type_width_map = {}
+    for d in cuda_data_type_enum_class:
+        m = cuda_data_type_pattern.match(d.name)
+
+        is_complex = m.group("cr").lower() == "c"
+        type_code = type_code_map[m.group("type").lower()]
+
+        # We'll generalize this if and when we support Gaussian integers.
+        if is_complex and type_code not in complex_types:
+            continue
+
+        width = int(m.group("width"))
+        if is_complex:
+            width *= 2
+            type_code = complex_types[type_code]
+
+        name = type_code + str(width)
+
+        # Handle narrow type kinds.
+        if width <= 8:
+            kind = m.group("kind").lower()
+            # Handle type qualifiers for narrow types.
+            kind += type_qualifier_map.get((width, kind), "")
+            if kind:
+                name += "_" + kind
+
+        cuda_data_type_map[name] = d
+        data_type_width_map[name] = width
+
+    return cuda_data_type_map, data_type_width_map
+
+
+def create_cuda_compute_type_map(cuda_compute_type_enum_class):
+    """
+    Map the data type name to the corresponding CUDA compute type.
+    """
+    cuda_compute_type_pattern = re.compile(r"COMPUTE_(?:(?P<width>\d+)(?P<type>F|I|U|BF)|(?P<tf32>TF32))")
+
+    type_code_map = {"i": "int", "u": "uint", "f": "float", "bf": "bfloat"}
+
+    cuda_compute_type_map = {}
+    for c in cuda_compute_type_enum_class:
+        if c.name == "COMPUTE_DEFAULT":
+            continue
+
+        m = cuda_compute_type_pattern.match(c.name)
+
+        if not m:
+            raise ValueError("Internal error - unexpected enum entry")
+
+        if m.group("tf32"):
+            continue
+
+        name = type_code_map[m.group("type").lower()] + m.group("width")
+        cuda_compute_type_map[name] = c
+
+    # Treat complex types as special case.
+    cuda_compute_type_map["bcomplex32"] = cuda_compute_type_enum_class.COMPUTE_16BF
+    cuda_compute_type_map["complex32"] = cuda_compute_type_enum_class.COMPUTE_16F
+    cuda_compute_type_map["complex64"] = cuda_compute_type_enum_class.COMPUTE_32F
+    cuda_compute_type_map["complex128"] = cuda_compute_type_enum_class.COMPUTE_64F
+
+    return cuda_compute_type_map
+
+
+NAME_TO_DATA_TYPE, NAME_TO_DATA_WIDTH = create_cuda_data_type_map(cudaDataType)
+NAME_TO_ITEM_SIZE = {k: v // 8 for k, v in NAME_TO_DATA_WIDTH.items() if v % 8 == 0}
+DATA_TYPE_TO_NAME = {v: k for k, v in NAME_TO_DATA_TYPE.items()}
+NAME_TO_COMPUTE_TYPE = create_cuda_compute_type_map(ComputeType)
+COMPUTE_TYPE_TO_NAME = {v: k for k, v in NAME_TO_COMPUTE_TYPE.items()}
+
+
+FFTW_SUPPORTED_SINGLE = (cudaDataType.CUDA_R_32F, cudaDataType.CUDA_C_32F)
+FFTW_SUPPORTED_DOUBLE = (cudaDataType.CUDA_R_64F, cudaDataType.CUDA_C_64F)
+FFTW_SUPPORTED_FLOAT = (cudaDataType.CUDA_R_32F, cudaDataType.CUDA_R_64F)
+FFTW_SUPPORTED_COMPLEX = (cudaDataType.CUDA_C_32F, cudaDataType.CUDA_C_64F)
+FFTW_SUPPORTED_TYPES = sorted(FFTW_SUPPORTED_SINGLE + FFTW_SUPPORTED_DOUBLE)