nvmath-python 0.5.0__cp313-cp313-win_amd64.whl

nvmath/device/random_states.py

@@ -0,0 +1,189 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import operator
+
+from nvmath.device import curand_kernel
+from nvmath.device.common_mathdx import CURAND_HOME  # noqa: F401
+
+from numba import cuda, types
+from numba.extending import models, register_model, typeof_impl
+from numba.cuda.cudadecl import register_global
+from numba.cuda.cudaimpl import lower
+from numba.core.typing.templates import AbstractTemplate, signature
+import numpy as np
+
+from llvmlite import ir
+
+xorwow_dtype = np.dtype(
+    [
+        ("d", np.uint32),
+        ("v", np.uint32, (5,)),
+        ("boxmuller_flag", np.int32),
+        ("boxmuller_flag_double", np.int32),
+        ("boxmuller_extra", np.float32),
+        ("boxmuller_extra_double", np.float64),
+    ],
+    align=True,
+)
+test_dtype = np.dtype([("v", np.uint32)], align=True)
+mrg32k3a_dtype = np.dtype(
+    [
+        ("s1", np.uint32, (3,)),
+        ("s2", np.uint32, (3,)),
+        ("boxmuller_flag", np.int32),
+        ("boxmuller_flag_double", np.int32),
+        ("boxmuller_extra", np.float32),
+        ("boxmuller_extra_double", np.float64),
+    ],
+    align=True,
+)
+sobol32_dtype = np.dtype(
+    [("i", np.uint32), ("x", np.uint32), ("c", np.uint32), ("direction_vectors", np.int32, (32,))], align=True
+)
+scrambled_sobol32_dtype = np.dtype(
+    [("i", np.uint32), ("x", np.uint32), ("c", np.uint32), ("direction_vectors", np.int32, (32,))], align=True
+)
+sobol64_dtype = np.dtype(
+    [("i", np.uint64), ("x", np.uint64), ("c", np.uint64), ("direction_vectors", np.int64, (64,))], align=True
+)
+scrambled_sobol64_dtype = np.dtype(
+    [("i", np.uint64), ("x", np.uint64), ("c", np.uint64), ("direction_vectors", np.int64, (64,))], align=True
+)
+philox4_32_10_dtype = np.dtype(
+    [
+        ("ctr", np.uint32, (4,)),
+        ("output", np.uint32, (4,)),
+        ("key", np.uint32, (2,)),
+        ("STATE", np.uint32),
+        ("boxmuller_flag", np.int32),
+        ("boxmuller_flag_double", np.int32),
+        ("boxmuller_extra", np.float32),
+        ("boxmuller_extra_double", np.float64),
+    ],
+    align=True,
+)
+
+NP_DTYPES = {
+    "curandStateXORWOW": xorwow_dtype,
+    "curandStateTest": test_dtype,
+    "curandStateMRG32k3a": mrg32k3a_dtype,
+    "curandStateSobol32": sobol32_dtype,
+    "curandStateScrambledSobol32": scrambled_sobol32_dtype,
+    "curandStateSobol64": sobol64_dtype,
+    "curandStateScrambledSobol64": scrambled_sobol64_dtype,
+    "curandStatePhilox4_32_10": philox4_32_10_dtype,
+}
+
+
+def make_curand_states(name: str):
+    """Create a boxed curand states object to represent an array of curand state.
+    Should invoke after creating the numba types for each curand state.
+    """
+
+    curand_state_name = name
+
+    state_ty = getattr(curand_kernel, f"_type_{curand_state_name}")
+
+    curand_states_name = curand_state_name.replace("State", "States")
+
+    # cuRAND state type as a NumPy dtype - this mirrors the state defined in
+    # curand_kernel.h. Can be used to inspect the state through the device array
+    # held by CurandStates.
+
+    curandStateDtype = NP_DTYPES[curand_state_name]
+
+    # Hold an array of cuRAND states - somewhat analogous to a curandState* in
+    # C/C++.
+
+    class CurandStates:
+        def __init__(self, n):
+            self._array = cuda.device_array(n, dtype=curandStateDtype)
+
+        @property
+        def data(self):
+            return self._array.__cuda_array_interface__["data"][0]
+
+    CurandStates.__name__ = curand_states_name
+
+    class CurandStatePointer(types.Type):
+        def __init__(self, name):
+            self.dtype = state_ty
+            super().__init__(name=name)
+
+    curand_state_pointer = CurandStatePointer(curand_states_name + "*")
+
+    @typeof_impl.register(CurandStates)
+    def typeof_curand_states(val, c):
+        return curand_state_pointer
+
+    @register_global(operator.getitem)
+    class CurandStatesPointerGetItem(AbstractTemplate):
+        def generic(self, args, kws):
+            assert not kws
+            [ptr, idx] = args
+            if ptr == curand_state_pointer:
+                return signature(types.CPointer(state_ty), curand_state_pointer, types.int64)
+
+    register_model(CurandStatePointer)(models.PointerModel)
+
+    @lower(operator.getitem, curand_state_pointer, types.int64)
+    def lower_curand_states_getitem(context, builder, sig, args):
+        [ptr, idx] = args
+
+        # Working out my own GEP
+        ptrint = builder.ptrtoint(ptr, ir.IntType(64))
+        itemsize = curandStateDtype.itemsize
+        offset = builder.mul(idx, context.get_constant(types.int64, itemsize))
+        ptrint = builder.add(ptrint, offset)
+        ptr = builder.inttoptr(ptrint, ptr.type)
+        return ptr
+
+    # Argument handling. When a CurandStatePointer is passed into a kernel, we
+    # really only need to pass the pointer to the data, not the whole underlying
+    # array structure. Our handler here transforms these arguments into a uint64
+    # holding the pointer.
+
+    class CurandStateArgHandler:
+        def prepare_args(self, ty, val, **kwargs):
+            if isinstance(val, CurandStates):
+                assert ty == curand_state_pointer
+                return types.uint64, val.data
+            else:
+                return ty, val
+
+    curand_state_arg_handler = CurandStateArgHandler()
+
+    return CurandStates, curand_state_arg_handler
+
+
+curandStates = [
+    "curandStateTest",
+    "curandStateXORWOW",
+    "curandStateMRG32k3a",
+    "curandStateSobol32",
+    "curandStateScrambledSobol32",
+    "curandStateSobol64",
+    "curandStateScrambledSobol64",
+    "curandStatePhilox4_32_10",
+    # Require additional type parsing of mtgp32_params_fast and mtgp32_kernel_params
+    # "curandStateMtgp32",
+]
+
+numpy_curand_states = []
+states_arg_handlers = []
+
+for state in curandStates:
+    states_obj, arg_handler = make_curand_states(state)
+    numpy_curand_states.append(states_obj)
+    states_arg_handlers.append(arg_handler)
+
+
+# Export
+globals().update({s.__name__: s for s in numpy_curand_states})
+globals().update({"states_arg_handlers": states_arg_handlers})
+
+__all__ = list(  # noqa: F822
+    {s.__name__ for s in numpy_curand_states} | {"states_arg_handlers"}
+)

nvmath/device/types.py

@@ -0,0 +1,11 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+
+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]

nvmath/device/vector_types_numba.py

@@ -0,0 +1,203 @@
+# Copyright (c) 2024-2025, 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 numba.core import types
+
+from numba.extending import overload, models, register_model, lower_builtin, as_numba_type, type_callable, lower_cast
+from numba.core.typeconv import Conversion
+from llvmlite import ir
+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
+
+
+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,117 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+try:
+    import mpi4py  # noqa: F401
+except ImportError as e:
+    # TODO: point to documentation with ways to install mpi4py
+    raise ImportError("nvmath.distributed requires mpi4py for bootstrapping. See [LINK] for installation guide.") from e
+
+import atexit
+import re
+from dataclasses import dataclass
+from threading import Lock
+
+from ._internal import nvshmem
+from ._utils import allocate_symmetric_memory, free_symmetric_memory
+
+from nvmath.distributed import fft
+from nvmath.distributed import reshape
+
+__all__ = [
+    "initialize",
+    "finalize",
+    "get_context",
+    "allocate_symmetric_memory",
+    "free_symmetric_memory",
+    "fft",
+    "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.
+        communicator: MPI communicator of participating processes.
+    """
+
+    device_id: int
+    communicator: mpi4py.MPI.Comm
+
+
+def initialize(device_id: int, communicator: mpi4py.MPI.Comm | None = None) -> None:
+    """Initialize nvmath.distributed. This is required before any distributed operations can
+    be performed. **Note that this is a collective operation and must be called by all
+    processes.**
+
+    Args:
+        device_id: CUDA device ID to associate with the nvmath.distributed runtime.
+        communicator: MPI communicator specifying the participating processes. If None, will
+                      use MPI.COMM_WORLD.
+    """
+    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}."
+        )
+
+    with _initialize_mutex:
+        global _atexit_registered, _ctx
+
+        if _ctx is not None:
+            raise RuntimeError("nvmath.distributed has already been initialized")
+
+        try:
+            # This initializes MPI if it hasn't been initialized yet.
+            from mpi4py import MPI
+        except RuntimeError as e:
+            if re.search(r"[Cc]annot load MPI library", str(e)):
+                raise RuntimeError("mpi4py could not find the MPI library. See [LINK] for installation tips") from e
+            raise
+
+        if not _atexit_registered:
+            # mpi4py's atexit handler is registered during `import MPI`, and to finalize
+            # NVSHMEM *before* MPI, we need to make sure that we register our exit handler
+            # *after* mpi4py's.
+            atexit.register(finalize)
+            _atexit_registered = True
+
+        if communicator is None:
+            communicator = MPI.COMM_WORLD
+        elif not isinstance(communicator, mpi4py.MPI.Comm):
+            raise TypeError(
+                "The provided communicator object should be of type mpi4py.MPI communicator, but "
+                f"got object of type {type(communicator)}."
+            )
+
+        nvshmem.initialize(device_id, communicator)
+
+        _ctx = DistributedContext(device_id=device_id, communicator=communicator)
+
+
+def finalize() -> None:
+    """Finalize nvmath.distributed runtime. **Note that this is a collective operation and
+    must be called by all processes.**"""
+    global _ctx
+    with _initialize_mutex:
+        if _ctx is None:
+            return
+        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