nvmath-python 0.5.0__cp313-cp313-win_amd64.whl

nvmath/device/common_cuda.py

@@ -0,0 +1,156 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = ["current_device_lto", "ComputeCapability", "CodeType", "ISAVersion", "Code", "Dim3", "MAX_SUPPORTED_CC"]
+
+from typing import NamedTuple
+from cuda.bindings import runtime as cudart, driver as cudadrv
+import logging
+from collections import namedtuple
+
+
+# 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: The buffer pointer.
+    """
+
+    pass
+
+
+# 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 current GPU compute capability.
+    """
+
+    pass
+
+
+# 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.
+    """
+
+    major: int
+    minor: int
+
+    @property
+    def integer(self) -> int:
+        """Integer representation of the ISAVersion"""
+        return self.major * 100 + self.minor * 10
+
+    def __str__(self):
+        """String representation of the ComputeCapability"""
+        return f"{self.major}.{self.minor}"
+
+    pass
+
+
+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).
+    """
+
+    pass
+
+
+# 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_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}")
+    return ComputeCapability(major, minor)
+
+
+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.")
+
+
+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()

nvmath/device/common_mathdx.py

@@ -0,0 +1,97 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from importlib.metadata import files, PackageNotFoundError
+import os
+import platform
+import re
+import sys
+import warnings
+
+
+CUDA_HOME = None
+CURAND_HOME = None
+
+
+PLATFORM_LINUX = sys.platform.startswith("linux")
+PLATFORM_WIN = sys.platform.startswith("win32")
+
+
+def conda_get_target_name():
+    if PLATFORM_LINUX:
+        plat = platform.processor()
+        if plat == "aarch64":
+            return "sbsa-linux"
+        else:
+            return f"{plat}-linux"
+    elif PLATFORM_WIN:
+        return "x64"
+    else:
+        raise AssertionError()
+
+
+def check_cuda_home():
+    # We need some CUDA headers for compiling mathDx headers.
+    # We assume users properly managing their local envs (ex: no mix-n-match).
+    global CUDA_HOME
+    global CURAND_HOME
+
+    # Try wheel
+    try:
+        # We need CUDA 12+ for device API support
+        cudart = files("nvidia-cuda-runtime-cu12")
+        cccl = files("nvidia-cuda-cccl-cu12")
+        curand = files("nvidia-curand-cu12")
+        # use cuda_fp16.h (which we need) as a proxy
+        cudart = [f for f in cudart if "cuda_fp16.h" in str(f)][0]
+        cudart = os.path.join(os.path.dirname(cudart.locate()), "..")
+        # use cuda/std/type_traits as a proxy
+        cccl = min([f for f in cccl if re.match(r".*cuda\/std\/type_traits.*", str(f))], key=lambda x: len(str(x)))
+        cccl = os.path.join(os.path.dirname(cccl.locate()), "../../..")
+        curand = [f for f in curand if "curand_kernel.h" in str(f)][0]
+        curand = os.path.dirname(curand.locate())
+    except PackageNotFoundError:
+        pass
+    except ValueError:
+        # cccl wheel is buggy (headers missing), skip using wheels
+        pass
+    else:
+        CUDA_HOME = (cudart, cccl)
+        CURAND_HOME = curand
+        return
+
+    # Try conda
+    if "CONDA_PREFIX" in os.environ:
+        if PLATFORM_LINUX:
+            conda_include = os.path.join(os.environ["CONDA_PREFIX"], "targets", f"{conda_get_target_name()}", "include")
+        elif PLATFORM_WIN:
+            conda_include = os.path.join(os.environ["CONDA_PREFIX"], "Library", "include")
+        else:
+            raise AssertionError()
+        if os.path.isfile(os.path.join(conda_include, "cuda_fp16.h")) and os.path.isfile(
+            os.path.join(conda_include, "cuda/std/type_traits")
+        ):
+            CUDA_HOME = (os.path.join(conda_include, ".."),)
+            CURAND_HOME = os.path.join(CUDA_HOME[0], "include")
+            return
+
+    # Try local
+    CUDA_PATH = os.environ.get("CUDA_PATH", None)
+    CUDA_HOME = os.environ.get("CUDA_HOME", None)
+    if CUDA_PATH is None and CUDA_HOME is None:
+        raise RuntimeError(
+            "cudart headers not found. Depending on how you install nvmath-python and other CUDA packages,\n"
+            "you may need to perform one of the steps below:\n"
+            "  - conda install -c conda-forge cuda-cudart-dev cuda-cccl cuda-version=12\n"
+            "  - export CUDA_HOME=/path/to/CUDA/Toolkit"
+        )
+    elif CUDA_PATH is not None and CUDA_HOME is None:
+        CUDA_HOME = CUDA_PATH
+    elif CUDA_PATH is not None and CUDA_HOME is not None and CUDA_HOME != CUDA_PATH:
+        warnings.warn("Both CUDA_HOME and CUDA_PATH are set but not consistent. Ignoring CUDA_PATH...")
+    CUDA_HOME = (CUDA_HOME,)
+    CURAND_HOME = os.path.join(CUDA_HOME[0], "include")
+
+
+check_cuda_home()

nvmath/device/common_numba.py

@@ -0,0 +1,159 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from llvmlite import ir
+from numba import types
+from numba.core import cgutils
+from numba.extending import models, overload_attribute
+import numpy as np
+
+from .vector_types_numba import float16x2_type, float16x4_type, float32x2_type, float64x2_type
+from .types import np_float16x2, np_float16x4
+from .common_opaque_tensor import OpaqueTensorType
+
+NP_TYPES_TO_NUMBA_FE_TYPES = {
+    np.float16: np.float16,
+    np.float32: np.float32,
+    np.float64: np.float64,
+    np_float16x2: float16x2_type,
+    np_float16x4: float16x4_type,
+    np.complex64: float32x2_type,
+    np.complex128: float64x2_type,
+    np.int8: np.int8,
+    np.int16: np.int16,
+    np.int32: np.int32,
+    np.int64: np.int64,
+    np.uint8: np.uint8,
+    np.uint16: np.uint16,
+    np.uint32: np.uint32,
+    np.uint64: np.uint64,
+}
+
+NUMBA_FE_TYPES_TO_NUMBA_IR = {
+    np.float16: types.float16,
+    np.float32: types.float32,
+    np.float64: types.float64,
+    float16x2_type: float16x2_type,
+    float16x4_type: float16x4_type,
+    float32x2_type: float32x2_type,
+    float64x2_type: float64x2_type,
+    np.int8: types.int8,
+    np.int16: types.int16,
+    np.int32: types.int32,
+    np.int64: types.int64,
+    np.uint8: types.uint8,
+    np.uint16: types.uint16,
+    np.uint32: types.uint32,
+    np.uint64: types.uint64,
+}
+
+
+def make_dx_codegen_one_arg(context, builder: ir.IRBuilder, type_, arg):
+    # mathdx expects everything to be passed by pointers.
+    # Arrays are lowered by passing a pointer to the data
+    # arg can be None, in which case a null pointer is passed
+    if isinstance(type_, OpaqueTensorType):
+        ptrTy = ir.PointerType(ir.IntType(8))
+        ldTy = ir.IntType(64)
+
+        opaque_tensor = cgutils.create_struct_proxy(type_)(context, builder, arg)
+        ptr = cgutils.create_struct_proxy(type_.buffer_type)(context, builder, opaque_tensor.buffer).data
+
+        # Future release of numba-cuda may have support for address spaces.
+        # It is not supported to pass a non generic pointer to device function
+        # call.
+        if ptr.type.addrspace != 0:
+            ptr = builder.addrspacecast(ptr, ir.PointerType(ptr.type.pointee), "generic")
+
+        ptr = builder.bitcast(ptr, ptrTy)
+
+        layout = cgutils.create_struct_proxy(type_.layout)(context, builder, opaque_tensor.layout)
+
+        member_values = [ptr]
+        member_types = [ptrTy]
+        if type_.layout.dynamic_ld:
+            member_types += [ldTy]
+            ld = builder.bitcast(layout.leading_dimension, ldTy)
+            member_values += [ld]
+
+        structTy = ir.LiteralStructType(member_types)
+        val = cgutils.pack_struct(builder, member_values)
+
+        return (structTy, val)
+    elif isinstance(type_, types.Array):
+        dtype = type_.dtype
+        valueTy = context.get_value_type(dtype)
+        ptrTy = ir.PointerType(valueTy)
+        if arg is None:
+            ptr = ptrTy(None)
+        else:
+            ptr = cgutils.create_struct_proxy(type_)(context, builder, arg).data
+
+        # Future release of numba-cuda may have support for address spaces.
+        # It is not supported to pass a non generic pointer to device function
+        # call.
+        if ptr.type.addrspace != 0:
+            ptr = builder.addrspacecast(ptr, ir.PointerType(ptr.type.pointee), "generic")
+
+        return (ptrTy, ptr)
+    # Integers are passed as-pointers
+    # arg can be None, in which case a 0 is passed
+    elif isinstance(type_, types.Integer):
+        intTy = context.get_value_type(type_)
+        if arg is None:
+            val = intTy(0)
+        else:
+            val = arg
+        ptrTy = ir.PointerType(intTy)
+        ptr = cgutils.alloca_once_value(builder, val)
+        return (ptrTy, ptr)
+    # Floats and Complex are passed by reference (pointer) This is because some CUDA C++
+    # types, such as __half2 are non-trivial, and those must be passed by reference. For
+    # consistency we pass everything by reference.
+    elif type_ in [float16x2_type, float16x4_type, float32x2_type, float64x2_type] or isinstance(  # noqa: UP038
+        type_, (types.Float, types.Complex)
+    ):
+        assert arg is not None
+        valueTy = context.get_value_type(type_)
+        ptrTy = ir.PointerType(valueTy)
+        ptr = cgutils.alloca_once_value(builder, arg)
+        return (ptrTy, ptr)
+    else:
+        raise RuntimeError(f"Unsupported lowering for type {type_} for arg {arg}")
+
+
+def make_function_call(symbol):
+    def codegen(context, builder, sig, args):
+        assert len(sig.args) == len(args)
+        argsTyAndArgs = [make_dx_codegen_one_arg(context, builder, t, a) for (t, a) in zip(sig.args, args, strict=True)]
+        argsTy = [t for (t, _) in argsTyAndArgs]
+        args = [v for (_, v) in argsTyAndArgs]
+        retTy = context.get_value_type(sig.return_type)
+        fnTy = ir.FunctionType(retTy, argsTy)
+        fn = cgutils.get_or_insert_function(builder.module, fnTy, symbol)
+        builder.call(fn, args)
+
+    return codegen
+
+
+class EmptyStructModel(models.StructModel):
+    """Data model that does not take space. Intended to be used with types that
+    are presented only at typing stage and not represented in memory."""
+
+    def __init__(self, dmm, fe_type):
+        members = []
+        super().__init__(dmm, fe_type, members)
+
+
+def overload_type_attribute(numba_type, attribute_base, attribute):
+    """Make type attribute available inside jitted code."""
+    assert issubclass(numba_type, types.Type)
+
+    @overload_attribute(numba_type, attribute, inline="always", target="cuda")
+    def ol_blas_attribute(blas_numba):
+        tp = blas_numba
+        if attribute_base != "":
+            tp = getattr(tp, attribute_base)
+        val = getattr(tp, attribute)
+        return lambda blas_numba: val

nvmath/device/common_opaque_tensor.py

@@ -0,0 +1,128 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import abstractmethod
+from functools import cached_property
+
+from numba.extending import types, models, type_callable, lower_builtin, overload_attribute
+from numba.cuda.models import register_model
+
+from nvmath.device.common import OpaqueTensor
+from numba.cuda.extending import make_attribute_wrapper
+
+from numba.core import cgutils
+
+from llvmlite.ir import IRBuilder
+from numba.core.base import BaseContext
+
+
+class LayoutType(types.Type):
+    """
+    Type class associated with opaque tensor layouts.
+    """
+
+    @property
+    @abstractmethod
+    def uid(self) -> int:
+        pass
+
+    @property
+    @abstractmethod
+    def dtype(self) -> str:
+        pass
+
+    @property
+    @abstractmethod
+    def size(self) -> int:
+        pass
+
+    @property
+    @abstractmethod
+    def dynamic_ld(self) -> bool:
+        pass
+
+
+@register_model(LayoutType)
+class LayoutModel(models.StructModel):
+    def __init__(self, dmm, fe_type: LayoutType):
+        members = []
+        if fe_type.dynamic_ld:
+            members += [("leading_dimension", types.int64)]
+        models.StructModel.__init__(self, dmm, fe_type, members)
+
+
+make_attribute_wrapper(LayoutType, "leading_dimension", "leading_dimension")
+
+
+@overload_attribute(LayoutType, "size", inline="always", strict=False)
+def ol_layout_size(layout: LayoutType):
+    assert isinstance(layout, LayoutType)
+    size = layout.size
+    return lambda _: size
+
+
+class OpaqueTensorType(types.Type):
+    def __init__(self, buffer_type: types.Array, layout: LayoutType):
+        assert layout.dtype == buffer_type.dtype
+        super().__init__(f"OpaqueTensor(layout={layout})")
+        self._buffer_type = buffer_type
+        self._layout = layout
+
+    @property
+    def uid(self):
+        return self._layout.uid
+
+    @property
+    def buffer_type(self):
+        return self._buffer_type
+
+    @property
+    def layout(self) -> LayoutType:
+        return self._layout
+
+    @cached_property
+    def ndim(self):
+        return 2
+
+    @cached_property
+    def dtype(self):
+        return self.buffer_type.dtype
+
+
+@register_model(OpaqueTensorType)
+class OpaqueTensorModel(models.StructModel):
+    def __init__(self, dmm, fe_type: OpaqueTensorType):
+        members = [
+            ("buffer", fe_type.buffer_type),
+            ("layout", fe_type.layout),
+        ]
+        models.StructModel.__init__(self, dmm, fe_type, members)
+
+
+make_attribute_wrapper(OpaqueTensorType, "buffer", "buffer")
+make_attribute_wrapper(OpaqueTensorType, "layout", "layout")
+
+
+@type_callable(OpaqueTensor)
+def type_callable_OpaqueTensor(context):
+    def typer(buffer_ty, layout_ty):
+        if not isinstance(buffer_ty, types.Array):
+            return
+        if not isinstance(layout_ty, LayoutType):
+            return
+        return OpaqueTensorType(buffer_ty, layout_ty)
+
+    return typer
+
+
+@lower_builtin(OpaqueTensor, types.Array, LayoutType)
+def impl_interval(context: BaseContext, builder: IRBuilder, sig, args):
+    typ = sig.return_type
+    buffer = args[0]
+    layout = args[1]
+    opaque_tensor = cgutils.create_struct_proxy(typ)(context, builder)
+    opaque_tensor.buffer = buffer
+    opaque_tensor.layout = layout
+
+    return opaque_tensor._getvalue()