nvmath-python 0.9.0__cp314-cp314-win_amd64.whl

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: The buffer pointer.
+    """
+
+    __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 current GPU compute capability.
+    """
+
+    __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 ISAVersion"""
+        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()

nvmath/device/common_numba.py

@@ -0,0 +1,310 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from collections.abc import Callable
+from typing import Any
+
+import llvmlite.ir as llvmir
+import numpy as np
+from llvmlite import ir
+from numba import types
+from numba.core import cgutils, typing
+from numba.core.base import BaseContext
+from numba.core.errors import TypingError
+from numba.cuda.cudaimpl import lower_constant
+from numba.cuda.extending import intrinsic, overload, overload_attribute, register_model
+from numba.extending import models, typeof_impl
+
+from nvmath.bindings import mathdx
+
+from .llvm_array import LLVMArray
+from .types import Complex, Vector, complex32, complex64, complex128, half2, half4, np_float16x2, np_float16x4
+from .vector_types_numba import float16x2_type, float16x4_type, float32x2_type, float64x2_type
+
+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,
+    complex32: float16x2_type,
+    complex64: float32x2_type,
+    complex128: float64x2_type,
+    half2: float16x2_type,
+    half4: float16x4_type,
+    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,
+}
+
+
+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")
+    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
+
+
+@typeof_impl.register(Complex)
+def typeof_complex(val: Complex, c: typing.Context) -> Any:
+    if val.real_dtype == np.float16:
+        return types.NumberClass(float16x2_type)
+    elif val.real_dtype == np.float32:
+        return types.NumberClass(float32x2_type)
+    elif val.real_dtype == np.float64:
+        return types.NumberClass(float64x2_type)
+
+    raise RuntimeError(f"Unsupported complex real dtype {val.real_dtype}")
+
+
+@typeof_impl.register(Vector)
+def typeof_vector(val: Vector, c: typing.Context) -> Any:
+    if val.real_dtype != np.float16 or val.size not in (2, 4):
+        raise RuntimeError(f"Unsupported vector type {val.real_dtype}x{val.size}")
+
+    return types.NumberClass(float16x2_type if val.size == 2 else float16x4_type)
+
+
+@intrinsic
+def get_array_ptr(typingctx: typing.Context, arr):
+    """Get raw pointer to the data of a Numba array."""
+    assert isinstance(arr, types.Array)
+
+    sig = typing.signature(types.CPointer(arr.dtype), arr)
+
+    def codegen(context: BaseContext, builder, sig, args):
+        arrTy = sig.args[0]
+        arr = args[0]
+
+        dtype = arrTy.dtype
+        valueTy = context.get_value_type(dtype)
+        ptrTy = ir.PointerType(valueTy)
+        if arr is None:
+            ptr = ptrTy(None)
+        else:
+            ptr = cgutils.create_struct_proxy(arrTy)(context, builder, arr).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 ptr
+
+    return sig, codegen
+
+
+def build_get_value_ptr(target_type=None):
+    @intrinsic
+    def get_value_ptr_impl(typingctx: typing.Context, value):
+        """Get raw pointer to the value."""
+        if value not in [float16x2_type, float16x4_type, float32x2_type, float64x2_type] and not isinstance(  # noqa: UP038
+            value, (types.Float, types.Complex, types.Integer)
+        ):
+            raise TypingError(f"get_value_ptr does not support type {value}")
+
+        return_type = target_type if target_type is not None else value
+        sig = typing.signature(types.CPointer(return_type), value)
+
+        def codegen(context: BaseContext, builder, sig, args):
+            if target_type is not None and sig.args[0] != target_type:
+                value = context.cast(builder, args[0], sig.args[0], target_type)
+            else:
+                value = args[0]
+            return cgutils.alloca_once_value(builder, value)
+
+        return sig, codegen
+
+    return get_value_ptr_impl
+
+
+get_value_ptr = build_get_value_ptr()
+get_uint32_value_ptr = build_get_value_ptr(target_type=types.uint32)
+
+
+class MathdxOpaqueTensorType(types.Type):
+    def __init__(self, dynamic_strides_size=0, dynamic_shape_size=0):
+        self.dynamic_strides_size = dynamic_strides_size
+        self.dynamic_shape_size = dynamic_shape_size
+        super().__init__(f"MathdxOpaqueTensor({dynamic_strides_size}, {dynamic_shape_size})")
+
+
+@register_model(MathdxOpaqueTensorType)
+class MathdxOpaqueTensorModel(models.StructModel):
+    def __init__(self, dmm, fe_type: MathdxOpaqueTensorType):
+        members = [("ptr", types.voidptr)]
+
+        if fe_type.dynamic_strides_size > 0:
+            array_type = LLVMArray(types.int64, fe_type.dynamic_strides_size)
+            members += [("strides", array_type)]
+
+        if fe_type.dynamic_shape_size > 0:
+            array_type = LLVMArray(types.int64, fe_type.dynamic_shape_size)
+            members += [("shape", array_type)]
+
+        models.StructModel.__init__(self, dmm, fe_type, members)
+
+    def get_value_type(self):
+        if self._value_type is None:
+            name = f"struct.libmathdx_tensor_{self._fe_type.dynamic_strides_size}s_{self._fe_type.dynamic_shape_size}s"
+            t = llvmir.global_context.get_identified_type(name)
+            if t.is_opaque:
+                t.set_body(*[m.get_value_type() for m in self._models])
+            self._value_type = t
+        return self._value_type
+
+
+class MathdxPipelineType(types.Type):
+    def __init__(self):
+        super().__init__("MathdxPipeline()")
+
+
+@register_model(MathdxPipelineType)
+class MathdxPipelineModel(models.StructModel):
+    def __init__(self, dmm, fe_type: MathdxPipelineType):
+        members = [("ptr", types.voidptr)]
+        models.StructModel.__init__(self, dmm, fe_type, members)
+
+    def get_value_type(self):
+        if self._value_type is None:
+            name = "struct.libmathdx_pipeline"
+            t = llvmir.global_context.get_identified_type(name)
+            if t.is_opaque:
+                t.set_body(*[m.get_value_type() for m in self._models])
+            self._value_type = t
+        return self._value_type
+
+
+@intrinsic
+def cast_to_void_pointer(typingctx: typing.Context, ptr_ty):
+    if not isinstance(ptr_ty, types.CPointer):
+        raise TypingError("cast_to_void_pointer expects a CPointer type")
+    i8_ptr_ty = types.CPointer(types.uint8)
+    sig = i8_ptr_ty(ptr_ty)
+
+    def codegen(context: BaseContext, builder, sig, args):
+        return builder.bitcast(args[0], context.get_value_type(i8_ptr_ty))
+
+    return sig, codegen
+
+
+_cabi_device_registry: dict[str | tuple[str, typing.Signature], Callable] = {}
+
+
+def declare_cabi_device(symbol: str, sig: typing.Signature, link=None):
+    """declare_cabi_device is an analog of cuda.declare_device but uses C ABI
+    calling convention instead of Numba ABI calling convention. It means that
+    the first argument is not a return value pointer."""
+    key: str | tuple[str, typing.Signature] = symbol
+    if mathdx.get_version_ex() < (0, 3, 0):
+        key = (symbol, sig)
+    device_func = _cabi_device_registry.get(key)
+    if device_func is None:
+        device_func = _declare_cabi_device(symbol, sig, link)
+        _cabi_device_registry[key] = device_func
+
+    return device_func
+
+
+def _declare_cabi_device(symbol: str, sig: typing.Signature, link=None):
+    intrinsic_sig = sig.return_type(types.Tuple(sig.args))
+
+    @intrinsic
+    def call_device(typingctx: typing.Context, args):
+        def codegen(context: BaseContext, builder, sig: typing.Signature, args):
+            if link is not None:
+                context.active_code_library.add_linking_file(link)
+
+            args = cgutils.unpack_tuple(builder, args[0])
+            argTypes = sig.args[0]
+
+            assert len(args) == len(argTypes)
+
+            # CUDA target maps types.void to i8*, but we need void for abi correctness.
+            if sig.return_type == types.void:
+                retTy = ir.VoidType()
+            else:
+                retTy = context.get_value_type(sig.return_type)
+            fnTy = ir.FunctionType(retTy, [context.get_value_type(argTy) for argTy in argTypes])
+            fn = cgutils.get_or_insert_function(builder.module, fnTy, symbol)
+            builder.call(fn, args)
+
+        return intrinsic_sig, codegen
+
+    def device_func():
+        pass
+
+    @overload(device_func, jit_options={"forceinline": True})
+    def ol_device_func(*args):
+        if len(args) != len(sig.args):
+            raise RuntimeError(
+                f"Invalid number of arguments for device function {symbol}: expected {len(sig.args)}, got {len(args)}"
+            )
+        for expected_t, provided_t in zip(sig.args, args, strict=True):
+            if expected_t == provided_t:
+                continue
+            raise RuntimeError(f"Invalid argument type for device function {symbol}: expected {expected_t}, got {provided_t}")
+        return lambda *args: call_device(tuple(args))
+
+    return device_func
+
+
+def register_dummy_numba_type(numba_type, base_type, name, attributes):
+    """Register and lower a dummy Numba type that does not take space in memory.
+    Intended for types that are present only at typing stage and not represented in memory,
+    but are allowed to be accessed inside kernel."""
+    for attribute in attributes:
+        overload_type_attribute(numba_type, name, attribute)
+
+    @typeof_impl.register(base_type)
+    def typeof_numba(val, context):
+        return numba_type(val)
+
+    register_model(numba_type)(EmptyStructModel)
+
+    @lower_constant(numba_type)
+    def constant_dummy(context, builder, typ, pyval):
+        struct_ptr = cgutils.create_struct_proxy(typ)(context, builder)
+        return struct_ptr._getvalue()

nvmath/device/common_opaque_tensor.py

@@ -0,0 +1,201 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from functools import cached_property
+
+import numpy
+from llvmlite import ir as llir
+from numba.core import cgutils
+from numba.core.base import BaseContext
+from numba.cuda.extending import lower_builtin, make_attribute_wrapper, type_callable
+from numba.cuda.models import register_model
+from numba.extending import models, types
+
+from nvmath.device.common import Layout, OpaqueTensor
+from nvmath.device.common_numba import NUMBA_FE_TYPES_TO_NUMBA_IR, overload_type_attribute
+from nvmath.device.llvm_array import LLVMArray
+
+_LIBMATHDX_RUNTIME = -9223372036854775808
+
+
+class OpaqueLayout(Layout):
+    """
+    A generic tensor layout that can represent arbitrary shapes and strides.
+    It is the same as libmathdx opaque tensor, but without the pointer
+    to data.
+    """
+
+    _dtype: numpy.number
+    _shape: tuple[int, ...]
+    _strides: tuple[int, ...]
+
+    def __init__(self, shape: tuple[int, ...], strides: tuple[int, ...], dtype: numpy.number):
+        assert len(shape) == len(strides)
+        self._shape = shape
+        self._strides = strides
+        self._dtype = dtype
+
+    @cached_property
+    def dynamic_strides_size(self) -> int:
+        return self._strides.count(_LIBMATHDX_RUNTIME)
+
+    @cached_property
+    def dynamic_shape_size(self) -> int:
+        return self._shape.count(_LIBMATHDX_RUNTIME)
+
+    @property
+    def dtype(self) -> numpy.number:
+        return self._dtype
+
+    @cached_property
+    def ndim(self) -> int:
+        return len(self._shape)
+
+    @property
+    def shape(self) -> tuple[int, ...]:
+        return self._shape
+
+    @property
+    def strides(self) -> tuple[int, ...]:
+        return self._strides
+
+
+class LayoutType(types.Type):
+    """
+    Type class associated with opaque tensor layouts.
+    """
+
+    def __init__(
+        self,
+        layout: OpaqueLayout,
+    ):
+        self._layout = layout
+        super().__init__(
+            f"Layout(dynamic_shape_size={layout.dynamic_shape_size}, dynamic_strides_size={layout.dynamic_strides_size})"
+        )
+
+    @property
+    def layout(self) -> OpaqueLayout:
+        return self._layout
+
+    @cached_property
+    def dtype(self) -> types.Number:
+        return NUMBA_FE_TYPES_TO_NUMBA_IR[self._layout.dtype]
+
+
+@register_model(LayoutType)
+class LayoutModel(models.StructModel):
+    def __init__(self, dmm, fe_type: LayoutType):
+        members = []
+
+        strides_size = fe_type.layout.dynamic_strides_size
+        if strides_size > 0:
+            array_type = LLVMArray(types.int64, strides_size)
+            members += [("strides", array_type)]
+
+        shape_size = fe_type.layout.dynamic_shape_size
+        if shape_size > 0:
+            array_type = LLVMArray(types.int64, shape_size)
+            members += [("shape", array_type)]
+
+        models.StructModel.__init__(self, dmm, fe_type, members)
+
+
+make_attribute_wrapper(LayoutType, "strides", "strides")
+make_attribute_wrapper(LayoutType, "shape", "shape")
+
+for attribute in ("uid", "alignment", "storage_bytes"):
+    overload_type_attribute(LayoutType, "", attribute)
+
+
+class OpaqueTensorType(types.Type):
+    def __init__(self, layout: LayoutType):
+        super().__init__(f"OpaqueTensor(layout={layout})")
+        self._layout = layout
+
+    @property
+    def uid(self):
+        return self._layout.layout.uid
+
+    @property
+    def layout(self) -> LayoutType:
+        return self._layout
+
+    @cached_property
+    def ndim(self):
+        return 2
+
+    @cached_property
+    def dtype(self):
+        return self._layout.dtype
+
+
+class OpaqueTensorModel(models.StructModel):
+    def __init__(self, dmm, fe_type: OpaqueTensorType):
+        members = [
+            ("ptr", types.voidptr),
+        ]
+        layout = fe_type.layout.layout
+
+        strides_size = layout.dynamic_strides_size
+        if strides_size > 0:
+            array_type = LLVMArray(types.int64, strides_size)
+            members += [("strides", array_type)]
+
+        shape_size = layout.dynamic_shape_size
+        if shape_size > 0:
+            array_type = LLVMArray(types.int64, shape_size)
+            members += [("shape", array_type)]
+
+        models.StructModel.__init__(self, dmm, fe_type, members)
+
+
+register_model(OpaqueTensorType)(OpaqueTensorModel)
+
+
+@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
+        assert buffer_ty.dtype == layout_ty.dtype
+        return OpaqueTensorType(layout_ty)
+
+    return typer
+
+
+@lower_builtin(OpaqueTensor, types.Array, LayoutType)
+def impl_opaque_tensor(context: BaseContext, builder: llir.IRBuilder, sig, args):
+    tensor_ty: OpaqueTensorType = sig.return_type
+    buffer, buffer_ty = args[0], sig.args[0]
+    layout, layout_ty = args[1], sig.args[1]
+
+    ptr = cgutils.create_struct_proxy(buffer_ty)(context, builder, 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, llir.PointerType(ptr.type.pointee), "generic")
+
+    voidPtrTy = llir.PointerType(llir.IntType(8))
+    ptr = builder.bitcast(ptr, voidPtrTy)
+
+    opaque_tensor = cgutils.create_struct_proxy(tensor_ty)(context, builder)
+    opaque_tensor.ptr = ptr
+
+    layout_struct = cgutils.create_struct_proxy(layout_ty)(context, builder, layout)
+
+    if tensor_ty.layout.layout.dynamic_shape_size > 0:
+        opaque_tensor.shape = layout_struct.shape
+
+    if tensor_ty.layout.layout.dynamic_strides_size > 0:
+        opaque_tensor.strides = layout_struct.strides
+
+    return opaque_tensor._getvalue()
+
+
+make_attribute_wrapper(OpaqueTensorType, "ptr", "ptr")
+overload_type_attribute(OpaqueTensorType, "layout", "layout")