nvmath-python 1.0.0__cp314-cp314t-win_amd64.whl

nvmath/distributed/process_group.py

@@ -0,0 +1,408 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import math
+import pickle
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from enum import IntEnum
+from typing import Literal, TypeAlias, TypeVar
+
+import numpy as np
+
+__all__ = [
+    "ReductionOp",
+    "ProcessGroup",
+    "MPIProcessGroup",
+    "TorchProcessGroup",
+]
+
+
+class ReductionOp(IntEnum):
+    SUM = 0
+    MIN = 1
+    MAX = 2
+
+
+T = TypeVar("T")
+Reducer: TypeAlias = Callable[[T, T], T]
+
+
+class ProcessGroup(ABC):
+    """A ProcessGroup represents a set of processes collectively running
+    ``nvmath.distributed`` operations. A ProcessGroup implements a small set of
+    communication collectives used by nvmath-python to initialize the distributed
+    runtime and set up distributed math operations. It is not used for compute
+    (each math library uses its chosen communication backend(s), like NCCL or NVSHMEM).
+    """
+
+    @property
+    @abstractmethod
+    def rank(self) -> int:
+        "Index of this process in the group."
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def nranks(self) -> int:
+        """Number of processes in the group."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def broadcast_buffer(self, array: np.ndarray, *, root: int = 0) -> None:
+        """Broadcast an array from one process to every process.
+
+        Args:
+            array: input (on root) and output of the collective.
+
+            root: rank of sending process.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def allreduce_buffer(self, array: np.ndarray, *, op: ReductionOp) -> None:
+        """Allreduce an array.
+
+        Args:
+            array: Input and output of the collective. The function operates in-place.
+
+            op: One of the values from :class:`ReduceOp` enum. Specifies an operation for
+                element-wise reductions.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def allreduce_object(self, obj: T, *, op: Reducer[T]) -> T:
+        """Reduces all Python objects contributed by members of the group. The result
+        is a single reduced object which is returned on every process.
+
+        Args:
+            obj: object contributed by this process.
+
+            op: A Python function that takes two objects and returns a single (reduced)
+                object.
+        """
+        raise NotImplementedError
+
+
+class MPIProcessGroup(ProcessGroup):
+    """ProcessGroup implemented on mpi4py.
+
+    Args:
+        mpi_comm: mpi4py communicator.
+    """
+
+    # Map from mpi4py reduction ops to this module's ReductionOp.
+    _reduction_op_map = {}  # type: ignore
+
+    def __init__(self, mpi_comm):
+        if not MPIProcessGroup._reduction_op_map:
+            from mpi4py import MPI
+
+            MPIProcessGroup._reduction_op_map = {
+                ReductionOp.MIN: MPI.MIN,
+                ReductionOp.MAX: MPI.MAX,
+                ReductionOp.SUM: MPI.SUM,
+            }
+        self._mpi_comm = mpi_comm
+
+    @property
+    def rank(self) -> int:
+        return self._mpi_comm.Get_rank()
+
+    @property
+    def nranks(self) -> int:
+        return self._mpi_comm.Get_size()
+
+    def broadcast_buffer(self, array: np.ndarray, *, root: int = 0) -> None:
+        self._mpi_comm.Bcast(array, root=root)
+
+    def allreduce_buffer(self, array: np.ndarray, *, op: ReductionOp) -> None:
+        mpi_reduction_op = MPIProcessGroup._reduction_op_map[op]
+        from mpi4py import MPI
+
+        self._mpi_comm.Allreduce(MPI.IN_PLACE, array, mpi_reduction_op)
+
+    def allreduce_object(self, obj: T, *, op: Reducer[T]) -> T:
+        return self._mpi_comm.allreduce(obj, op=op)
+
+
+def _init_binomial_tree(rank: int, nranks: int) -> tuple[int | None, list[int]]:
+    """Return parent and children of process with given rank in a binomial
+    tree of nranks processes."""
+    max_height = math.ceil(math.log2(nranks))
+    max_size = 1 << max_height  # max number of nodes in tree (2^max_height)
+
+    # Compute order/height of this subtree (in a full binomial tree of max_size nodes).
+    # Relabel this node for the computation (so 0 is deepest leaf instead of root).
+    p = label = max_size - 1 - rank
+    order = 0
+    while p > 0:
+        if p % 2 == 0:
+            break
+        p //= 2
+        order += 1
+
+    if rank == 0:
+        parent = None
+    else:
+        parent = label + (1 << order)  # parent = label + 2^order
+        parent = max_size - 1 - parent  # convert back to original numbering scheme
+
+    children = []
+    # Max number of children is the order of this subtree.
+    for i in range(order):
+        child = label - (1 << i)  # child = label - 2^i
+        child = max_size - 1 - child  # convert back to original numbering scheme
+        if child <= nranks - 1:
+            children.append(child)
+
+    return parent, children
+
+
+class MsgSizeError(Exception):
+    def __init__(self, size: int):
+        self.size = size
+
+
+class TorchProcessGroup(ProcessGroup):
+    """ProcessGroup implemented on ``torch.distributed``.
+
+    Args:
+        device_id: Device used by the ``torch.distributed`` process group backend.
+
+        torch_process_group: ``torch.distributed`` process group handle (e.g. returned
+            by ``torch.distributed.new_group()``), or None to use the default torch
+            process group.
+    """
+
+    # Map from torch.distributed reduction ops to this module's ReductionOp.
+    _reduction_op_map = {}  # type: ignore
+    # The minimum buffer size (in bytes) for allreduce_object()
+    MIN_ALL_REDUCE_OBJ_BUFFER_SIZE = 128
+
+    def __init__(
+        self,
+        *,
+        device_id: int | Literal["cpu"],
+        torch_process_group=None,
+        logger=None,
+    ):
+        import torch
+        import torch.distributed as dist
+
+        if not isinstance(device_id, int) and device_id != "cpu":
+            raise ValueError("device_id must be int or 'cpu'")
+
+        if not dist.is_initialized():
+            raise RuntimeError("torch.distributed is not initialized")
+
+        if logger is None:
+            logger = logging.getLogger()
+        self.logger = logger
+
+        # Get communication backend used by this torch process group.
+        backend = dist.get_backend(torch_process_group)
+        if backend in ("nccl", "xccl") and device_id == "cpu":
+            raise ValueError(f"This process group is using a GPU-only backend ({backend}), but the provided device_id is 'cpu'")
+        if backend == "mpi":
+            logger.warning(
+                "You're using torch.distributed with MPI backend. In this case it's "
+                "recommended to initialize nvmath.distributed with "
+                "nvmath.distributed.process_group.MPIProcessGroup instead of "
+                "TorchProcessGroup."
+            )
+
+        self._device_id = device_id
+        self._device = torch.device(f"cuda:{device_id}") if device_id != "cpu" else torch.device("cpu")
+        self._torch_process_group = torch_process_group
+
+        if not TorchProcessGroup._reduction_op_map:
+            TorchProcessGroup._reduction_op_map = {
+                ReductionOp.MIN: dist.ReduceOp.MIN,
+                ReductionOp.MAX: dist.ReduceOp.MAX,
+                ReductionOp.SUM: dist.ReduceOp.SUM,
+            }
+
+        if backend == "nccl":
+            # Ensure a NCCL communicator is formed for the whole group. Otherwise if a
+            # parent NCCL communicator is not yet formed when batched P2P operations are
+            # called on a subset of ranks (see allreduce_object() below), they could fail.
+            with torch.cuda.device(self._device):
+                a = torch.ones(1, device=self._device)
+                dist.all_reduce(a, op=dist.ReduceOp.SUM, group=self._torch_process_group)
+                if a[0] != self.nranks:
+                    raise RuntimeError("torch.distributed incorrect all_reduce output")
+
+        # A serialized MsgSizeError must fit in MIN_ALL_REDUCE_OBJ_BUFFER_SIZE
+        assert len(pickle.dumps(MsgSizeError(1 << 30))) < self.MIN_ALL_REDUCE_OBJ_BUFFER_SIZE, "Internal error."
+        # Initial buffer size for allreduce_object (in bytes)
+        self.allreduce_obj_buffer_size = 2048
+
+        # Determine my parent and children for allreduce_object() implementation.
+        # A binomial tree is good for small message sizes.
+        self._parent, self._children = _init_binomial_tree(self.rank, self.nranks)
+
+    @property
+    def device_id(self) -> int | Literal["cpu"]:
+        """Device used by the communication backend of this torch process group."""
+        return self._device_id
+
+    @property
+    def rank(self) -> int:
+        import torch.distributed as dist
+
+        return dist.get_rank(group=self._torch_process_group)
+
+    @property
+    def nranks(self) -> int:
+        import torch.distributed as dist
+
+        return dist.get_world_size(group=self._torch_process_group)
+
+    def broadcast_buffer(self, array: np.ndarray, *, root: int = 0) -> None:
+        import torch
+
+        if self._device_id == "cpu":
+            torch.distributed.broadcast(torch.from_numpy(array), group=self._torch_process_group, group_src=root)
+        else:
+            tensor_gpu = torch.from_numpy(array).to(self._device)
+            torch.distributed.broadcast(tensor_gpu, group=self._torch_process_group, group_src=root)
+            array[:] = tensor_gpu.cpu().numpy()
+
+    def allreduce_buffer(self, array: np.ndarray, *, op: ReductionOp) -> None:
+        import torch
+
+        torch_reduction_op = TorchProcessGroup._reduction_op_map[op]
+        if self._device_id == "cpu":
+            torch.distributed.all_reduce(torch.from_numpy(array), op=torch_reduction_op, group=self._torch_process_group)
+        else:
+            tensor_gpu = torch.from_numpy(array).to(self._device)
+            torch.distributed.all_reduce(tensor_gpu, op=torch_reduction_op, group=self._torch_process_group)
+            array[:] = tensor_gpu.cpu().numpy()
+
+    @property
+    def allreduce_obj_buffer_size(self):
+        """Current buffer size for allreduce_object() (in bytes)"""
+        return self._allreduce_obj_buffer_size
+
+    @allreduce_obj_buffer_size.setter
+    def allreduce_obj_buffer_size(self, value):
+        """**NOTE: Same size needs to be set on every process**"""
+        if not isinstance(value, int) or value < self.MIN_ALL_REDUCE_OBJ_BUFFER_SIZE:
+            raise ValueError(f"Buffer size must be >= {self.MIN_ALL_REDUCE_OBJ_BUFFER_SIZE}")
+        self._allreduce_obj_buffer_size = value
+
+    def allreduce_object(self, obj: T, *, op: Reducer[T]) -> T:
+        import torch
+
+        NUM_TRIES = 5
+        for _ in range(NUM_TRIES):
+            try:
+                if self._device_id != "cpu":
+                    # It's recommended to set torch cuda device for batch_isend_irecv
+                    with torch.cuda.device(self._device):
+                        return self._allreduce_object(obj, op)
+                else:
+                    return self._allreduce_object(obj, op)
+            except MsgSizeError as e:
+                new_size = int(e.size * 1.2)
+                self.logger.warning(
+                    "allreduce_obj_buffer_size is too small for this allreduce: "
+                    f"current size is {self.allreduce_obj_buffer_size} and the required "
+                    f"size is {e.size}. Increasing to {new_size}. "
+                    "You can set the size manually with 'allreduce_obj_buffer_size' property "
+                    " to avoid dynamic resizing."
+                )
+                self.allreduce_obj_buffer_size = new_size
+        raise RuntimeError(
+            f"Could not find buffer size for allreduce_object() after {NUM_TRIES} tries. "
+            "Please set the limit manually with 'allreduce_obj_buffer_size' property."
+        )
+
+    def _allreduce_object(self, obj: T, reducer: Reducer[T]) -> T:
+        import torch
+        import torch.distributed as dist
+
+        def to_tensor(obj, out_tensor):
+            """Serializes Python object into out_tensor buffer.
+            If the tensor buffer is not large enough to accommodate the serialized object,
+            insert a serialized MsgSizeError (with required size) instead."""
+            pickled_obj = pickle.dumps(obj)
+            pickled_obj_arr = np.frombuffer(pickled_obj, dtype=np.uint8).copy()
+            if pickled_obj_arr.nbytes > out_tensor.nbytes:
+                pickled_obj = pickle.dumps(MsgSizeError(pickled_obj_arr.nbytes))
+                pickled_obj_arr = np.frombuffer(pickled_obj, dtype=np.uint8).copy()
+            out_tensor[: pickled_obj_arr.nbytes] = torch.from_numpy(pickled_obj_arr)
+            return out_tensor
+
+        def from_tensor(tensor):
+            """Deserializes Python object from tensor buffer."""
+            return pickle.loads(tensor.cpu().numpy().tobytes())
+
+        num_children = len(self._children)
+
+        if num_children > 0:
+            # 1) Get partially reduced objects from children.
+
+            # Post a recv for each child.
+            # Use batch_isend_irecv because, when using the NCCL backend, individual
+            # (unbatched) point-to-point (P2P) send or receive operations can trigger
+            # the creation of new, separate NCCL communicators for each pair of
+            # communicating processes, which can have significant overhead or use
+            # large amounts of GPU memory. Instead, we want the p2p operations to use
+            # the parent NCCL communicator whenever possible.
+            data = [
+                torch.empty(self._allreduce_obj_buffer_size, dtype=torch.uint8, device=self._device)
+                for _ in range(num_children)
+            ]
+            recv_ops = []
+            for i, child in enumerate(self._children):
+                # The tag is to match the recv against the remote send (in case the sender
+                # is sending multiple different messages).
+                recv_op = dist.P2POp(dist.irecv, data[i], group=self._torch_process_group, tag=2355, group_peer=child)
+                recv_ops.append(recv_op)
+
+            reqs = dist.batch_isend_irecv(recv_ops)
+            for req in reqs:
+                req.wait()
+
+            # 2) Apply the reducer on my object and objects received from children.
+            reduced_obj = obj
+            errors = []
+            for i in range(num_children):
+                child_obj = from_tensor(data[i])
+                if isinstance(child_obj, MsgSizeError):
+                    errors.append(child_obj)
+                elif not errors:
+                    reduced_obj = reducer(reduced_obj, child_obj)
+            if errors:
+                # Max reduction of all received errors to get the maximum required buffer
+                # size for this subtree.
+                reduced_obj = MsgSizeError(max(e.size for e in errors))  # type: ignore
+        else:
+            reduced_obj = obj
+
+        data = [torch.empty(self._allreduce_obj_buffer_size, dtype=torch.uint8, device=self._device)]
+        if self._parent is not None:
+            # Send reduced object to parent.
+            send_op = dist.P2POp(
+                dist.isend, to_tensor(reduced_obj, data[0]), group=self._torch_process_group, tag=2355, group_peer=self._parent
+            )
+            dist.batch_isend_irecv([send_op])[0].wait()
+
+            # Wait for broadcast of reduction from root.
+            dist.broadcast(data[0], group=self._torch_process_group, group_src=0)
+            result = from_tensor(data[0])
+        else:
+            # This is the root: broadcast the result.
+            assert self.rank == 0
+            dist.broadcast(to_tensor(reduced_obj, data[0]), group=self._torch_process_group, group_src=0)
+            result = reduced_obj
+
+        if isinstance(result, MsgSizeError):
+            raise result
+
+        return result

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,189 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = ["FFTDirection", "FFTOptions", "DeviceCallable", "ExecutionCUDA", "ExecutionCPU"]
+
+from dataclasses import dataclass
+from enum import IntEnum
+from logging import Logger
+from typing import Literal
+
+from nvmath._internal import templates
+from nvmath.memory import BaseCUDAMemoryManager
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class ExecutionCUDA(templates.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`.
+    """
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class ExecutionCPU(templates.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`.
+    """
+
+
+@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`.
+
+            .. experimental:: attribute
+
+        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
+            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.
+
+        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"
+    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'.")
+
+
+@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,82 @@
+# 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
+
+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