nvmath-python 0.9.0__cp314-cp314-win_amd64.whl

nvmath/sparse/ust/tensor_jit.py

@@ -0,0 +1,212 @@
+# 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__ = ["launch_kernel"]
+
+import re
+from collections.abc import Sequence
+
+from nvmath.internal import utils
+
+# TODO: support cuda.core >=0.5.0
+try:
+    from cuda.core import (
+        LaunchConfig,
+        Linker,
+        LinkerOptions,
+        ObjectCode,
+        Program,
+        ProgramOptions,
+        launch,
+    )
+except ImportError:
+    from cuda.core.experimental import (
+        LaunchConfig,
+        Linker,
+        LinkerOptions,
+        ObjectCode,
+        Program,
+        ProgramOptions,
+        launch,
+    )
+
+
+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)
+
+
+# TODO: take grid instead of problem size.
+def launch_kernel(kernel, parameters, problem_size, *, device_id, stream_holder, blocking=True):
+    """
+    Launch the kernel on the specified device and stream.
+
+    Args:
+        kernel: the CUDA kernel to launch.
+        parameters: the kernel arguments as a sequence or None.
+        problem_size: the size of the problem (work to divide among threads).
+        device_id: the device ID on which to launch the kernel.
+        stream_holder: a `nvmath.internal.utils.StreamHolder` object encapsulating the
+            stream to use.
+        blocking: if True, the call blocks.
+    """
+
+    if problem_size == 0:
+        return
+
+    # Prepare grid.
+    threads_per_block = min(1024, problem_size)
+    batch_size = (problem_size + threads_per_block - 1) // threads_per_block
+    grid = LaunchConfig(grid=batch_size, block=threads_per_block)
+
+    # Launch the kernel.
+    with utils.device_ctx(device_id):
+        launch(stream_holder.obj, grid, kernel, *parameters)
+
+    # TODO: make this non-blocking by default.
+    if blocking:
+        stream_holder.obj.sync()

nvmath/sparse/ust/tensor_kernel.py

@@ -0,0 +1,348 @@
+# 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__ = [
+    "KernelGen",
+]
+
+import time
+
+from cuda import pathfinder
+
+from ._utils import COMPLEX_PRECISION, Cache, type_str
+from .semiring_ops import (
+    epilog_impl,
+    prolog_a_impl,
+    prolog_b_impl,
+    prolog_c_impl,
+    prolog_epilog_decl,
+    semiring_add_impl,
+    semiring_atomic_add_impl,
+    semiring_mul_impl,
+    semiring_ops_decl,
+)
+from .tensor_emitter import emit_matmul
+from .tensor_jit import compile_cpp_and_link
+
+# Maps matmul(a,b,c) to pre-defined kernel source code. All kernels should adhere to the
+# parameter passing conventions for each of the three UST parameters.
+_tensor_matmul_kernels = {
+    # (format_a, format_b, format_c, transpose_a, transpose_b)
+    ("DIAI", "DenseVector", "DenseVector", False, False): """
+extern "C" __global__ void matmul(
+    const POS* __restrict__ pos0,
+    const CRD* __restrict__ diags,
+    const VAL* __restrict__ A,
+    unsigned long m,
+    unsigned long n,
+    unsigned long l0,
+    unsigned long l1,
+    const VAL *__restrict__ B,
+    VAL *__restrict__ C) {
+  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 = prolog_c(static_cast<CTP>(C[i]));
+    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 = add(acc, mul(prolog_a(static_cast<CTP>(A[di])),
+                           prolog_b(static_cast<CTP>(B[j]))));
+      }
+    }
+    C[i] = static_cast<VAL>(epilog(acc));
+  }
+}
+""",
+    ("DIAJ", "DenseVector", "DenseVector", False, False): """
+extern "C" __global__ void matmul(
+    const POS* __restrict__ pos0,
+    const CRD* __restrict__ diags,
+    const VAL* __restrict__ A,
+    unsigned long m,
+    unsigned long n,
+    unsigned long l0,
+    unsigned long l1,
+    const VAL *__restrict__ B,
+    VAL *__restrict__ C) {
+  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 = prolog_c(static_cast<CTP>(C[i]));
+    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 = add(acc, mul(static_cast<CTP>(prolog_a(A[di + j])),
+                           static_cast<CTP>(prolog_b(B[j]))));
+      }
+    }
+    C[i] = static_cast<VAL>(epilog(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,
+        epilog=None,
+        semiring=None,
+        transpose_a=False,
+        transpose_b=False,
+        arch=None,
+    ):
+        """
+        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 semiring is None:
+            user_mul, user_add, user_atomic_add = None, None, None
+        else:
+            user_mul, user_add, user_atomic_add = semiring.get("mul"), semiring.get("add"), semiring.get("atomic_add")
+        if prologs is None:
+            user_prolog_a, user_prolog_b, user_prolog_c = None, None, None
+        else:
+            user_prolog_a, user_prolog_b, user_prolog_c = prologs.get("a"), prologs.get("b"), prologs.get("c")
+        # Try cache lookup.
+        key = (
+            format_a,
+            format_b,
+            format_c,
+            vtp,
+            itp,
+            compute_type,
+            user_mul,
+            user_add,
+            user_atomic_add,
+            user_prolog_a,
+            user_prolog_b,
+            user_prolog_c,
+            epilog,
+            transpose_a,
+            transpose_b,
+            arch,
+        )
+        kernel = self._matmul_kernels_cache.get(key)
+        if kernel is None:
+            # Nothing in cache, try kernel lookup.
+            typeless_key = (format_a, format_b, format_c, transpose_a, transpose_b)
+            code = _tensor_matmul_kernels.get(typeless_key)
+
+            # Prepare prolog.
+            ctp = type_str(compute_type)
+            if a.dtype in ["float8_e4m3fn", "float8_e5m2"]:
+                prolog = f"#include <cuda_fp8.h>\nusing ATP = {ctp};\nusing CTP = {ctp};\n"  # TODO: ATP=vtp
+            elif a.dtype == "float16":
+                prolog = f"#include <cuda_fp16.h>\nusing ATP = {vtp};\nusing CTP = {ctp};\n"
+            elif a.dtype == "bfloat16":
+                prolog = f"#include <cuda_bf16.h>\nusing ATP = {vtp};\nusing CTP = {ctp};\n"
+            elif a.dtype in ["complex64", "complex128"]:
+                prolog = f"#include <cuda/std/complex>\nusing ATP = {COMPLEX_PRECISION[a.dtype]};\nusing CTP = {ctp};\n"
+            else:
+                prolog = f"using ATP = {vtp};\nusing 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}\nusing VAL = {vtp};\nusing POS = {itp};\n"
+                    f"using CRD = {itp};\n{semiring_ops_decl}{prolog_epilog_decl}{code}"
+                )
+            else:
+                if user_prolog_c is not None or epilog is not None:
+                    raise NotImplementedError(
+                        "The code generation (emitter) path does not currently support \
+prolog for the `c` operand or epilog."
+                    )
+                # (2) Sparse emitter. Emits proper types.
+                source = emit_matmul(a, b, c, ctp, transpose_a, transpose_b)
+
+            # The source and object code lists.
+            src_code_list, obj_code_list = [source], []
+
+            # User-provided or default semiring mul operation.
+            src_code_list, obj_code_list = _set_user_or_default_code(
+                user_code=user_mul,
+                default_code=semiring_mul_impl,
+                default_prolog=prolog,
+                src_code_list=src_code_list,
+                obj_code_list=obj_code_list,
+            )
+
+            # User-provided or default semiring add operation.
+            src_code_list, obj_code_list = _set_user_or_default_code(
+                user_code=user_add,
+                default_code=semiring_add_impl,
+                default_prolog=prolog,
+                src_code_list=src_code_list,
+                obj_code_list=obj_code_list,
+            )
+
+            # User-provided or default semiring atomic add operation.
+            src_code_list, obj_code_list = _set_user_or_default_code(
+                user_code=user_atomic_add,
+                default_code=semiring_atomic_add_impl,
+                default_prolog=prolog,
+                src_code_list=src_code_list,
+                obj_code_list=obj_code_list,
+            )
+
+            # User-provided or default prolog operation on A.
+            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,
+            )
+
+            # User-provided or default prolog operation on B.
+            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,
+            )
+
+            # User-provided or default prolog operation on C.
+            src_code_list, obj_code_list = _set_user_or_default_code(
+                user_code=user_prolog_c,
+                default_code=prolog_c_impl,
+                default_prolog=prolog,
+                src_code_list=src_code_list,
+                obj_code_list=obj_code_list,
+            )
+
+            # User-provided or default epilog operation.
+            src_code_list, obj_code_list = _set_user_or_default_code(
+                user_code=epilog,
+                default_code=epilog_impl,
+                default_prolog=prolog,
+                src_code_list=src_code_list,
+                obj_code_list=obj_code_list,
+            )
+
+            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 = compile_cpp_and_link(
+                src_code=src_code_list,
+                object_code=obj_code_list,
+                function_name="matmul",
+                compiler_options=compiler_options,
+                linker_options=linker_options,
+            )
+            assert kernel is not None, "Internal error."
+            self._matmul_kernels_cache[key] = kernel
+
+            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