nvmath-python 1.0.0__cp314-cp314t-win_amd64.whl

nvmath/sparse/ust/_jit.py

@@ -0,0 +1,188 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+This module defines a JIT runtime for the universal sparse tensor (UST).
+"""
+
+__all__ = []
+
+import re
+from collections.abc import Sequence
+
+from cuda.core import (
+    LaunchConfig,
+    Linker,
+    LinkerOptions,
+    ObjectCode,
+    Program,
+    ProgramOptions,
+    launch,
+)
+
+from nvmath.internal import utils
+
+
+def _check_numba_available():
+    try:
+        import numba  # noqa: F401
+    except ModuleNotFoundError as e:
+        raise RuntimeError("Numba is required to compile user-defined Python functions provided to the UST APIs.") from e
+
+
+def compile_python_function(
+    function,
+    function_name,
+    signature,
+    *,
+    representation="ltoir",
+    compute_capability=None,
+):
+    """
+    Compile the Python function to an intermediate representation using Numba.
+
+    Args:
+        function: The Python function to compile.
+        function_name: The symbol to use for the function name.
+        signature: the function signature as a sequence (return_type, *parameter_types). The
+            type names should correspond to the C type names recognized by Numba.
+        representation: the intermediate representation ("ltoir", "ptx").
+        compute_capability: the device compute capability. If not specified, the
+            current devices CC will be used.
+
+    Returns:
+        The compiled IR.
+    """
+
+    _check_numba_available()
+
+    import numba
+    import numba.cuda
+
+    user_return_type, *user_arg_types = signature
+
+    return_type = getattr(numba.cuda.types.types, user_return_type)
+
+    arg_types = []
+    for arg_type in user_arg_types:
+        m = re.match(r"(\w+)((\[:\])|( *(\*)))?", arg_type)
+        assert m is not None, "Internal error."
+        t = getattr(numba.cuda.types.types, m.group(1))
+        if m.group(2) is not None:
+            if m.group(3) == "[:]":
+                t = t[:]
+            elif m.group(5) == "*":
+                t = numba.cuda.types.types.voidptr
+            else:
+                raise AssertionError("Internal error.")
+        arg_types.append(t)
+
+    signature = numba.core.typing.signature(return_type, *arg_types)
+
+    return numba.cuda.compile(
+        function,
+        signature,
+        debug=None,
+        lineinfo=False,
+        device=True,
+        fastmath=False,
+        cc=compute_capability,
+        opt=True,
+        abi="c",
+        abi_info={"abi_name": function_name},
+        output=representation,
+        forceinline=False,
+        launch_bounds=None,
+    )[0]
+
+
+def compile_cpp_and_link(
+    *,
+    src_code=None,
+    object_code=None,
+    function_name,
+    compiler_options=None,
+    linker_options=None,
+    intermediate_type="ltoir",
+    target_type="cubin",
+):
+    """
+    Compile the source code (if provided) and link with the intermediate code (if provided).
+
+    Args:
+        src_code: A string or sequence of strings containing the C++ source code.
+        object_code: A byte-string or sequence of byte-string objects to be linked
+            along with the source code.
+        function_name: symbol name to use for the compiled and linked function.
+        compiler_options: the compiler options to use as a
+            :class:`cuda.core.ProgramOptions` object. Alternatively, a dictionary containing
+            the options can be provided.
+        linker_options: the linker options to use as a
+            :class:`cuda.core.LinkerOptions` object. Alternatively, a dictionary containing
+            the options can be provided.
+
+    Returns:
+        The linked kernel.
+    """
+
+    assert not (src_code is None and object_code is None), "Internal error."
+
+    compiler_options = utils.check_or_create_options(ProgramOptions, compiler_options, "Compiler Options")
+    linker_options = utils.check_or_create_options(LinkerOptions, linker_options, "Linker Options")
+
+    if isinstance(src_code, str):
+        src_code = [src_code]
+    else:
+        if src_code is not None and not isinstance(src_code, Sequence):
+            raise TypeError("The source code type 'type(src_code)' must be a string or a sequence of strings.")
+
+    # Handle the case where no source code is provided.
+    if src_code is None:
+        src_code = []
+
+    # Compile source code to LTO-IR.
+    compiled_code = [Program(s, "c++", options=compiler_options).compile(intermediate_type) for s in src_code]
+
+    if isinstance(object_code, bytes):
+        object_code = [object_code]
+    else:
+        if object_code is not None and not isinstance(object_code, Sequence):
+            raise TypeError(
+                "The intermediate code type 'type(object_code)' must be a byte-string or a sequence of byte-strings."
+            )
+
+    assert intermediate_type == "ltoir", "Internal error."
+
+    # Handle the case where no object code is provided.
+    if object_code is None:
+        object_code = []
+
+    object_code = [ObjectCode.from_ltoir(o) for o in object_code]
+
+    # Link all objects together.
+    linker = Linker(*(compiled_code + object_code), options=linker_options)
+    linked_code = linker.link(target_type=target_type)
+
+    # Extract kernel entry by function name.
+    return linked_code.get_kernel(function_name)
+
+
+def launch_kernel(kernel, grid, block, params, *, device_id, stream_holder, blocking=True):
+    """
+    Launch the kernel on the specified device and stream.
+
+    Args:
+        kernel: the CUDA kernel to launch.
+        grid, block: the CUDA grid and block dimensions for launch.
+        params: the kernel arguments as a sequence or None.
+        device_id: the device ID on which to launch the kernel.
+        stream_holder: a `nvmath.internal.utils.StreamHolder` object for the stream to use.
+        blocking: if True, the call blocks.
+    """
+    config = LaunchConfig(grid=grid, block=block)
+    with utils.device_ctx(device_id):
+        launch(stream_holder.obj, config, kernel, *params)
+    # TODO: make this non-blocking by default.
+    if blocking:
+        stream_holder.obj.sync()

nvmath/sparse/ust/_kernel.py

@@ -0,0 +1,282 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+This module defines a kernel generator for various sparse tensor operations.
+"""
+
+__all__ = []
+
+import time
+
+from cuda import pathfinder
+
+from ._cpp import prolog_a_impl, prolog_b_impl, prolog_decl
+from ._emitter import Index, emit_matmul_kernel
+from ._jit import compile_cpp_and_link
+from ._utils import Cache, type_str
+
+# Maps matmul(a,b,c) to pre-defined kernel source code.
+_tensor_matmul_kernels = {
+    # (format_a, format_b, format_c, transpose_a, transpose_b)
+    ("DIAI", "DenseVector", "DenseVector", False, False): """
+extern "C" __global__ void matmul_kernel(
+    const POS* __restrict__ pos0,
+    const CRD* __restrict__ diags,
+    const VAL* __restrict__ A,
+    const VAL *__restrict__ B,
+    VAL *__restrict__ C,
+    unsigned long m,
+    unsigned long n,
+    unsigned long l0,
+    unsigned long l1) {
+  const unsigned long tidx = threadIdx.x;
+  const unsigned long bidx = blockIdx.x;
+  const unsigned long i = bidx * blockDim.x + tidx;
+  if (i < m) {
+    CTP acc = 0;
+    const POS numDiags = pos0[1];
+    for (unsigned long d = 0, di = i; d < numDiags; d++, di += m) {
+      const long j = i + diags[d];
+      if (0 <= j && j < n) {
+        acc += prolog_a(static_cast<CTP>(A[di])) *
+               prolog_b(static_cast<CTP>(B[j]));
+      }
+    }
+    C[i] += static_cast<VAL>(acc);
+  }
+}
+""",
+    ("DIAJ", "DenseVector", "DenseVector", False, False): """
+extern "C" __global__ void matmul_kernel(
+    const POS* __restrict__ pos0,
+    const CRD* __restrict__ diags,
+    const VAL* __restrict__ A,
+    const VAL *__restrict__ B,
+    VAL *__restrict__ C,
+    unsigned long m,
+    unsigned long n,
+    unsigned long l0,
+    unsigned long l1) {
+  const unsigned long tidx = threadIdx.x;
+  const unsigned long bidx = blockIdx.x;
+  const unsigned long i = bidx * blockDim.x + tidx;
+  if (i < m) {
+    CTP acc = 0;
+    const POS numDiags = pos0[1];
+    for (unsigned long d = 0, di = 0; d < numDiags; d++, di += n) {
+      const long j = i + diags[d];
+      if (0 <= j && j < n) {
+        acc += static_cast<CTP>(prolog_a(A[di + j])) *
+               static_cast<CTP>(prolog_b(B[j]));
+      }
+    }
+    C[i] += static_cast<VAL>(acc);
+  }
+}
+""",
+}
+
+
+def _set_user_or_default_code(*, user_code, default_code, default_prolog, src_code_list, obj_code_list):
+    """
+    Args:
+        user_code: The user-provided LTO-IR as a bytes object or CUDA C++ source code
+            string.
+        default_code: The CUDA C++ source code that will be used if the user-provided
+            code is None.
+        default_prolog: The prolog to the CUDA C++ source code that will be used if the
+            user-provided code is None.
+        src_code_list: The list to which the source code will be appended to.
+        obj_code_list: The list to which the LTO-IR code will be appended to.
+
+    Returns:
+        The pair (src_code_list, obj_code_list)
+    """
+
+    if user_code is not None:
+        if isinstance(user_code, str):
+            target = src_code_list
+        elif isinstance(user_code, bytes):
+            target = obj_code_list
+        else:
+            raise AssertionError("Internal error.")
+        target.append(user_code)
+        return src_code_list, obj_code_list
+
+    assert default_code is not None and isinstance(default_code, str), "Internal error."
+
+    source_code = f"{default_prolog} {default_code}"
+    src_code_list.append(source_code)
+
+    return src_code_list, obj_code_list
+
+
+class KernelGen:
+    """
+    This class defines a kernel generator for various sparse tensor operations.
+    The generated kernels can be JITted on-demand. Caching avoids recompilation
+    on subsequent uses.
+
+    The generation uses:
+    (1) kernel lookup from a collection of pre-defined kernels
+    (2) kernel generated by a sparse emitter
+
+    Clients can easily add new pre-defined kernels in the table above (for example,
+    by hand optimizing the automatically generated cuda code from the sparse emitter).
+    """
+
+    def __init__(self, debug=False):
+        self._matmul_kernels_cache = Cache()
+        self._debug = debug
+
+    def generate_matmul(
+        self,
+        a,
+        b,
+        c,
+        compute_type,
+        prologs=None,
+        transpose_a=False,
+        transpose_b=False,
+        arch=None,
+        kernel=0,
+    ):
+        """
+        Matmul semantics:
+
+        Dot Product: If both inputs are 1-dimensional, it computes the dot
+        product (inner product).
+
+        Matrix-Vector Product: If the first input is a 2D tensor (matrix) and
+        the second is a 1D tensor (vector), it performs a matrix-vector
+        multiplication.
+
+        Vector-Matrix Product: If the first input is a 1D tensor and the second
+        is a 2D tensor, it performs a vector-matrix multiplication.
+
+        Matrix-Matrix Product: If both inputs are 2D tensors, it performs
+        standard matrix multiplication, requiring the number of columns in
+        the first matrix to equal the number of rows in the second.
+
+        Batched Matrix Product: If either input has more than two dimensions,
+        the extra leading dimensions are treated as "batch dimensions" and it
+        performs a batched matrix multiplication over the last two dimensions.
+        """
+
+        format_a = a.tensor_format.name
+        format_b = b.tensor_format.name
+        format_c = c.tensor_format.name
+        vtp = type_str(a.dtype)
+        itp = type_str(a.index_type)
+
+        # Get user code.
+        if prologs is None:
+            user_prolog_a, user_prolog_b = None, None
+        else:
+            user_prolog_a, user_prolog_b, user_prolog_c = prologs.get("a"), prologs.get("b"), prologs.get("c")
+            if user_prolog_c is not None:
+                raise NotImplementedError("The code generation currently does not support prolog for C.")
+
+        # Try cache lookup.
+        key = (
+            format_a,
+            format_b,
+            format_c,
+            vtp,
+            itp,
+            compute_type,
+            user_prolog_a,
+            user_prolog_b,
+            transpose_a,
+            transpose_b,
+            arch,
+            kernel,
+        )
+        entry = self._matmul_kernels_cache.get(key)
+        if entry is not None:
+            return entry
+
+        # Nothing in cache, try kernel lookup. Note that kernel==0 overrides codegen
+        # with this lookup mechanism, all subsequent choices go to actual codegen.
+        typeless_key = (format_a, format_b, format_c, transpose_a, transpose_b)
+        code = _tensor_matmul_kernels.get(typeless_key) if kernel == 0 else None
+
+        # Prepare prolog.
+        ctp = type_str(compute_type)
+        if a.dtype in ["float8_e4m3fn", "float8_e5m2"]:
+            prolog = f"#include <cuda_fp8.h>\n\nusing CTP = {ctp};\n"
+        elif a.dtype == "float16":
+            prolog = f"#include <cuda_fp16.h>\n\nusing CTP = {ctp};\n"
+        elif a.dtype == "bfloat16":
+            prolog = f"#include <cuda_bf16.h>\n\nusing CTP = {ctp};\n"
+        elif a.dtype in ["complex64", "complex128"]:
+            prolog = f"#include <cuda/std/complex>\n\nusing CTP = {ctp};\n"
+        else:
+            prolog = f"using CTP = {ctp};\n"
+
+        # Use kernel lookup or emitter.
+        if code is not None:
+            # (1) Kernel lookup successful. Inject types into the typeless kernel.
+            source = f"{prolog}using VAL = {vtp};\nusing POS = {itp};\nusing CRD = {itp};\n{prolog_decl}{code}"
+            grid_iter = [(Index.FREE, 0)], True
+        else:
+            # (2) Sparse emitter. Emits proper types.
+            source, grid_iter = emit_matmul_kernel(
+                a, b, c, ctp, transpose_a=transpose_a, transpose_b=transpose_b, kernel=kernel
+            )
+
+        # The source and object code lists.
+        src_code_list, obj_code_list = [source], []
+
+        # User-provided or default prolog operation on A and B.
+        src_code_list, obj_code_list = _set_user_or_default_code(
+            user_code=user_prolog_a,
+            default_code=prolog_a_impl,
+            default_prolog=prolog,
+            src_code_list=src_code_list,
+            obj_code_list=obj_code_list,
+        )
+        src_code_list, obj_code_list = _set_user_or_default_code(
+            user_code=user_prolog_b,
+            default_code=prolog_b_impl,
+            default_prolog=prolog,
+            src_code_list=src_code_list,
+            obj_code_list=obj_code_list,
+        )
+
+        # TODO: this should use logger.debug(...)
+        if self._debug:
+            print("\n==> Compiling matmul\n")
+            for s in src_code_list:
+                print(s)
+                print("/////////////////////\n")
+            t1 = time.perf_counter_ns()
+
+        # Compile and link matmul function. Cache the result.
+        complex_path = pathfinder.find_nvidia_header_directory("cccl")
+        narrow_prec_path = pathfinder.find_nvidia_header_directory("cudart")
+        compiler_options = {
+            "std": "c++17",
+            "link_time_optimization": True,
+            "arch": arch,
+            "include_path": [complex_path, narrow_prec_path],
+        }
+        linker_options = {"link_time_optimization": True}
+        kernel_obj = compile_cpp_and_link(
+            src_code=src_code_list,
+            object_code=obj_code_list,
+            function_name="matmul_kernel",
+            compiler_options=compiler_options,
+            linker_options=linker_options,
+        )
+        assert kernel_obj is not None, "Internal error."
+        self._matmul_kernels_cache[key] = kernel_obj, grid_iter
+
+        if self._debug:
+            t2 = time.perf_counter_ns()
+            tt = (t2 - t1) / 1000.0  # ns -> us
+            print(f"<== Compiled matmul in {tt:10.2f} us.\n")
+
+        return kernel_obj, grid_iter

nvmath/sparse/ust/_utils.py

@@ -0,0 +1,149 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+This module has internal utilities used in implementing the UST and related APIs.
+"""
+
+__all__ = []
+
+import threading
+from contextlib import nullcontext
+
+import numpy as np
+
+from nvmath.internal import utils
+from nvmath.internal.ndbuffer import NDBuffer
+from nvmath.internal.package_ifc import StreamHolder
+
+
+class Cache(dict):
+    """
+    Cache for sharing kernels and matmuls.
+    """
+
+    def __init__(self):
+        self.lock = threading.Lock()
+
+    def free(self):
+        with self.lock:
+            for obj in self.values():
+                obj.free()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args, **kwargs):
+        self.free()
+
+
+class LevelMap(dict):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    # TODO: for "subarray" like pos() and crd(), need to allocate one chunk and offset.
+    # Also need to keep track of the offsets, or recompute? Add pos.base, crd.base?
+    def to(self, device_id, stream_holder):
+        level_map = LevelMap(self)
+        for k in level_map:
+            level_map[k] = self[k].to(device_id, stream_holder)
+        return level_map
+
+    def empty_like(self, stream_holder):
+        level_map = LevelMap(self)
+        for k in level_map:
+            level_map[k] = utils.create_empty_tensor(
+                self[k].__class__, self[k].shape, self[k].dtype, self[k].device_id, stream_holder, verify_strides=False
+            )
+        return level_map
+
+    def copy_(self, src, stream_holder):
+        for k in self:
+            self[k].copy_(src[k], stream_holder)
+
+
+_CTPS = {
+    "complex64": "cuda::std::complex<float>",
+    "complex128": "cuda::std::complex<double>",
+    "float8_e4m3fn": "__nv_fp8_e4m3",
+    "float8_e5m2": "__nv_fp8_e5m2",
+    "bfloat16": "__nv_bfloat16",
+    "float16": "__half",
+    "float32": "float",
+    "float64": "double",
+    "int64": "long long",
+    "int32": "int",
+    "int16": "int",
+    "int8": "int",
+}
+
+_NP_ENVELOPE = {
+    "complex32": np.complex64,
+    "complex64": np.complex64,
+    "complex128": np.complex128,
+    "float8_e4m3fn": np.float16,
+    "float8_e5m2": np.float16,
+    "bfloat16": np.float32,
+    "float16": np.float16,
+    "float32": np.float32,
+    "float64": np.float64,
+    "int64": np.int64,
+    "int32": np.int32,
+    "int16": np.int16,
+    "int8": np.int8,
+}
+
+
+def type_str(tp):
+    return _CTPS[tp]
+
+
+def np_enveloping_type(tp):
+    return _NP_ENVELOPE[tp]
+
+
+def resolve_stream(stream, tensor, target_device_id=None):
+    """
+    Get a default package stream (based on the dense tensor package)
+    or wrap user-provided stream.
+
+    Similar logic is implemented inside statefull instances for matmul,
+    ffts etc, but ust.Tensor is a standalone public API, so we need
+    to implement the logic separately.
+    For internal usages from stateful objects, we don't interfere
+    and just use the provided StreamHolder.
+    """
+    # For internal use, we accept StreamHolder so that UST `to` has a
+    # consistent interface with `TensorHolder.to`.
+    if isinstance(stream, StreamHolder):
+        return stream
+    device_id = tensor.device_id
+    if device_id == "cpu":
+        if target_device_id is None or target_device_id == "cpu":
+            return None
+        device_id = target_device_id
+    package = tensor._dense_tensorholder_type.device_tensor_class.name
+    return utils.get_or_create_stream(device_id, stream, package)
+
+
+def stream_context(stream_holder):
+    if stream_holder is None:
+        return nullcontext()
+    return stream_holder.ctx
+
+
+def as_external_tensor(tensor, stream_holder):
+    """
+    Return unchanged external package tensor (numpy, torch, etc.)
+    or convert a NDBuffer to a host tensor if necessary.
+    Ensures the tensor can be subsequently used with
+    ufuncs, reductions, etc. (which are not available with NDBuffer).
+    """
+    if type(tensor) is not NDBuffer:
+        return tensor
+    if tensor.device_id == "cpu":
+        return tensor.as_numpy()
+    tensor_host = NDBuffer.empty_like(tensor, device_id="cpu")
+    tensor_host.copy_(tensor, stream=stream_holder)
+    return tensor_host.as_numpy()