nvmath-python 1.0.0__cp314-cp314t-win_amd64.whl

nvmath/internal/enum_utils.py

@@ -0,0 +1,142 @@
+# Copyright (c) 2024-2026, 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
+import re
+from collections.abc import Callable
+from enum import IntEnum
+from typing import Any, ClassVar
+
+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, slots=True)
+    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-2026, 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-2026, 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/memory.pxd

@@ -0,0 +1,13 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+cimport cython
+from libc.stdint cimport int64_t, intptr_t, uint64_t
+
+cpdef get_device_memory_resource(int device_id)
+cpdef get_pinned_async_memory_resource(int device_id)
+cpdef get_legacy_pinned_memory_resource(int device_id)
+cpdef allocate_from_mr(mr, int64_t size, stream, int device_id, logger=*)
+cpdef free_reserved_memory(bint sync=*)
+cdef get_default_stream(int device_id)

nvmath/internal/memory.pyi

@@ -0,0 +1,50 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import _cython_3_2_5
+from typing import Any, ClassVar
+
+__pyx_capi__: dict
+__test__: dict
+allocate_from_mr: _cython_3_2_5.cython_function_or_method
+free_reserved_memory: _cython_3_2_5.cython_function_or_method
+get_device_memory_resource: _cython_3_2_5.cython_function_or_method
+get_legacy_pinned_memory_resource: _cython_3_2_5.cython_function_or_method
+get_pinned_async_memory_resource: _cython_3_2_5.cython_function_or_method
+
+class _MemoryPointer:
+    """
+    Temporary internal NDBuffer allocation adapter class. NDBuffer expects
+    custom allocator to return a cuda.core.Buffer instance. Until all supported
+    cuda.core versions have unified support for wrapping external allocations
+    with Buffer.from_handle, this class servers as an adapter/workaround.
+
+    WARNING: This is internal tool subject to change/removal without notice.
+
+    Internally, it is used conditionally if any of the following is needed:
+        * wrap external allocations that don't come as cuda.core.Buffer (e.g. from cupy).
+        This is needed because prior to cuda.core 0.5.0, it's not possible to
+        pass reference to externall RAII object as ``owner`` parameter to Buffer.from_handle.
+        * provide debug logging on deallocation. In the future, we can use
+        Buffer.from_handle(owner=...) to inject deallocation callback (cuda.core >= 0.5.0)
+        or weakref.finalize for Buffer (cuda.core >= 0.6.0).
+
+    The only publicly exposed field is the handle - a base pointer to the allocated memory.
+    """
+    from_handle: ClassVar[method] = ...
+    __pyx_vtable__: ClassVar[PyCapsule] = ...
+    handle: handle
+    owner: owner
+    @classmethod
+    def __init__(cls, *args, **kwargs) -> None:
+        """Create and return a new object.  See help(type) for accurate signature."""
+    def close(self, stream) -> Any:
+        """_MemoryPointer.close(self, stream)"""
+    def __del__(self, *args, **kwargs) -> None: ...
+    def __reduce__(self):
+        """_MemoryPointer.__reduce_cython__(self)"""
+    def __reduce_cython__(self) -> Any:
+        """_MemoryPointer.__reduce_cython__(self)"""
+    def __setstate_cython__(self, __pyx_state) -> Any:
+        """_MemoryPointer.__setstate_cython__(self, __pyx_state)"""

nvmath/internal/ndbuffer/__init__.pxd

@@ -0,0 +1,3 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0

nvmath/internal/ndbuffer/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from nvmath.internal._layout import StridedLayout
+
+from ._ndbuffer import NDBuffer
+
+__all__ = ["NDBuffer", "StridedLayout"]

nvmath/internal/ndbuffer/_copy_kernel.pxd

@@ -0,0 +1,10 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+cimport cython
+from libc.stdint cimport intptr_t
+from .._layout._layout cimport StridedLayout
+
+
+cdef int launch_copy_kernel(StridedLayout dst, StridedLayout src, intptr_t dst_ptr, intptr_t src_ptr, int device_id, intptr_t stream_ptr, object logger) except -1

nvmath/internal/ndbuffer/_jit.pxd

@@ -0,0 +1,7 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from libc.stdint cimport intptr_t
+
+cdef intptr_t get_kernel(str kernel_code, str kernel_name, int device_id, include_path, object logger=*) except? 0

nvmath/internal/ndbuffer/_ndbuffer.pxd

@@ -0,0 +1,40 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+cimport cython
+from libc.stdint cimport intptr_t
+from .._layout._layout cimport StridedLayout
+
+# cdef class Exporting
+
+@cython.final
+cdef class NDBuffer:
+    cdef readonly StridedLayout layout
+    cdef readonly object data
+    cdef int data_device_id
+    cdef readonly intptr_t data_ptr
+    cdef readonly str dtype_name
+
+    cdef readonly NDBuffer base
+
+    # We use numpy_view internally to dispatch host memory
+    # operations to numpy (e.g. h2h copy).
+    # If the dtype_name is not supported by numpy, we still
+    # create a numpy view for internal use with any dtype
+    # of the corresponding itemsize, but such a view
+    # should not be returned to the user
+    cdef object numpy_view
+    cdef bint has_supported_dtype
+
+    # lazy evaluated properties
+    cdef str prop_device   # "cpu" or "cuda"
+    cdef object prop_device_id   # "cpu" or int
+
+    cdef inline int init_from_data(NDBuffer self, StridedLayout layout, object data, intptr_t data_ptr, str dtype_name, int data_device_id) except -1:
+        self.layout = layout
+        self.data = data
+        self.data_ptr = data_ptr
+        self.data_device_id = data_device_id
+        self.dtype_name = dtype_name
+        return 0