nvmath-python 0.5.0__cp313-cp313-win_amd64.whl

nvmath/memory.py

@@ -0,0 +1,266 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Interface for pluggable memory handlers."""
+
+__all__ = ["BaseCUDAMemoryManager", "MemoryPointer"]
+
+from abc import abstractmethod
+from collections.abc import Callable
+from typing import Protocol, runtime_checkable
+import logging
+import weakref
+
+import cuda.core.experimental as ccx
+
+from nvmath.internal import utils
+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()
+
+
+@runtime_checkable
+class BaseCUDAMemoryManager(Protocol):
+    """
+    Protocol for memory manager plugins.
+
+    .. seealso:: :class:`numba.cuda.BaseCUDAMemoryManager`
+    """
+
+    @abstractmethod
+    def __init__(self, device_id: int, logger: logging.Logger):
+        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: logging.Logger):
+        raise NotImplementedError
+
+    @abstractmethod
+    def memalloc_async(self, size: int, stream: ccx.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: logging.Logger):
+        """
+        __init__(device_id)
+        """
+        self.device_id = device_id
+        self.logger = logger
+
+    def memalloc_async(self, size: int, stream: ccx.Stream) -> MemoryPointer:
+        with utils.device_ctx(self.device_id) as device:
+            buffer = device.allocate(size=size, stream=stream)
+            device_ptr = int(buffer.handle)
+
+        self.logger.debug(
+            "_RawCUDAMemoryManager (allocate memory): size = %d, ptr = %d, device_id = %d, stream = %s",
+            size,
+            device_ptr,
+            self.device_id,
+            stream,
+        )
+
+        def finalizer():
+            nonlocal buffer, stream, device_ptr
+            self.logger.debug(
+                "_RawCUDAMemoryManager (release memory): ptr = %d, device_id = %d, stream = %s",
+                device_ptr,
+                self.device_id,
+                stream,
+            )
+            with utils.device_ctx(self.device_id):
+                buffer.close(stream=stream)
+
+        return MemoryPointer(device_ptr, size, finalizer=finalizer)
+
+
+_MEMORY_MANAGER: dict[str, type[BaseCUDAMemoryManager] | type[BaseCUDAMemoryManagerAsync]] = {
+    "_raw": _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: logging.Logger):
+            """
+            __init__(device_id)
+            """
+            self.device_id = device_id
+            self.logger = logger
+
+        def memalloc_async(self, size: int, stream) -> MemoryPointer:
+            stream_ctx = CupyPackage.to_stream_context(
+                CupyPackage.create_external_stream(self.device_id, CUDAPackage.to_stream_pointer(stream))
+            )
+            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: logging.Logger):
+            """
+            __init__(device_id)
+            """
+            self.device_id = device_id
+            self.logger = logger
+
+        def memalloc_async(self, size: int, stream: ccx.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,18 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from nvmath._utils import force_loading_cudss
+
+try:
+    force_loading_cudss("12")
+except RuntimeError:
+    pass
+
+del force_loading_cudss
+
+from . import advanced  # noqa: E402
+
+__all__ = [
+    "advanced",
+]

nvmath/sparse/_internal/__init__.py

@@ -0,0 +1,21 @@
+# Copyright (c) 2025, 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,81 @@
+# Copyright (c) 2024-2025, 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.utils import infer_object_package
+from nvmath.sparse._internal.sparse_format_helpers import PACKAGE_HELPER_MAP
+
+
+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.
+    """
+
+    helper = PACKAGE_HELPER_MAP[infer_object_package(native_operand)]
+    sparse_format = helper.sparse_format_name(native_operand)
+    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()
+    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