nvmath-python 1.0.0__cp314-cp314t-win_amd64.whl

nvmath/sparse/_internal/cudss_data_ifc.py

@@ -0,0 +1,399 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Interface classes to encapsulate low-level calls to get or set cuDSS "data" attributes.
+"""
+
+__all__ = ["FactorizationInfo", "PlanInfo", "memory_estimates_dtype"]
+
+import re
+import threading
+from typing import TypeAlias
+
+import numpy as np
+
+from nvmath.bindings import cudss
+from nvmath.internal import utils
+
+DataParamEnum: TypeAlias = cudss.DataParam  # type: ignore[name-defined, attr-defined]
+
+# Future-proofing - cuDSS is currently not thread-safe.
+# https://docs.nvidia.com/cuda/cudss/general.html#thread-safety
+_tls = threading.local()
+_tls.size_written = np.empty((1,), dtype=np.uint64)
+
+
+# TODO: factor out common utilities (taken from fft.py).
+def complex_to_real_equivalent(name):
+    assert "complex" in name, f"Internal Error ({name=})"
+    m = name.split("complex")
+    assert len(m) in (1, 2)
+    size = int(m[-1]) // 2
+    if len(m) == 1:
+        return f"float{size}"
+    else:
+        return f"{m[0]}float{size}"
+
+
+def _get_attribute(handle, data_ptr, name, attribute, length=1):
+    """
+    name      = cudss enumerator for the attribute.
+    attribute = numpy ndarray object into which the value is stored.
+    """
+    cudss.data_get(
+        handle, data_ptr, name, attribute.ctypes.data, length * attribute.dtype.itemsize, _tls.size_written.ctypes.data
+    )
+    assert _tls.size_written[0] <= length * attribute.dtype.itemsize, "Internal error."
+
+
+def _check_valid_solver(info):
+    """
+    Check if the configuration points to a valid DirectSolver object.
+    """
+    if not info._solver.valid_state:
+        m = re.match(r"<nvmath\..*\.(.*Info) object at (.*)>$", str(info))
+        assert m is not None, "Internal error."
+        name = m.group(1)
+        address = m.group(2)
+
+        raise RuntimeError(f"The {name} object at {address} cannot be used after it's solver object is free'd.")
+
+
+memory_estimates_dtype = np.dtype(
+    [
+        ("permanent_device_memory", "<u8"),
+        ("peak_device_memory", "<u8"),
+        ("permanent_host_memory", "<u8"),
+        ("peak_host_memory", "<u8"),
+        ("hybrid_min_device_memory", "<u8"),
+        ("hybrid_max_device_memory", "<u8"),
+        ("reserved", "<u8", (10,)),
+    ]
+)
+
+
+class PlanInfo:
+    """
+    An interface to query information returned by
+    :meth:`nvmath.sparse.advanced.DirectSolver.plan`.
+    """
+
+    def __init__(self, solver):
+        """
+        ctor for internal use only.
+        """
+        self._solver = solver
+        self._handle = self._solver.handle
+        self._data_ptr = self._solver.data_ptr
+        assert self._data_ptr is not None, "Internal error."
+
+        self._N = self._solver._N
+        self._batched = self._solver.batched
+        self._index_type = self._solver.index_type
+        self._value_type = self._solver.value_type
+        self._nsuperpanels = np.zeros((1,), dtype=self._index_type)
+
+        # Allocate permutation arrays lazily, and only if not batched.
+        self._perm_reorder_col = None
+        self._perm_reorder_row = None
+        self._perm_matching = None
+
+        self._memory_estimates = np.zeros((), dtype=memory_estimates_dtype).view(np.recarray)
+
+    def _check_valid_solver_wrapper(self, *args, **kwargs):
+        _check_valid_solver(self)
+
+    @property
+    @utils.precondition(_check_valid_solver_wrapper)
+    def memory_estimates(self):
+        """
+        Query the memory estimates. See the
+        `cuDSS documentation
+        <https://docs.nvidia.com/cuda/cudss/types.html#cudssdataparam-t>`_
+        for more information.
+        """
+        _get_attribute(self._handle, self._data_ptr, DataParamEnum.MEMORY_ESTIMATES, self._memory_estimates, length=1)
+
+        return self._memory_estimates
+
+    @property
+    @utils.precondition(_check_valid_solver_wrapper)
+    def col_permutation(self):
+        """
+        Query the column permutation after planning (reordering). See the
+        `cuDSS documentation
+        <https://docs.nvidia.com/cuda/cudss/types.html#cudssdataparam-t>`_
+        for more information.
+        """
+        if self._batched:
+            raise RuntimeError("Column permutation is not available for batched systems.")
+
+        if self._perm_reorder_col is None:
+            self._perm_reorder_col = np.empty((self._N,), dtype=self._index_type)
+
+        _get_attribute(self._handle, self._data_ptr, DataParamEnum.PERM_REORDER_COL, self._perm_reorder_col, length=self._N)
+
+        return self._perm_reorder_col
+
+    @property
+    @utils.precondition(_check_valid_solver_wrapper)
+    def row_permutation(self):
+        """
+        Query the row permutation after planning (reordering). See the
+        `cuDSS documentation
+        <https://docs.nvidia.com/cuda/cudss/types.html#cudssdataparam-t>`_
+        for more information.
+        """
+        if self._batched:
+            raise RuntimeError("Row permutation is not available for batched systems.")
+
+        if self._perm_reorder_row is None:
+            self._perm_reorder_row = np.empty((self._N,), dtype=self._index_type)
+
+        _get_attribute(self._handle, self._data_ptr, DataParamEnum.PERM_REORDER_ROW, self._perm_reorder_row, length=self._N)
+
+        return self._perm_reorder_row
+
+    @property
+    @utils.precondition(_check_valid_solver_wrapper)
+    def matching_col_permutation(self):
+        """
+        Query the matching (column) permutation after planning (reordering). See the
+        `cuDSS documentation
+        <https://docs.nvidia.com/cuda/cudss/types.html#cudssdataparam-t>`_
+        for more information.
+        """
+        if self._batched:
+            raise RuntimeError("Matching (column) permutation is not available for batched systems.")
+
+        if self._solver.plan_config.matching_algorithm == cudss.MatchingAlg.NONE:
+            # documented in https://github.com/NVIDIA/CUDALibrarySamples/blob/main/cuDSS/simple_get_set/simple_get_set.cpp#L269-L274  # noqa: E501, W505
+            raise RuntimeError("Matching (column) permutation is not available if matching is not enabled.")
+
+        if self._perm_matching is None:
+            self._perm_matching = np.empty((self._N,), dtype=self._index_type)
+
+        _get_attribute(self._handle, self._data_ptr, DataParamEnum.PERM_MATCHING, self._perm_matching, length=self._N)
+
+        return self._perm_matching
+
+    @property
+    @utils.precondition(_check_valid_solver_wrapper)
+    def num_superpanels(self):
+        """
+        Query the number of number of superpanels after planning (symbolic factorization).
+        See the `cuDSS documentation
+        <https://docs.nvidia.com/cuda/cudss/types.html#cudssdataparam-t>`_
+        for more information.
+        """
+        _get_attribute(self._handle, self._data_ptr, DataParamEnum.NSUPERPANELS, self._nsuperpanels)
+
+        return self._nsuperpanels.item()
+
+
+class FactorizationInfo:
+    """
+    An interface to query information returned by
+    :meth:`nvmath.sparse.advanced.DirectSolver.factorize`.
+    """
+
+    def __init__(self, solver):
+        """
+        ctor for internal use only.
+        """
+        self._solver = solver
+        self._handle = self._solver.handle
+        self._data_ptr = self._solver.data_ptr
+        assert self._data_ptr is not None, "Internal error"
+
+        self._N = self._solver._N
+        self._batched = self._solver.batched
+        self._index_type = self._solver.index_type
+        self._value_type = self._solver.value_type
+
+        get_dtype = cudss.get_data_param_dtype
+
+        self._info = np.zeros((1,), dtype=get_dtype(DataParamEnum.INFO))
+        self._lu_nnz = np.zeros((1,), dtype=get_dtype(DataParamEnum.LU_NNZ))
+
+        # Allocate arrays that depend on the index or value type lazily.
+        self._npivots = None
+        self._inertia = None
+
+        # Allocate permutation and diagonal arrays lazily, and only if not batched.
+        self._perm_col = None
+        self._perm_row = None
+        self._diag = None
+
+        # Allocate matching row and col scale arrays lazily, and only if not batched.
+        self._scale_row = None
+        self._scale_col = None
+
+    def _check_valid_solver_wrapper(self, *args, **kwargs):
+        _check_valid_solver(self)
+
+    @property
+    @utils.precondition(_check_valid_solver_wrapper)
+    def info(self):
+        """
+        Query the error info after factorization. See the
+        `cuDSS documentation
+        <https://docs.nvidia.com/cuda/cudss/types.html#cudssdataparam-t>`_
+        for more information.
+        """
+        _get_attribute(self._handle, self._data_ptr, DataParamEnum.INFO, self._info)
+
+        return self._info.item()
+
+    @property
+    @utils.precondition(_check_valid_solver_wrapper)
+    def lu_nnz(self):
+        """
+        Query the number of non-zeros  in the LU factorization. See the
+        `cuDSS documentation
+        <https://docs.nvidia.com/cuda/cudss/types.html#cudssdataparam-t>`_
+        for more information.
+        """
+        _get_attribute(self._handle, self._data_ptr, DataParamEnum.LU_NNZ, self._lu_nnz)
+
+        return self._lu_nnz.item()
+
+    @property
+    @utils.precondition(_check_valid_solver_wrapper)
+    def npivots(self):
+        """
+        Query the number of pivots encountered in the LU factorization. See the
+        `cuDSS documentation
+        <https://docs.nvidia.com/cuda/cudss/types.html#cudssdataparam-t>`_
+        for more information.
+        """
+        if self._npivots is None:
+            self._npivots = np.empty((1,), dtype=self._index_type)
+
+        _get_attribute(self._handle, self._data_ptr, DataParamEnum.NPIVOTS, self._npivots)
+
+        return self._npivots.item()
+
+    @property
+    @utils.precondition(_check_valid_solver_wrapper)
+    def inertia(self):
+        """
+        Query the inertia (number of positive and negative pivots) encountered in the
+        LU factorization. See the
+        `cuDSS documentation
+        <https://docs.nvidia.com/cuda/cudss/types.html#cudssdataparam-t>`_
+        for more information.
+        """
+        if self._inertia is None:
+            self._inertia = np.empty((2,), dtype=self._index_type)
+
+        _get_attribute(self._handle, self._data_ptr, DataParamEnum.INERTIA, self._inertia, length=2)
+
+        return self._inertia
+
+    @property
+    @utils.precondition(_check_valid_solver_wrapper)
+    def col_permutation(self):
+        """
+        Query the column permutation after factorization. See the
+        `cuDSS documentation
+        <https://docs.nvidia.com/cuda/cudss/types.html#cudssdataparam-t>`_
+        for more information.
+        """
+        if self._batched:
+            raise RuntimeError("Column permutation is not available for batched systems.")
+
+        if self._perm_col is None:
+            self._perm_col = np.empty((self._N,), dtype=self._index_type)
+
+        _get_attribute(self._handle, self._data_ptr, DataParamEnum.PERM_COL, self._perm_col, length=self._N)
+
+        return self._perm_col
+
+    @property
+    @utils.precondition(_check_valid_solver_wrapper)
+    def row_permutation(self):
+        """
+        Query the row permutation after factorization. See the
+        `cuDSS documentation
+        <https://docs.nvidia.com/cuda/cudss/types.html#cudssdataparam-t>`_
+        for more information.
+        """
+        if self._batched:
+            raise RuntimeError("Row permutation is not available for batched systems.")
+
+        if self._perm_row is None:
+            self._perm_row = np.empty((self._N,), dtype=self._index_type)
+
+        _get_attribute(self._handle, self._data_ptr, DataParamEnum.PERM_ROW, self._perm_row, length=self._N)
+
+        return self._perm_row
+
+    @property
+    @utils.precondition(_check_valid_solver_wrapper)
+    def diag(self):
+        """
+        Query the diagonal of the factorized system. See the
+        `cuDSS documentation
+        <https://docs.nvidia.com/cuda/cudss/types.html#cudssdataparam-t>`_
+        for more information.
+        """
+        if self._batched:
+            raise RuntimeError("The factorized system's diagonal is not available for batched systems.")
+
+        if self._diag is None:
+            self._diag = np.empty((self._N,), dtype=self._value_type)
+
+        _get_attribute(self._handle, self._data_ptr, DataParamEnum.DIAG, self._diag, length=self._N)
+
+        return self._diag
+
+    @property
+    @utils.precondition(_check_valid_solver_wrapper)
+    def row_scale_factors(self):
+        """
+        Query the scale factors for the rows of the factorized system, if matching was
+        used. See the `cuDSS documentation
+        <https://docs.nvidia.com/cuda/cudss/types.html#cudssdataparam-t>`_
+        for more information.
+        """
+        if self._batched:
+            raise RuntimeError("The factorized system's row scale factors is not available for batched systems.")
+
+        if self._solver.plan_config.matching_algorithm == cudss.MatchingAlg.NONE:
+            raise RuntimeError("Row scale factors is only available if matching is enabled.")
+
+        if self._scale_row is None:
+            vtype = self._value_type
+            dtype = complex_to_real_equivalent(vtype) if "complex" in vtype else vtype
+            self._scale_row = np.empty((self._N,), dtype=dtype)
+
+        _get_attribute(self._handle, self._data_ptr, DataParamEnum.SCALE_ROW, self._scale_row, length=self._N)
+
+        return self._scale_row
+
+    @property
+    @utils.precondition(_check_valid_solver_wrapper)
+    def col_scale_factors(self):
+        """
+        Query the scale factors for the columns of the factorized system, if matching was
+        used. See the `cuDSS documentation
+        <https://docs.nvidia.com/cuda/cudss/types.html#cudssdataparam-t>`_
+        for more information.
+        """
+        if self._batched:
+            raise RuntimeError("The factorized system's column scale factors is not available for batched systems.")
+
+        if self._solver.plan_config.matching_algorithm == cudss.MatchingAlg.NONE:
+            raise RuntimeError("Column scale factors is only available if matching is enabled.")
+
+        if self._scale_col is None:
+            vtype = self._value_type
+            dtype = complex_to_real_equivalent(vtype) if "complex" in vtype else vtype
+            self._scale_col = np.empty((self._N,), dtype=dtype)
+
+        _get_attribute(self._handle, self._data_ptr, DataParamEnum.SCALE_COL, self._scale_col, length=self._N)
+
+        return self._scale_col