nvmath-python 0.9.0__cp314-cp314-win_amd64.whl

nvmath/device/types.py

@@ -0,0 +1,147 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = [
+    "Complex",
+    "Vector",
+    "complex32",
+    "complex64",
+    "complex128",
+    "half2",
+    "half4",
+    "np_float16x2",
+    "np_float16x4",
+    "REAL_NP_TYPES",
+    "INT_NP_TYPES",
+]
+
+import warnings
+
+import numpy as np
+
+from ._deprecated import deprecated
+
+np_float16x2 = np.dtype([("x", np.float16), ("y", np.float16)], align=True)
+np_float16x4 = np.dtype([("x", np.float16), ("y", np.float16), ("z", np.float16), ("w", np.float16)], align=True)
+
+REAL_NP_TYPES: list = [np.float16, np.float32, np.float64]
+INT_NP_TYPES: list = [np.int8, np.int16, np.int32, np.int64, np.uint8, np.uint16, np.uint32, np.uint64]
+
+_alignment_warning_msg = (
+    "You are using the host counterpart of a dtype that is under-aligned "
+    "compared to what the device function expects. In most cases this will "
+    "work, since device memory is typically allocated with at least 256-byte "
+    "alignment. For details, please review the alignment guidelines at "
+    "https://nvidia.com/docs/nvmath/alignment"
+)
+
+
+class Complex:
+    """
+    Complex type that can be used to represent complex numbers both
+    on host and device side. Numpy does not provide a built-in complex type with
+    16-bit real and imaginary parts, so we define our own dtype for that case.
+    For 32-bit and 64-bit complex numbers, we can use the built-in numpy dtypes.
+    However on device side we expect those types to be aligned to the full size
+    of the complex type, so the array defined on host and device side will have
+    different type and alignment. :py:const:`np_float16x2`,
+    :py:const:`numpy.dtype(numpy.complex64)` and
+    :py:const:`numpy.dtype(numpy.complex128)` are the host side dtypes and
+    :py:const:`float16x2_type`, :py:const:`float32x2_type` and
+    :py:const:`float64x2_type` are the device side types.
+    """
+
+    def __init__(self, real_dtype):
+        self._real_dtype = real_dtype
+
+    @property
+    def real_dtype(self):
+        return self._real_dtype
+
+    @property
+    def dtype(self):
+        warnings.warn(_alignment_warning_msg, UserWarning, stacklevel=2)
+        if self._real_dtype == np.float16:
+            return np_float16x2
+        elif self._real_dtype == np.float32:
+            return np.dtype(np.complex64)
+        assert self._real_dtype == np.float64
+        return np.dtype(np.complex128)
+
+    @property
+    def _numba_type(self):
+        from .vector_types_numba import float16x2_type, float32x2_type, float64x2_type
+
+        if self.real_dtype == np.float16:
+            return float16x2_type
+        if self.real_dtype == np.float32:
+            return float32x2_type
+        assert self.real_dtype == np.float64
+        return float64x2_type
+
+    @property
+    @deprecated("This is a numba fallback behavior and will be removed in future releases, please use numba types directly")
+    def make(self):
+        return self._numba_type.make
+
+
+complex32 = Complex(np.float16)
+complex64 = Complex(np.float32)
+complex128 = Complex(np.float64)
+
+
+class Vector:
+    """
+    Vector type that can be used to represent vector numbers both
+    on host and device side. Host side representation uses numpy structured
+    dtypes to represent the vector components, while device side representation
+    uses custom numba types. This difference is necessary because device
+    functions expect alignment of the vector types to be the same as the size of
+    the vector, which is not the case for numpy structured dtypes.
+    :py:const:`np_float16x2` and :py:const:`np_float16x4` are the host side
+    dtypes and :py:const:`float16x2_type` and :py:const:`float16x4_type` are the
+    device side types.
+    """
+
+    def __init__(self, real_dtype, size):
+        if size not in (2, 4):
+            raise ValueError(f"Unsupported vector size {size}, only 2 and 4 are supported")
+        if real_dtype != np.float16:
+            raise ValueError(f"Unsupported vector real dtype {real_dtype}, only float16 is supported")
+        self._real_dtype = real_dtype
+        self._size = size
+
+    @property
+    def real_dtype(self):
+        return self._real_dtype
+
+    @property
+    def size(self):
+        return self._size
+
+    @property
+    def dtype(self):
+        warnings.warn(_alignment_warning_msg, UserWarning, stacklevel=2)
+        if self._size == 2:
+            return np_float16x2
+        assert self._size == 4
+        return np_float16x4
+
+    @property
+    def _numba_type(self):
+        from .vector_types_numba import float16x2_type, float16x4_type
+
+        if self._size == 2:
+            return float16x2_type
+        assert self._size == 4
+        return float16x4_type
+
+    @property
+    @deprecated("This is a numba fallback behavior and will be removed in future releases, please use numba types directly")
+    def make(self):
+        return self._numba_type.make
+
+
+half2 = Vector(np.float16, 2)
+half4 = Vector(np.float16, 4)

nvmath/device/vector_types_numba.py

@@ -0,0 +1,203 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = [
+    "float16x2",
+    "float16x4",
+    "float32x2",
+    "float64x2",
+    "float16x2_type",
+    "float16x4_type",
+    "float32x2_type",
+    "float64x2_type",
+]
+
+import typing
+
+from llvmlite import ir
+from numba.core import types
+from numba.core.typeconv import Conversion
+from numba.core.typing.templates import AttributeTemplate
+from numba.cuda.cudadecl import registry as cuda_registry
+from numba.cuda.cudaimpl import lower_attr as cuda_lower_attr
+from numba.cuda.extending import as_numba_type, lower_builtin, lower_cast, overload, register_model, type_callable
+from numba.extending import models
+
+
+def make_vector(
+    float_bitwidth: int, numba_float_type: types.scalars.Float, llvm_float_type, vector_length: typing.Literal[2, 4]
+):
+    assert vector_length == 2 or vector_length == 4
+    vector_bitwidth = vector_length * float_bitwidth
+
+    # User visible types and functions
+    class vector:
+        def __init__(self, *args):
+            pass
+
+    def create(input):
+        pass
+
+    # FE types
+    class vectorType(types.Number):
+        def __init__(self):
+            super().__init__(name=f"vector({numba_float_type.name}x{vector_length})")
+            self.dtype = numba_float_type
+            self.count = vector_length
+            self.bitwidth = vector_bitwidth
+            self.make = vector
+
+        def can_convert_from(self, typingctx, other):
+            if vector_length == 2 and isinstance(other, types.Complex):
+                if other.bitwidth == vector_bitwidth:
+                    return Conversion.exact
+                elif other.bitwidth < vector_bitwidth:
+                    return Conversion.safe
+                elif other.bitwidth > vector_bitwidth:
+                    return Conversion.unsafe
+
+        def can_convert_to(self, typingctx, other):
+            if vector_length == 2 and isinstance(other, types.Complex):
+                if other.bitwidth == vector_bitwidth:
+                    return Conversion.exact
+                elif other.bitwidth < vector_bitwidth:
+                    return Conversion.unsafe
+                elif other.bitwidth > vector_bitwidth:
+                    return Conversion.safe
+
+        def __name__(self):
+            return f"vector({numba_float_type.name}x{vector_length})"
+
+    vector_type = vectorType()
+
+    types.number_domain = types.number_domain | frozenset([vector_type])
+
+    # Type inference (Python --> FE type)
+    as_numba_type.register(vector, vector_type)
+
+    # How to build a vectorType from individual elements
+    @type_callable(vector)
+    def type(context):
+        def typer(x, y=None, z=None, w=None):
+            if (
+                (vector_length == 2 and isinstance(x, types.Complex) and y is None and z is None and w is None)
+                or (vector_length == 2 and all(isinstance(v, types.Float) for v in [x, y]) and z is None and w is None)
+                or (vector_length == 4 and all(isinstance(v, types.Float) for v in [x, y, z, w]))
+            ):
+                return vector_type
+
+        return typer
+
+    # How to represent vectorType in Memory?
+    @register_model(vectorType)
+    class vectorTypeModel(models.PrimitiveModel):
+        def __init__(self, dmm, fe_type):
+            be_type = ir.IntType(vector_bitwidth)
+            super().__init__(dmm, fe_type, be_type)
+
+    # How to build a vectorType from individual float values?
+    @lower_builtin(vector, *(vector_length * [types.Float]))  # FIXME
+    def vector_ctor(context, builder, sig, args):
+        typ = sig.return_type
+        args = list(args)
+        assert len(args) == vector_length
+        for i in range(vector_length):
+            args[i] = context.cast(builder, args[i], sig.args[i], numba_float_type)
+
+        vt = ir.VectorType(llvm_float_type, vector_length)
+        val = ir.Constant(vt, ir.Undefined)
+
+        for i in range(vector_length):
+            val = builder.insert_element(val, args[i], context.get_constant(types.int32, i))
+
+        return builder.bitcast(val, context.get_value_type(typ))
+
+    # How to access .x, .y, .z, .w
+    @cuda_registry.register_attr
+    class complex_attrs(AttributeTemplate):
+        key = vectorType
+
+        def resolve_x(self, mod):
+            return numba_float_type
+
+        def resolve_y(self, mod):
+            return numba_float_type
+
+        def resolve_z(self, mod):
+            return numba_float_type
+
+        def resolve_w(self, mod):
+            return numba_float_type
+
+    def complex_get(context, builder, typ, val, i):
+        vt = ir.VectorType(llvm_float_type, vector_length)
+        vec = builder.bitcast(val, vt)
+        index = context.get_constant(types.int32, i)
+        return builder.extract_element(vec, index)
+
+    # This cannot be done with a loop
+    @cuda_lower_attr(vectorType, "x")
+    def complex_get_x(context, builder, typ, val):
+        return complex_get(context, builder, typ, val, 0)
+
+    @cuda_lower_attr(vectorType, "y")
+    def complex_get_y(context, builder, typ, val):
+        return complex_get(context, builder, typ, val, 1)
+
+    if vector_length == 4:
+
+        @cuda_lower_attr(vectorType, "z")
+        def complex_get_z(context, builder, typ, val):
+            return complex_get(context, builder, typ, val, 2)
+
+        @cuda_lower_attr(vectorType, "w")
+        def complex_get_w(context, builder, typ, val):
+            return complex_get(context, builder, typ, val, 3)
+
+    # Conversions from types.Complex
+    if vector_length == 2:
+
+        @overload(create)
+        def make_complex(cplx):
+            if isinstance(cplx, types.Complex):
+
+                def impl(cplx):
+                    return vector(cplx.real, cplx.imag)
+
+                return impl
+
+        @lower_cast(types.Complex, vector_type)
+        def np_complex_to_complex(context, builder, fromty, toty, val):
+            src = context.make_complex(builder, fromty, value=val)
+            src_float_ty = fromty.underlying_float
+            ctor_sig = toty(src_float_ty, src_float_ty)
+            ctor_args = (src.real, src.imag)
+            return vector_ctor(context, builder, ctor_sig, ctor_args)
+
+        @lower_cast(vector_type, types.Complex)
+        def complex_to_np_complex(context, builder, fromty, toty, val):
+            dst_float_ty = toty.underlying_float
+            x = complex_get(context, builder, fromty, val, 0)
+            y = complex_get(context, builder, fromty, val, 1)
+            x = context.cast(builder, x, numba_float_type, dst_float_ty)
+            y = context.cast(builder, y, numba_float_type, dst_float_ty)
+            vt = context.get_value_type(toty)
+            val = ir.Constant(vt, ir.Undefined)
+            val = builder.insert_value(val, x, [0])
+            val = builder.insert_value(val, y, [1])
+            return val
+
+    vector.__name__ = f"float{float_bitwidth}x{vector_length}"
+    vector.__doc__ = f"""
+    float{float_bitwidth}x{vector_length}({"x, y" if vector_length == 2 else "x, y, z, w"})
+
+    Create a Numba compliant vector object for float{float_bitwidth} with vector length {vector_length}.
+    """
+    return vector, vector_type
+
+
+float16x2, float16x2_type = make_vector(16, types.float16, ir.IntType(16), 2)
+float16x4, float16x4_type = make_vector(16, types.float16, ir.IntType(16), 4)
+float32x2, float32x2_type = make_vector(32, types.float32, ir.FloatType(), 2)
+float64x2, float64x2_type = make_vector(64, types.float64, ir.DoubleType(), 2)

nvmath/distributed/__init__.py

@@ -0,0 +1,205 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import sys
+
+if sys.platform != "linux":
+    raise ImportError("nvmath.distributed is only supported on Linux.")
+
+import atexit
+from collections.abc import Sequence
+from dataclasses import dataclass
+from threading import Lock
+from typing import TYPE_CHECKING, Any, Literal
+
+import numpy as np
+
+try:
+    from cuda.core import Device
+except ImportError:
+    from cuda.core.experimental import Device
+
+if TYPE_CHECKING:
+    import mpi4py
+
+from nvmath.distributed import fft, linalg, reshape  # noqa: E402
+from nvmath.internal.utils import device_ctx
+
+from . import distribution
+from ._internal import nvshmem
+from ._utils import allocate_symmetric_memory, free_symmetric_memory
+from .process_group import MPIProcessGroup, ProcessGroup, TorchProcessGroup
+
+__all__ = [
+    "initialize",
+    "finalize",
+    "get_context",
+    "ProcessGroup",
+    "MPIProcessGroup",
+    "TorchProcessGroup",
+    "allocate_symmetric_memory",
+    "free_symmetric_memory",
+    "distribution",
+    "fft",
+    "linalg",
+    "reshape",
+]
+
+_initialize_mutex = Lock()
+_atexit_registered = False
+_ctx = None
+
+
+@dataclass(frozen=True)
+class DistributedContext:
+    """
+    Context of initialized ``nvmath.distributed`` runtime.
+
+    Attributes:
+        device_id: CUDA device ID associated with the distributed runtime
+            on this process.
+
+        process_group: ``nvmath.distributed`` participating processes.
+
+        nvshmem_available: True if NVSHMEM backend was selected at initialization.
+
+        nccl_comm: nccl4py communicator if NCCL backend was selected at
+            initialization, None otherwise.
+    """
+
+    device_id: int
+    process_group: ProcessGroup
+    nvshmem_available: bool
+    nccl_comm: Any | None
+
+
+def initialize(
+    device_id: int,
+    process_group: ProcessGroup | mpi4py.MPI.Comm,
+    backends: Sequence[Literal["nvshmem", "nccl"]],
+) -> None:
+    """Initialize ``nvmath.distributed`` runtime. This is required before any distributed
+    operations can be performed. **Note that this is a collective operation and must be
+    called by all processes.**
+
+    If the runtime is already initialized this function will raise an error. If you need
+    to reinitialize the runtime (for example with different backends) you have to finalize
+    it first.
+
+    Note: NCCL doesn't allow assigning more than one process to the same GPU.
+
+    Args:
+        device_id: CUDA device ID to associate with the ``nvmath.distributed`` runtime
+            on this process.
+
+        process_group: ProcessGroup (or mpi4py communicator) specifying the participating
+            processes. This is used for setup and not for communication during compute.
+
+        backends: Communication backends to use in distributed computations. Valid values
+            are "nvshmem" and "nccl". Note that specific libraries (cuFFTMp, cuBLASMp, ...)
+            have specific required backends.
+    """
+    if not isinstance(device_id, int):
+        raise TypeError(
+            "The device ID used to initialize the nvmath.distributed module "
+            f"must be an integer. The provided device ID is {device_id}."
+        )
+
+    valid_backends = ("nvshmem", "nccl")
+    for backend in backends:
+        if backend not in valid_backends:
+            raise ValueError(f"backend must be one of {valid_backends}, got {backend}")
+
+    if len(backends) == 0:
+        raise ValueError(f"Need to specify at least one of {valid_backends} communication backends")
+
+    if process_group is None:
+        raise ValueError("process_group cannot be None")
+
+    with _initialize_mutex:
+        global _atexit_registered, _ctx
+
+        if _ctx is not None:
+            raise RuntimeError("nvmath.distributed has already been initialized")
+
+        if not _atexit_registered:
+            atexit.register(finalize)
+            _atexit_registered = True
+
+        if not isinstance(process_group, ProcessGroup):
+            valid_process_group = False
+            if "mpi4py" in sys.modules:
+                import mpi4py.MPI
+
+                if isinstance(process_group, mpi4py.MPI.Comm):
+                    process_group = MPIProcessGroup(process_group)
+                    valid_process_group = True
+            if not valid_process_group:
+                raise TypeError(
+                    f"Unrecognized process group type ({process_group}). "
+                    "Need nvmath.distributed.ProcessGroup or mpi4py communicator."
+                )
+
+        assert isinstance(process_group, ProcessGroup)  # for type-checker
+        rank = process_group.rank
+        nranks = process_group.nranks
+
+        # Set the device for NVSHMEM and NCCL initialization, but also need to make sure
+        # that a CUDA context has been created. We can't rely on `device_ctx` to do it
+        # since it's not guaranteed to make a runtime API call.
+        old_device = Device()
+        device = Device(device_id)
+        device.set_current()
+
+        try:
+            nvshmem_available = False
+            if "nvshmem" in backends:
+                nvshmem.initialize(device_id, process_group)
+                nvshmem_available = True
+
+            nccl_comm = None
+            if "nccl" in backends:
+                # Don't want to import this at the global scope yet because
+                # `import nccl.core` imports cupy and torch (with nccl4py 0.1.1)
+                import nccl.core as nccl  # type: ignore
+
+                # Create NCCL communicator.
+                unique_id = nccl.get_unique_id()
+                process_group.broadcast_buffer(unique_id.as_ndarray.view(np.int8), root=0)
+                nccl_comm = nccl.Communicator.init(nranks=nranks, rank=rank, unique_id=unique_id)
+
+            _ctx = DistributedContext(
+                device_id=device_id, process_group=process_group, nvshmem_available=nvshmem_available, nccl_comm=nccl_comm
+            )
+        finally:
+            old_device.set_current()
+
+
+def finalize() -> None:
+    """Finalize ``nvmath.distributed`` runtime (this is called automatically at exit
+    if the runtime is initialized). **Note that this is a collective operation and
+    must be called by all processes.**"""
+    global _ctx
+    with _initialize_mutex:
+        if _ctx is None:
+            return
+
+        linalg.advanced.matmulmod._grid_cache.clear()
+
+        if _ctx.nccl_comm is not None:
+            with device_ctx(_ctx.device_id):
+                _ctx.nccl_comm.finalize()
+                _ctx.nccl_comm.destroy()
+
+        if _ctx.nvshmem_available:
+            nvshmem.finalize(_ctx.device_id)
+
+        _ctx = None
+
+
+def get_context() -> DistributedContext | None:
+    """Return the distributed runtime's context or None if not initialized."""
+    return _ctx