nvmath-python 1.0.0__cp314-cp314t-win_amd64.whl

nvmath/memory.py

@@ -0,0 +1,279 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Interface for pluggable memory handlers."""
+
+__all__ = ["BaseCUDAMemoryManager", "MemoryPointer"]
+
+import weakref
+from abc import abstractmethod
+from collections.abc import Callable
+from typing import Protocol, runtime_checkable
+
+from cuda.core import Stream
+
+from nvmath._internal.utils import LoggerLike
+from nvmath.internal import utils
+from nvmath.internal.memory import allocate_from_mr, get_device_memory_resource
+from nvmath.internal.package_ifc_cuda import CUDAPackage
+
+
+class MemoryPointer:
+    """
+    An RAII class for a device memory buffer.
+
+    Args:
+        device_ptr: The address of the device memory buffer.
+        size: The size of the memory buffer in bytes.
+        finalizer: A nullary callable that will be called when the buffer is to be freed.
+
+    .. seealso:: :class:`numba.cuda.MemoryPointer`
+    """
+
+    def __init__(self, device_ptr: int, size: int, finalizer: None | Callable[[], None]):
+        self.device_ptr = device_ptr
+        self.size = size
+        self._finalizer: weakref.finalize | None
+        if finalizer is not None:
+            self._finalizer = weakref.finalize(self, finalizer)
+        else:
+            self._finalizer = None
+
+    def free(self):
+        """
+        "Frees" the memory buffer by calling the finalizer.
+        """
+        if self._finalizer is None:
+            return
+
+        if not self._finalizer.alive:
+            raise RuntimeError("The buffer has already been freed.")
+        self._finalizer()
+
+
+class _UnmanagedMemoryPointer:
+    """
+    An alternative to :class:`MemoryPointer` to wrap objects that already implements
+    its own RAII semantics.
+
+    Args:
+        device_ptr: The address of the device memory buffer.
+        size: The size of the memory buffer in bytes.
+        owner: An object that owns the memory buffer and releases the allocation
+               on object deletion.
+
+    .. seealso:: :class:`MemoryPointer`
+    """
+
+    __slots__ = ("device_ptr", "size", "_owner")
+
+    def __init__(self, device_ptr: int, size: int, owner: object):
+        self.device_ptr = device_ptr
+        self.size = size
+        self._owner = owner
+
+    def free(self):
+        """
+        "Frees" by removing the reference to the owner.
+        """
+        self._owner = None
+
+
+@runtime_checkable
+class BaseCUDAMemoryManager(Protocol):
+    """
+    Protocol for memory manager plugins.
+
+    .. seealso:: :class:`numba.cuda.BaseCUDAMemoryManager`
+    """
+
+    @abstractmethod
+    def __init__(self, device_id: int, logger: LoggerLike):
+        raise NotImplementedError
+
+    @abstractmethod
+    def memalloc(self, size: int) -> MemoryPointer:
+        """
+        Allocate device memory synchronously or on the current stream.
+
+        Args:
+            size: The size of the memory buffer in bytes.
+
+        Returns:
+            An object that owns the allocated memory and is responsible for releasing it (to
+            the OS or a pool). The object must have an attribute named ``device_ptr``,
+            ``device_pointer``, or ``ptr`` specifying the pointer to the allocated memory
+            buffer. See :class:`MemoryPointer` for an example interface.
+
+        Note:
+            Objects of type :class:`numba.cuda.MemoryPointer` as well as
+            :class:`cupy.cuda.MemoryPointer` meet the requirements listed above for the
+            device memory pointer object.
+        """
+        raise NotImplementedError
+
+
+@runtime_checkable
+class BaseCUDAMemoryManagerAsync(Protocol):
+    """
+    Protocol for async memory manager plugins.
+
+    .. seealso:: :class:`BaseCUDAMemoryManager`
+    """
+
+    @abstractmethod
+    def __init__(self, device_id: int, logger: LoggerLike):
+        raise NotImplementedError
+
+    @abstractmethod
+    def memalloc_async(self, size: int, stream: Stream) -> MemoryPointer:
+        """
+        Allocate device memory asynchronously on the provided stream.
+
+        Args:
+            size: The size of the memory buffer in bytes.
+            stream: A cuda.core.Stream object on which the allocation will be performed.
+
+        Returns:
+            An object that owns the allocated memory and is responsible for releasing it (to
+            the OS or a pool). The object must have an attribute named ``device_ptr``,
+            ``device_pointer``, or ``ptr`` specifying the pointer to the allocated memory
+            buffer. See :class:`MemoryPointer` for an example interface.
+
+        Note:
+            Objects of type :class:`numba.cuda.MemoryPointer` as well as
+            :class:`cupy.cuda.MemoryPointer` meet the requirements listed above for the
+            device memory pointer object.
+        """
+        raise NotImplementedError
+
+
+class _RawCUDAMemoryManager(BaseCUDAMemoryManagerAsync):
+    """
+    Raw device 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: int, logger: LoggerLike):
+        """
+        __init__(device_id)
+        """
+        self.device_id = device_id
+        self.logger = logger
+        self._pool = get_device_memory_resource(device_id)
+
+    def memalloc_async(self, size: int, stream: Stream) -> MemoryPointer:
+        buffer = allocate_from_mr(self._pool, size, stream, self.device_id, self.logger)
+        device_ptr = int(buffer.handle)
+
+        return _UnmanagedMemoryPointer(device_ptr, size, buffer)  # type: ignore[return-value]
+
+
+_MEMORY_MANAGER: dict[str, type[BaseCUDAMemoryManager] | type[BaseCUDAMemoryManagerAsync]] = {
+    "_raw": _RawCUDAMemoryManager,
+    "cuda": _RawCUDAMemoryManager,
+}
+
+
+def lazy_load_cupy():
+    global _MEMORY_MANAGER
+    import cupy as cp
+
+    from nvmath.internal.package_ifc_cupy import CupyPackage
+
+    class _CupyCUDAMemoryManager(BaseCUDAMemoryManagerAsync):
+        """
+        CuPy device 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: int, logger: LoggerLike):
+            """
+            __init__(device_id)
+            """
+            self.device_id = device_id
+            self.logger = logger
+
+        def memalloc_async(self, size: int, stream: Stream) -> MemoryPointer:
+            stream_ptr = CUDAPackage.to_stream_pointer(stream)
+            external = CupyPackage.create_external_stream(self.device_id, stream_ptr)
+            stream_ctx = CupyPackage.to_stream_context(external)
+            with utils.device_ctx(self.device_id), stream_ctx:
+                cp_mem_ptr = cp.cuda.alloc(size)
+                device_ptr = cp_mem_ptr.ptr
+
+            self.logger.debug(
+                "_CupyCUDAMemoryManager (allocate memory): size = %d, ptr = %d, device_id = %d, stream = %s",
+                size,
+                device_ptr,
+                self.device_id,
+                stream,
+            )
+
+            def finalizer():
+                # The cupy MemoryPointer object is RAII, so we keep a reference to it
+                # until we don't need it anymore.
+                nonlocal cp_mem_ptr, device_ptr
+                self.logger.debug("_CupyCUDAMemoryManager (release memory): ptr = %d", device_ptr)
+                del cp_mem_ptr
+
+            return MemoryPointer(device_ptr, size, finalizer=finalizer)
+
+    _MEMORY_MANAGER["cupy"] = _CupyCUDAMemoryManager
+
+
+def lazy_load_torch():
+    global _MEMORY_MANAGER
+
+    from torch.cuda import caching_allocator_alloc, caching_allocator_delete
+
+    from nvmath.internal.package_ifc_torch import TorchPackage
+
+    class _TorchCUDAMemoryManager(BaseCUDAMemoryManagerAsync):
+        """
+        Torch caching 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: int, logger: LoggerLike):
+            """
+            __init__(device_id)
+            """
+            self.device_id = device_id
+            self.logger = logger
+
+        def memalloc_async(self, size: int, stream: Stream) -> MemoryPointer:
+            torch_stream = TorchPackage.create_external_stream(self.device_id, CUDAPackage.to_stream_pointer(stream))
+            device_ptr = caching_allocator_alloc(size, device=self.device_id, stream=torch_stream)
+
+            self.logger.debug(
+                "_TorchCUDAMemoryManager (allocate memory): size = %d, ptr = %d, device_id = %d, stream = %s",
+                size,
+                device_ptr,
+                self.device_id,
+                stream,
+            )
+
+            def finalizer():
+                nonlocal device_ptr, stream
+                self.logger.debug(
+                    "_TorchCUDAMemoryManager (release memory): ptr = %d, device_id = %d, stream = %s",
+                    device_ptr,
+                    self.device_id,
+                    stream,
+                )
+                caching_allocator_delete(device_ptr)
+
+            return MemoryPointer(device_ptr, size, finalizer=finalizer)
+
+    _MEMORY_MANAGER["torch"] = _TorchCUDAMemoryManager

nvmath/sparse/__init__.py

@@ -0,0 +1,38 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from . import advanced, ust  # noqa: E402
+from .generic import ComputeType, ExecutionCUDA, Matmul, MatmulOptions, matmul
+from .generic import (
+    compile_add as compile_matmul_add,
+)
+from .generic import (
+    compile_atomic_add as compile_matmul_atomic_add,
+)
+from .generic import (
+    compile_epilog as compile_matmul_epilog,
+)
+from .generic import (
+    compile_mul as compile_matmul_mul,
+)
+from .generic import (
+    compile_prolog as compile_matmul_prolog,
+)
+from .generic import matrix_qualifiers_dtype as matmul_matrix_qualifiers_dtype
+
+__all__ = [
+    "advanced",
+    "ust",
+    "ExecutionCUDA",
+    "MatmulOptions",
+    "matmul_matrix_qualifiers_dtype",
+    "compile_matmul_prolog",
+    "compile_matmul_epilog",
+    "compile_matmul_add",
+    "compile_matmul_atomic_add",
+    "compile_matmul_mul",
+    "Matmul",
+    "matmul",
+    "ComputeType",
+]

nvmath/sparse/_internal/__init__.py

@@ -0,0 +1,21 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from . import common_utils as common_utils
+from . import cudss_config_ifc as cudss_config_ifc
+from . import cudss_data_ifc as cudss_data_ifc
+from . import cudss_utils as cudss_utils
+from . import sparse_csr_ifc as sparse_csr_ifc
+from . import sparse_format_helpers as sparse_format_helpers
+from . import sparse_tensor_ifc as sparse_tensor_ifc
+
+__all__ = [
+    "common_utils",
+    "cudss_config_ifc",
+    "cudss_data_ifc",
+    "cudss_utils",
+    "sparse_csr_ifc",
+    "sparse_format_helpers",
+    "sparse_tensor_ifc",
+]

nvmath/sparse/_internal/common_utils.py

@@ -0,0 +1,147 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+A collection of sparse utilities.
+"""
+
+import importlib
+from collections.abc import Sequence
+
+from nvmath.internal.ndbuffer import NDBuffer
+from nvmath.internal.utils import infer_object_package
+from nvmath.sparse._internal.sparse_format_helpers import PACKAGE_HELPER_MAP
+
+RECOGNIZED_NAMED_FORMATS = {"BSR", "BSC", "CSR", "CSC", "COO", "DIA"}
+
+
+def check_supported_package(native_operands):
+    """
+    Check if the operands belong to one of the supported packages.
+    """
+    package = {infer_object_package(o) for o in native_operands}
+    if len(package) != 1:
+        raise ValueError(f"The sparse operands do not belong to the same package: {package}.")
+
+    package = package.pop()
+    if package not in PACKAGE_HELPER_MAP:
+        raise TypeError(f"The sparse operand package {package} is not supported.")
+
+    return package
+
+
+def wrap_sparse_operand(native_operand):
+    """
+    Wrap one "native" sparse operand so that package-agnostic API can be used.
+    """
+
+    package = infer_object_package(native_operand)
+    helper = PACKAGE_HELPER_MAP[package]
+    sparse_format = helper.sparse_format_name(native_operand)
+
+    # The sparse format should be a recognized named format or an UST.
+    if sparse_format not in RECOGNIZED_NAMED_FORMATS:
+        if package == "nvmath":
+            from nvmath.sparse._internal.sparse_ust_ifc import USTensorHolder
+
+            return USTensorHolder(native_operand)
+
+        raise TypeError(f"The sparse format {sparse_format} is not currently supported.")
+
+    # Recognized named format.
+    attr_name_map = getattr(helper, f"{sparse_format.lower()}_attribute_names")()
+
+    module_name = f"nvmath.sparse._internal.sparse_{sparse_format.lower()}_ifc"
+    module = importlib.import_module(module_name)
+    SparseHolder = getattr(module, f"{sparse_format}TensorHolder")
+
+    return SparseHolder.create_from_tensor(native_operand, attr_name_map=attr_name_map)
+
+
+def wrap_sparse_operands(native_operands):
+    """
+    Wrap each "native" sparse operand so that package-agnostic API can be used.
+    """
+
+    if not isinstance(native_operands, Sequence):
+        return wrap_sparse_operand(native_operands)
+
+    package = check_supported_package(native_operands)
+    helper = PACKAGE_HELPER_MAP[package]
+
+    sparse_format = {helper.sparse_format_name(o) for o in native_operands}
+    if len(sparse_format) != 1:
+        raise ValueError(f"The sparse operands do not have the same sparse format: {sparse_format}")
+    sparse_format = sparse_format.pop()
+
+    # The sparse format should be a recognized named format or an UST.
+    if sparse_format not in RECOGNIZED_NAMED_FORMATS:
+        if package == "nvmath":
+            from nvmath.sparse._internal.sparse_ust_ifc import USTensorHolder
+
+            return [USTensorHolder(o) for o in native_operands]
+
+        raise TypeError(f"The sparse format {sparse_format} is not currently supported.")
+
+    attr_name_map = getattr(helper, f"{sparse_format.lower()}_attribute_names")()
+
+    module_name = f"nvmath.sparse._internal.sparse_{sparse_format.lower()}_ifc"
+    module = importlib.import_module(module_name)
+    SparseHolder = getattr(module, f"{sparse_format}TensorHolder")
+
+    return [SparseHolder.create_from_tensor(o, attr_name_map=attr_name_map) for o in native_operands]
+
+
+# TODO: Unify all the get_attribute functions in internal.utils
+# as get(operands, attribute, attribute_desc).
+def get_operands_index_type(operands):
+    """
+    Return the index type name of the tensors.
+    """
+    index_type = operands[0].index_type
+    if not all(operand.index_type == index_type for operand in operands):
+        index_types = {operand.index_type for operand in operands}
+        raise ValueError(f"All operands must have the same index type. The index types found = {index_types}.")
+    return index_type
+
+
+def sparse_or_dense(operand):
+    """
+    Determine if the operand is sparse or dense.
+    """
+
+    dense_tensor_modules = {"cupy", "numpy", "torch"}
+    sparse_tensor_modules = set(PACKAGE_HELPER_MAP.keys())
+
+    module = infer_object_package(operand)
+
+    if module == "torch":
+        import torch
+
+        if not torch.is_tensor(operand):
+            raise TypeError(f"The operand '{operand}' is not a tensor object.")
+
+        if operand.layout == torch.strided:
+            return "dense"
+
+        return "sparse"
+
+    # TODO: use the level spec instead of format name to infer if dense.
+    if module == "nvmath":
+        if type(operand) is NDBuffer:
+            return "dense"
+        name = operand.tensor_format.name
+        if name == "Scalar" or "Dense" in name:
+            return "dense"
+        return "sparse"
+
+    # We'll check if the object is a tensor later, we'll assume here that it is
+    # solely based on the package.
+    if module in sparse_tensor_modules:
+        return "sparse"
+
+    if module in dense_tensor_modules:
+        return "dense"
+
+    raise TypeError(f"The package '{module}' is not supported.")