nvmath-python 1.0.0__cp314-cp314t-win_amd64.whl

nvmath/device/common.py

@@ -0,0 +1,315 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import importlib.util
+import os
+import tempfile
+from abc import abstractmethod
+from collections.abc import Sequence
+from typing import Any
+
+import numpy as np
+
+from .common_cuda import MAX_SUPPORTED_CC, MIN_SUPPORTED_CC, CodeType, ComputeCapability, get_current_device_cc
+
+_HAS_NUMBA = importlib.util.find_spec("numba.cuda") is not None
+_HAS_NUMBA_CUDA_MLIR = importlib.util.find_spec("numba_cuda_mlir") is not None
+
+__all__ = [
+    "make_tensor",
+    "make_fragment_like",
+    "OpaqueTensor",
+    "Layout",
+    "axpby",
+    "copy",
+    "copy_fragment",
+    "clear",
+    "copy_wait",
+]
+
+SHARED_DEVICE_DOCSTRINGS = {
+    "precision": """\
+The computation precision specified as a numpy float dtype, currently supports ``numpy.float16``, ``numpy.float32`` and
+``numpy.float64``.""".replace("\n", " "),
+    #
+    "code_type": "The target GPU code and compute-capability.",
+    "sm": "Target mathdx compute-capability.",
+    #
+    "execution": "A string specifying the execution method, can be ``'Block'`` or ``'Thread'``.",
+}
+
+
+# TODO: maybe pre-compile regular expression
+def make_binary_tempfile(content, suffix: str) -> tempfile._TemporaryFileWrapper:
+    """Write `content` to a temporary file with the given `suffix`.
+
+    A closed file object returned; it is the user's responsibility to delete the file when
+    finished.
+
+    .. seealso:: :py:func:`delete_binary_tempfiles`
+
+    """
+    with tempfile.NamedTemporaryFile(mode="w+b", suffix=suffix, delete=False) as tmp:
+        tmp.write(content)
+        tmp.flush()
+    return tmp
+
+
+def delete_binary_tempfiles(filenames: list[str]):
+    for name in filenames:
+        if os.path.isfile(name):
+            os.remove(name)
+
+
+def check_in(name, value, coll, format="{name} must be in {coll_str} ; got {name} = {value}"):
+    if value not in coll:
+        coll_str = ", ".join(f'"{t}"' for t in coll)
+        msg = format.format(name=name, value=value, coll_str=coll_str)
+        raise ValueError(msg)
+
+
+def check_not_in(name, value, coll, format="{name} must not be any of these values: {coll_str} ; got {name} = {value}"):
+    if value in coll:
+        coll_str = ", ".join(f'"{t}"' for t in coll)
+        msg = format.format(name=name, value=value, coll_str=coll_str)
+        raise ValueError(msg)
+
+
+def check_contains(coll, key):
+    check_in("", key, coll, "missing required key {value!r}; got keys {coll_str}")
+
+
+def parse_sm(sm: Any) -> ComputeCapability:
+    if sm is None:
+        sm = get_current_device_cc()
+    else:
+        if not isinstance(sm, ComputeCapability) and not isinstance(sm, int):
+            raise ValueError(f"sm should be a ComputeCapability or an int; got sm = {sm}")
+        if isinstance(sm, int):
+            sm = ComputeCapability(sm // 10, sm % 10)
+
+    return sm
+
+
+def check_sm(sm, library_name: str, var_name: str = "sm"):
+    if not isinstance(sm, ComputeCapability):
+        raise ValueError(f"{var_name} should be an instance of ComputeCapability ; got {var_name} = {sm}")
+    if sm < MIN_SUPPORTED_CC:
+        raise RuntimeError(f"Minimum compute capability {MIN_SUPPORTED_CC} is required by {library_name}, got {sm}")
+    if sm > MAX_SUPPORTED_CC:
+        raise RuntimeError(f"The maximum compute capability currently supported by device APIs is {MAX_SUPPORTED_CC}, got {sm}")
+    if sm.minor < 0:
+        raise ValueError(f"{var_name}.minor must be >= 0 ; got {var_name}.minor = {sm.minor}")
+
+
+def parse_code_type(code_type: Any) -> CodeType:
+    if not isinstance(code_type, Sequence) or len(code_type) != 2:
+        raise ValueError(f"code_type should be an instance of CodeType or a 2-tuple ; got code_type = {code_type}")
+
+    return CodeType(code_type[0], ComputeCapability(*code_type[1]))
+
+
+def check_code_type(code_type, library_name):
+    if not isinstance(code_type, CodeType):
+        raise ValueError(f"code_type should be an instance of CodeType ; got code_type = {code_type}")
+    check_sm(code_type.cc, library_name, "code_type.cc")
+    check_in("code_type.kind", code_type.kind, ["lto"])
+
+
+def pad_or_truncate(list, target_len):
+    return list[:target_len] + [0] * (target_len - len(list))
+
+
+class Layout:
+    """
+    Layout for the :py:class:`nvmath.device.OpaqueTensor`.
+
+    .. note:: Do not create directly, use appropriate method from
+        :py:class:`nvmath.device.Matmul`. Refer to
+        :cublasdx_doc:`api/other_tensors.html#imported-tensor-utilities`
+        for guidance on which method to use.
+    """
+
+    def __init__(self):
+        raise RuntimeError("Layout should not be called directly.")
+
+    @property
+    @abstractmethod
+    def size(self) -> int:
+        """
+        Number of valid elements in a tensor. This is simply a product of all
+        shape dimensions.
+
+        Refer to the cuBLASDx documentation for more details on how to use this attribute:
+        :cublasdx_doc:`api/other_tensors.html#imported-tensor-utilities`
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def cosize(self) -> int:
+        """
+        Returns a distance from last element of a tensor to its first element.
+        It describes how many elements the argument layout spans.
+
+        Refer to the cuBLASDx documentation for more details on how to use this attribute:
+        :cublasdx_doc:`api/other_tensors.html#imported-tensor-utilities`
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def alignment(self) -> int:
+        """
+        Returns the required alignment (in bytes) for the tensor data buffer.
+
+        Refer to the cuBLASDx documentation for more details on how to use this attribute:
+        :cublasdx_doc:`api/other_tensors.html#imported-tensor-utilities`
+        """
+        pass
+
+
+class OpaqueTensor:
+    """
+    Abstraction over the cuBLASDx tensor type (an alias of the CuTe tensor type).
+    The CuTe tensor layout is powerful and supports layouts not provided by NumPy,
+    so this is a bridge to add this functionality to Python.
+
+    .. note:: Do not create directly, use :py:func:`nvmath.device.make_tensor`.
+
+    Refer to the cuBLASDx documentation for more details on how to use this class:
+    :cublasdx_doc:`api/other_tensors.html#tensors`
+    """
+
+    buffer: np.ndarray
+    layout: Layout
+    leading_dimension: int | None
+
+    def __init__(self, *args):
+        raise RuntimeError("OpaqueTensor should not be called directly outside of a jitted kernel.")
+
+
+def make_tensor(array: np.ndarray, layout: Layout) -> OpaqueTensor:
+    """
+    make_tensor is a helper function for creating
+    :py:class:`nvmath.device.OpaqueTensor` objects.
+
+    Args:
+        array (np.ndarray): The input array to be wrapped as an OpaqueTensor.
+        layout (Layout): The layout of the tensor, which describes how the data is
+            organized in memory.
+
+    Refer to the cuBLASDx documentation for more details on how to use this function:
+    :cublasdx_doc:`api/other_tensors.html#create-tensor-other-label`
+    """
+    raise RuntimeError("make_tensor should not be called directly outside of a jitted kernel.")
+
+
+def make_fragment_like(tensor: OpaqueTensor, dtype) -> OpaqueTensor:
+    """
+    make_fragment_like is a helper function for creating register fragments with
+    the same layout as input tensor, but different dtype.
+
+    Args:
+        tensor (OpaqueTensor): The input tensor to be used as a template for the fragment.
+        dtype: The data type of the fragment to be created.
+
+    Refer to the cuBLASDx documentation for more details on how to use this function:
+    :cublasdx_doc:`api/other_tensors.html#imported-tensor-utilities`
+    """
+    raise RuntimeError("make_fragment_like should not be called directly outside of a jitted kernel.")
+
+
+def axpby(alpha: float, x_tensor: OpaqueTensor, beta: float, y_tensor: OpaqueTensor) -> None:
+    """
+    AXPBY operation: y = alpha * x + beta * y
+
+    Args:
+        alpha (float): Scalar multiplier for x_tensor.
+        x_tensor (OpaqueTensor): Input tensor x.
+        beta (float): Scalar multiplier for y_tensor.
+        y_tensor (OpaqueTensor): Input/output tensor y, which will be updated
+            with the result.
+
+    Refer to the cuBLASDx documentation for more details on how to use this function:
+    :cublasdx_doc:`api/other_tensors.html#imported-tensor-utilities`
+    """
+    raise RuntimeError("axpby should not be called directly outside of a jitted kernel.")
+
+
+def copy(src: OpaqueTensor, dst: OpaqueTensor, alignment=None):
+    """
+    Copies data from the source tensor to the destination tensor.
+
+    Args:
+        src (OpaqueTensor): The source tensor to copy from.
+        dst (OpaqueTensor): The destination tensor to copy to.
+        alignment (int, optional): The alignment (in bytes) for the copy operation.
+
+    Refer to the cuBLASDx documentation for more details on how to use this function:
+    :cublasdx_doc:`api/other_tensors.html#cooperative-global-shared-copying`
+    """
+    raise RuntimeError("copy should not be called directly outside of a jitted kernel.")
+
+
+def copy_fragment(src: OpaqueTensor, dst: OpaqueTensor, alignment=None):
+    """
+    A bidirectional copying method to copy data between register fragments and
+    global memory tensors.
+
+    Args:
+        src (OpaqueTensor): The source tensor to copy from.
+        dst (OpaqueTensor): The destination tensor to copy to.
+        alignment (int, optional): The alignment (in bytes) for the copy operation.
+
+    Refer to the cuBLASDx documentation for more details on how to use this function:
+    :cublasdx_doc:`api/other_tensors.html#copying-registers-tensors`
+    """
+    raise RuntimeError("copy_fragment should not be called directly outside of a jitted kernel.")
+
+
+def clear(arr: OpaqueTensor):
+    """
+    Clears the contents of the given tensor by setting all elements to zero.
+
+    Args:
+        arr (OpaqueTensor): The tensor to be cleared.
+
+    Refer to the cuBLASDx documentation for more details on how to use this function:
+    :cublasdx_doc:`api/other_tensors.html#imported-tensor-utilities`
+    """
+    raise RuntimeError("clear should not be called directly outside of a jitted kernel.")
+
+
+def copy_wait():
+    """
+    Creates a synchronization point. It has to be called after :py:func:`nvmath.device.copy`
+    to ensure that the copy operation has completed before any subsequent
+    operations are executed.
+
+    Refer to the cuBLASDx documentation for more details on how to use this function:
+    :cublasdx_doc:`api/other_tensors.html#cooperative-global-shared-copying`
+    """
+    raise RuntimeError("copy_wait should not be called directly outside of a jitted kernel.")
+
+
+def check_positive_integer_sequence(arg, arg_name, min_len, max_len):
+    if not isinstance(arg, Sequence):
+        raise ValueError(f'Parameter "{arg_name}" must be a sequence. Got: {type(arg).__name__}.')
+
+    for i, x in enumerate(arg):
+        if not isinstance(x, (int, np.integer)):
+            raise ValueError(
+                f'Parameter "{arg_name}" must be a sequence of positive integers. '
+                + f"Element {arg_name}[{i}] is {type(x).__name__}."
+            )
+        if x <= 0:
+            raise ValueError(f'Parameter "{arg_name}" values must be positive integers. Got {arg_name}[{i}] = {x}.')
+
+    if min_len == max_len:
+        if len(arg) != min_len:
+            raise ValueError(f'Parameter "{arg_name}" must be a sequence of length {min_len}. Got length: {len(arg)}.')
+    else:
+        if len(arg) > max_len or len(arg) < min_len:
+            raise ValueError(f'Parameter "{arg_name}" must define {min_len} to {max_len} values. Got {len(arg)} value(s).')

nvmath/device/common_backend.py

@@ -0,0 +1,131 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import weakref
+from collections.abc import Callable
+
+import numpy as np
+
+from nvmath.bindings import mathdx
+from nvmath.device.common_cuda import ComputeCapability, ISAVersion
+
+from .types import complex32, complex64, complex128, half2, half4
+
+MATHDX_TYPES_TO_NP = {
+    mathdx.CommondxValueType.R_16F: np.float16,
+    mathdx.CommondxValueType.R_16F2: half2,
+    mathdx.CommondxValueType.R_32F: np.float32,
+    mathdx.CommondxValueType.R_64F: np.float64,
+    mathdx.CommondxValueType.C_16F: complex32,
+    mathdx.CommondxValueType.C_16F2: half4,
+    mathdx.CommondxValueType.C_32F: complex64,
+    mathdx.CommondxValueType.C_64F: complex128,
+    mathdx.CommondxValueType.R_8I: np.int8,
+    mathdx.CommondxValueType.R_16I: np.int16,
+    mathdx.CommondxValueType.R_32I: np.int32,
+    mathdx.CommondxValueType.R_64I: np.int64,
+    mathdx.CommondxValueType.R_8UI: np.uint8,
+    mathdx.CommondxValueType.R_16UI: np.uint16,
+    mathdx.CommondxValueType.R_32UI: np.uint32,
+    mathdx.CommondxValueType.R_64UI: np.uint64,
+}
+
+
+NP_TYPES_TO_MATHDX_TYPES = {np_type: mathdx_type for mathdx_type, np_type in MATHDX_TYPES_TO_NP.items()}
+
+NP_TYPES_TO_MATHDX_PRECISION = {
+    np.float16: mathdx.CommondxPrecision.F16,
+    np.float32: mathdx.CommondxPrecision.F32,
+    np.float64: mathdx.CommondxPrecision.F64,
+    np.int8: mathdx.CommondxPrecision.I8,
+    np.int16: mathdx.CommondxPrecision.I16,
+    np.int32: mathdx.CommondxPrecision.I32,
+    np.int64: mathdx.CommondxPrecision.I64,
+    np.uint8: mathdx.CommondxPrecision.UI8,
+    np.uint16: mathdx.CommondxPrecision.UI16,
+    np.uint32: mathdx.CommondxPrecision.UI32,
+    np.uint64: mathdx.CommondxPrecision.UI64,
+}
+
+EXECUTION_STR_TO_MATHDX = {
+    "Block": mathdx.CommondxExecution.BLOCK,
+    "Thread": mathdx.CommondxExecution.THREAD,
+}
+
+NVARG_GEN_OPT_LTO = "-gen-opt-lto"
+
+
+class DescriptorWrapper:
+    """
+    A smart pointer to a descriptor.
+
+    Lives as long as we need the descriptor (due to lru cache usage). Once the
+    descriptor is no longer needed/cached, it will trigger proper resource
+    deallocation.
+    """
+
+    def __init__(self, descriptor, destructor):
+        self.descriptor = descriptor
+        self._destructor = destructor
+        if destructor is not None:
+            weakref.finalize(self, self._destructor, self.descriptor)
+
+
+def get_lto(code_descriptor: int) -> bytes:
+    """Get lto binary from the mathdx common code descriptor."""
+    lto_size = mathdx.commondx_get_code_ltoir_size(code_descriptor)
+
+    lto_fn = bytearray(lto_size)
+    mathdx.commondx_get_code_ltoir(code_descriptor, lto_size, lto_fn)
+    return bytes(lto_fn)
+
+
+def get_isa_version(code_descriptor: int) -> ISAVersion:
+    """Parse isa version from the mathdx common code descriptor."""
+    isa = mathdx.commondx_get_code_option_int64(code_descriptor, mathdx.CommondxOption.CODE_ISA)
+    return ISAVersion.from_integer(isa)
+
+
+def build_get_int_traits(get_trait_int64s: Callable[[int, int, int, int], int]):
+    """Generate function that returns mathdx tuple of int traits."""
+
+    def get_int_traits(handle: int, trait_type: int, size: int) -> tuple:
+        int_buffer = np.zeros(size, np.int64)
+        get_trait_int64s(handle, trait_type, size, int_buffer.ctypes.data)
+        return tuple(map(int, int_buffer))
+
+    return get_int_traits
+
+
+def build_get_str_trait(
+    get_trait_str_size: Callable[[int, int], int],
+    get_trait_str: Callable[[int, int, int, bytearray], bytearray],
+):
+    """Generate function that returns mathdx string trait."""
+
+    def get_str_trait(handle: int, trait_type: int) -> str:
+        symbol_size = get_trait_str_size(handle, trait_type)
+
+        symbol = bytearray(symbol_size)
+
+        get_trait_str(handle, trait_type, symbol_size, symbol)
+
+        # terminate trailing 0 (indicator of c-string end)
+        symbol_str = symbol[:-1].decode()
+
+        return symbol_str
+
+    return get_str_trait
+
+
+def get_mathdx_sm(version: ComputeCapability) -> list:
+    sm = [version.integer]
+    if version.arch == "a":
+        sm.append(mathdx.CommondxArchModifier.ARCH_SPECIFIC)
+    return sm
+
+
+def set_code_target_sm(code: int, version: ComputeCapability):
+    sm = get_mathdx_sm(version)
+    mathdx.commondx_set_code_option_int64s(code, mathdx.CommondxOption.TARGET_SM, len(sm), sm)

nvmath/device/common_cuda.py

@@ -0,0 +1,201 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = [
+    "current_device_lto",
+    "current_device_sm",
+    "ComputeCapability",
+    "CodeType",
+    "ISAVersion",
+    "Code",
+    "Dim3",
+    "MAX_SUPPORTED_CC",
+]
+
+import logging
+from collections import namedtuple
+from typing import NamedTuple
+
+from cuda.bindings import driver as cudadrv
+from cuda.bindings import runtime as cudart
+
+
+# Code = CodeType + ISAVersion + buffer
+class Code(namedtuple("Code", ("code_type", "isa_version", "data"))):
+    """
+    A namedtuple class that encapsulates code type, version, and buffer.
+
+    Attributes:
+        code_type (CodeType): The underlying code type.
+        isa_version (ISAVersion): The instruction set architecture version for the code.
+        data (bytes): The code buffer.
+    """
+
+    __slots__ = ()
+
+
+# CodeType = type + CC
+class CodeType(namedtuple("CodeType", ("kind", "cc"))):
+    """
+    A namedtuple class that encapsulates code kind and compute capability.
+
+    Attributes:
+        kind (str): A string denoting the nature of the code, e.g, ``'lto'``.
+        cc (ComputeCapability): The target compute capability of the code.
+    """
+
+    __slots__ = ()
+
+
+# CC = e.g. SM 9.0 (Hopper)
+class ComputeCapability(NamedTuple):
+    """
+    A namedtuple class that encapsulates the major and minor compute capability.
+
+    Attributes:
+        major (int): The major compute capability.
+        minor (int): The minor compute capability.
+        arch (str): An optional string representing the GPU accelerated
+            architecture. Must be either empty or one of 'f' (for Family),
+            'a' (for Architecture).
+    """
+
+    major: int
+    minor: int
+    arch: str = ""
+
+    @property
+    def integer(self) -> int:
+        """Integer representation of the ComputeCapability"""
+        return self.major * 100 + self.minor * 10
+
+    def __str__(self):
+        """String representation of the ComputeCapability"""
+        return f"{self.major}.{self.minor}{self.arch}"
+
+
+MIN_SUPPORTED_CC = ComputeCapability(7, 0)
+MAX_SUPPORTED_CC = ComputeCapability(12, 1)
+
+
+class Dim3(namedtuple("Dim3", ("x", "y", "z"), defaults=(1, 1, 1))):
+    """
+    A namedtuple class that encapsulates the dimensions for grids and blocks.
+
+    Attributes:
+        x (int): The dimension in the x direction (default 1).
+        y (int): The dimension in the y direction (default 1).
+        z (int): The dimension in the z direction (default 1).
+    """
+
+    __slots__ = ()
+
+
+# ISAVersion = e.g. 12.3 (from CUDA 12.3)
+class ISAVersion(NamedTuple):
+    """
+    A namedtuple class that encapsulates the code version.
+
+    Attributes:
+        major (int): The major code version.
+        minor (int): The minor code version.
+    """
+
+    major: int
+    minor: int
+
+    @classmethod
+    def from_integer(cls, isa: int):
+        major = isa // 1000
+        minor = (isa % 1000) // 10
+
+        return cls(major, minor)
+
+
+def CHECK_CUDART(err):
+    if err != cudart.cudaError_t.cudaSuccess:
+        err2, str = cudart.cudaGetErrorString(err)
+        raise RuntimeError(f"CUDArt Error: {str} ({err})")
+
+
+def CHECK_CUDA(err):
+    if err != cudadrv.CUresult.CUDA_SUCCESS:
+        err2, str = cudadrv.cuGetErrorName(err)
+        raise RuntimeError(f"CUDA Error: {str} ({err})")
+
+
+def get_current_device() -> int:
+    (err,) = cudadrv.cuInit(0)
+    CHECK_CUDA(err)
+    # Check if a context exist
+    err, pctx = cudadrv.cuCtxGetCurrent()
+    CHECK_CUDA(err)
+    if int(pctx) == 0:
+        # If not, return the CC of device 0
+        device = 0
+    else:
+        err, device = cudart.cudaGetDevice()
+        CHECK_CUDART(err)
+
+    return device
+
+
+def get_current_device_cc():
+    (err,) = cudadrv.cuInit(0)
+    CHECK_CUDA(err)
+    # Check if a context exist
+    err, pctx = cudadrv.cuCtxGetCurrent()
+    CHECK_CUDA(err)
+    if int(pctx) == 0:
+        # If not, return the CC of device 0
+        device = 0
+    else:
+        err, device = cudart.cudaGetDevice()
+        CHECK_CUDART(err)
+    err, prop = cudart.cudaGetDeviceProperties(device)
+    CHECK_CUDART(err)
+    major, minor = prop.major, prop.minor
+    if (major, minor) > MAX_SUPPORTED_CC:
+        logging.info(
+            "The current device supports compute capability "
+            f"{prop.major}.{prop.minor}, but the generated LTO version is "
+            f"capped at {MAX_SUPPORTED_CC}."
+        )
+        major, minor = MAX_SUPPORTED_CC
+    logging.info(f"Using device {device} for default compute capability, found cc = {prop.major}.{prop.minor}")
+    arch = ""
+    if (major, minor) >= (9, 0):
+        arch = "a"
+    return ComputeCapability(major, minor, arch)
+
+
+def get_default_code_type() -> CodeType:
+    try:
+        return CodeType("lto", get_current_device_cc())
+    except RuntimeError as e:
+        raise ValueError(f"Failed to get the current GPU compute capability (got {e}). Set code_type explicitly.") from e
+
+
+def current_device_lto():
+    """
+    A helper function to get the default code type for link time optimization (LTO) on the
+    current device.
+
+    Returns:
+        A :class:`CodeType` object representing the default LTO code type for the current
+        device.
+    """
+    return get_default_code_type()
+
+
+def current_device_sm() -> ComputeCapability:
+    """
+    A helper function to get the default SM for mathdx types on the
+    current device.
+
+    Returns:
+        A :class:`ComputeCapability` object representing the default SM for
+        mathdx types on the current device.
+    """
+    return get_current_device_cc()