nvmath-python 1.0.0__cp314-cp314t-win_amd64.whl

nvmath/fft/_helpers.py

@@ -0,0 +1,237 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = ["compile_prolog", "compile_epilog"]
+
+import functools
+
+from nvmath.bindings import cufft as _cufft  # type: ignore
+
+# Currently supported element data types for FFT callback functions.
+VALID_DTYPES = {"float32": "Real", "float64": "DoubleReal", "complex64": "Complex", "complex128": "DoubleComplex"}
+
+# Python <> CUDA magled types.
+TYPE_MAP = {
+    "float32": "f",  # float
+    "float64": "d",  # double
+    "complex64": "6float2",  # float2
+    "complex128": "7double2",  # double2
+}
+
+
+def _check_numba_available():
+    try:
+        import numba  # noqa: F401
+    except ModuleNotFoundError as e:
+        raise RuntimeError("Numba is required to compile FFT prolog and epilog functions.") from e
+
+
+def _compile(
+    function, element_dtype, user_info_dtype, *, name=None, phase=None, representation="ltoir", compute_capability=None
+):
+    _check_numba_available()
+
+    import numba
+    import numba.cuda
+
+    assert phase in ["Load", "Store"], "Internal error."
+
+    if element_dtype not in VALID_DTYPES:
+        raise ValueError(
+            f"The specified operand data type '{element_dtype}' is not currently supported. "
+            f"It must be one of {VALID_DTYPES.keys()}."
+        )
+
+    data_type = getattr(numba.cuda.types.types, element_dtype)
+
+    if user_info_dtype in VALID_DTYPES:
+        info_type = getattr(numba.cuda.types.types, user_info_dtype)
+    elif isinstance(user_info_dtype, numba.types.Type):
+        info_type = user_info_dtype
+    else:
+        raise ValueError(
+            f"The specified user information data type '{user_info_dtype}' is not supported. "
+            f"It must be a Numba custom type or one of {VALID_DTYPES.keys()}."
+        )
+
+    dataptr_type = numba.types.CPointer(data_type)
+    infoptr_type = numba.types.CPointer(info_type)
+    offset_type = numba.types.uint64
+    smemptr_type = dataptr_type
+
+    if phase == "Load":
+        insert = ""
+        return_type = data_type
+        signature = numba.core.typing.signature(return_type, dataptr_type, offset_type, infoptr_type, smemptr_type)
+    else:
+        insert = TYPE_MAP[element_dtype]
+        return_type = numba.types.void
+        signature = numba.core.typing.signature(return_type, dataptr_type, offset_type, data_type, infoptr_type, smemptr_type)
+
+    if name is None:
+        snippet = VALID_DTYPES[element_dtype]
+        length = 16 + len(phase) + len(snippet)
+        # the signature of callbacks differs - the offset argument
+        # changed type from size_t to unsigned long long int
+        offset_dtype = "m" if _cufft.get_version() < 11300 else "y"
+        name = f"_Z{length}cufftJITCallback{phase}{snippet}Pv{offset_dtype}{insert}S_S_"
+
+    if compute_capability is None:
+        compute_capability = numba.cuda.get_current_device().compute_capability
+    else:
+        if not isinstance(compute_capability, str):
+            raise ValueError(
+                f"The compute capability must be specified as a string ('80', '89', ...). "
+                f"The provided value {compute_capability} is invalid."
+            )
+        cc_number = int(compute_capability)
+        compute_capability = (cc_number // 10, cc_number % 10)
+
+    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": name},
+        output=representation,
+        forceinline=False,
+        launch_bounds=None,
+    )[0]
+
+
+compile_prolog = functools.wraps(_compile)(functools.partial(_compile, phase="Load"))
+compile_epilog = functools.wraps(_compile)(functools.partial(_compile, phase="Store"))
+
+SHARED_FFT_HELPER_DOCUMENTATION = {
+    "element_dtype": """The data type of the ``data_in`` argument, one of ``['float32', 'float64', 'complex64',
+            'complex128']``. It must have the same data type as that of the FFT operand for prolog functions or the FFT
+            result for epilog functions.""",
+    "user_info_dtype": """The data type of the ``user_info`` argument. It must be one of ``['float32', 'float64',
+            'complex64', 'complex128']`` or an object of type :class:`numba.types.Type`. The offset is computed based on
+            the memory layout (shape and strides) of the operand (input for prolog, output for epilog). If the user
+            would like to pass additional tensor as `user_info` and access it based on the offset, it is crucial to know
+            memory layout of the operand. Please note, the actual layout of the input tensor may differ from the layout
+            of the tensor passed to fft call. To learn the memory layout of the input or output, please use stateful FFT
+            API and :meth:`nvmath.fft.FFT.get_input_layout` :meth:`nvmath.fft.FFT.get_output_layout` respectively.
+
+            .. note::
+                Currently, in the callback, the position of the element in the input and output operands are described
+                with a single flat offset, even if the original operand is multi-dimensional tensor.
+
+            """,
+    "compute_capability": """The target compute capability, specified as a string (``'80'``, ``'89'``, ...). The
+            default is the compute capability of the current device. For cuFFT in CUDA Toolkit 12.8, 12.9
+            on Blackwell devices (compute capability 12.0), the default
+            compute capability of the current device may fail to compile. As a workaround, ``'50'``
+            can be specified.""",
+}
+
+compile_prolog.__doc__ = """
+    compile_prolog(prolog_fn, element_dtype, user_info_dtype, *, compute_capability=None)
+
+    Compile a Python function to LTO-IR to provide as a prolog function for
+    :func:`~nvmath.fft.fft` and :meth:`~nvmath.fft.FFT.plan`.
+
+    Args:
+        prolog_fn: The prolog function to be compiled to LTO-IR. It must have the signature:
+            ``prolog_fn(data_in, offset, user_info, reserved_for_future_use)``, and it
+            essentially returns transformed ``data_in`` at ``offset``.
+
+        element_dtype: {element_dtype}
+
+        user_info_dtype: {user_info_dtype}
+
+        compute_capability: {compute_capability}
+
+    Returns:
+        The function compiled to LTO-IR as `bytes` object.
+
+    .. seealso::
+        :func:`~nvmath.fft.fft`, :meth:`~nvmath.fft.FFT.plan`,
+        :func:`~nvmath.fft.compile_epilog`.
+
+    Notes:
+        - The user must ensure that the specified argument types meet the requirements
+          listed above.
+""".format(**SHARED_FFT_HELPER_DOCUMENTATION)
+compile_prolog.__name__ = "compile_prolog"
+
+compile_epilog.__doc__ = """
+    compile_epilog(epilog_fn, element_dtype, user_info_dtype, *, compute_capability=None)
+
+    Compile a Python function to LTO-IR to provide as an epilog function for
+    :func:`~nvmath.fft.fft` and :meth:`~nvmath.fft.FFT.plan`.
+
+    Args:
+        epilog_fn: The epilog function to be compiled to LTO-IR. It must have the signature:
+            ``epilog_fn(data_out, offset, data, user_info, reserved_for_future_use)``, and
+            it essentially stores transformed ``data`` into ``data_out`` at ``offset``.
+
+        element_dtype: {element_dtype}
+
+        user_info_dtype: {user_info_dtype}
+
+        compute_capability: {compute_capability}
+
+    Returns:
+        The function compiled to LTO-IR as `bytes` object.
+
+    .. seealso::
+        :func:`~nvmath.fft.fft`, :meth:`~nvmath.fft.FFT.plan`,
+        :func:`~nvmath.fft.compile_prolog`.
+
+    Examples:
+
+        The cuFFT library expects the end user to manage scaling of the outputs, so in order
+        to replicate the ``norm`` option found in `other Python FFT libraries
+        <https://numpy.org/doc/stable/reference/routines.fft.html#normalization>`_ we can
+        define an epilog which performs the scaling.
+
+        >>> import cupy as cp
+        >>> import nvmath
+        >>> import math
+
+        Create the data for a batched 1-D FFT.
+
+        >>> B, N = 256, 1024
+        >>> a = cp.random.rand(B, N, dtype=cp.float64) + 1j * cp.random.rand(B, N, dtype=cp.float64)
+
+        Compute a normalization factor that will create unitary transforms.
+
+        >>> norm_factor = 1.0 / math.sqrt(N)
+
+        Define the epilog function for the FFT.
+
+        >>> def rescale(data_out, offset, data, user_info, unused):
+        ...     data_out[offset] = data * norm_factor
+
+        Compile the epilog to LTO-IR. In a system with GPUs that have different compute
+        capability, the `compute_capability` option must be specified to the
+        `compile_prolog` or `compile_epilog` helpers. Alternatively, the epilog can be
+        compiled in the context of the device where the FFT to which the epilog is provided
+        is executed. In this case we use the current device context, where the operands have
+        been created.
+
+        >>> with cp.cuda.Device():
+        ...     epilog = nvmath.fft.compile_epilog(rescale, "complex128", "complex128")
+
+        Perform the forward FFT, applying the rescaling as a epilog.
+
+        >>> r = nvmath.fft.fft(a, axes=[-1], epilog=dict(ltoir=epilog))
+
+        Test that the fused FFT run result matches the result of other libraries.
+
+        >>> s = cp.fft.fftn(a, axes=[-1], norm="ortho")
+        >>> assert cp.allclose(r, s)
+
+    Notes:
+        - The user must ensure that the specified argument types meet the requirements
+          listed above.
+""".format(**SHARED_FFT_HELPER_DOCUMENTATION)
+compile_epilog.__name__ = "compile_epilog"