nvmath-python 1.0.0__cp314-cp314t-win_amd64.whl

nvmath/device/cublasdx_numba_cuda_mlir.py

@@ -0,0 +1,208 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from numba_cuda_mlir import types
+from numba_cuda_mlir.compiler import declare_device
+from numba_cuda_mlir.extending import (
+    overload,
+    overload_method,
+    typing_registry,
+)
+from numba_cuda_mlir.lowering_utilities.type_conversions import to_numba_type
+from numba_cuda_mlir.numba_cuda.cudadrv.linkable_code import LTOIR
+from numba_cuda_mlir.numba_cuda.extending import typeof_impl
+
+from nvmath.device.common import (
+    axpby,
+    clear,
+    copy,
+    copy_fragment,
+    copy_wait,
+    make_fragment_like,
+    make_tensor,
+)
+from nvmath.device.common_cuda import get_default_code_type
+from nvmath.device.common_numba_cuda_mlir import (
+    HostDescriptorBase,
+    get_array_pointer,
+    get_value_ptr,
+    int_to_uint32_ptr,
+    register_dummy_type,
+)
+from nvmath.device.cublasdx import (
+    DevicePipeline,
+    Matmul,
+    TilePipeline,
+    _BlasMatmulLayout,
+    _BlasMatmulLikeLayout,
+    compile_blas_execute,
+)
+
+from .cublasdx_backend import (
+    _BLAS_COMPILED_ARGS,
+    _BLAS_DEFINITION_ARGS,
+)
+
+# ==============================================
+# Helpers
+# ==============================================
+
+
+def _check_array_value_types(MM: Matmul, a: types.Array, b: types.Array, c: types.Array) -> bool:
+    expected = tuple(to_numba_type(vt) for vt in MM._traits.value_types)
+    return (a.dtype, b.dtype, c.dtype) == expected
+
+
+_NOT_SUPPORTED = (
+    "Advanced matmul and opaque tensor APIs are not yet supported by the numba-cuda-mlir backend."
+    " Only Matmul.execute with array operands is available."
+)
+
+
+def _unsupported(*args, **kwargs):
+    raise NotImplementedError(_NOT_SUPPORTED)
+
+
+# ==============================================
+# Matmul Type
+# ==============================================
+
+
+class BlasType(HostDescriptorBase):
+    """
+    Type class associated with the `cublasdx.Matmul`.
+    """
+
+    def __init__(self, blas: Matmul) -> None:
+        assert isinstance(blas, Matmul)
+        super().__init__(blas, _BLAS_DEFINITION_ARGS, "BlasNumba")
+
+
+register_dummy_type(BlasType, Matmul, _BLAS_DEFINITION_ARGS + _BLAS_COMPILED_ARGS)
+
+
+# ==============================================
+# Matmul executes
+# ==============================================
+
+
+@overload_method(BlasType, "execute", strict=False, typing_registry=typing_registry)
+def ol_blas_execute(blas_numba: BlasType, _arg1, _arg2, _arg3, _arg4=None, _arg5=None, _arg6=None, _arg7=None, _arg8=None):
+    if not isinstance(blas_numba, BlasType):
+        return
+
+    none_set = {None, types.Omitted(None)}
+    MM = blas_numba.host_descriptor
+
+    if {_arg4, _arg5, _arg6, _arg7, _arg8} <= none_set:
+        _unsupported()
+
+    if {_arg6, _arg7, _arg8} <= none_set:
+        if not all(isinstance(x, types.Array) for x in (_arg2, _arg3, _arg5)):
+            _unsupported()
+        return _build_basic_impl(MM, _arg1, _arg2, _arg3, _arg4, _arg5)
+
+    if not all(isinstance(x, types.Array) for x in (_arg2, _arg4, _arg7)):
+        _unsupported()
+    return _build_ldabc_impl(MM, _arg1, _arg2, _arg3, _arg4, _arg5, _arg6, _arg7, _arg8)
+
+
+def _build_basic_impl(MM: Matmul, alpha, a, b, beta, c):
+    if (
+        not all(isinstance(x, types.Number) for x in (alpha, beta))
+        or not all(isinstance(x, types.Array) for x in (a, b, c))
+        or not _check_array_value_types(MM, a, b, c)
+    ):
+        return None
+
+    code, symbol = compile_blas_execute(
+        MM,
+        code_type=get_default_code_type(),
+        execute_api="static_leading_dimensions",
+    )
+
+    c_value_type = MM.c_value_type
+    c_ptr = types.CPointer(c.dtype)
+    sig = types.void(c_ptr, types.CPointer(a.dtype), types.CPointer(b.dtype), c_ptr, c_ptr)
+    blas_device_func = declare_device(symbol, sig, link=LTOIR(code.data), abi="c")
+
+    def impl(_, alpha, a, b, beta, c, _arg6=None, _arg7=None, _arg8=None):
+        aptr = get_array_pointer(a)
+        bptr = get_array_pointer(b)
+        cptr = get_array_pointer(c)
+        alpha_ptr = get_value_ptr(c_value_type(alpha))
+        beta_ptr = get_value_ptr(c_value_type(beta))
+
+        blas_device_func(alpha_ptr, aptr, bptr, beta_ptr, cptr)
+
+    return impl
+
+
+def _build_ldabc_impl(MM: Matmul, alpha, a, lda, b, ldb, beta, c, ldc):
+    if (
+        not all(isinstance(x, types.Number) for x in (alpha, beta))
+        or not all(isinstance(x, types.Array) for x in (a, b, c))
+        or not all(isinstance(x, types.Integer) for x in (lda, ldb, ldc))
+        or not _check_array_value_types(MM, a, b, c)
+    ):
+        return None
+
+    code, symbol = compile_blas_execute(
+        MM,
+        code_type=get_default_code_type(),
+        execute_api="dynamic_leading_dimensions",
+    )
+
+    c_value_type = MM.c_value_type
+    ld_ptr = types.CPointer(types.uint32)
+    c_ptr = types.CPointer(c.dtype)
+    sig = types.void(c_ptr, types.CPointer(a.dtype), ld_ptr, types.CPointer(b.dtype), ld_ptr, c_ptr, c_ptr, ld_ptr)
+    blas_device_func = declare_device(symbol, sig, link=LTOIR(code.data), abi="c")
+
+    def impl(_, alpha, a, lda, b, ldb, beta, c, ldc):
+        aptr = get_array_pointer(a)
+        bptr = get_array_pointer(b)
+        cptr = get_array_pointer(c)
+        alpha_ptr = get_value_ptr(c_value_type(alpha))
+        beta_ptr = get_value_ptr(c_value_type(beta))
+        lda_ptr = int_to_uint32_ptr(lda)
+        ldb_ptr = int_to_uint32_ptr(ldb)
+        ldc_ptr = int_to_uint32_ptr(ldc)
+
+        blas_device_func(alpha_ptr, aptr, lda_ptr, bptr, ldb_ptr, beta_ptr, cptr, ldc_ptr)
+
+    return impl
+
+
+# ==============================================
+# Unsupported API surface
+# ==============================================
+
+_UNSUPPORTED_BLAS_METHODS = (
+    "get_accumulator",
+    "suggest_accumulator",
+    "get_layout_smem_a",
+    "get_layout_smem_b",
+    "get_layout_smem_c",
+    "get_layout_gmem_a",
+    "get_layout_gmem_b",
+    "get_layout_gmem_c",
+    "suggest_layout_smem_a",
+    "suggest_layout_smem_b",
+    "suggest_layout_smem_c",
+    "suggest_layout_rmem_c",
+    "get_layout_rmem_c",
+    "_get_accumulator_c",
+    "_suggest_accumulator_c",
+)
+for _method in _UNSUPPORTED_BLAS_METHODS:
+    overload_method(BlasType, _method, strict=False, typing_registry=typing_registry)(_unsupported)
+
+_UNSUPPORTED_FREE_FUNCTIONS = (make_tensor, make_fragment_like, copy, copy_fragment, clear, copy_wait, axpby)
+for _func in _UNSUPPORTED_FREE_FUNCTIONS:
+    overload(_func, strict=False, typing_registry=typing_registry)(_unsupported)
+
+_UNSUPPORTED_TYPES = (DevicePipeline, TilePipeline, _BlasMatmulLayout, _BlasMatmulLikeLayout)
+for _cls in _UNSUPPORTED_TYPES:
+    typeof_impl.register(_cls)(_unsupported)

nvmath/device/cufftdx.py

@@ -0,0 +1,373 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = ["FFT", "compile_fft_execute"]
+from functools import cached_property
+from typing import Any
+
+from nvmath.bindings import mathdx
+from nvmath.internal.utils import docstring_decorator
+
+from .common import (
+    SHARED_DEVICE_DOCSTRINGS,
+    check_code_type,
+    parse_code_type,
+    parse_sm,
+)
+from .common_backend import MATHDX_TYPES_TO_NP, get_isa_version, get_lto
+from .common_cuda import (
+    Code,
+    Dim3,
+)
+from .cufftdx_backend import (
+    generate_code,
+    generate_FFT,
+    get_data_type_trait,
+    get_int_trait,
+    get_int_traits,
+    get_knobs,
+    get_str_trait,
+    validate,
+    validate_execute_api,
+)
+
+CUFFTDX_DATABASE = None
+
+FFTDX_DOCSTRING = SHARED_DEVICE_DOCSTRINGS.copy()
+FFTDX_DOCSTRING.update(
+    {
+        "size": """\
+The size of the FFT to calculate.""".replace("\n", " "),
+        #
+        "fft_type": """\
+A string specifying the type of FFT operation, can be ``'c2c'``, ``'c2r'`` or ``'r2c'``.""".replace("\n", " "),
+        #
+        "direction": """\
+A string specifying the direction of FFT, can be ``'forward'`` or ``'inverse'``. If not provided, will be ``'forward'``
+if real-to-complex FFT is specified and ``'inverse'`` if complex-to-real FFT is specified.""".replace("\n", " "),
+        #
+        "ffts_per_block": """\
+The number of FFTs calculated per CUDA block, optional. The default is 1. Alternatively, if provided as ``'suggested'``,
+will be set to a suggested value.""".replace("\n", " "),
+        #
+        "elements_per_thread": """\
+The number of elements per thread, optional. If not provided, the default value selected by cuFFTDx is used.
+Alternatively, if provided as ``'suggested'``, will be set to a suggested value.""".replace("\n", " "),
+        #
+        "real_fft_options": """\
+A dictionary specifying the options for real FFT operation, optional.""".replace("\n", " "),
+    }
+)
+
+
+@docstring_decorator(FFTDX_DOCSTRING, skip_missing=False)
+class FFT:
+    """
+    A class that encapsulates a partial FFT device function. A partial device function can
+    be queried for available or optimal values for some knobs (such as `ffts_per_block`
+    or `elements_per_thread`).
+
+    .. versionchanged:: 0.7.0
+        `FFT` has replaced `FFTOptions` and `FFTOptionsComplete`.
+
+    Args:
+        size (int): {size}
+
+        precision (type[np.floating]): {precision}
+
+        fft_type (str): {fft_type}
+
+        execution (str): {execution}
+
+        sm (ComputeCapability): {sm}
+
+        direction (str): {direction}
+
+        ffts_per_block (int): {ffts_per_block}
+
+        elements_per_thread (int): {elements_per_thread}
+
+        real_fft_options (dict): {real_fft_options} User may specify the following options
+            in the dictionary:
+
+            - ``'complex_layout'``, currently supports ``'natural'``, ``'packed'``, and
+              ``'full'``.
+            - ``'real_mode'``, currently supports ``'normal'`` and ``'folded'``.
+
+    .. seealso::
+        The attributes of this class provide a 1:1 mapping with the CUDA C++ cuFFTDx APIs.
+        For further details, please refer to :cufftdx_doc:`cuFFTDx documentation
+        <index.html>`.
+    """
+
+    def __init__(
+        self,
+        size,
+        precision,
+        fft_type,
+        execution,
+        *,
+        sm=None,
+        direction=None,
+        ffts_per_block=None,
+        elements_per_thread=None,
+        real_fft_options=None,
+    ):
+        sm = parse_sm(sm)
+
+        #
+        # Check that the knobs are, individually, valid
+        #
+
+        validate(
+            size=size,
+            precision=precision,
+            fft_type=fft_type,
+            sm=sm,
+            execution=execution,
+            direction=direction,
+            ffts_per_block=ffts_per_block,
+            elements_per_thread=elements_per_thread,
+            real_fft_options=real_fft_options,
+        )
+
+        if direction is None and fft_type == "r2c":
+            direction = "forward"
+        elif direction is None and fft_type == "c2r":
+            direction = "inverse"
+
+        #
+        # Traits set by input
+        #
+
+        self._precision = precision
+        self._fft_type = fft_type
+        self._direction = direction
+        self._size = size
+        self._execution = execution
+        self._ffts_per_block = ffts_per_block
+        self._elements_per_thread = elements_per_thread
+        self._real_fft_options = real_fft_options
+        self._sm = sm
+
+        #
+        # Update suggested traits
+        #
+
+        if elements_per_thread == "suggested":
+            self._elements_per_thread = self._suggested_elements_per_thread
+
+        if ffts_per_block == "suggested":
+            self._ffts_per_block = self._suggested_ffts_per_block
+
+    @cached_property
+    def _traits(self):
+        return _FFTTraits(self)
+
+    @property
+    def elements_per_thread(self):
+        if self._elements_per_thread is None:
+            return self._traits.elements_per_thread
+        return self._elements_per_thread
+
+    @property
+    def precision(self):
+        return self._precision
+
+    @property
+    def ffts_per_block(self):
+        if self._ffts_per_block is None:
+            return self._traits.ffts_per_block
+        return self._ffts_per_block
+
+    @property
+    def fft_type(self):
+        return self._fft_type
+
+    @property
+    def direction(self):
+        return self._direction
+
+    @property
+    def size(self):
+        return self._size
+
+    @property
+    def execution(self):
+        return self._execution
+
+    @property
+    def real_fft_options(self):
+        return self._real_fft_options
+
+    @property
+    def sm(self):
+        return self._sm
+
+    #
+    # Extensions
+    #
+
+    def valid(self, *knobs):
+        if not (set(knobs) <= {"ffts_per_block", "elements_per_thread"}):
+            raise ValueError(f"Unsupported knob. Only valid knobs are ffts_per_block and elements_per_thread but got {knobs}")
+
+        return self._get_knobs(*knobs)
+
+    @property
+    def value_type(self):
+        return self._traits.value_type
+
+    @property
+    def input_type(self):
+        return self._traits.input_type
+
+    @property
+    def output_type(self):
+        return self._traits.output_type
+
+    @property
+    def storage_size(self):
+        return self._traits.storage_size
+
+    @property
+    def shared_memory_size(self):
+        return self._traits.shared_memory_size
+
+    @property
+    def stride(self):
+        return self._traits.stride
+
+    @property
+    def block_dim(self):
+        return self._traits.block_dim
+
+    @property
+    def implicit_type_batching(self):
+        return self._traits.implicit_type_batching
+
+    @property
+    def extensions(self):
+        raise NotImplementedError("Extensions not supported yet")
+
+    def execute(*args):
+        raise RuntimeError("execute should not be called directly outside of a jitted kernel.")
+
+    #
+    # Private implementations
+    #
+
+    def _suggested(self, what):
+        # Generate full PTX
+        h = generate_FFT(
+            size=self._size,
+            precision=self._precision,
+            fft_type=self._fft_type,
+            direction=self._direction,
+            ffts_per_block=(None if self._ffts_per_block == "suggested" else self._ffts_per_block),
+            elements_per_thread=(None if self._elements_per_thread == "suggested" else self._elements_per_thread),
+            real_fft_options=frozenset(self._real_fft_options.items()) if self._real_fft_options else None,
+            sm=self._sm,
+            execution=self._execution,
+        )
+
+        if what == "elements_per_thread":
+            return get_int_trait(h.descriptor, mathdx.CufftdxTraitType.ELEMENTS_PER_THREAD)
+
+        if what == "suggested_ffts_per_block":
+            return get_int_trait(h.descriptor, mathdx.CufftdxTraitType.SUGGESTED_FFTS_PER_BLOCK)
+
+        raise Exception(f"Unknown suggested option '{what}'")
+
+    @cached_property
+    def _suggested_ffts_per_block(self):
+        return self._suggested("suggested_ffts_per_block")
+
+    @cached_property
+    def _suggested_elements_per_thread(self):
+        return self._suggested("elements_per_thread")
+
+    def _get_knobs(self, *knobs):
+        if not (set(knobs) <= {"ffts_per_block", "elements_per_thread"}):
+            raise ValueError(f"Unsupported knob. Only valid knobs are ffts_per_block and elements_per_thread but got {knobs}")
+
+        h = generate_FFT(
+            size=self._size,
+            precision=self._precision,
+            fft_type=self._fft_type,
+            direction=self._direction,
+            ffts_per_block=self._ffts_per_block,
+            elements_per_thread=self._elements_per_thread,
+            real_fft_options=frozenset(self._real_fft_options.items()) if self._real_fft_options else None,
+            sm=self._sm,
+            execution=self._execution,
+        )
+
+        return get_knobs(h.descriptor, knobs)
+
+
+class _FFTTraits:
+    def __init__(self, FFT: FFT):
+        h = generate_FFT(
+            size=FFT._size,
+            precision=FFT._precision,
+            fft_type=FFT._fft_type,
+            direction=FFT._direction,
+            sm=FFT._sm,
+            execution=FFT._execution,
+            ffts_per_block=FFT._ffts_per_block if FFT._execution == "Block" else None,
+            elements_per_thread=FFT._elements_per_thread if FFT._execution == "Block" else None,
+            real_fft_options=frozenset(FFT._real_fft_options.items()) if FFT._real_fft_options else None,
+        ).descriptor
+
+        self.value_type = MATHDX_TYPES_TO_NP[get_data_type_trait(h, mathdx.CufftdxTraitType.VALUE_TYPE)]
+        self.input_type = MATHDX_TYPES_TO_NP[get_data_type_trait(h, mathdx.CufftdxTraitType.INPUT_TYPE)]
+        self.output_type = MATHDX_TYPES_TO_NP[get_data_type_trait(h, mathdx.CufftdxTraitType.OUTPUT_TYPE)]
+
+        self.storage_size = get_int_trait(h, mathdx.CufftdxTraitType.STORAGE_SIZE)
+        self.stride = get_int_trait(h, mathdx.CufftdxTraitType.STRIDE)
+        self.elements_per_thread = get_int_trait(h, mathdx.CufftdxTraitType.ELEMENTS_PER_THREAD)
+        self.implicit_type_batching = get_int_trait(h, mathdx.CufftdxTraitType.IMPLICIT_TYPE_BATCHING)
+
+        if FFT.execution == "Block":
+            self.block_dim: Dim3 | None = Dim3(*get_int_traits(h, mathdx.CufftdxTraitType.BLOCK_DIM, 3))
+            self.shared_memory_size: int | None = get_int_trait(h, mathdx.CufftdxTraitType.SHARED_MEMORY_SIZE)
+            self.ffts_per_block: int | None = get_int_trait(h, mathdx.CufftdxTraitType.FFTS_PER_BLOCK)
+        else:
+            self.block_dim = None
+            self.shared_memory_size = None
+            self.ffts_per_block = None
+
+
+def compile_fft_execute(
+    fft: FFT,
+    code_type: Any,
+    execute_api: str | None = None,
+) -> tuple[Code, str]:
+    code_type = parse_code_type(code_type)
+    check_code_type(code_type, "cuFFTDx")
+    validate_execute_api(fft.execution, execute_api)
+
+    h = generate_FFT(
+        size=fft._size,
+        precision=fft._precision,
+        fft_type=fft._fft_type,
+        direction=fft._direction,
+        sm=fft._sm,
+        execution=fft._execution,
+        ffts_per_block=fft._ffts_per_block if fft._execution == "Block" else None,
+        elements_per_thread=fft._elements_per_thread if fft._execution == "Block" else None,
+        real_fft_options=frozenset(fft._real_fft_options.items()) if fft._real_fft_options else None,
+        execute_api=execute_api,
+    ).descriptor
+
+    code = generate_code(h, code_type.cc)
+
+    # Compile
+    lto_fn = get_lto(code.descriptor)
+    isa_version = get_isa_version(code.descriptor)
+
+    symbol = get_str_trait(h, mathdx.CufftdxTraitType.SYMBOL_NAME)
+
+    return Code(code_type, isa_version, lto_fn), symbol