nvmath-python 1.0.0__cp314-cp314t-win_amd64.whl

nvmath/device/cufftdx_backend.py

@@ -0,0 +1,220 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from collections.abc import Sequence
+from functools import lru_cache
+from typing import Protocol
+
+import numpy
+
+from nvmath.bindings import mathdx
+from nvmath.device.common_backend import DescriptorWrapper
+from nvmath.device.common_cuda import ComputeCapability
+
+from .common import check_contains, check_in, check_not_in, check_sm
+from .common_backend import (
+    EXECUTION_STR_TO_MATHDX,
+    NP_TYPES_TO_MATHDX_PRECISION,
+    NVARG_GEN_OPT_LTO,
+    build_get_int_traits,
+    build_get_str_trait,
+    set_code_target_sm,
+)
+from .types import REAL_NP_TYPES
+
+_FFT_TYPE_TO_MATHDX = {t.name.lower(): t for t in mathdx.CufftdxType}
+
+_FFT_DIRECTION_TO_MATHDX = {d.name.lower(): d for d in mathdx.CufftdxDirection}
+
+_FFT_COMPLEX_LAYOUT_TO_MATHDX = {cl.name.lower(): cl for cl in mathdx.CufftdxComplexLayout}
+
+_FFT_REAL_MODE_TO_MATHDX = {m.name.lower(): m for m in mathdx.CufftdxRealMode}
+
+_FFT_API_STR_TO_MATHDX = {
+    "register_memory": mathdx.CufftdxApi.LMEM,
+    "shared_memory": mathdx.CufftdxApi.SMEM,
+}
+
+_FFT_KNOB_TYPE_TO_MATHDX = {t.name.lower(): t for t in mathdx.CufftdxKnobType}
+
+
+_FFT_DEFINITION_ARGS = [
+    "size",
+    "precision",
+    "fft_type",
+    "execution",
+    "sm",
+    "direction",
+    "ffts_per_block",
+    "elements_per_thread",
+    "real_fft_options",
+]
+
+_FFT_COMPILED_ARGS = [
+    "value_type",
+    "input_type",
+    "output_type",
+    "storage_size",
+    "shared_memory_size",
+    "stride",
+    "block_dim",
+    "implicit_type_batching",
+]
+
+
+class CallableGetIntTraits(Protocol):
+    def __call__(self, handle: int, trait_type: mathdx.CufftdxTraitType, size: int) -> tuple: ...
+
+
+class CallableGetStrTrait(Protocol):
+    def __call__(self, handle: int, trait_type: mathdx.CufftdxTraitType) -> str: ...
+
+
+get_int_traits: CallableGetIntTraits = build_get_int_traits(mathdx.cufftdx_get_trait_int64s)
+
+get_str_trait: CallableGetStrTrait = build_get_str_trait(mathdx.cufftdx_get_trait_str_size, mathdx.cufftdx_get_trait_str)
+
+
+def get_int_trait(handle: int, trait_type: mathdx.CufftdxTraitType) -> int:
+    return int(mathdx.cufftdx_get_trait_int64(handle, trait_type))
+
+
+def get_data_type_trait(handle: int, trait_type: mathdx.CufftdxTraitType) -> mathdx.CommondxValueType:
+    return mathdx.CommondxValueType(mathdx.cufftdx_get_trait_commondx_data_type(handle, trait_type))
+
+
+def validate(
+    size,
+    precision,
+    fft_type,
+    execution,
+    direction,
+    ffts_per_block,
+    elements_per_thread,
+    real_fft_options,
+    sm,
+):
+    if size <= 0:
+        raise ValueError(f"size must be > 0. Got {size}")
+    check_in("precision", precision, REAL_NP_TYPES)
+    check_in("fft_type", fft_type, ["c2c", "c2r", "r2c"])
+    check_in("execution", execution, ["Block", "Thread"])
+    if direction is not None:
+        check_in("direction", direction, ["forward", "inverse"])
+    if ffts_per_block in (None, "suggested"):
+        pass
+    else:
+        if ffts_per_block <= 0:
+            raise ValueError(
+                f"ffts_per_block must be None, 'suggested' or a positive integer ; got ffts_per_block = {ffts_per_block}"
+            )
+    if elements_per_thread in (None, "suggested"):
+        pass
+    else:
+        if elements_per_thread <= 0:
+            raise ValueError(
+                f"elements_per_thread must be None, 'suggested' or a positive integer ; "
+                f"got elements_per_thread = {elements_per_thread}"
+            )
+    if real_fft_options is None:
+        pass
+    else:
+        check_contains(real_fft_options, "complex_layout")
+        check_contains(real_fft_options, "real_mode")
+        check_in("real_fft_options['complex_layout']", real_fft_options["complex_layout"], ["natural", "packed", "full"])
+        check_in("real_fft_options['real_mode']", real_fft_options["real_mode"], ["normal", "folded"])
+    check_sm(sm, "cuFFTDx")
+
+
+def validate_execute_api(execution: str, execute_api: str | None):
+    """
+    Validate the execute_api argument.
+    """
+    if execution == "Block":
+        check_in("execute_api", execute_api, list(_FFT_API_STR_TO_MATHDX.keys()) + [None])
+    else:
+        if execute_api is not None:
+            raise ValueError(f"execute_api may be set only for 'Block' execution; got execution = {execution}")
+
+
+@lru_cache
+def generate_FFT(
+    size,
+    precision,
+    fft_type,
+    direction,
+    sm,
+    execution,
+    ffts_per_block,
+    elements_per_thread,
+    real_fft_options,
+    execute_api=None,
+):
+    check_not_in("ffts_per_block", ffts_per_block, ["suggested"])
+    check_not_in("elements_per_thread", elements_per_thread, ["suggested"])
+
+    h = mathdx.cufftdx_create_descriptor()
+
+    # TODO: remove after migrating to libmathdx 0.2.4+
+    if execute_api is None:
+        execute_api = "register_memory"
+
+    if execute_api is not None:
+        mathdx.cufftdx_set_operator_int64(h, mathdx.CufftdxOperatorType.API, _FFT_API_STR_TO_MATHDX[execute_api])
+
+    mathdx.cufftdx_set_operator_int64(h, mathdx.CufftdxOperatorType.SM, sm.major * 100 + sm.minor * 10)
+    mathdx.cufftdx_set_operator_int64(h, mathdx.CufftdxOperatorType.SIZE, size)
+    mathdx.cufftdx_set_operator_int64(h, mathdx.CufftdxOperatorType.PRECISION, NP_TYPES_TO_MATHDX_PRECISION[precision])
+    mathdx.cufftdx_set_operator_int64(h, mathdx.CufftdxOperatorType.TYPE, _FFT_TYPE_TO_MATHDX[fft_type])
+    mathdx.cufftdx_set_operator_int64(h, mathdx.CufftdxOperatorType.DIRECTION, _FFT_DIRECTION_TO_MATHDX[direction])
+
+    if real_fft_options:
+        real_fft_options = dict(real_fft_options)
+        ll = real_fft_options["complex_layout"]
+        mm = real_fft_options["real_mode"]
+        mathdx.cufftdx_set_operator_int64s(
+            h, mathdx.CufftdxOperatorType.REAL_FFT_OPTIONS, 2, [_FFT_COMPLEX_LAYOUT_TO_MATHDX[ll], _FFT_REAL_MODE_TO_MATHDX[mm]]
+        )
+
+    mathdx.cufftdx_set_operator_int64(h, mathdx.CufftdxOperatorType.EXECUTION, EXECUTION_STR_TO_MATHDX[execution])
+
+    if ffts_per_block:
+        mathdx.cufftdx_set_operator_int64(h, mathdx.CufftdxOperatorType.FFTS_PER_BLOCK, ffts_per_block)
+
+    if elements_per_thread:
+        mathdx.cufftdx_set_operator_int64(h, mathdx.CufftdxOperatorType.ELEMENTS_PER_THREAD, elements_per_thread)
+
+    return DescriptorWrapper(h, mathdx.cufftdx_destroy_descriptor)
+
+
+@lru_cache
+def generate_code(handle, version: ComputeCapability):
+    code = mathdx.commondx_create_code()
+
+    set_code_target_sm(code, version)
+    mathdx.commondx_set_code_option_str(code, mathdx.CommondxOption.EXTRA_NVRTC_ARGS, NVARG_GEN_OPT_LTO)
+    mathdx.cufftdx_finalize_code(code, handle)
+
+    return DescriptorWrapper(code, mathdx.commondx_destroy_code)
+
+
+@lru_cache
+def get_knobs(handle, knobs: Sequence[str]):
+    knobs_dx: list[mathdx.CufftdxKnobType] = [_FFT_KNOB_TYPE_TO_MATHDX[k] for k in knobs]
+    knobs_size = mathdx.cufftdx_get_knob_int64size(handle, len(knobs), knobs_dx)
+
+    knobs_result_dx = numpy.empty(knobs_size, dtype=numpy.int64)
+    mathdx.cufftdx_get_knob_int64s(
+        handle,
+        len(knobs),
+        knobs_dx,
+        knobs_size,
+        knobs_result_dx.ctypes.data,
+    )
+
+    knobs_result = [
+        tuple(int(knobs_result_dx[j]) for j in range(i, i + len(knobs))) for i in range(0, len(knobs_result_dx), len(knobs))
+    ]
+
+    return knobs_result

nvmath/device/cufftdx_numba.py

@@ -0,0 +1,140 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from numba import cuda
+from numba.core import cgutils, typing
+from numba.cuda.cudaimpl import lower_constant
+from numba.cuda.extending import overload_method, typeof_impl
+from numba.cuda.models import register_model
+from numba.extending import types
+
+from nvmath.device.common_cuda import get_default_code_type
+from nvmath.device.cufftdx import FFT, compile_fft_execute
+
+from .common_numba import (
+    NUMBA_FE_TYPES_TO_NUMBA_IR,
+    EmptyStructModel,
+    declare_cabi_device,
+    get_array_ptr,
+    overload_type_attribute,
+)
+from .cufftdx_backend import _FFT_COMPILED_ARGS, _FFT_DEFINITION_ARGS
+
+
+class FFTType(types.Type):
+    """
+    Type class associated with the `cufftdx.FFT`.
+    """
+
+    def __init__(self, fft: FFT):
+        assert isinstance(fft, FFT)
+        self._fft = fft
+        attributes = [f"{attr}={getattr(fft, attr)}" for attr in _FFT_DEFINITION_ARGS if getattr(fft, attr)]
+        attributes.sort()
+
+        self.name = "FFT(" + ",".join(attributes) + ")"
+
+    @property
+    def fft(self) -> FFT:
+        return self._fft
+
+
+register_model(FFTType)(EmptyStructModel)
+
+
+@lower_constant(FFTType)
+def constant_dummy(context, builder, typ, pyval):
+    struct_ptr = cgutils.create_struct_proxy(typ)(context, builder)
+    return struct_ptr._getvalue()
+
+
+@typeof_impl.register(FFT)
+def typeof_fft_numba(val: FFT, c: typing.Context) -> FFTType:
+    return FFTType(val)
+
+
+for attribute in _FFT_DEFINITION_ARGS + _FFT_COMPILED_ARGS:
+    overload_type_attribute(FFTType, "fft", attribute)
+
+
+# Numba does not support method overload or variadic arguments, so we using
+# default values as a workaround
+# https://github.com/numba/numba/issues/9980
+# https://github.com/numba/numba/issues/9979
+# https://github.com/numba/numba/issues/10143
+@overload_method(FFTType, "execute", jit_options={"forceinline": True}, strict=False)
+def ol_fft_numba_execute(fft_numba: FFTType, _arg1, _arg2=None):
+    if _arg2 in {None, types.Omitted(None), types.none}:
+        return ol_fft_type_execute_rmem(fft_numba, _arg1)
+    return ol_fft_type_execute_smem(fft_numba, _arg1, _arg2)
+
+
+def ol_fft_type_execute_rmem(
+    fft_numba: FFTType,
+    thread_data: types.Array,
+):
+    if not isinstance(fft_numba, FFTType):
+        return
+    if not isinstance(thread_data, types.Array):
+        return
+
+    FFT = fft_numba.fft
+    value_type = NUMBA_FE_TYPES_TO_NUMBA_IR[FFT.value_type]
+    if thread_data.dtype != value_type:
+        return
+
+    code, symbol = compile_fft_execute(
+        FFT,
+        code_type=get_default_code_type(),
+        execute_api="shared_memory" if FFT.execution == "Block" else None,
+    )
+
+    lto = cuda.LTOIR(code.data)
+
+    sig = types.void(types.CPointer(value_type))
+    fft_device_func = declare_cabi_device(symbol, sig, link=lto)
+
+    def impl(_, thread_data, _arg2=None):
+        tptr = get_array_ptr(thread_data)
+        fft_device_func(tptr)
+
+    return impl
+
+
+def ol_fft_type_execute_smem(
+    fft_numba: FFTType,
+    thread_data: types.Array,
+    smem: types.Array,
+):
+    if not isinstance(fft_numba, FFTType):
+        return
+    if not isinstance(thread_data, types.Array):
+        return
+    if not isinstance(smem, types.Array):
+        return
+    FFT = fft_numba.fft
+    value_type = NUMBA_FE_TYPES_TO_NUMBA_IR[FFT.value_type]
+    if smem.dtype != value_type:
+        return
+    if thread_data.dtype != value_type:
+        return
+
+    code, symbol = compile_fft_execute(
+        FFT,
+        code_type=get_default_code_type(),
+        execute_api="register_memory" if FFT.execution == "Block" else None,
+    )
+
+    lto = cuda.LTOIR(code.data)
+
+    value_type = NUMBA_FE_TYPES_TO_NUMBA_IR[FFT.value_type]
+    sig = types.void(types.CPointer(value_type), types.CPointer(value_type))
+    fft_device_func = declare_cabi_device(symbol, sig, link=lto)
+
+    def impl(_, thread_data, smem):
+        tptr = get_array_ptr(thread_data)
+        sptr = get_array_ptr(smem)
+        fft_device_func(tptr, sptr)
+
+    return impl

nvmath/device/cufftdx_numba_cuda_mlir.py

@@ -0,0 +1,79 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from numba_cuda_mlir.compiler import declare_device
+from numba_cuda_mlir.extending import overload_method, typing_registry
+from numba_cuda_mlir.lowering_utilities.type_conversions import to_numba_type
+from numba_cuda_mlir.numba_cuda import types
+from numba_cuda_mlir.numba_cuda.cudadrv.linkable_code import LTOIR
+
+from nvmath.device.common_cuda import get_default_code_type
+from nvmath.device.common_numba_cuda_mlir import HostDescriptorBase, get_array_pointer, register_dummy_type
+from nvmath.device.cufftdx import FFT, compile_fft_execute
+
+from .cufftdx_backend import _FFT_COMPILED_ARGS, _FFT_DEFINITION_ARGS
+
+
+class FFTType(HostDescriptorBase):
+    """
+    Type class associated with the `cufftdx.FFT`.
+    """
+
+    def __init__(self, fft: FFT) -> None:
+        assert isinstance(fft, FFT)
+        super().__init__(fft, _FFT_DEFINITION_ARGS, "FFT")
+
+
+register_dummy_type(FFTType, FFT, _FFT_DEFINITION_ARGS + _FFT_COMPILED_ARGS)
+
+
+@overload_method(FFTType, "execute", strict=False, typing_registry=typing_registry)
+def ol_execute(fft_numba, thread_data, smem=None):
+    if not isinstance(fft_numba, FFTType):
+        return
+    if not isinstance(thread_data, types.Array):
+        return
+    if smem is not None and not isinstance(smem, (types.Array, types.Omitted)):
+        return
+
+    fft = fft_numba.host_descriptor
+
+    value_type = to_numba_type(fft.value_type)
+
+    if thread_data.dtype != value_type:
+        raise ValueError(f"Thread data dtype {thread_data.dtype} does not match value type {value_type}")
+
+    is_smem = smem is not None and not isinstance(smem, types.Omitted)
+    if is_smem and smem.dtype != value_type:
+        raise ValueError(f"Smem dtype {smem.dtype} does not match value type {value_type}")
+
+    execute_api = (
+        "register_memory" if fft.execution == "Block" and is_smem else "shared_memory" if fft.execution == "Block" else None
+    )
+
+    code, symbol = compile_fft_execute(
+        fft,
+        code_type=get_default_code_type(),
+        execute_api=execute_api,
+    )
+
+    if not is_smem:
+        sig = types.void(types.CPointer(value_type))
+        device_func = declare_device(symbol, sig, link=LTOIR(code.data), abi="c")
+
+        def impl(fft_numba, thread_data, smem=None):
+            ptr = get_array_pointer(thread_data)
+            device_func(ptr)
+
+        return impl
+    else:
+        sig = types.void(types.CPointer(value_type), types.CPointer(value_type))
+        device_func = declare_device(symbol, sig, link=LTOIR(code.data), abi="c")
+
+        def impl(fft_numba, thread_data, smem=None):
+            tptr = get_array_pointer(thread_data)
+            sptr = get_array_pointer(smem)
+            device_func(tptr, sptr)
+
+        return impl