nvmath-python 0.5.0__cp313-cp313-win_amd64.whl

nvmath/distributed/_internal/nvshmem.py

@@ -0,0 +1,243 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = ["initialize", "finalize", "is_initialized", "nvshmem_empty_dlpack", "free", "NvshmemMemoryManager"]
+
+import logging
+import numpy as np
+import cuda.core.experimental as ccx
+
+from nvmath import memory
+from nvmath.bindings import nvshmem  # type: ignore
+from nvmath.internal.utils import device_ctx
+
+# Indicates if this module has initialized NVSHMEM
+_nvshmem_initialized_here = False
+
+
+def initialize(device_id: int, mpi_comm) -> None:
+    """Initialize NVSHMEM runtime if not initialized, otherwise do nothing."""
+
+    global _nvshmem_initialized_here
+
+    rank = mpi_comm.Get_rank()
+    nranks = mpi_comm.Get_size()
+
+    # Here we set the device for NVSHMEM initialization, but we also need to make sure that
+    # a CUDA context has been created before initializing NVSHMEM. We can't rely on
+    # `device_ctx` to do it since it's not guaranteed to make a runtime API call.
+    old_device = ccx.Device()
+    ccx.Device(device_id).set_current()
+
+    try:
+        status = nvshmem.init_status()
+        if status > nvshmem.STATUS_IS_BOOTSTRAPPED:
+            # NVSHMEM is already initialized
+            # NOTE: We assume that the user has passed the same communicator and device_id
+            # on which NVSHMEM was initialized (we have no way of checking here).
+            assert nvshmem.n_pes() == nranks
+            return
+        elif status == nvshmem.STATUS_IS_BOOTSTRAPPED:
+            # A value of 0 indicates an initialization that is similar to when nvshmem_init
+            # is used.
+            nvshmem.hostlib_init_attr(0, 0)
+            assert nvshmem.n_pes() == nranks
+            # Sanity check, can eventually remove
+            assert nvshmem.init_status() > nvshmem.STATUS_IS_BOOTSTRAPPED
+            _nvshmem_initialized_here = True
+            return
+
+        attr = nvshmem.InitAttr()
+        unique_id = nvshmem.UniqueId()
+        # PE 0 queries the unique ID
+        if rank == 0:
+            nvshmem.get_uniqueid(unique_id.ptr)
+        # PE 0 broadcasts the unique ID
+        mpi_comm.Bcast(unique_id._data.view(np.int8), root=0)
+        nvshmem.set_attr_uniqueid_args(rank, nranks, unique_id.ptr, attr.ptr)
+        nvshmem.hostlib_init_attr(nvshmem.Flags.INIT_WITH_UNIQUEID, attr.ptr)
+
+        # sanity check
+        assert nvshmem.init_status() > nvshmem.STATUS_IS_BOOTSTRAPPED
+        _nvshmem_initialized_here = True
+    finally:
+        old_device.set_current()
+
+
+def finalize(device_id: int) -> None:
+    """Finalize NVSHMEM runtime if initialized"""
+    global _nvshmem_initialized_here
+    if not _nvshmem_initialized_here or (nvshmem.init_status() < nvshmem.STATUS_IS_INITIALIZED):
+        return
+
+    with device_ctx(device_id):
+        nvshmem.hostlib_finalize()
+    _nvshmem_initialized_here = False
+
+
+def is_initialized() -> bool:
+    return nvshmem.init_status() >= nvshmem.STATUS_IS_INITIALIZED
+
+
+def _check_initialized():
+    if nvshmem.init_status() < nvshmem.STATUS_IS_INITIALIZED:
+        raise RuntimeError("NVSHMEM is not initialized. Please initialize nvmath.distributed")
+
+
+# Keeps track of memory allocated with nvshmem_empty_dlpack. This is used to report memory
+# leaks and double free errors. This functionality might be handled by nvshmem4py in the
+# future. More precisely, how it works is that when a buffer allocated via
+# nvshmem_empty_dlpack is garbage-collected and found to not have been freed explicitly
+# with the `free` function in this module, we report it.
+_resource_registry = {}
+
+
+class _NvshmemResource(ccx.MemoryResource):
+    def __init__(self, device):
+        self.device = device
+
+    def allocate(self, size, stream=None) -> ccx.Buffer:
+        # NOTE: setting the device is left to the caller
+        ptr = nvshmem.malloc(size)
+        if ptr == 0 and size != 0:
+            raise MemoryError("nvshmem_malloc returned NULL")
+        self.freed = False
+        return ccx.Buffer.from_handle(ptr=ptr, size=size, mr=self)
+
+    def deallocate(self, ptr, size, stream=None, manual=False):
+        # NOTE: setting the device is left to the caller
+        if not manual:
+            if self.freed:
+                return  # this is fine
+            else:
+                # We can't call nvshmem_free when deallocate is triggered by the GC, since
+                # the GC has non-deterministic behavior and nvshmem_free is a collective
+                # call.
+                raise RuntimeError("Symmetric heap memory needs to be deallocated explicitly")
+        if self.freed:
+            raise RuntimeError("This memory resource was already deallocated")
+        nvshmem.free(ptr)
+        self.freed = True
+
+    @property
+    def is_device_accessible(self) -> bool:
+        return True
+
+    @property
+    def is_host_accessible(self) -> bool:
+        return False
+
+    @property
+    def device_id(self) -> int:
+        return self.device.device_id
+
+
+def nvshmem_empty_dlpack(size, device_id, comm, make_symmetric=False, logger=None):
+    """Return uninitialized DLPack buffer of given size in bytes, allocated using
+    nvshmem_malloc (which makes this a *collective* call). Note that the DLPack
+    buffer currently does not include any shape, dtype, or stride information.
+
+    IMPORTANT: device_id must be the one with which NVSHMEM was initialized, and
+    setting the device is left to the caller.
+    """
+
+    _check_initialized()
+
+    global _resource_registry
+
+    logger = logger if logger is not None else logging.getLogger()
+
+    from mpi4py import MPI
+
+    max_size = np.array([-size, size], dtype=np.int64)
+    comm.Allreduce(MPI.IN_PLACE, max_size, MPI.MAX)
+    if -max_size[0] != max_size[1]:
+        # The buffer size is not the same on all processes.
+        if not make_symmetric:
+            raise ValueError(
+                "The buffer size for symmetric memory allocation is not the same on all processes. "
+                "Consider using make_symmetric=True if you have uneven data distribution."
+            )
+        else:
+            logger.info(
+                "Symmetric memory allocator: the buffer will be padded on some processes to "
+                f"satisfy symmetric requirement (make_symmetric=True), size={size} max_size={max_size[1]}."
+            )
+    else:
+        logger.info(f"Symmetric memory allocator: the requested buffer size ({size}) is the same on all processes.")
+    # Sizes are equal or make_symmetric=True.
+    size = max_size[1]
+
+    mem = _NvshmemResource(ccx.Device(device_id))
+    mem_buffer = mem.allocate(size)
+    pointer = mem_buffer._mnff.ptr
+    assert pointer not in _resource_registry
+    _resource_registry[pointer] = (mem, size)
+    return mem_buffer
+
+
+def free(pointer):
+    """This is a *collective* call (invokes nvshmem_free), used to free buffers
+    allocated with nvshmem_empty_dlpack.
+    Setting the device is left to the caller"""
+
+    _check_initialized()
+
+    global _resource_registry
+
+    try:
+        resource, size = _resource_registry.pop(pointer)
+    except KeyError:
+        raise RuntimeError(
+            "Unknown pointer to free. Possible causes:\n"
+            " - This memory was not allocated with nvmath.distributed helpers.\n"
+            " - This memory was already freed. Possible causes:\n"
+            "     - Free was called multiple times on the same tensor.\n"
+            "     - You have multiple tensors sharing the same symmetric memory buffer,\n"
+            "       e.g. as a result of inplace operations, or tensor operations that\n"
+            "       result in views such as slicing."
+        )
+    resource.deallocate(pointer, size, manual=True)
+
+
+class NvshmemMemoryManager(memory.BaseCUDAMemoryManager):
+    """
+    Nvshmem memory allocator.
+
+    Args:
+        device_id: The ID (int) of the device on which memory is to be allocated.
+        logger (logging.Logger): Python Logger object.
+    """
+
+    def __init__(self, device_id, logger):
+        """
+        __init__(device_id)
+        """
+        _check_initialized()
+        self.device_id = device_id
+        self.logger = logger
+
+    def memalloc(self, size):
+        """This is a *collective* call (invokes nvshmem_malloc)"""
+        mem = _NvshmemResource(ccx.Device(self.device_id))
+        mem_buffer = mem.allocate(size)
+        pointer = mem_buffer._mnff.ptr
+        assert pointer not in _resource_registry
+        _resource_registry[pointer] = (mem, size)
+
+        self.logger.debug(
+            f"NvshmemMemoryManager (allocate memory): size = {size}, pointer = {pointer}, device_id = {self.device_id}"
+        )
+
+        return SymmetricMemoryPointer(mem_buffer)
+
+
+class SymmetricMemoryPointer(memory.MemoryPointer):
+    def __init__(self, mem_buffer):
+        super().__init__(mem_buffer._mnff.ptr, mem_buffer.size, finalizer=None)
+        self.mem_buffer = mem_buffer
+
+    def free(self):
+        """This is a *collective* call (invokes nvshmem_free)"""
+        free(self.device_ptr)

nvmath/distributed/_internal/tensor_ifc_cupy.py

@@ -0,0 +1,111 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Interface to seamlessly use distributed Cupy ndarray objects.
+"""
+
+from __future__ import annotations  # allows typehint of class methods to return the self class
+
+__all__ = ["CupyDistributedTensor"]
+
+import math
+
+try:
+    import cupy
+except ImportError:
+    cupy = None
+
+from collections.abc import Sequence
+
+import nvmath.distributed
+from nvmath.internal.tensor_ifc_cupy import CupyTensor
+from nvmath.internal.utils import device_ctx
+from nvmath.bindings import nvshmem  # type: ignore
+from nvmath.distributed._internal.nvshmem import nvshmem_empty_dlpack
+
+from .tensor_ifc_numpy import NumpyDistributedTensor
+
+
+# Most methods aren't redefined, because they simply act on the local array
+class CupyDistributedTensor(CupyTensor):
+    """
+    Tensor wrapper for distributed cupy ndarrays.
+    """
+
+    def __init__(self, tensor):
+        super().__init__(tensor)
+        if nvshmem.ptr(tensor.data.ptr, nvshmem.my_pe()) == 0:
+            raise TypeError(
+                "Operand must be on the symmetric heap. Consider allocating it "
+                "with nvmath.distributed.allocate_symmetric_memory()."
+            )
+
+    @classmethod
+    def empty(cls, shape, device_id="cpu", *, dtype="float32", strides=None, **context) -> CupyDistributedTensor:
+        """
+        Create an empty tensor of the specified shape and data type.
+
+        Note, that the strides, if specified, MUST correspond to a dense (possibly permuted)
+        tensor, otherwise the created tensor may be corrupted.
+        """
+        dtype = CupyTensor.name_to_dtype[dtype]
+
+        from nvmath.distributed._utils import calculate_strides
+
+        ctx = nvmath.distributed.get_context()
+        assert ctx is not None, "nvmath.distributed has not been initialized"
+
+        make_symmetric = context.get("make_symmetric", False)
+        logger = context.get("logger")
+
+        order = "C"
+        if strides is not None:
+            if list(strides) == calculate_strides(shape, reversed(range(len(shape)))):
+                order = "C"
+            elif list(strides) == calculate_strides(shape, range(len(shape))):
+                order = "F"
+            else:
+                raise ValueError("CupyDistributedTensor.empty() only supports 'C' or 'F' order")
+
+        with device_ctx(device_id):
+            # TODO: ideally strides should be set in DLPack, but cuda.core doesn't support
+            # ndarray yet and instead returns a flat buffer.
+            size = math.prod(shape, start=dtype.itemsize)
+            dlpack_buf = nvshmem_empty_dlpack(size, device_id, ctx.communicator, make_symmetric=make_symmetric, logger=logger)
+            tensor = cupy.from_dlpack(dlpack_buf)
+            # Buffer may be padded if make_symmetric=True.
+            tensor = tensor[:size].view(dtype).reshape(shape, order=order)
+            # assert tensor is not a copy
+            assert tensor.base is not None
+
+        return cls(tensor)
+
+    def to(self, device_id, stream_holder) -> NumpyDistributedTensor | CupyDistributedTensor:
+        """
+        In addition to the base class semantics:
+          - Source or target device must be the one used to initialize NVSHMEM on this
+            process. This implies that copy from one CUDA device to another is not allowed.
+          - Memory layout is preserved.
+          - Strides must be dense non-overlapping.
+        """
+        if not (device_id == "cpu" or isinstance(device_id, int)):
+            raise ValueError(f"The device must be specified as an integer or 'cpu', not '{device_id}'.")
+
+        if device_id == "cpu":
+            # NOTE: not using self.numpy() because it doesn't preserve memory layout.
+            np_tensor = NumpyDistributedTensor.empty(self.shape, dtype=self.dtype, strides=self.strides)
+            np_tensor.copy_(self, stream_holder)
+            return np_tensor
+
+        if device_id != self.device_id:
+            raise ValueError("Cannot copy distributed tensor to a different CUDA device")
+
+        with device_ctx(device_id), stream_holder.ctx:
+            return CupyDistributedTensor(cupy.asarray(self.tensor))
+
+    def reshape(self, shape: Sequence[int], *, copy: bool | None = None) -> CupyDistributedTensor:
+        if copy:
+            raise NotImplementedError("reshape with copy=True is not supported for CupyDistributedTensor")
+        return super().reshape(shape, copy=copy)

nvmath/distributed/_internal/tensor_ifc_numpy.py

@@ -0,0 +1,53 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Interface to seamlessly use distributed Numpy ndarray objects.
+"""
+
+from __future__ import annotations  # allows typehint of class methods to return the self class
+
+__all__ = ["NumpyDistributedTensor"]
+
+from nvmath.internal.tensor_ifc_numpy import NumpyTensor
+from nvmath.internal.utils import device_ctx
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    # Can't import CupyDistributedTensor at runtime here due to circular import, but mypy
+    # needs it for type checking.
+    from .tensor_ifc_cupy import CupyDistributedTensor
+
+
+# Most methods aren't redefined, because they simply act on the local array
+class NumpyDistributedTensor(NumpyTensor):
+    """
+    Tensor wrapper for distributed numpy ndarrays.
+    """
+
+    def __init__(self, tensor):
+        super().__init__(tensor)
+
+    def to(self, device_id, stream_holder) -> NumpyDistributedTensor | CupyDistributedTensor:
+        """
+        In addition to the base class semantics:
+          - Target device must be the one used to initialize NVSHMEM on this process.
+          - Memory layout is preserved.
+          - Strides must be dense non-overlapping.
+        """
+        if not (device_id == "cpu" or isinstance(device_id, int)):
+            raise ValueError(f"The device must be specified as an integer or 'cpu', not '{device_id}'.")
+
+        if device_id == "cpu":
+            return NumpyDistributedTensor(self.tensor)
+
+        from .tensor_ifc_cupy import CupyDistributedTensor
+
+        with device_ctx(device_id), stream_holder.ctx:
+            tensor_device = CupyDistributedTensor.empty(
+                self.shape, dtype=self.dtype, device_id=device_id, strides=self.strides, make_symmetric=True
+            )
+            tensor_device.copy_(self, stream_holder)
+            return tensor_device

nvmath/distributed/_internal/tensor_ifc_torch.py

@@ -0,0 +1,108 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Interface to seamlessly use distributed Torch tensor objects.
+"""
+
+from __future__ import annotations  # allows typehint of class methods to return the self class
+
+__all__ = ["TorchDistributedTensor"]
+
+import math
+
+try:
+    import torch
+except ImportError:
+    torch = None  # type: ignore
+
+from collections.abc import Sequence
+
+import nvmath.distributed
+from nvmath.internal.tensor_ifc_torch import TorchTensor
+from nvmath.internal.utils import device_ctx
+from nvmath.bindings import nvshmem  # type: ignore
+from nvmath.distributed._internal.nvshmem import nvshmem_empty_dlpack
+
+
+# Most methods aren't redefined, because they simply act on the local array
+class TorchDistributedTensor(TorchTensor):
+    """
+    TensorHolder for distributed torch tensors.
+    """
+
+    def __init__(self, tensor):
+        super().__init__(tensor)
+        if tensor.device.index is not None and nvshmem.ptr(tensor.data_ptr(), nvshmem.my_pe()) == 0:
+            raise TypeError(
+                "Operand must be on the symmetric heap. Consider allocating it "
+                "with nvmath.distributed.allocate_symmetric_memory()."
+            )
+
+    @classmethod
+    def empty(cls, shape, device_id="cpu", *, dtype="float32", strides=None, **context) -> TorchDistributedTensor:
+        """
+        Create an empty tensor of the specified shape and data type on the specified device
+        (None, 'cpu', or device id).
+
+        Note, that the strides, if specified, should correspond to a dense
+        (possibly permuted) tensor and MUST NOT overlap.
+        Otherwise, the behaviour is not defined.
+        """
+        if device_id == "cpu":
+            return super().empty(shape, device_id, dtype=dtype, strides=strides)
+
+        dtype = TorchTensor.name_to_dtype[dtype]
+
+        ctx = nvmath.distributed.get_context()
+        assert ctx is not None, "nvmath.distributed has not been initialized"
+
+        make_symmetric = context.get("make_symmetric", False)
+        logger = context.get("logger")
+
+        with device_ctx(device_id):
+            size = math.prod(shape, start=dtype.itemsize)
+            # TODO: ideally strides should be set in DLPack, but cuda.core doesn't support
+            # ndarray yet and instead returns a flat buffer.
+            dlpack_buf = nvshmem_empty_dlpack(size, device_id, ctx.communicator, make_symmetric=make_symmetric, logger=logger)
+            tensor = torch.from_dlpack(dlpack_buf)
+            # Buffer may be padded if make_symmetric=True.
+            tensor = tensor[:size]
+            if strides is None:
+                tensor = tensor.view(dtype).view(shape)
+            else:
+                tensor = torch.as_strided(tensor.view(dtype), shape, strides)
+
+        return cls(tensor)
+
+    def to(self, device_id, stream_holder) -> TorchDistributedTensor:
+        """
+        In addition to the base class semantics:
+          - Source or target device must be the one used to initialize NVSHMEM on this
+            process. This implies that copy from one CUDA device to another is not allowed.
+          - Memory layout is preserved.
+          - Strides must be non-overlapping.
+        """
+        if not (device_id == "cpu" or isinstance(device_id, int)):
+            raise ValueError(f"The device must be specified as an integer or 'cpu', not '{device_id}'.")
+
+        if device_id == "cpu" or self.device_id == device_id:
+            with stream_holder.ctx:
+                tensor = self.tensor.to(device=device_id, non_blocking=(device_id != "cpu"))
+            return TorchDistributedTensor(tensor)
+
+        if self.device_id != "cpu" and self.device_id != device_id:
+            raise ValueError("Cannot copy distributed tensor to a different CUDA device")
+
+        with stream_holder.ctx:
+            tensor_device = TorchDistributedTensor.empty(
+                self.shape, device_id=device_id, dtype=self.dtype, strides=self.strides, make_symmetric=True
+            )
+            tensor_device.tensor.copy_(self.tensor, non_blocking=True)
+            return tensor_device
+
+    def reshape(self, shape: Sequence[int], *, copy: bool | None = None) -> TorchDistributedTensor:
+        if copy and self.device_id != "cpu":
+            raise NotImplementedError("reshape with copy=True is not supported for TorchDistributedTensor on GPU")
+        return super().reshape(shape, copy=copy)

nvmath/distributed/_internal/tensor_wrapper.py

@@ -0,0 +1,59 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Entry point to using tensors from different libraries seamlessly.
+"""
+
+__all__ = ["wrap_operand", "copy_"]
+
+import warnings
+from collections.abc import Sequence
+
+from nvmath.internal.tensor_ifc import Tensor, TensorHolder
+from nvmath.internal.tensor_wrapper import infer_tensor_package
+
+from .tensor_ifc_numpy import NumpyDistributedTensor
+
+_TENSOR_TYPES: dict[str, type[TensorHolder]] = {"numpy": NumpyDistributedTensor}
+
+# Optional modules
+try:
+    from .tensor_ifc_cupy import CupyDistributedTensor
+
+    _TENSOR_TYPES["cupy"] = CupyDistributedTensor
+except ImportError:
+    pass
+
+try:
+    from .tensor_ifc_torch import TorchDistributedTensor
+
+    _TENSOR_TYPES["torch"] = TorchDistributedTensor
+except ImportError:
+    pass
+
+
+def wrap_operand(native_operand: Tensor) -> TensorHolder[Tensor]:
+    """
+    Wrap one "native" operand so that package-agnostic API can be used.
+    """
+    if isinstance(native_operand, TensorHolder):
+        msg = (
+            "wrap_operand() is being called unnecessarily because the input is already a TensorHolder. "
+            "Only public facing APIs should call wrap_operand(). "
+            "Internal APIs should assume the operands are TensorHolder already."
+            "Trying to wrap a TensorHolder will become an error in the future."
+        )
+        warnings.warn(msg, DeprecationWarning)
+        return native_operand
+    wrapped_operand = _TENSOR_TYPES[infer_tensor_package(native_operand)](native_operand)
+    return wrapped_operand
+
+
+def copy_(src: Sequence[TensorHolder], dest: Sequence[TensorHolder], stream_holder):
+    """
+    Copy the wrapped operands in dest to the corresponding wrapped operands in src.
+    """
+    for s, d in zip(src, dest, strict=True):
+        d.copy_(s, stream_holder=stream_holder)