nvmath-python 0.5.0__cp313-cp313-win_amd64.whl

nvmath/distributed/_utils.py

@@ -0,0 +1,152 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+__all__ = ["allocate_symmetric_memory", "free_symmetric_memory"]
+
+from logging import Logger
+from collections.abc import Iterable, Sequence
+from types import ModuleType
+from typing import Literal
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch
+
+import numpy as np
+from numpy.typing import DTypeLike
+
+import nvmath.distributed
+from nvmath.internal.utils import device_ctx
+from nvmath.distributed._internal import tensor_wrapper
+
+from ._internal.nvshmem import free as nvshmem_free_wrapper
+
+# Supported packages for tensors backed by symmetric memory.
+_SUPPORTED_PACKAGES = ("cupy", "torch")
+
+
+# TODO: do a general refactor of this function in nvmath
+def calculate_strides(shape: Sequence[int], axis_order: Iterable[int]):
+    """
+    Calculate the strides for the provided shape and axis order.
+    """
+    strides = [0] * len(shape)
+
+    stride = 1
+    for axis in axis_order:
+        strides[axis] = stride
+        stride *= shape[axis]
+
+    return strides
+
+
+def allocate_symmetric_memory(
+    shape: int | Sequence[int],
+    package: ModuleType,
+    *,
+    dtype: DTypeLike | torch.dtype | None = None,
+    axis_order: Literal["C", "F"] | Sequence[int] | None = None,
+    make_symmetric: bool = False,
+    logger: Logger | None = None,
+):
+    """Return uninitialized tensor of given shape and type, allocated from the symmetric
+    heap, on the device on which nvmath.distributed was initialized. The tensor type is
+    determined by the provided package (e.g. cupy, torch).
+    **This is a collective operation and must be called by all processes**.
+    Note that the buffer size must be the same on all processes, or you can
+    use ``make_symmetric=True`` to pad all buffers to the same max size.
+
+    Args:
+        shape: Shape of the tensor to allocate.
+
+        package: Python package determining the tensor type (e.g. cupy, torch).
+
+        dtype: Tensor dtype.
+
+        axis_order: Axis order.
+
+        make_symmetric: If buffer sizes do not match across processes, will allocate
+            the maximum size on every process to ensure the allocation is symmetric.
+
+        logger (logging.Logger): Python Logger object. The root logger will be used if a
+            logger object is not provided.
+    """
+
+    if package.__name__ not in _SUPPORTED_PACKAGES:
+        raise ValueError(f"The package must be one of {_SUPPORTED_PACKAGES}. Got {package}.")
+
+    distributed_ctx = nvmath.distributed.get_context()
+    if distributed_ctx is None:
+        raise RuntimeError("nvmath.distributed has not been initialized")
+
+    device_id = distributed_ctx.device_id
+
+    if isinstance(shape, int):
+        shape = (shape,)
+
+    strides = None
+    if axis_order is not None:
+        if axis_order == "C":
+            strides = calculate_strides(shape, reversed(range(len(shape))))
+        elif axis_order == "F":
+            strides = calculate_strides(shape, range(len(shape)))
+        else:
+            strides = calculate_strides(shape, axis_order)
+
+    if package.__name__ == "cupy":
+        from ._internal.tensor_ifc_cupy import CupyDistributedTensor
+
+        if dtype is None:
+            # This mimics numpy and cupy
+            dtype = np.float64
+
+        dtype = np.dtype(dtype).name  # type: ignore
+        return CupyDistributedTensor.empty(
+            shape, dtype=dtype, device_id=device_id, strides=strides, make_symmetric=make_symmetric, logger=logger
+        ).tensor
+    elif package.__name__ == "torch":
+        from ._internal.tensor_ifc_torch import TorchDistributedTensor
+
+        if dtype is None:
+            import torch
+
+            dtype = torch.get_default_dtype()
+
+        dtype = str(dtype).split(".")[1]
+        return TorchDistributedTensor.empty(
+            shape, dtype=dtype, device_id=device_id, strides=strides, make_symmetric=make_symmetric, logger=logger
+        ).tensor
+
+
+def free_symmetric_memory(*tensors) -> None:
+    """Frees tensors' data buffer where the buffer was allocated on the symmetric heap.
+    Note that this is only meant to be called on tensors returned by
+    ``allocate_symmetric_memory()``.
+
+    **This is a collective operation and must be called by all processes, with tensors
+    in the same order**."""
+    for tensor in tensors:
+        package = _get_tensor_package(tensor)
+        if package not in _SUPPORTED_PACKAGES:
+            raise ValueError(
+                f"The tensor package must be one of {_SUPPORTED_PACKAGES}. Got {type(tensor)} from package {package}."
+            )
+
+    for tensor in tensors:
+        wrapped_tensor = tensor_wrapper.wrap_operand(tensor)
+        if not isinstance(wrapped_tensor.device_id, int):
+            raise ValueError("Tensor must be on GPU symmetric memory")
+        with device_ctx(wrapped_tensor.device_id):
+            nvshmem_free_wrapper(wrapped_tensor.data_ptr)
+
+
+def _get_tensor_package(tensor):
+    if issubclass(tensor.__class__, np.ndarray):
+        return "numpy"
+    module = tensor.__class__.__module__
+    package = module.split(".")[0]
+    return package

nvmath/distributed/fft/__init__.py

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

nvmath/distributed/fft/_configuration.py

@@ -0,0 +1,86 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = ["FFTDirection", "FFTOptions", "Slab"]
+
+from dataclasses import dataclass
+from enum import IntEnum
+from logging import Logger
+from typing import Literal
+from nvmath.bindings import cufftMp  # type: ignore
+
+
+@dataclass
+class FFTOptions:
+    """
+    A data class for providing options to the :class:`FFT` object and the family of wrapper
+    functions :func:`fft` and :func:`ifft`.
+
+    Attributes:
+        fft_type: The type of FFT to perform, available options include ``'C2C'``.
+
+        reshape: Reshape the output distribution to the same slab distribution used by the
+            input. This only applies when using a Slab distribution. The default is `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. The execution methods always block when the input tensor is on
+            the CPU, to ensure that the user doesn't inadvertently use the result before it
+            becomes available. The default is ``"auto"``.
+
+    See Also:
+        :class:`FFT`, :func:`fft` and :func:`ifft`.
+    """
+
+    fft_type: Literal["C2C"] | None = None
+    reshape: bool = True
+    logger: Logger | None = None
+    blocking: Literal[True, "auto"] = "auto"
+
+    def __post_init__(self):
+        valid_fft_types = [None, "C2C"]
+        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.reshape, bool):
+            raise ValueError("The value specified for 'reshape' must be of type bool (True or False).")
+
+        if self.blocking not in (True, "auto"):
+            raise ValueError("The value specified for 'blocking' must be either True or 'auto'.")
+
+
+class FFTDirection(IntEnum):
+    """An IntEnum class specifying the direction of the transform.
+
+    See Also:
+        :meth:`FFT.execute`, :func:`fft`
+    """
+
+    FORWARD = -1
+    INVERSE = 1
+
+
+class Slab(IntEnum):
+    """An IntEnum class to specify a cuFFTMp Slab distribution.
+
+    Given an array of size X * Y * Z distributed over n GPUs, there are two possible slab
+    distributions depending on whether the data is partitioned on the X or Y axis:
+
+    * X axis partitioning: the first X % n GPUs each own (X/n+1) * Y * Z elements and
+      the remaining GPUs each own (X/n) * Y * Z elements.
+
+    * Y axis partitioning: the first Y % n GPUs each own X * (Y/n+1) * Z elements and
+      the remaining GPUs each own X * (Y/n) * Z elements.
+
+    See Also:
+        :class:`FFT`, :func:`fft`
+    """
+
+    X = cufftMp.XtSubFormat.FORMAT_INPLACE
+    Y = cufftMp.XtSubFormat.FORMAT_INPLACE_SHUFFLED