nvmath-python 0.9.0__cp314-cp314-win_amd64.whl

nvmath/fft/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from ._configuration import *  # noqa: F403
+from ._helpers import *  # noqa: F403
+from .fft import *  # noqa: F403

nvmath/fft/_configuration.py

@@ -0,0 +1,208 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = ["FFTDirection", "FFTOptions", "DeviceCallable", "ExecutionCUDA", "ExecutionCPU"]
+
+import warnings
+from dataclasses import dataclass, field
+from enum import IntEnum
+from logging import Logger
+from typing import Literal
+
+from nvmath.memory import BaseCUDAMemoryManager
+
+
+@dataclass
+class ExecutionCUDA:
+    """
+    A data class for providing GPU execution options to the :class:`FFT` object and the
+    family of wrapper functions :func:`fft`, :func:`ifft`, :func:`rfft`, and :func:`irfft`.
+
+    Attributes:
+        device_id: CUDA device ordinal (only used if the operand resides on the CPU). The
+            default value is 0.
+
+    .. seealso::
+       :class:`ExecutionCPU`, :class:`FFT`, :func:`fft`, :func:`ifft`, :func:`rfft`, and
+       :func:`irfft`.
+    """
+
+    name: Literal["cuda"] = field(default="cuda", init=False)
+    # If not specified, it defaults to the deprecated FFTOptions.device_id or 0
+    # Keep None as an option to differentiate between the user-provided 0 and the default 0.
+    device_id: int | None = None
+
+
+@dataclass
+class ExecutionCPU:
+    """
+    A data class for providing CPU execution options to the :class:`FFT` object and the
+    family of wrapper functions :func:`fft`, :func:`ifft`, :func:`rfft`, and :func:`irfft`.
+
+    Attributes:
+        num_threads: The number of CPU threads used to execute the plan.
+                     If not specified, defaults to the number of CPU cores available to the
+                     process.
+
+    .. seealso::
+       :class:`ExecutionCUDA`, :class:`FFT`, :func:`fft`, :func:`ifft`, :func:`rfft`, and
+       :func:`irfft`.
+    """
+
+    name: Literal["cpu"] = field(default="cpu", init=False)
+    num_threads: int | None = None
+
+
+@dataclass
+class FFTOptions:
+    """
+    A data class for providing options to the :class:`FFT` object and the family of wrapper
+    functions :func:`fft`, :func:`ifft`, :func:`rfft`, and :func:`irfft`.
+
+    Attributes:
+        fft_type: The type of FFT to perform, available options include ``'C2C'``,
+            ``'C2R'``, and ``'R2C'``. The default is ``'C2C'`` for complex input and
+            ``'R2C'`` for real input.
+
+        inplace: Specify if the operation is in-place (`True` or `False`). The operand is
+            overwritten by the result if ``inplace`` is `True`. The default is `False`.
+
+        last_axis_parity: For complex-to-real FFT (corresponding to ``fft_type='C2R'``),
+            specify whether the size of the last axis in the result should be even or odd.
+            The even size is calculated as :math:`2 * (m - 1)`, where :math:`m` is the the
+            size of the last axis of the operand, and the odd size is calculated as :math:`2
+            * (m - 1) + 1`. The specified value should be either ``'even'`` or ``'odd'``,
+            with the default being ``'even'``.
+
+        result_layout: The layout to use for the result, either ``'natural'`` or
+            ``'optimized'``. For the ``'natural'`` option, the result layout is the same as
+            that of the operand. The default is ``'optimized'``, which generally provides
+            much better performance and should be used if the user doesn't care about the
+            result layout matching the operand layout. However in rare cases, depending on
+            the device type, shape and strides of the operand, and the FFT dimensions, the
+            ``'natural'`` layout may perform better. This option is ignored if ``inplace``
+            is specified to be True.
+
+        device_id: CUDA device ordinal (used if the operand resides on the CPU).
+            Device 0 will be used if not specified.
+
+            .. deprecated:: 0.2.0
+                The ``device_id`` should be specified as ``execution`` option,
+                see :attr:`ExecutionCUDA.device_id`.
+
+
+        logger (logging.Logger): Python Logger object. The root logger will be used if a
+            logger object is not provided.
+
+        blocking: A flag specifying the behavior of the execution functions and methods,
+            such as :func:`fft` and :meth:`FFT.execute`. When ``blocking`` is `True`, the
+            execution methods do not return until the operation is complete. When
+            ``blocking`` is ``"auto"``, the methods return immediately when the input tensor
+            is on the GPU and ``execution`` is set to 'CUDA'. The execution methods always
+            block when the input tensor is on the CPU or ``execution`` is specified to 'CPU'
+            to ensure that the user doesn't inadvertently use the result before it becomes
+            available. The default is ``"auto"``.
+
+        allocator: An object that supports the :class:`BaseCUDAMemoryManager` protocol, used
+            to draw device memory. If an allocator is not provided, a memory allocator from
+            the library package will be used (:func:`torch.cuda.caching_allocator_alloc` for
+            PyTorch operands, :func:`cupy.cuda.alloc` otherwise). Currently, CPU execution
+            does not allow specifying a memory pool.
+
+    .. seealso::
+        :class:`ExecutionCUDA`, :class:`ExecutionCPU`, :class:`FFT`, :func:`fft`,
+        :func:`ifft`, :func:`rfft`, and :func:`irfft`.
+    """
+
+    fft_type: Literal["C2C", "C2R", "R2C"] | None = None
+    inplace: bool = False
+    last_axis_parity: Literal["even", "odd"] = "even"
+    result_layout: Literal["natural", "optimized"] = "optimized"
+    device_id: int | None = None
+    logger: Logger | None = None
+    blocking: Literal[True, "auto"] = "auto"
+    allocator: BaseCUDAMemoryManager | None = None
+
+    def __post_init__(self):
+        valid_fft_types = [None, "C2C", "C2R", "R2C"]
+        if self.fft_type not in valid_fft_types:
+            raise ValueError(f"The value specified for 'fft_type' must be one of {valid_fft_types}.")
+
+        if not isinstance(self.inplace, bool):
+            raise ValueError("The value specified for 'inplace' must be of type bool (True or False).")
+
+        valid_last_axis_parity = ["even", "odd"]
+        if self.last_axis_parity not in valid_last_axis_parity:
+            raise ValueError(f"The value specified for 'last_axis_parity' must be one of {valid_last_axis_parity}.")
+
+        valid_result_layout_options = ["natural", "optimized"]
+        if self.result_layout not in valid_result_layout_options:
+            raise ValueError(f"The value specified for 'result_layout' must be one of {valid_result_layout_options}.")
+
+        if self.blocking not in (True, "auto"):
+            raise ValueError("The value specified for 'blocking' must be either True or 'auto'.")
+
+        if self.device_id is not None:
+            msg = "`FFTOptions.device_id` is deprecated. Use `ExecutionCUDA.device_id` instead."
+            warnings.warn(msg, DeprecationWarning)
+
+
+@dataclass
+class DeviceCallable:
+    """A data class capturing LTO-IR callables.
+
+    Attributes:
+        ltoir: A device-callable function in LTO-IR format, which can be provided as as
+            either as a :class:`bytes` object or as a pointer to the LTO-IR as Python
+            :class:`int`.
+
+        size: The size of the LTO-IR callable. If not specified and a :class:`bytes` object
+            is passed for ``ltoir``, the size is calculated from it. If a pointer is
+            provided for the LTO-IR, `size` must be specified.
+
+        data:  A device pointer to user data used in the callback. The default is None,
+            which means a null pointer will be used in the callback.
+
+    .. seealso::
+        :meth:`FFT.plan`, :func:`fft`, :func:`ifft`, :func:`rfft`, and :func:`irfft`.
+    """
+
+    ltoir: int | bytes | None = None
+    size: int | None = None
+    data: int | None = None
+
+    def __post_init__(self):
+        if self.ltoir is None:
+            return
+        if not isinstance(self.ltoir, int | bytes):
+            raise ValueError(
+                "The LTO-IR code must be provided as a bytes object or as a Python int "
+                "representing the pointer to the LTO-IR code."
+            )
+        if isinstance(self.ltoir, int) and self.size is None:
+            raise ValueError(
+                "The size of the LTO-IR code specified as a pointer must be explicitly provided using the 'size' option."
+            )
+        if isinstance(self.ltoir, bytes):
+            self.size = len(self.ltoir)
+        if not isinstance(self.size, int):
+            raise ValueError(f"Invalid size value: {self.size}.")
+        if self.data is None:
+            self.data = 0
+        if not isinstance(self.data, int):
+            raise ValueError("The 'data' attribute must be a Python int representing a device pointer to user data.")
+
+
+KeywordArgType = dict
+
+
+class FFTDirection(IntEnum):
+    """An IntEnum class specifying the direction of the transform.
+
+    .. seealso::
+        :meth:`FFT.execute`, :func:`fft`
+    """
+
+    FORWARD = -1
+    INVERSE = 1

nvmath/fft/_exec_utils.py

@@ -0,0 +1,109 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+
+from nvmath.bindings import cufft  # type: ignore
+from nvmath.bindings._internal import utils as _bindings_utils  # type: ignore
+from nvmath.fft import _configuration
+
+IS_EXEC_GPU_AVAILABLE = False
+IS_EXEC_CPU_AVAILABLE = False
+NUM_THREADS_DEFAULT = None
+
+
+def _get_num_threads_default():
+    global NUM_THREADS_DEFAULT
+    if NUM_THREADS_DEFAULT is None:
+        if os.name == "posix":
+            num_threads = len(os.sched_getaffinity(0))
+        else:
+            # `sched_getaffinity` is not supported on Windows
+            num_threads = os.cpu_count() or 1
+        num_threads = max(1, num_threads // 2)
+        NUM_THREADS_DEFAULT = num_threads
+    return NUM_THREADS_DEFAULT
+
+
+def _check_init_cufft():
+    global IS_EXEC_GPU_AVAILABLE
+    if not IS_EXEC_GPU_AVAILABLE:
+        try:
+            cufft.get_version()
+        except (
+            _bindings_utils.FunctionNotFoundError,
+            _bindings_utils.NotSupportedError,
+            RuntimeError,
+        ) as e:
+            raise RuntimeError(
+                "The FFT CUDA execution is not available. "
+                "Please check if CUDA toolkit and cuFFT are installed and visible to nvmath."
+            ) from e
+
+        IS_EXEC_GPU_AVAILABLE = True
+
+
+def _check_init_fftw():
+    global IS_EXEC_CPU_AVAILABLE
+    if not IS_EXEC_CPU_AVAILABLE:
+        from nvmath.bindings.nvpl import fft as _fftw
+        from nvmath.bindings.nvpl._internal import fft as _internal_fftw
+
+        env_lib_name = os.environ.get("NVMATH_FFT_CPU_LIBRARY")
+
+        if env_lib_name:
+            _internal_fftw._set_lib_so_names((env_lib_name,))
+        try:
+            _internal_fftw._inspect_function_pointers()
+            loaded_lib_name = _internal_fftw._get_current_lib_so_name()
+        except (
+            _bindings_utils.FunctionNotFoundError,
+            _bindings_utils.NotSupportedError,
+            RuntimeError,
+        ) as e:
+            raise RuntimeError(
+                "The FFT CPU execution is not available, because no FFTW-compatible "
+                "library was found. Please make sure to install optional dependency required "
+                "for FFT execution (NVPL on aarch64 platforms or MKL on x86_64 platforms). "
+                "Users can use other FFTW3-compatible libraries. In order to do that, please "
+                "specify the library name as "
+                "`NVMATH_FFT_CPU_LIBRARY=some_fftw3_compatible_lib.so` and make sure "
+                "the library can be found within current `LD_LIBRARY_PATH`"
+            ) from e
+
+        if env_lib_name:
+            assert env_lib_name == loaded_lib_name
+
+        lib_nvpl = "libnvpl_fftw.so"
+        if not loaded_lib_name.startswith(lib_nvpl):
+            _fftw.init_threads_float()
+            _fftw.init_threads_double()
+
+        IS_EXEC_CPU_AVAILABLE = True
+
+
+def _cross_setup_execution_and_options(
+    options: _configuration.FFTOptions,
+    execution: _configuration.ExecutionCUDA | _configuration.ExecutionCPU,
+):
+    if execution.name == "cpu":
+        _check_init_fftw()
+        if options.device_id is not None:
+            raise ValueError("The 'device_id' is not a valid option when 'execution' is specified to be 'cpu'.")
+        if execution.num_threads is None:
+            execution.num_threads = _get_num_threads_default()
+        if not isinstance(execution.num_threads, int) or execution.num_threads <= 0:
+            raise ValueError("The 'num_threads' must be a positive integer")
+    else:
+        assert execution.name == "cuda"
+        _check_init_cufft()
+        if execution.device_id is None:
+            execution.device_id = options.device_id or 0
+        elif options.device_id is not None and execution.device_id != options.device_id:
+            raise ValueError(
+                f"Got conflicting 'device_id' passed in 'execution' ({execution.device_id}) "
+                f"and 'options' ({options.device_id}). It should be passed as 'execution' "
+                f"configuration only."
+            )
+    return options, execution

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"