nvmath-python 0.5.0__cp313-cp313-win_amd64.whl

nvmath/internal/__init__.py

@@ -0,0 +1,5 @@
+"""APIs in this module are for NVIDIA-use only.
+
+APIs in internal are not intended for third-party developers, but are shared between NVIDIA
+packages.
+"""

nvmath/internal/enum_utils.py

@@ -0,0 +1,142 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Factories for create options dataclasses, as well as utilities to add docstring to enum
+classes.
+"""
+
+import dataclasses
+from enum import IntEnum
+import re
+from typing import Any, ClassVar
+from collections.abc import Callable
+
+import numpy
+
+
+def create_options_class_from_enum(
+    options_class_name: str,
+    enum_class: IntEnum,
+    get_attr_dtype: Callable,
+    description: str,
+    filter_re: str = r"(?P<option_name>.*)",
+):
+    """
+    Create an options dataclass from a Python enum class. Names can be filtered if desired.
+
+    Args:
+       options_class_name: Name of the dataclass that will be created.
+
+       enum_class: The IntEnum class that contains the options for the dataclass.
+
+       get_attr_dtype: A callable that takes in an enum value as the argument and returns
+            the size in bytes of the cuTensorNet.
+
+       filter_re: A re definition that defines the match named 'option_name'.
+    """
+    if r"(?P<option_name>" not in filter_re:
+        message = """Incorrect re.
+The re for the filter must contain the named group 'option_name'."""
+        raise ValueError(message)
+
+    # Helper vars for creating attribute docstring.
+    doc = f"""A data class for capturing the {description} options.
+
+    Attributes:
+"""
+    indent = " " * 8
+    prefix = determine_enum_prefix(enum_class, "_ATTRIBUTE")
+
+    filter_pattern = re.compile(filter_re)
+    option_to_enum: dict[str, IntEnum] = {}
+    option_to_dtype: dict[str, Any] = {}
+    # mypy claims that IntEnum is not iterable
+    for e in enum_class:  # type: ignore
+        m = filter_pattern.match(e.name)
+        if not m:
+            continue
+        option_name = m.group("option_name").lower()
+        option_to_enum[option_name] = e
+        option_to_dtype[option_name] = get_attr_dtype(e)
+
+        # Add docstring for this attribute.
+        doc += indent + option_name + ":" + f" See `{prefix + '_' + m.group(0)}`.\n"
+
+    field: tuple[str, Any, None | dict]
+    fields: list[tuple[str, Any, None | dict]] = []
+    for option_name, dtype in option_to_dtype.items():
+        if numpy.issubdtype(dtype, numpy.integer):
+            field = option_name, int | None, dataclasses.field(default=None)
+        else:
+            field = option_name, Any | None, dataclasses.field(default=None)
+        fields.append(field)
+
+    # Add class attributes.
+
+    field = "option_to_enum", ClassVar[dict], dataclasses.field(default=option_to_enum)
+    fields.append(field)
+
+    field = "option_to_dtype", ClassVar[dict], dataclasses.field(default=option_to_dtype)
+    fields.append(field)
+
+    # Create the options class.
+    options_class = dataclasses.make_dataclass(options_class_name, fields)
+    options_class.__doc__ = doc
+
+    return options_class
+
+
+def snake_to_camel(names):
+    name = ""
+    for i, sub_name in enumerate(names):
+        if i == 0:
+            name += sub_name.lower()
+        else:
+            name += sub_name[0].upper() + sub_name[1:]
+    name += "_t"
+    return name
+
+
+def camel_to_snake(name, upper=True):
+    """
+    Convert string from camel case to snake style.
+    """
+
+    def transformer(m):
+        prefix = iter(("", "_", ""))
+        for i in 1, 3, 5:
+            first, second = i, i + 1
+            s = next(prefix)
+            if m.group(second):
+                if m.group(first):  # If the second group doesn't exist, the first won't either by the design of the RE.
+                    s += m.group(first).lower() + "_"
+                s += m.group(second).lower()
+                break
+        return s
+
+    name = re.sub(r"^([A-Z]*)([A-Z])|(?<!_)([A-Z]*)([A-Z])|([A-Z]*)([A-Z])", transformer, name)
+    if upper:
+        name = name.upper()
+    return name
+
+
+def determine_enum_prefix(enum_class, chomp):
+    """
+    This function assumes that the convention used to translate C enumerators to Python enum
+    names holds.
+    """
+
+    prefix = enum_class.__module__.split(".")[-1].upper()
+    prefix += "_" + camel_to_snake(enum_class.__name__)
+    prefix = re.sub(chomp, "", prefix)
+    return prefix
+
+
+def add_enum_class_doc(enum_class, chomp):
+    """
+    Add docstring to enum classes.
+    """
+    for e in enum_class:
+        e.__doc__ = f"See `{determine_enum_prefix(enum_class, chomp) + '_' + e.name.upper()}`."

nvmath/internal/formatters.py

@@ -0,0 +1,87 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Formatters for printing data.
+"""
+
+import numpy as np
+
+
+class MemoryStr:
+    """
+    A simple type to pretty-print memory-like values.
+    """
+
+    def __init__(self, memory, base_unit="B"):
+        self.memory = memory
+        self.base_unit = base_unit
+        self.base = 1024
+
+    def __str__(self):
+        """
+        Convert large values to powers of 1024 for readability.
+        """
+
+        base, base_unit, memory = self.base, self.base_unit, self.memory
+
+        if memory < base:
+            value, unit = memory, base_unit
+        elif memory < base**2:
+            value, unit = memory / base, f"Ki{base_unit}"
+        elif memory < base**3:
+            value, unit = memory / base**2, f"Mi{base_unit}"
+        else:
+            value, unit = memory / base**3, f"Gi{base_unit}"
+
+        return f"{value:0.2f} {unit}"
+
+
+class FLOPSStr:
+    """
+    A simple type to pretty-print FLOP count and FLOP/s-like values.
+    """
+
+    def __init__(self, flops, base_unit):
+        self.flops = flops
+        self.base_unit = base_unit
+        self.base = 1000
+
+    def __str__(self):
+        """
+        Convert large values to powers of 1000 for readability.
+        """
+
+        base, base_unit, flops = self.base, self.base_unit, self.flops
+
+        if flops < base:
+            value, unit = flops, base_unit
+        elif flops < base**2:
+            value, unit = flops / base, f"K{base_unit}"
+        elif flops < base**3:
+            value, unit = flops / base**2, f"M{base_unit}"
+        elif flops < base**4:
+            value, unit = flops / base**3, f"G{base_unit}"
+        elif flops < base**5:
+            value, unit = flops / base**4, f"T{base_unit}"
+        else:
+            value, unit = flops / base**5, f"P{base_unit}"
+
+        return f"{value:0.3f} {unit}"
+
+
+def array2string(array_like):
+    """
+    String representation of an array-like object with possible truncation of "interior"
+    values to limit string size.
+
+    The NumPy function "set_printoptions" can be used to control the display of the array.
+    """
+
+    return np.array2string(
+        np.asanyarray(array_like, dtype=object),
+        separator=", ",
+        # NumPy hates empty strings so we print 'None' instead.
+        formatter={"object": lambda s: s if s != "" else "None"},
+    )

nvmath/internal/mem_limit.py

@@ -0,0 +1,51 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Memory specification regular expression.
+"""
+
+__all__ = ["MEM_LIMIT_RE_PCT", "MEM_LIMIT_RE_VAL", "MEM_LIMIT_DOC", "check_memory_str"]
+
+import re
+
+MEM_LIMIT_RE_PCT = re.compile(r"(?P<value>[+]?[0-9]*\.?[0-9]+(?:[eE][-+]?[0-9]+)?)\s*%\s*$")
+MEM_LIMIT_RE_VAL = re.compile(
+    r"(?P<value>[+]?[0-9]*\.?[0-9]+(?:[eE][-+]?[0-9]+)?)\s*(?P<units>[kmg])?(?P<binary>(?<=[kmg])i)?b\s*$",
+    re.IGNORECASE,
+)
+MEM_LIMIT_DOC = """The {kind} must be specified in one of the following forms:
+  (1) A number (int or float). If the number is a float between 0 and 1 inclusive, the
+      {kind} is interpreted as a fraction of the total device memory; otherwise it is
+      interpreted as the number of bytes of memory, with float value being cast to int.
+      Examples: 0.75, 50E6, 50000000, ...
+  (2) A string containing a positive value followed by B, kB, MB, or GB for powers of 1000.
+      Examples: "0.05 GB", "50 MB", "50000000 B" ...
+  (3) A string containing a positive value followed by kiB, MiB, or GiB for powers of 1024.
+      Examples:  "0.05 GiB", "51.2 MiB", "53687091 B" ...
+  (4) A string with value in the range [0, 100] followed by a % symbol.
+      Examples: "26%","82%", ...
+
+  Whitespace between values and units is optional.
+
+The provided {kind} is "{value}".
+"""
+
+
+def check_memory_str(value, kind):
+    """
+    Check if the memory specification string is valid.
+
+    value = the memory specification string.
+    kind  = a string denoting the type of memory being checked, used in error messages.
+    """
+    if not isinstance(value, int | float):
+        m1 = MEM_LIMIT_RE_PCT.match(value)
+        if m1:
+            factor = float(m1.group("value"))
+            if factor < 0 or factor > 100:
+                raise ValueError(f"The {kind} percentage must be in the range [0, 100].")
+        m2 = MEM_LIMIT_RE_VAL.match(value)
+        if not (m1 or m2):
+            raise ValueError(MEM_LIMIT_DOC.format(kind=kind, value=value))

nvmath/internal/package_ifc.py

@@ -0,0 +1,119 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+An abstract interface to certain package-provided operations.
+
+Even though cuda.core does not have the concept of a "current stream" or a "stream context",
+we promise to honor third-party CUDA stream contexts if they match the input operand's
+package. The interface class defined in this module wraps around those stream concepts to
+make them uniform.
+
+The strategy is to use cuda.core everywhere internally except for context managers which
+will need to wrap back around to the external implementation.
+"""
+
+__all__ = ["Package", "StreamHolder"]
+
+from abc import ABC, abstractmethod
+from collections.abc import Hashable
+from contextlib import AbstractContextManager
+from dataclasses import dataclass
+from typing import Generic, Protocol, TypeVar
+
+import cuda.core.experimental as ccx
+
+
+class AnyStream(Hashable, Protocol):
+    """Any supported Stream object such as a Stream from cuda.core, CuPy, or PyTorch."""
+
+    # 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.
+    pass
+
+
+"""
+A generic type for the third-party Stream which a given Package implementation wraps
+around.
+"""
+Stream = TypeVar("Stream")
+
+
+class Package(ABC, Generic[Stream]):
+    @staticmethod
+    @abstractmethod
+    def get_current_stream(device_id: int) -> Stream:
+        """
+        Obtain the current stream on the device.
+
+        Args:
+            device_id: The id (ordinal) of the device.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def to_stream_pointer(stream: Stream) -> int:
+        """
+        Obtain the stream pointer.
+
+        Args:
+            stream: The stream object.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def to_stream_context(stream: Stream) -> AbstractContextManager[Stream]:
+        """
+        Create a context manager from the stream.
+
+        Args:
+            stream: The stream object.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def create_external_stream(device_id: int, stream_ptr: int) -> Stream:
+        """
+        Wrap a stream pointer into an external stream object.
+
+        Args:
+            device_id: The id (ordinal) of the device.
+            stream: The stream pointer (int) to be wrapped.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def create_stream(cls, external: Stream) -> ccx.Stream:
+        """
+        Wrap an external Stream object into a cuda.core.Stream.
+
+        Args:
+            external: The external Stream object.
+        """
+        return ccx.Stream.from_handle(cls.to_stream_pointer(external))
+
+
+@dataclass
+class StreamHolder(Generic[Stream]):
+    """A data class for easing CUDA stream manipulation.
+
+    Attributes:
+        ctx: A context manager for using the specified stream.
+        device_id (int): The device ID where the encapsulated stream locates.
+        external: A foreign object that holds the stream alive.
+        obj: The cuda.core Stream object wrapping external.
+        package (str): The name of the package to which the external stream belongs.
+        ptr (int): The address of the underlying ``cudaStream_t`` object.
+    """
+
+    ctx: AbstractContextManager[Stream]
+    device_id: int
+    external: Stream
+    obj: ccx.Stream
+    package: str
+    ptr: int

nvmath/internal/package_ifc_cuda.py

@@ -0,0 +1,43 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Interface to cuda.core operations.
+"""
+
+__all__ = ["CUDAPackage"]
+
+import contextlib
+import typing
+
+import cuda.core.experimental as ccx
+
+from .package_ifc import Package
+
+
+class CUDAPackage(Package[ccx.Stream]):
+    @staticmethod
+    def get_current_stream(device_id: int):
+        message = "cuda.core has no concept of a current stream or a stream context."
+        raise NotImplementedError(message)
+
+    @staticmethod
+    def to_stream_pointer(stream: ccx.Stream) -> int:  # type: ignore[override]
+        return int(stream.handle)
+
+    @staticmethod
+    def to_stream_context(stream: ccx.Stream):  # type: ignore[override]
+        @contextlib.contextmanager
+        def stream_context() -> typing.Iterator[ccx.Stream]:
+            yield stream
+
+        return stream_context
+
+    @staticmethod
+    def create_external_stream(device_id: int, stream_ptr: int) -> ccx.Stream:
+        return ccx.Stream.from_handle(stream_ptr)
+
+    @classmethod
+    def create_stream(cls, external: ccx.Stream) -> ccx.Stream:  # type: ignore[override]
+        return external

nvmath/internal/package_ifc_cupy.py

@@ -0,0 +1,49 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Interface to CuPy operations.
+"""
+
+__all__ = ["CupyPackage"]
+
+import cupy as cp
+
+from . import utils
+from .package_ifc import Package
+
+# Using the functional API is faster than setting a device context
+if int(cp.__version__.split(".")[0]) >= 13:
+    _get_current_stream = cp.cuda.get_current_stream
+else:
+
+    def _get_current_stream(device_id: int):
+        with utils.device_ctx(device_id):
+            stream = cp.cuda.get_current_stream()
+        return stream
+
+
+# Monkey patch older versions of CuPy, so that Streams are hashable
+# NOTE: We choose not to patch the BaseStream/_BaseStream class because of name change
+if int(cp.__version__.split(".")[0]) < 11:
+    cp.cuda.Stream.__hash__ = lambda self: hash(self.ptr)
+    cp.cuda.ExternalStream.__hash__ = lambda self: hash(self.ptr)
+
+
+class CupyPackage(Package[cp.cuda.Stream]):
+    @staticmethod
+    def get_current_stream(device_id: int):
+        return _get_current_stream(device_id)
+
+    @staticmethod
+    def to_stream_pointer(stream: cp.cuda.Stream) -> int:
+        return stream.ptr
+
+    @staticmethod
+    def to_stream_context(stream: cp.cuda.Stream):
+        return stream
+
+    @staticmethod
+    def create_external_stream(device_id: int, stream_ptr: int) -> cp.cuda.ExternalStream:
+        return cp.cuda.ExternalStream(stream_ptr)

nvmath/internal/package_ifc_torch.py

@@ -0,0 +1,31 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Interface to Torch operations.
+"""
+
+__all__ = ["TorchPackage"]
+
+import torch
+
+from .package_ifc import Package
+
+
+class TorchPackage(Package[torch.cuda.Stream]):
+    @staticmethod
+    def get_current_stream(device_id: int) -> torch.cuda.Stream:
+        return torch.cuda.current_stream(device=device_id)
+
+    @staticmethod
+    def to_stream_pointer(stream: torch.cuda.Stream) -> int:  # type: ignore[override]
+        return stream.cuda_stream
+
+    @staticmethod
+    def to_stream_context(stream: torch.cuda.Stream) -> torch.cuda.StreamContext:  # type: ignore[override]
+        return torch.cuda.stream(stream)
+
+    @staticmethod
+    def create_external_stream(device_id: int, stream_ptr: int):
+        return torch.cuda.ExternalStream(stream_ptr, device=device_id)

nvmath/internal/package_wrapper.py

@@ -0,0 +1,14 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Create map from package names to package interface objects.
+"""
+
+__all__ = ["PACKAGE", "AnyStream", "StreamHolder"]
+
+from .package_ifc import Package, AnyStream, StreamHolder
+from .package_ifc_cuda import CUDAPackage
+
+PACKAGE: dict[str, type[Package]] = {"cuda": CUDAPackage}

nvmath/internal/tensor_ifc.py

@@ -0,0 +1,143 @@
+# Copyright (c) 2024-2025, 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 typing import Any, Literal, Generic, Protocol, TypeVar
+from types import ModuleType
+
+from . import typemaps
+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
+
+    @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]:
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def size(self) -> int:
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def strides(self) -> Sequence[int]:
+        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 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