nvmath-python 0.9.0__cp314-cp314-win_amd64.whl

nvmath/sparse/_internal/sparse_tensor_ifc.py

@@ -0,0 +1,121 @@
+# Copyright (c) 2026, 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
+
+    @property
+    @abstractmethod
+    def values(self):
+        raise NotImplementedError
+
+    @property
+    def dense_tensorholder_type(self):
+        """The tensor holder type for the dense constituent tensors."""
+        return self._dense_tensorholder_type
+
+    @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
+
+    @abstractmethod
+    def copy_(self, src: SparseTensorHolder, stream_holder: StreamHolder | None) -> None:
+        """Overwrite the sparse tensor (in-place) with a copy of src."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def to_ust(self, *, stream):
+        """Create an UST from the named representation. This is a zero-copy operation."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def to_package(self):
+        """
+        This will create a sparse tensor for the original package from which this
+        interface was created.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def release(self):
+        """
+        This method will release the wrapped tensor and any format-specific data
+        by setting them to None.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def reset_unchecked(self, tensor):
+        """
+        This method will reset the wrapped tensor to the specified one, and update
+        any format-specific data accordingly. It assumes that all attributes
+        like the device, shape, etc are consistent between the existing and new tensor.
+        """
+        raise NotImplementedError

nvmath/sparse/_internal/sparse_ust_ifc.py

@@ -0,0 +1,146 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Interface to wrap an UST that encapsulates an uncommon or novel format.
+"""
+
+from __future__ import annotations  # allows typehint of class methods to return the self class
+
+__all__ = ["USTensorHolder"]
+
+from typing import Literal, TypeVar
+
+from nvmath.internal.package_ifc import StreamHolder
+from nvmath.internal.tensor_ifc import TensorHolder
+
+from .sparse_tensor_ifc import SparseTensorHolder
+
+"""
+A generic wrapper type for a UST.
+"""
+USTensor = TypeVar("USTensor")
+
+
+class USTensorHolder(SparseTensorHolder):
+    """
+    A simple wrapper type for UST objects to enable them to be used in the same way as
+    the other abstractions like CSRTensorHolder etc.
+    """
+
+    _format_name: str = "UST"
+
+    def __init__(self, tensor):
+        # tensor is the native tensor.
+        self.tensor = tensor
+
+        # The tensor holder type for the dense constituent tensors.
+        self._dense_tensorholder_type = tensor._dense_tensorholder_type
+
+    @classmethod
+    def create_from_tensor(cls, tensor, *, attr_name_map=None):
+        # TODO: check that tensor is UST.
+        return USTensorHolder(tensor=tensor)
+
+    @property
+    def attr_name_map(self):
+        return {}
+
+    @property
+    def device(self):
+        return self.tensor.device
+
+    @property
+    def device_id(self):
+        return self.tensor.device_id
+
+    @property
+    def dtype(self):
+        return self.tensor.dtype
+
+    @property
+    def index_type(self):
+        return self.tensor.index_type
+
+    @property
+    def format_name(self) -> str:
+        return USTensorHolder._format_name
+
+    @property
+    def num_dimensions(self):
+        return self.tensor.num_dimensions
+
+    @property
+    def shape(self):
+        return self.tensor.shape
+
+    @property
+    def pos(self):
+        """The LevelMap object wrapping the pos array."""
+        return self.tensor._pos
+
+    @property
+    def crd(self):
+        """The LevelMap object wrapping the crd array."""
+        return self.tensor._crd
+
+    @property
+    def val(self) -> TensorHolder:
+        """The TensorHolder object wrapping the val array."""
+        return self.tensor._val
+
+    @property
+    def values(self) -> TensorHolder:
+        """The TensorHolder object wrapping the val array. We use the values
+        attribute name to get the NNZ etc."""
+        return self.tensor._val
+
+    def to(self, device_id: int | Literal["cpu"], stream_holder: StreamHolder | None):
+        """Copy the UST representation to a different device.
+
+        No copy is performed if the UST is already on the requested device.
+        """
+
+        # TODO: create an internal function that directly uses the stream holder.
+        stream = stream_holder.ptr if stream_holder is not None else None
+        target = self.tensor.to(device_id, stream)
+
+        return USTensorHolder(target)
+
+    def copy_(self, src, stream_holder: StreamHolder | None) -> None:
+        """Overwrite self.tensor (in-place) with a copy of src (USTensorHolder)."""
+
+        # TODO: create an internal function that directly uses the stream holder.
+        stream = stream_holder.ptr if stream_holder is not None else None
+        self.tensor.copy_(src.tensor, stream)
+
+    def to_ust(self, *, stream):
+        return self.tensor
+
+    def to_package(self):
+        """
+        This will create a sparse tensor for the original package from which this
+        interface was created.
+        """
+        raise NotImplementedError
+
+    def release(self):
+        """
+        This method will release the wrapped tensor and any format-specific data
+        by setting them to None.
+        """
+        # The tensor reference.
+        self.tensor = None
+
+    def reset_unchecked(self, tensor):
+        """
+        This method will reset the wrapped tensor to the specified one, and update
+        any format-specific data accordingly. It assumes that all attributes
+        like the device, shape, etc are consistent between the existing and new tensor.
+
+        Args:
+            tensor: The native tensor to wrap.
+        """
+        # Update the tensor reference.
+        self.tensor = tensor

nvmath/sparse/_internal/utils.py

@@ -0,0 +1,56 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import sys
+import typing
+
+from nvmath.internal.tensor_wrapper import wrap_operand
+
+
+def synchronize_dense_tensor(wrapped_tensor, *, stream_ptr=None):
+    """
+    This utility synchronizes the provided wrapped tensor for operations performed
+    on it in the stream specified by ``stream_ptr``.
+
+    The returned object is also a wrapped tensor.
+    """
+    if stream_ptr is None or wrapped_tensor.device == "cpu":
+        return wrapped_tensor
+
+    package = sys.modules[wrapped_tensor.name]
+    capsule = wrapped_tensor.tensor.__dlpack__(stream=stream_ptr)
+    tensor = package.from_dlpack(capsule)
+
+    wrapped_tensor = wrap_operand(tensor)
+    return wrapped_tensor
+
+
+def axis_order_in_memory(strides):
+    """
+    Compute the order in which the axes appear in memory.
+    """
+    if len(strides) == 0:
+        return ()
+
+    _, axis_order = zip(*sorted(zip(strides, range(len(strides)), strict=True)), strict=True)
+
+    return axis_order
+
+
+def calculate_strides(shape: typing.Sequence[int], axis_order: typing.Sequence[int], min_stride: int = 1):
+    """
+    Calculate the strides for the provided shape and axis order.
+    """
+    assert len(axis_order) == len(shape), f"axis_order length ({len(axis_order)}) must equal shape length ({len(shape)})"
+    assert len(set(axis_order)) == len(axis_order), f"axis_order must not contain duplicates: {axis_order}"
+    assert set(axis_order) == set(range(len(shape))), f"axis_order must be permutation of range({len(shape)}): {axis_order}"
+
+    strides: list[None | int] = [None] * len(shape)
+
+    stride = min_stride
+    for axis in axis_order:
+        strides[axis] = stride
+        stride *= shape[axis]
+
+    return strides

nvmath/sparse/advanced/__init__.py

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

nvmath/sparse/advanced/_configuration.py

@@ -0,0 +1,230 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = [
+    "DirectSolverOptions",
+    "DirectSolverMatrixType",
+    "DirectSolverMatrixViewType",
+    "ExecutionCUDA",
+    "ExecutionHybrid",
+    "HybridMemoryModeOptions",
+    "DirectSolverPlanPreferences",
+    "DirectSolverFactorizationPreferences",
+    "DirectSolverSolutionPreferences",
+]
+
+
+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`.
+
+    .. seealso::
+       :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`.
+
+    .. seealso::
+       :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.
+
+    .. seealso::
+       :class:`ExecutionCUDA`, :class:`DirectSolver`, :func:`direct_solver`.
+    """
+
+    name: ClassVar[Literal["hybrid"]] = "hybrid"
+    device_id: int = 0
+    num_threads: int | None = None
+
+
+@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, which depends on how nvmath-python is
+            :ref:`installed <nvmath installation>`. For example, if you install using
+            ``pip install nvmath-python[cu13]`` on Linux, the multithreading library can
+            be found in ``<install_path>/nvidia/cu13/lib/libcudss_mtlayer_gomp.so.0``,
+            where  ``<install_path>`` is the path shown by ``pip show nvidia-cudss-cu13``.
+
+        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.
+
+    .. seealso::
+        :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)
+
+
+@dataclass(slots=True, kw_only=True)
+class DirectSolverPlanPreferences:
+    """
+    A data class for providing plan preferences to the :func:`direct_solver` function.
+
+    Attributes:
+        host_nthreads: The number of host threads to use.
+        reordering_algorithm: The reordering algorithm to use.
+        pivot_type: The pivot type to use.
+        pivot_threshold: The pivot threshold to use.
+        max_nnz: The maximum number of non-zeros in the LU factorization.
+        use_matching: Whether to use matching.
+        matching_algorithm: The matching algorithm to use.
+        nd_min_levels: The number of levels in the nested dissection.
+        use_superpanels: Whether to use superpanels.
+
+    .. seealso::
+        :func:`direct_solver`
+        :class:`DirectSolverPlanConfig`
+        `cuDSS configuration documentation <https://docs.nvidia.com/cuda/cudss/types.html#cudssconfigparam-t>`_
+    """
+
+    host_nthreads: int | None = None
+    reordering_algorithm: cudss.AlgType | None = None
+    pivot_type: int | None = None
+    pivot_threshold: float | None = None
+    max_nnz: int | None = None
+    use_matching: bool | None = None
+    matching_algorithm: cudss.AlgType | None = None
+    nd_min_levels: int | None = None
+    use_superpanels: bool | None = None
+
+
+@dataclass(slots=True, kw_only=True)
+class DirectSolverFactorizationPreferences:
+    """
+    A data class for providing factorization preferences to
+    the :func:`direct_solver` function.
+
+    Attributes:
+        factorization_algorithm: The factorization algorithm to use.
+        pivot_eps_algorithm: The pivot epsilon algorithm to use.
+        pivot_eps: The pivot epsilon to use.
+
+    .. seealso::
+       :func:`direct_solver`
+       :class:`DirectSolverFactorizationConfig`
+       `cuDSS configuration documentation <https://docs.nvidia.com/cuda/cudss/types.html#cudssconfigparam-t>`_
+    """
+
+    factorization_algorithm: cudss.AlgType | None = None
+    pivot_eps_algorithm: cudss.AlgType | None = None
+    pivot_eps: float | None = None
+
+
+@dataclass(slots=True, kw_only=True)
+class DirectSolverSolutionPreferences:
+    """
+    A data class for providing solution preferences to the :func:`direct_solver` function.
+
+    Attributes:
+        solution_algorithm: The solution algorithm to use.
+        ir_num_steps: The number of steps to use for iterative refinement.
+
+    .. seealso::
+       :func:`direct_solver`
+       :class:`DirectSolverSolutionConfig`
+       `cuDSS configuration documentation <https://docs.nvidia.com/cuda/cudss/types.html#cudssconfigparam-t>`_
+    """
+
+    solution_algorithm: cudss.AlgType | None = None
+    ir_num_steps: int | None = None