nvmath-python 0.9.0__cp314-cp314-win_amd64.whl

nvmath/internal/ndbuffer/copy_kernel.pxd

@@ -0,0 +1,9 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+cimport cython
+from libc.stdint cimport intptr_t
+from .data_layout cimport Layout
+
+cdef int launch_copy_kernel(Layout dst, Layout src, intptr_t dst_ptr, intptr_t src_ptr, int device_id, intptr_t stream_ptr, object logger=*) except -1 nogil

nvmath/internal/ndbuffer/copy_kernel.pyi

@@ -0,0 +1,9 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import _cython_3_2_4
+
+__pyx_capi__: dict
+__test__: dict
+get_include_path: _cython_3_2_4.cython_function_or_method

nvmath/internal/ndbuffer/data_layout.pxd

@@ -0,0 +1,72 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+cimport cython
+from libc.stdint cimport int64_t, uint32_t, intptr_t
+from libcpp cimport vector
+
+ctypedef int64_t extent_t
+ctypedef int64_t stride_t
+ctypedef int axis_t
+
+ctypedef uint32_t axes_mask_t
+
+ctypedef vector.vector[extent_t] shape_t
+ctypedef vector.vector[stride_t] strides_t
+ctypedef vector.vector[axis_t] axis_order_t
+
+
+ctypedef fused vector_type:
+    shape_t
+    strides_t
+    axis_order_t
+
+
+cdef enum OrderFlag:
+    C_ORDER = 0
+    F_ORDER = 1
+    CUSTOM_PERMUTATION = 2
+
+
+@cython.final
+cdef class Layout:
+    cdef readonly shape_t shape
+    cdef readonly strides_t strides
+    cdef readonly int64_t volume
+    cdef readonly int ndim
+    cdef readonly int itemsize
+
+
+@cython.overflowcheck(True)
+cdef inline int64_t size_in_bytes(Layout layout) except? -1 nogil:
+    return layout.volume * layout.itemsize
+
+
+cdef int tuple2vec(vector_type &vec, object t) except -1
+cdef int set_strides_tuple(Layout layout, object strides, bint strides_in_bytes) except -1
+cdef int64_t set_strides_in_order(strides_t& strides, shape_t &shape, OrderFlag order_flag, axis_order_t *axis_order=*) except -1 nogil
+cdef int64_t overflow_checked_volume(shape_t& shape) except? -1 nogil
+cdef int zero_strides(strides_t& strides, int ndim) except -1 nogil
+cdef int divide_strides(strides_t &strides, int ndim, int itemsize) except -1 nogil
+cdef tuple get_strides_in_bytes_tuple(Layout layout)
+cdef Layout create_layout_without_strides(object shape, int itemsize)
+cdef Layout create_layout(object shape, object strides, int itemsize, bint strides_in_bytes)
+cdef Layout copy_layout(Layout src)
+cdef Layout empty_layout_with_dtype_like(Layout src)
+cdef bint is_overlapping_layout(Layout sorted_layout) except -1 nogil
+cdef bint is_overlapping_layout_in_order(Layout layout, axis_order_t& axis_order) except -1 nogil
+cdef int64_t transpose_squeeze_zeros_ones_layout(Layout out_layout, Layout in_layout, axis_order_t& axis_order) except -1 nogil
+cdef int transpose_layout(Layout layout, axis_order_t& axis_order) except -1 nogil
+cdef bint is_c_contiguous_layout(Layout sorted_layout) except -1 nogil
+cdef bint is_f_contiguous_layout(Layout sorted_layout) except -1 nogil
+cdef bint is_contiguous_layout_in_order(Layout layout, axis_order_t& axis_order) except -1 nogil
+cdef int squeeze_layouts_together(Layout layout_a, Layout layout_b, int ndim) except -1 nogil
+cdef int squeeze_layout(shape_t& out_shape, strides_t& out_strides, Layout in_layout) except? -1 nogil
+cdef bint split_strides(Layout new_layout, shape_t& old_shape, strides_t& old_strides) except -1 nogil
+cdef int vectorize_together(Layout layout_a, intptr_t ptr_a, Layout layout_b, intptr_t ptr_b, int max_vec_size=*, int max_itemsize=*) except -1 nogil
+cdef int get_axis_order(axis_order_t& axis_order, Layout layout) except -1 nogil
+cdef axes_mask_t get_contiguous_axes_up_to_vol(int64_t &suffix_vol, axes_mask_t forbidden_axes, int64_t max_volume, Layout layout, int* axis_order=*) except? -1 nogil
+cdef int parse_py_axis_order(OrderFlag& order_flag, axis_order_t& axis_order_vec, Layout other, object axis_order_arg) except -1
+cdef bint is_c_or_f(OrderFlag& order_flag, shape_t& shape, strides_t& strides, int ndim) except -1 nogil
+cdef bint is_c_or_f_layout(OrderFlag& order_flag, Layout layout) except -1 nogil

nvmath/internal/ndbuffer/data_layout.pyi

@@ -0,0 +1,23 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import _cython_3_2_4
+from _typeshed import Incomplete
+
+__pyx_capi__: dict
+__reduce_cython__: _cython_3_2_4.cython_function_or_method
+__setstate_cython__: _cython_3_2_4.cython_function_or_method
+__test__: dict
+
+class Layout:
+    itemsize: Incomplete
+    ndim: Incomplete
+    shape: Incomplete
+    strides: Incomplete
+    volume: Incomplete
+    @classmethod
+    def __init__(cls, *args, **kwargs) -> None:
+        """Create and return a new object.  See help(type) for accurate signature."""
+    def __reduce__(self):
+        """Layout.__reduce_cython__(self)"""

nvmath/internal/ndbuffer/jit.pxd

@@ -0,0 +1,7 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from libc.stdint cimport intptr_t
+
+cdef intptr_t get_kernel(str kernel_code, str kernel_name, int device_id, include_path, object logger=*) except? 0

nvmath/internal/ndbuffer/jit.pyi

@@ -0,0 +1,6 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+__pyx_capi__: dict
+__test__: dict

nvmath/internal/ndbuffer/ndbuffer.pxd

@@ -0,0 +1,36 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+cimport cython
+from libc.stdint cimport intptr_t
+from .data_layout cimport Layout, axis_order_t, OrderFlag
+
+
+@cython.final
+cdef class NDBuffer:
+    cdef Layout layout
+    cdef readonly object data
+    cdef int data_device_id
+    cdef int flags
+    cdef readonly intptr_t data_ptr
+    cdef readonly str dtype_name
+
+    # possibly lazy evaluated properties
+    # accessible publicly in python
+    cdef prop_strides
+    cdef prop_shape
+    cdef prop_device
+    cdef prop_device_id
+    cdef prop_strides_in_bytes
+
+
+cdef NDBuffer _no_data_dense_like(NDBuffer other, axis_order_t* axis_order_vec, OrderFlag order_flag)
+cdef NDBuffer _no_data_like(NDBuffer other, bint copy_data)
+cdef int _set_flags(NDBuffer ndbuffer, bint is_wrapping_tensor=*) except -1 nogil
+
+cpdef NDBuffer wrap_external(data, intptr_t ptr, str dtype_name, object shape, object strides, int device_id, int itemsize, bint strides_in_bytes=*)
+cpdef NDBuffer empty(object shape, int device_id, str dtype_name, int itemsize, object axis_order=*, object strides=*, object host_memory_pool=*, object device_memory_pool=*, object stream=*, bint strides_in_bytes=*, object logger=*)
+cpdef NDBuffer empty_like(NDBuffer other, object axis_order=*, object device_id=*, object stream=*, object host_memory_pool=*, object device_memory_pool=*, object logger=*)
+cpdef int copy_into(NDBuffer dst, NDBuffer src, object stream, object host_memory_pool=*, object device_memory_pool=*, object logger=*) except -1
+cpdef NDBuffer reshaped_view(NDBuffer other, object shape, object logger=*)

nvmath/internal/ndbuffer/ndbuffer.pyi

@@ -0,0 +1,41 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import _cython_3_2_4
+from _typeshed import Incomplete
+from typing import Any
+
+CPU_DEVICE_ID: int
+__pyx_capi__: dict
+__test__: dict
+copy_into: _cython_3_2_4.cython_function_or_method
+empty: _cython_3_2_4.cython_function_or_method
+empty_like: _cython_3_2_4.cython_function_or_method
+reshaped_view: _cython_3_2_4.cython_function_or_method
+wrap_external: _cython_3_2_4.cython_function_or_method
+
+class NDBuffer:
+    data: Incomplete
+    data_ptr: Incomplete
+    device: Incomplete
+    device_id: Incomplete
+    dtype_name: Incomplete
+    itemsize: Incomplete
+    ndim: Incomplete
+    shape: Incomplete
+    size: Incomplete
+    size_in_bytes: Incomplete
+    strides: Incomplete
+    strides_in_bytes: Incomplete
+    @classmethod
+    def __init__(cls, *args, **kwargs) -> None:
+        """Create and return a new object.  See help(type) for accurate signature."""
+    def cf_order(self) -> Any:
+        """NDBuffer.cf_order(self)"""
+    def __reduce__(self):
+        """NDBuffer.__reduce_cython__(self)"""
+    def __reduce_cython__(self) -> Any:
+        """NDBuffer.__reduce_cython__(self)"""
+    def __setstate_cython__(self, __pyx_state) -> Any:
+        """NDBuffer.__setstate_cython__(self, __pyx_state)"""

nvmath/internal/ndbuffer/package_utils.pxd

@@ -0,0 +1,11 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from .ndbuffer cimport NDBuffer
+
+
+cpdef NDBuffer empty_numpy_like(NDBuffer other, object axis_order=*)
+cpdef NDBuffer wrap_numpy_array(object array)
+cpdef NDBuffer wrap_cupy_array(object array)
+cpdef str is_c_or_f(object shape, object strides)

nvmath/internal/ndbuffer/package_utils.pyi

@@ -0,0 +1,12 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import _cython_3_2_4
+
+__pyx_capi__: dict
+__test__: dict
+empty_numpy_like: _cython_3_2_4.cython_function_or_method
+is_c_or_f: _cython_3_2_4.cython_function_or_method
+wrap_cupy_array: _cython_3_2_4.cython_function_or_method
+wrap_numpy_array: _cython_3_2_4.cython_function_or_method

nvmath/internal/package_ifc.py

@@ -0,0 +1,147 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+An abstract interface to certain package-provided operations.
+
+Even though cuda.core does not have the concept of a "current stream" or a "stream context",
+we promise to honor third-party CUDA stream contexts if they match the input operand's
+package. The interface class defined in this module wraps around those stream concepts to
+make them uniform.
+
+The strategy is to use cuda.core everywhere internally except for context managers which
+will need to wrap back around to the external implementation.
+"""
+
+__all__ = ["Package", "StreamHolder"]
+
+from abc import ABC, abstractmethod
+from collections.abc import Hashable
+from contextlib import AbstractContextManager
+from dataclasses import dataclass
+from typing import Generic, Protocol, TypeVar
+
+try:
+    from cuda.core import Stream
+except ImportError:
+    from cuda.core.experimental import Stream
+
+from ._device_utils import get_device
+
+
+class AnyStream(Hashable, Protocol):
+    """Any supported Stream object such as a Stream from cuda.core, CuPy, or PyTorch."""
+
+    # This empty protocol class is just a placeholder to make type checking more helpful.
+    # Since the class is empty, it doesn't assume anything is true about the implementation
+    # classes.
+    pass
+
+
+"""
+A generic type for the third-party Stream which a given Package implementation wraps
+around.
+"""
+S = TypeVar("S")
+
+
+class _cuda_core_stream_holder:
+    """
+    Dummy class implementing ``__cuda_stream__`` protocol for ``cuda.core.Stream``.
+
+    Calling ``cuda.core.Device.create_stream(_cuda_core_stream_holder(handle, external))``
+    is similar to ``Stream.from_handle(handle)``, but additionally makes sure
+    to extend the ``external`` reference lifetime as long as the created
+    ``cuda.core.Stream`` object is alive.
+    """
+
+    def __init__(self, handle: int, external: S):
+        self.handle = handle
+        self.external = external
+
+    def __cuda_stream__(self):
+        return (0, self.handle)
+
+
+class Package(ABC, Generic[S]):
+    @staticmethod
+    @abstractmethod
+    def get_current_stream(device_id: int) -> S:
+        """
+        Obtain the current stream on the device.
+
+        Args:
+            device_id: The id (ordinal) of the device.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def to_stream_pointer(stream: S) -> int:
+        """
+        Obtain the stream pointer.
+
+        Args:
+            stream: The stream object.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def to_stream_context(stream: S) -> AbstractContextManager[S]:
+        """
+        Create a context manager from the stream.
+
+        Args:
+            stream: The stream object.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def create_external_stream(device_id: int, stream_ptr: int) -> S:
+        """
+        Wrap a stream pointer into an external stream object.
+
+        Args:
+            device_id: The id (ordinal) of the device.
+            stream: The stream pointer (int) to be wrapped.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def create_stream(cls, external: S, device_id: int) -> Stream:
+        """
+        Wrap an external Stream object into a cuda.core.Stream.
+
+        Args:
+            external: The external Stream object.
+        """
+        # use get_device to ensure the initial set_current is called
+        device = get_device(device_id)
+        # the stream holder ensures we tie the external reference to the
+        # cuda.core stream object, extending its lifetime.
+        holder = _cuda_core_stream_holder(cls.to_stream_pointer(external), external)
+        return device.create_stream(holder)
+
+
+@dataclass
+class StreamHolder(Generic[S]):
+    """A data class for easing CUDA stream manipulation.
+
+    Attributes:
+        ctx: A context manager for using the specified stream.
+        device_id (int): The device ID where the encapsulated stream locates.
+        external: A foreign object that holds the stream alive.
+        obj: The cuda.core Stream object wrapping external.
+        package (str): The name of the package to which the external stream belongs.
+        ptr (int): The address of the underlying ``cudaStream_t`` object.
+    """
+
+    ctx: AbstractContextManager[S]
+    device_id: int
+    external: S
+    obj: Stream
+    package: str
+    ptr: int

nvmath/internal/package_ifc_cuda.py

@@ -0,0 +1,49 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Interface to cuda.core operations.
+"""
+
+__all__ = ["CUDAPackage"]
+
+import contextlib
+
+try:
+    from cuda.core import Stream
+except ImportError:
+    from cuda.core.experimental import Stream
+
+from ._device_utils import get_device
+from .package_ifc import Package
+
+
+class CUDAPackage(Package[Stream]):
+    @staticmethod
+    def get_current_stream(device_id: int):
+        # Use get_device to ensure the cuda.core
+        # device has been initialized.
+        # In cuda.core 0.5.0, Stream.__hash__ requires context
+        # to be set.
+        # See https://github.com/NVIDIA/cuda-python/issues/1480
+        # TODO(ktokarski): Once we drop support for cuda.core 0.5.0,
+        # we can remove this precaution.
+        device = get_device(device_id)
+        return device.default_stream
+
+    @staticmethod
+    def to_stream_pointer(stream: Stream) -> int:  # type: ignore[override]
+        return int(stream.handle)
+
+    @staticmethod
+    def to_stream_context(stream: Stream):  # type: ignore[override]
+        return contextlib.nullcontext(stream)
+
+    @staticmethod
+    def create_external_stream(device_id: int, stream_ptr: int) -> Stream:
+        return Stream.from_handle(stream_ptr)
+
+    @classmethod
+    def create_stream(cls, external: Stream, device_id: int) -> Stream:  # type: ignore[override]
+        return external

nvmath/internal/package_ifc_cupy.py

@@ -0,0 +1,49 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Interface to CuPy operations.
+"""
+
+__all__ = ["CupyPackage"]
+
+import cupy as cp
+
+from . import utils
+from .package_ifc import Package
+
+# Using the functional API is faster than setting a device context
+if int(cp.__version__.split(".")[0]) >= 13:
+    _get_current_stream = cp.cuda.get_current_stream
+else:
+
+    def _get_current_stream(device_id: int):
+        with utils.device_ctx(device_id):
+            stream = cp.cuda.get_current_stream()
+        return stream
+
+
+# Monkey patch older versions of CuPy, so that Streams are hashable
+# NOTE: We choose not to patch the BaseStream/_BaseStream class because of name change
+if int(cp.__version__.split(".")[0]) < 11:
+    cp.cuda.Stream.__hash__ = lambda self: hash(self.ptr)
+    cp.cuda.ExternalStream.__hash__ = lambda self: hash(self.ptr)
+
+
+class CupyPackage(Package[cp.cuda.Stream]):
+    @staticmethod
+    def get_current_stream(device_id: int):
+        return _get_current_stream(device_id)
+
+    @staticmethod
+    def to_stream_pointer(stream: cp.cuda.Stream) -> int:
+        return stream.ptr
+
+    @staticmethod
+    def to_stream_context(stream: cp.cuda.Stream):
+        return stream
+
+    @staticmethod
+    def create_external_stream(device_id: int, stream_ptr: int) -> cp.cuda.ExternalStream:
+        return cp.cuda.ExternalStream(stream_ptr)

nvmath/internal/package_ifc_torch.py

@@ -0,0 +1,31 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Interface to Torch operations.
+"""
+
+__all__ = ["TorchPackage"]
+
+import torch
+
+from .package_ifc import Package
+
+
+class TorchPackage(Package[torch.cuda.Stream]):
+    @staticmethod
+    def get_current_stream(device_id: int) -> torch.cuda.Stream:
+        return torch.cuda.current_stream(device=device_id)
+
+    @staticmethod
+    def to_stream_pointer(stream: torch.cuda.Stream) -> int:  # type: ignore[override]
+        return stream.cuda_stream
+
+    @staticmethod
+    def to_stream_context(stream: torch.cuda.Stream) -> torch.cuda.StreamContext:  # type: ignore[override]
+        return torch.cuda.stream(stream)
+
+    @staticmethod
+    def create_external_stream(device_id: int, stream_ptr: int):
+        return torch.cuda.ExternalStream(stream_ptr, device=device_id)

nvmath/internal/package_wrapper.py

@@ -0,0 +1,14 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Create map from package names to package interface objects.
+"""
+
+__all__ = ["PACKAGE", "AnyStream", "StreamHolder"]
+
+from .package_ifc import AnyStream, Package, StreamHolder
+from .package_ifc_cuda import CUDAPackage
+
+PACKAGE: dict[str, type[Package]] = {"cuda": CUDAPackage}