nvmath-python 0.9.0__cp314-cp314-win_amd64.whl

nvmath/device/cufftdx.py

@@ -0,0 +1,510 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = ["fft", "FFT", "compile_fft_execute"]
+import warnings
+from functools import cached_property
+from typing import Any
+
+from nvmath.bindings import mathdx
+from nvmath.internal.utils import docstring_decorator
+
+from ._deprecated import deprecated
+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. The default is 1. 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", " "),
+        "execute_api": """\
+A string specifying the signature of the function that handles problems with input in register or in shared memory buffers.
+Could be ``'shared_memory'`` or ``'register_memory'``.""".replace("\n", " "),
+    }
+)
+
+_workspace_deprecation_warning = lambda: warnings.warn(
+    "Using workspaces is deprecated and will be removed in future release.", DeprecationWarning
+)
+
+
+@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`). It does not contain a compiled, ready-to-use,
+    device function until finalized using :meth:`create`.
+
+    .. versionchanged:: 0.7.0
+        `FFT` has replaced `FFTOptions` and `FFTOptionsComplete`.
+
+    Args:
+        size (int): {size}
+
+        precision (str): {precision}
+
+        fft_type (str): {fft_type}
+
+        sm (ComputeCapability): {sm}
+
+        execution (str): {execution}
+
+        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)
+
+    @deprecated("definition is deprecated and may be removed in future versions")
+    def definition(self):
+        """
+        .. deprecated:: 0.7.0
+        """
+        dd = {
+            "size": self.size,
+            "precision": self.precision,
+            "fft_type": self.fft_type,
+            "sm": self.sm,
+            "execution": self.execution,
+            "direction": self.direction,
+            "ffts_per_block": self.ffts_per_block,
+            "elements_per_thread": self.elements_per_thread,
+            "real_fft_options": self.real_fft_options,
+        }
+        return dd
+
+    @deprecated("create is deprecated and may be removed in future versions. Use `functools.partial` instead")
+    def create(self, **kwargs):
+        """
+        Creates a copy of the instance with provided arguments updated.
+
+        .. deprecated:: 0.7.0
+            Please use :py:func:`functools.partial` instead.
+        """
+        code_type = kwargs.pop("code_type", None)
+        if code_type is not None:
+            DeprecationWarning("code_type is deprecated and will be removed in future releases. It is no longer needed.")
+        compiler = kwargs.pop("compiler", None)
+        if compiler is not None:
+            DeprecationWarning("compiler is deprecated and will be removed in future releases. It is no longer needed.")
+        dd = self.definition()
+        dd.update(**kwargs)
+        return FFT(**dd)
+
+    @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 requires_workspace(self):
+        _workspace_deprecation_warning()
+        return False
+
+    @property
+    def workspace_size(self):
+        return self._traits.workspace_size
+
+    @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 is a device function and can not be called on host.")
+
+    @deprecated("Calling MM(...) directly is deprecated, please use MM.execute(...) method instead.")
+    def __call__(self, *args):
+        raise RuntimeError("__call__ is a device function and can not be called on host.")
+
+    @property
+    @deprecated("files is deprecated and is no longer required and will be removed in future releases.")
+    def files(self) -> list:
+        return []
+
+    #
+    # 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)
+
+        self.workspace_size = 0
+        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
+
+
+@docstring_decorator(FFTDX_DOCSTRING, skip_missing=False)
+def fft(*, compiler=None, code_type=None, execute_api=None, **kwargs):
+    """
+    Create an :class:`FFT` object that encapsulates a compiled and ready-to-use FFT
+    device function.
+
+    .. deprecated:: 0.7.0
+
+    Args:
+        size (int): {size}
+
+        precision (str): {precision}
+
+        fft_type (str): {fft_type}
+
+        sm (ComputeCapability): {sm}
+
+        execution (str): {execution}
+
+        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'``.
+
+        compiler: {compiler}
+
+            .. versionchanged:: 0.7.0
+                compiler is no longer needed and does not take effect. Use
+                :py:func:`nvmath.device.compile_fft_execute` to get device
+                function code.
+
+        code_type (CodeType): {code_type}
+
+            .. versionchanged:: 0.7.0
+                code_type should be used by
+                :py:func:`nvmath.device.compile_fft_execute` and no longer
+                needed for numba-cuda usage.
+
+        execute_api (str): {execute_api}
+
+            .. versionchanged:: 0.7.0
+                execute_api should be used by
+                :py:func:`nvmath.device.compile_fft_execute` and no longer
+                needed for numba-cuda usage.
+
+    .. seealso::
+        The attributes of :class:`FFT` provide a 1:1 mapping with the CUDA C++
+        cuFFTDx APIs. For further details, please refer to
+        :cufftdx_doc:`cuFFTDx documentation <index.html>`.
+
+    Examples:
+        Examples can be found in the `nvmath/examples/device
+        <https://github.com/NVIDIA/nvmath-python/tree/main/examples/device>`_ directory.
+    """
+    DeprecationWarning("fft is deprecated and will be removed in future releases. Please use FFT class directly.")
+    if code_type is not None:
+        DeprecationWarning("code_type is deprecated and will be removed in future releases. It is no longer needed.")
+    if compiler is not None:
+        DeprecationWarning("compiler is deprecated and will be removed in future releases. It is no longer needed.")
+    if execute_api is not None:
+        DeprecationWarning("execute_api is deprecated and will be removed in future releases. It is no longer needed.")
+    return FFT(**kwargs)

nvmath/device/cufftdx_backend.py

@@ -0,0 +1,196 @@
+# 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}
+
+
+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, "sm")
+
+
+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"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