nvmath-python 0.5.0__cp313-cp313-win_amd64.whl

nvmath/sparse/_internal/sparse_csr_ifc.py

@@ -0,0 +1,154 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Interface to seamlessly use CSR tensors from different libraries.
+"""
+
+from __future__ import annotations  # allows typehint of class methods to return the self class
+
+__all__ = ["CSRTensorHolder"]
+
+from typing import Literal, TypeVar
+
+from .sparse_tensor_ifc import SparseTensorHolder
+from nvmath.internal import tensor_wrapper
+from nvmath.internal.package_ifc import StreamHolder
+from nvmath.internal.tensor_ifc import TensorHolder
+
+
+"""
+A generic type for the third-party CSRTensor which a CSRTensorHolder implementation wraps
+around.
+"""
+CSRTensor = TypeVar("CSRTensor")
+
+
+class CSRTensorHolder(SparseTensorHolder):
+    """
+    A simple wrapper type for CSR sparse tensors to make the API package-agnostic. The
+    naming convention is sparse format name with the suffix "TensorHolder", which convention
+    is used in the sparse format helper.
+
+    Since the implementation is in terms of our dense tensor abstraction, package-specific
+    implementation is not needed.
+    """
+
+    _format_name: str = "CSR"
+
+    def __init__(self, crow_indices, col_indices, values, *, attr_name_map=None, shape=None, tensor=None):
+        # tensor is the native tensor.
+        self.tensor = tensor
+
+        # attr_name_map is the canonical attribute name to package attribute name map.
+        # It is set when constructing from a tensor, otherwise it is None.
+        self._attr_name_map = attr_name_map
+
+        assert (shape is not None) ^ (tensor is not None), "Internal error."
+
+        if shape is not None:
+            self._shape = shape
+        else:
+            self._shape = tuple(tensor.shape)
+
+        # Wrap constituent dense tensors, if required.
+        if not isinstance(crow_indices, TensorHolder):
+            crow_indices = tensor_wrapper.wrap_operand(crow_indices)
+
+        if not isinstance(col_indices, TensorHolder):
+            col_indices = tensor_wrapper.wrap_operand(col_indices)
+
+        if not isinstance(values, TensorHolder):
+            values = tensor_wrapper.wrap_operand(values)
+
+        # The wrapped constituent dense tensors.
+        self._crow_indices = crow_indices
+        self._col_indices = col_indices
+        self._values = values
+
+        # Ensure constituent dense tensors are all on the same device.
+        device_ids = {d.device_id for d in (self._crow_indices, self._col_indices, self._values)}
+        message = f"Internal error: the arrays defining the sparse format aren't on the same device {device_ids}."
+        assert len(device_ids) == 1, message
+        self._device_id = device_ids.pop()
+
+        self._device = self._values.device
+
+    @classmethod
+    def create_from_tensor(cls, tensor, *, attr_name_map=None):
+        assert attr_name_map is not None, "Internal error."
+
+        # tensor is the native tensor.
+        crow_indices = attr_name_map["crow_indices"](tensor)
+        col_indices = attr_name_map["col_indices"](tensor)
+        values = attr_name_map["values"](tensor)
+        return CSRTensorHolder(crow_indices, col_indices, values, attr_name_map=attr_name_map, tensor=tensor)
+
+    @property
+    def attr_name_map(self):
+        return self._attr_name_map
+
+    @property
+    def device(self):
+        return self._device
+
+    @property
+    def device_id(self):
+        return self._device_id
+
+    @property
+    def dtype(self):
+        return self._values.dtype
+
+    @property
+    def index_type(self):
+        return self._crow_indices.dtype
+
+    @property
+    def format_name(self) -> str:
+        return CSRTensorHolder._format_name
+
+    @property
+    def num_dimensions(self):
+        return len(self._shape)
+
+    @property
+    def shape(self):
+        return self._shape
+
+    @property
+    def crow_indices(self) -> TensorHolder:
+        """The TensorHolder object wrapping the crow_indices array."""
+        return self._crow_indices
+
+    @property
+    def col_indices(self) -> TensorHolder:
+        """The TensorHolder object wrapping the col_indices array."""
+        return self._col_indices
+
+    @property
+    def values(self) -> TensorHolder:
+        """The TensorHolder object wrapping the values array."""
+        return self._values
+
+    def to(self, device_id: int | Literal["cpu"], stream_holder: StreamHolder | None):
+        """Copy the CSRTensor representation to a different device.
+
+        No copy is performed if the CSRTensor is already on the requested device.
+        """
+
+        target_crow_indices = self.crow_indices.to(device_id=device_id, stream_holder=stream_holder)
+        target_col_indices = self.col_indices.to(device_id=device_id, stream_holder=stream_holder)
+        target_values = self.values.to(device_id=device_id, stream_holder=stream_holder)
+
+        return CSRTensorHolder(
+            target_crow_indices, target_col_indices, target_values, attr_name_map=self.attr_name_map, shape=self.shape
+        )
+
+    def copy_(self, src, stream_holder: StreamHolder | None) -> None:
+        """Overwrite self.tensor (in-place) with a copy of src (CSRTensor)."""
+
+        self.crow_indices.copy_(src=src.crow_indices, stream_holder=stream_holder)
+        self.col_indices.copy_(src=src.col_indices, stream_holder=stream_holder)
+        self.values.copy_(src=src.values, stream_holder=stream_holder)

nvmath/sparse/_internal/sparse_format_helpers.py

@@ -0,0 +1,88 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Package-agnostic interface for key sparse tensor attributes.
+"""
+
+__all__ = ["PACKAGE_HELPER_MAP", "CupySparseFormatHelper", "ScipySparseFormatHelper", "TorchSparseFormatHelper"]
+
+from abc import ABC, abstractmethod
+
+import re
+
+
+class SparseFormatHelper(ABC):
+    @staticmethod
+    @abstractmethod
+    def sparse_format_name(tensor):
+        raise NotImplementedError
+
+    # Attribute name map: canonical name <> package name.
+    # One for each named sparse format type.
+    @staticmethod
+    @abstractmethod
+    def csr_attribute_names():
+        raise NotImplementedError
+
+    # A ctor for each named sparse format type.
+    @staticmethod
+    @abstractmethod
+    def create_csr(shape, crow_indices, col_indices, values):
+        raise NotImplementedError
+
+
+class TorchSparseFormatHelper(SparseFormatHelper):
+    @staticmethod
+    def sparse_format_name(tensor):
+        name = re.sub(r"^torch\.sparse_", "", str(tensor.layout)).upper()
+        return name
+
+    @staticmethod
+    def csr_attribute_names():
+        return {
+            "crow_indices": lambda a: a.crow_indices(),
+            "col_indices": lambda a: a.col_indices(),
+            "values": lambda a: a.values(),
+        }
+
+    @staticmethod
+    def create_csr(shape, crow_indices, col_indices, values):
+        # import torch
+        # return torch.sparse_csr_tensor(crow_indices.tensor, col_indices.tensor,
+        #                                values.tensor, shape, device=values.device_id)
+        raise NotImplementedError
+
+
+class CupySparseFormatHelper(SparseFormatHelper):
+    @staticmethod
+    def sparse_format_name(tensor):
+        name = re.sub(r"_(array|matrix)$", "", type(tensor).__name__).upper()
+        return name
+
+    @staticmethod
+    def csr_attribute_names():
+        return {"crow_indices": lambda a: a.indptr, "col_indices": lambda a: a.indices, "values": lambda a: a.data}
+
+    @staticmethod
+    def create_csr(shape, crow_indices, col_indices, values):
+        raise NotImplementedError
+
+
+class ScipySparseFormatHelper(SparseFormatHelper):
+    @staticmethod
+    def sparse_format_name(tensor):
+        name = re.sub(r"_(array|matrix)$", "", type(tensor).__name__).upper()
+        return name
+
+    @staticmethod
+    def csr_attribute_names():
+        return {"crow_indices": lambda a: a.indptr, "col_indices": lambda a: a.indices, "values": lambda a: a.data}
+
+    @staticmethod
+    def create_csr(shape, crow_indices, col_indices, values):
+        raise NotImplementedError
+
+
+PACKAGE_HELPER_MAP = {"cupyx": CupySparseFormatHelper, "scipy": ScipySparseFormatHelper, "torch": TorchSparseFormatHelper}

nvmath/sparse/_internal/sparse_tensor_ifc.py

@@ -0,0 +1,80 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Interface to seamlessly use sparse tensors from different libraries.
+"""
+
+from __future__ import annotations  # allows typehint of class methods to return the self class
+
+__all__ = ["SparseTensorHolder"]
+
+
+from abc import ABC, abstractmethod
+from typing import Literal
+
+from nvmath.internal.package_ifc import StreamHolder
+
+
+class SparseTensorHolder(ABC):
+    """
+    A simple base type for attributes common to *all* sparse formats.
+    """
+
+    @classmethod
+    @abstractmethod
+    def create_from_tensor(cls, attr_name_map, tensor):
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def attr_name_map(self):
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def device(self):
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def device_id(self):
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def dtype(self):
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def index_type(self):
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def format_name(self) -> str:
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def num_dimensions(self):
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def shape(self):
+        raise NotImplementedError
+
+    @abstractmethod
+    def to(self, device_id: int | Literal["cpu"], stream_holder: StreamHolder | None):
+        """Copy the SparseTensor representation to a different device.
+
+        No copy is performed if the SparseTensor is already on the requested device.
+        """
+        raise NotImplementedError
+
+    def copy_(self, src: SparseTensorHolder, stream_holder: StreamHolder | None) -> None:
+        """Overwrite the sparse tensor (in-place) with a copy of src."""
+        raise NotImplementedError

nvmath/sparse/advanced/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from ._configuration import *  # noqa: F403
+from .direct_solver import *  # noqa: F403
+from .._internal.cudss_data_ifc import memory_estimates_dtype  # noqa: F401

nvmath/sparse/advanced/_configuration.py

@@ -0,0 +1,150 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = [
+    "DirectSolverOptions",
+    "DirectSolverMatrixType",
+    "DirectSolverMatrixViewType",
+    "ExecutionCUDA",
+    "ExecutionHybrid",
+    "HybridMemoryModeOptions",
+]
+
+
+from dataclasses import dataclass, field
+from logging import Logger
+from typing import ClassVar, Literal, TypeAlias
+
+from nvmath.bindings import cudss
+from nvmath.internal import mem_limit
+
+DirectSolverMatrixType: TypeAlias = cudss.MatrixType
+DirectSolverMatrixViewType: TypeAlias = cudss.MatrixViewType
+
+
+@dataclass
+class HybridMemoryModeOptions:
+    """
+    A data class for providing options related to the use of hybrid (CPU-GPU) memory to
+    those execution spaces that support it.
+
+    Attributes:
+        hybrid_memory_mode: If set to True, use CPU memory to store factors (default:
+            ``False``). See :attr:`nvmath.bindings.cudss.ConfigParam.HYBRID_MODE`.
+        device_memory_limit: The maximum device memory available for execution. It can
+            be specified as a value (with optional suffix like K[iB], M[iB], G[iB]) or as a
+            percentage. The default is based on internal heuristics. See
+            :attr:`nvmath.bindings.cudss.ConfigParam.HYBRID_DEVICE_MEMORY_LIMIT`.
+        register_cuda_memory: Specify whether to register memory using
+            ``cudaHostRegister()`` if hybrid memory mode is used. The default is
+            ``True``. See
+            :attr:`nvmath.bindings.cudss.ConfigParam.USE_CUDA_REGISTER_MEMORY`.
+
+    See Also:
+       :class:`ExecutionHybrid`, :class:`DirectSolver`, :func:`direct_solver`.
+    """
+
+    hybrid_memory_mode: bool = False
+    hybrid_device_memory_limit: int | str | None = None  # Internal heuristic.
+    register_cuda_memory: bool = True
+
+    def __post_init__(self):
+        if self.hybrid_device_memory_limit is not None:
+            mem_limit.check_memory_str(self.hybrid_device_memory_limit, "hybrid device memory limit")
+
+
+@dataclass
+class ExecutionCUDA:
+    """
+    A data class for providing GPU execution options to the :class:`DirectSolver`
+    object and the wrapper function :func:`direct_solver`.
+
+    Attributes:
+        device_id: CUDA device ordinal (only used if the operand resides on the CPU). The
+            default value is 0.
+        hybrid_memory_mode_options: Options controlling the use of hybrid (CPU-GPU) memory
+            as an object of type :class:`HybridMemoryModeOptions` or a `dict`.
+
+    See Also:
+       :class:`ExecutionHybrid`, :class:`HybridMemoryModeOptions`, :class:`DirectSolver`,
+       :func:`direct_solver`.
+    """
+
+    name: ClassVar[Literal["cuda"]] = "cuda"
+    device_id: int = 0
+    hybrid_memory_mode_options: object = field(default_factory=HybridMemoryModeOptions)
+
+
+@dataclass
+class ExecutionHybrid:
+    """
+    A data class for providing hybrid (CPU-GPU) execution options to the
+    :class:`DirectSolver` object and the wrapper function :func:`direct_solver`.
+
+    Attributes:
+        device_id: CUDA device ordinal (only used if the operand resides on the CPU). The
+            default value is 0.
+        num_threads: The number of CPU threads used to execute the plan.
+            If not specified, defaults to the number of CPU cores available to the process.
+
+    See Also:
+       :class:`ExecutionCUDA`, :class:`DirectSolver`, :func:`direct_solver`.
+    """
+
+    name: ClassVar[Literal["hybrid"]] = "hybrid"
+    device_id: int = 0
+    num_threads: int | None = None
+
+
+# TODO: docstring
+@dataclass
+class DirectSolverOptions:
+    """
+    A data class for providing options to the :class:`DirectSolver` object and the wrapper
+    function :func:`direct_solver`.
+
+    Attributes:
+        sparse_system_type (:class:`DirectSolverMatrixType`): The type of the sparse
+            system of equations (general, symmetric, symmetric positive definite, etc).
+            The default is ``DirectSolverMatrixType.GENERAL``.
+
+        sparse_system_view (:class:`DirectSolverMatrixViewType`): The desired view of the
+            sparse system of equations (full, upper, lower). The default is
+            ``DirectSolverMatrixViewType.FULL``.
+
+        multithreading_lib: The location (full path) to the library implementing the
+            threading layer interface. **TODO**: link to docs and default pip path.
+
+        logger (logging.Logger): Python Logger object. The root logger will be used if a
+            logger object is not provided.
+
+        blocking: A flag specifying the behavior of the execution functions and methods,
+            such as :func:`direct_solver` and :meth:`DirectSolver.solve`.
+            When ``blocking`` is `True`, the execution methods do not return until the
+            operation is complete. When
+            ``blocking`` is ``"auto"``, the methods return immediately when the input tensor
+            is on the GPU and ``execution`` is set to 'CUDA'. The execution methods always
+            block when the input tensor is on the CPU or ``execution`` is specified to 'CPU'
+            to ensure that the user doesn't inadvertently use the result before it becomes
+            available. The default is ``"auto"``.
+
+        handle: cuDSS library handle. A handle will be created if one is not provided.
+
+    See Also:
+        :class:`ExecutionCUDA`, :class:`ExecutionHybrid`, :class:`DirectSolver`, and
+        :func:`direct_solver`.
+    """
+
+    sparse_system_type: int = DirectSolverMatrixType.GENERAL
+    sparse_system_view: int = DirectSolverMatrixViewType.FULL
+    multithreading_lib: str | None = None
+    logger: Logger | None = None
+    blocking: Literal[True, "auto"] = "auto"
+    handle: int | None = None
+
+    def __post_init__(self):
+        if self.blocking not in (True, "auto"):
+            raise ValueError("The value specified for 'blocking' must be either True or 'auto'.")
+        self.sparse_system_type = DirectSolverMatrixType(self.sparse_system_type)
+        self.sparse_system_view = DirectSolverMatrixViewType(self.sparse_system_view)