nvmath-python 0.9.0__cp314-cp314-win_amd64.whl

nvmath/device/cufftdx_numba.py

@@ -0,0 +1,196 @@
+# 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.cudaimpl import registry as cuda_registry
+from numba.cuda.extending import overload, overload_method, typeof_impl
+from numba.cuda.models import register_model
+from numba.extending import types, utils
+
+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,
+)
+
+_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 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):
+    return ol_fft_numba(fft_numba, _arg1, _arg2)
+
+
+@overload_method(FFTType, "__call__", strict=False)
+def ol_fft_numba_call(fft_numba: FFTType, _arg1, _arg2=None):
+    return ol_fft_numba(fft_numba, _arg1, _arg2)
+
+
+def ol_fft_numba(fft_numba: FFTType, _arg1, _arg2=None):
+    if _arg2 in {None, types.Omitted(None)}:
+        return lambda _, smem, _arg2=None: _fft_type___call__(_, smem)
+    else:
+        return lambda _, thread_data, smem: _fft_type___call__(_, thread_data, smem)
+
+
+# TODO: use overload_method when supported
+def _fft_type___call__(*args):
+    raise Exception("Stub for overloads")
+
+
+@overload(_fft_type___call__, jit_options={"forceinline": True}, strict=False)
+def ol_fft_type___call___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):
+        tptr = get_array_ptr(thread_data)
+        fft_device_func(tptr)
+
+    return impl
+
+
+@overload(_fft_type___call__, jit_options={"forceinline": True}, strict=False)
+def ol_fft_type___call___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
+
+
+# __call__ overload is not supported by numba, however adding this overload
+# kind of activates proper behaviour and works like magic.
+# Issue reference: https://github.com/numba/numba/issues/5885
+# TODO: remove once supported
+@cuda_registry.lower(FFTType, FFTType, types.VarArg(types.Any))
+def method_impl(context, builder, sig, args):
+    typing_context = context.typing_context
+    fnty = typing_context.resolve_value_type(ol_fft_numba_call)
+    sig = fnty.get_call_type(typing_context, sig.args, {})
+    sig = sig.replace(pysig=utils.pysignature(ol_fft_numba_call))
+
+    call = context.get_function(fnty, sig)
+    # Link dependent library
+    context.add_linking_libs(getattr(call, "libs", ()))
+    return call(builder, args)