nvmath-python 1.0.0__cp314-cp314t-win_amd64.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- nvmath/__init__.pxd +0 -0
- nvmath/__init__.py +45 -0
- nvmath/_internal/__init__.py +0 -0
- nvmath/_internal/attribute_ifc_factory.py +330 -0
- nvmath/_internal/layout.py +70 -0
- nvmath/_internal/templates.py +130 -0
- nvmath/_internal/threadsafe.py +106 -0
- nvmath/_internal/utils.py +43 -0
- nvmath/_internal/workspace.py +490 -0
- nvmath/_utils.py +147 -0
- nvmath/bindings/__init__.py +60 -0
- nvmath/bindings/_internal/__init__.pxd +0 -0
- nvmath/bindings/_internal/__init__.py +0 -0
- nvmath/bindings/_internal/common_types.pxd +31 -0
- nvmath/bindings/_internal/cublas.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/_internal/cublas.pxd +530 -0
- nvmath/bindings/_internal/cublasLt.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/_internal/cublasLt.pxd +59 -0
- nvmath/bindings/_internal/cublasMp.pxd +52 -0
- nvmath/bindings/_internal/cudss.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/_internal/cudss.pxd +54 -0
- nvmath/bindings/_internal/cufft.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/_internal/cufft.pxd +70 -0
- nvmath/bindings/_internal/cufftMp.pxd +77 -0
- nvmath/bindings/_internal/curand.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/_internal/curand.pxd +42 -0
- nvmath/bindings/_internal/cusolver.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/_internal/cusolver.pxd +15 -0
- nvmath/bindings/_internal/cusolverDn.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/_internal/cusolverDn.pxd +406 -0
- nvmath/bindings/_internal/cusolverMp.pxd +71 -0
- nvmath/bindings/_internal/cusolverSp.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/_internal/cusolverSp.pxd +75 -0
- nvmath/bindings/_internal/cusparse.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/_internal/cusparse.pxd +471 -0
- nvmath/bindings/_internal/cusparseLt.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/_internal/cusparseLt.pxd +48 -0
- nvmath/bindings/_internal/cutensor.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/_internal/cutensor.pxd +58 -0
- nvmath/bindings/_internal/mathdx.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/_internal/mathdx.pxd +116 -0
- nvmath/bindings/_internal/nvshmem.pxd +29 -0
- nvmath/bindings/_internal/utils.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/_internal/utils.pxd +174 -0
- nvmath/bindings/_internal/utils.pyi +10 -0
- nvmath/bindings/cublas.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cublas.pxd +558 -0
- nvmath/bindings/cublas.pyi +812 -0
- nvmath/bindings/cublasLt.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cublasLt.pxd +109 -0
- nvmath/bindings/cublasLt.pyi +1461 -0
- nvmath/bindings/cublasMp.pxd +85 -0
- nvmath/bindings/cublasMp.pyi +267 -0
- nvmath/bindings/cudss.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cudss.pxd +98 -0
- nvmath/bindings/cudss.pyi +443 -0
- nvmath/bindings/cufft.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cufft.pxd +118 -0
- nvmath/bindings/cufft.pyi +301 -0
- nvmath/bindings/cufftMp.pxd +124 -0
- nvmath/bindings/cufftMp.pyi +326 -0
- nvmath/bindings/curand.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/curand.pxd +71 -0
- nvmath/bindings/curand.pyi +189 -0
- nvmath/bindings/cusolver.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cusolver.pxd +62 -0
- nvmath/bindings/cusolver.pyi +320 -0
- nvmath/bindings/cusolverDn.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cusolverDn.pxd +430 -0
- nvmath/bindings/cusolverDn.pyi +422 -0
- nvmath/bindings/cusolverMp.pxd +98 -0
- nvmath/bindings/cusolverMp.pyi +114 -0
- nvmath/bindings/cusolverSp.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cusolverSp.pxd +95 -0
- nvmath/bindings/cusolverSp.pyi +70 -0
- nvmath/bindings/cusparse.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cusparse.pxd +546 -0
- nvmath/bindings/cusparse.pyi +1017 -0
- nvmath/bindings/cusparseLt.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cusparseLt.pxd +99 -0
- nvmath/bindings/cusparseLt.pyi +252 -0
- nvmath/bindings/cutensor.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cutensor.pxd +98 -0
- nvmath/bindings/cutensor.pyi +324 -0
- nvmath/bindings/cycublas.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cycublas.pxd +664 -0
- nvmath/bindings/cycublasLt.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cycublasLt.pxd +1045 -0
- nvmath/bindings/cycublasMp.pxd +171 -0
- nvmath/bindings/cycudss.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cycudss.pxd +277 -0
- nvmath/bindings/cycufft.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cycufft.pxd +333 -0
- nvmath/bindings/cycufftMp.pxd +342 -0
- nvmath/bindings/cycurand.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cycurand.pxd +141 -0
- nvmath/bindings/cycusolver.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cycusolver.pxd +137 -0
- nvmath/bindings/cycusolverDn.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cycusolverDn.pxd +443 -0
- nvmath/bindings/cycusolverMp.pxd +107 -0
- nvmath/bindings/cycusolverSp.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cycusolverSp.pxd +93 -0
- nvmath/bindings/cycusparse.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cycusparse.pxd +679 -0
- nvmath/bindings/cycusparseLt.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cycusparseLt.pxd +135 -0
- nvmath/bindings/cycutensor.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cycutensor.pxd +189 -0
- nvmath/bindings/cymathdx.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cymathdx.pxd +552 -0
- nvmath/bindings/cynvshmem.pxd +118 -0
- nvmath/bindings/mathdx.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/mathdx.pxd +182 -0
- nvmath/bindings/mathdx.pyi +1562 -0
- nvmath/bindings/nvpl/__init__.pxd +0 -0
- nvmath/bindings/nvpl/__init__.py +13 -0
- nvmath/bindings/nvpl/_internal/__init__.pxd +0 -0
- nvmath/bindings/nvpl/_internal/__init__.py +0 -0
- nvmath/bindings/nvpl/_internal/blas.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/nvpl/_internal/blas.pxd +237 -0
- nvmath/bindings/nvpl/_internal/fft.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/nvpl/_internal/fft.pxd +36 -0
- nvmath/bindings/nvpl/blas.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/nvpl/blas.pxd +131 -0
- nvmath/bindings/nvpl/blas.pyi +168 -0
- nvmath/bindings/nvpl/cyblas.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/nvpl/cyblas.pxd +280 -0
- nvmath/bindings/nvpl/cyfft.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/nvpl/cyfft.pxd +93 -0
- nvmath/bindings/nvpl/fft.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/nvpl/fft.pxd +100 -0
- nvmath/bindings/nvpl/fft.pyi +168 -0
- nvmath/bindings/nvshmem.pxd +54 -0
- nvmath/bindings/nvshmem.pyi +191 -0
- nvmath/device/__init__.py +38 -0
- nvmath/device/_deprecated.py +33 -0
- nvmath/device/common.py +315 -0
- nvmath/device/common_backend.py +131 -0
- nvmath/device/common_cuda.py +201 -0
- nvmath/device/common_numba.py +300 -0
- nvmath/device/common_numba_cuda_mlir.py +202 -0
- nvmath/device/common_opaque_tensor.py +201 -0
- nvmath/device/cublasdx.py +1606 -0
- nvmath/device/cublasdx_backend.py +860 -0
- nvmath/device/cublasdx_numba.py +1534 -0
- nvmath/device/cublasdx_numba_cuda_mlir.py +208 -0
- nvmath/device/cufftdx.py +373 -0
- nvmath/device/cufftdx_backend.py +220 -0
- nvmath/device/cufftdx_numba.py +140 -0
- nvmath/device/cufftdx_numba_cuda_mlir.py +79 -0
- nvmath/device/curand_kernel.py +9147 -0
- nvmath/device/cusolverdx.py +2708 -0
- nvmath/device/cusolverdx_backend.py +440 -0
- nvmath/device/cusolverdx_numba.py +567 -0
- nvmath/device/cusolverdx_numba_cuda_mlir.py +604 -0
- nvmath/device/cusolverdx_overload_backend.py +1029 -0
- nvmath/device/llvm_array.py +29 -0
- nvmath/device/random.py +441 -0
- nvmath/device/random_helpers.py +23 -0
- nvmath/device/random_states.py +187 -0
- nvmath/device/types.py +138 -0
- nvmath/device/vector_types_numba.py +259 -0
- nvmath/distributed/__init__.py +200 -0
- nvmath/distributed/_internal/__init__.py +0 -0
- nvmath/distributed/_internal/nccl.py +86 -0
- nvmath/distributed/_internal/nvshmem.py +307 -0
- nvmath/distributed/_internal/symmetric_memory.py +35 -0
- nvmath/distributed/_internal/tensor_ifc.py +70 -0
- nvmath/distributed/_internal/tensor_ifc_cupy.py +68 -0
- nvmath/distributed/_internal/tensor_ifc_host_device.py +172 -0
- nvmath/distributed/_internal/tensor_ifc_numpy.py +46 -0
- nvmath/distributed/_internal/tensor_ifc_torch.py +162 -0
- nvmath/distributed/_internal/tensor_wrapper.py +81 -0
- nvmath/distributed/_utils.py +167 -0
- nvmath/distributed/distribution/__init__.py +30 -0
- nvmath/distributed/distribution/_configuration.py +39 -0
- nvmath/distributed/distribution/distributions.py +1024 -0
- nvmath/distributed/distribution/redistribute.py +1284 -0
- nvmath/distributed/fft/__init__.py +7 -0
- nvmath/distributed/fft/_configuration.py +82 -0
- nvmath/distributed/fft/fft.py +2742 -0
- nvmath/distributed/linalg/__init__.py +22 -0
- nvmath/distributed/linalg/_internal/__init__.py +3 -0
- nvmath/distributed/linalg/_internal/epilog_protocol.py +586 -0
- nvmath/distributed/linalg/_internal/matmul_desc_ifc.py +28 -0
- nvmath/distributed/linalg/advanced/__init__.py +8 -0
- nvmath/distributed/linalg/advanced/_configuration.py +171 -0
- nvmath/distributed/linalg/advanced/matmulmod.py +3573 -0
- nvmath/distributed/linalg/generic/__init__.py +8 -0
- nvmath/distributed/linalg/generic/_caching.py +66 -0
- nvmath/distributed/linalg/generic/_configuration.py +61 -0
- nvmath/distributed/linalg/generic/_factorization.py +172 -0
- nvmath/distributed/linalg/generic/_initialization.py +966 -0
- nvmath/distributed/linalg/generic/_problem_spec.py +511 -0
- nvmath/distributed/linalg/generic/solvermod.py +1368 -0
- nvmath/distributed/process_group.py +408 -0
- nvmath/fft/__init__.py +7 -0
- nvmath/fft/_configuration.py +189 -0
- nvmath/fft/_exec_utils.py +82 -0
- nvmath/fft/_helpers.py +237 -0
- nvmath/fft/fft.py +3122 -0
- nvmath/internal/__init__.pxd +3 -0
- nvmath/internal/__init__.py +10 -0
- nvmath/internal/_bindings.cp314t-win_amd64.pyd +0 -0
- nvmath/internal/_bindings.pxd +18 -0
- nvmath/internal/_device_utils.py +45 -0
- nvmath/internal/_layout/__init__.pxd +3 -0
- nvmath/internal/_layout/__init__.py +7 -0
- nvmath/internal/_layout/_layout.cp314t-win_amd64.pyd +0 -0
- nvmath/internal/_layout/_layout.pxd +1303 -0
- nvmath/internal/_layout/_layout.pyi +1145 -0
- nvmath/internal/enum_utils.py +142 -0
- nvmath/internal/formatters.py +87 -0
- nvmath/internal/mem_limit.py +51 -0
- nvmath/internal/memory.cp314t-win_amd64.pyd +0 -0
- nvmath/internal/memory.pxd +13 -0
- nvmath/internal/memory.pyi +50 -0
- nvmath/internal/ndbuffer/__init__.pxd +3 -0
- nvmath/internal/ndbuffer/__init__.py +9 -0
- nvmath/internal/ndbuffer/_copy_kernel.cp314t-win_amd64.pyd +0 -0
- nvmath/internal/ndbuffer/_copy_kernel.pxd +10 -0
- nvmath/internal/ndbuffer/_jit.cp314t-win_amd64.pyd +0 -0
- nvmath/internal/ndbuffer/_jit.pxd +7 -0
- nvmath/internal/ndbuffer/_ndbuffer.cp314t-win_amd64.pyd +0 -0
- nvmath/internal/ndbuffer/_ndbuffer.pxd +40 -0
- nvmath/internal/ndbuffer/_ndbuffer.pyi +463 -0
- nvmath/internal/ndbuffer/copy_kernel/args.h +34 -0
- nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/array_view.h +52 -0
- nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/elementwise.h +68 -0
- nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/grid_indexer.h +69 -0
- nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/transposed.h +242 -0
- nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/type_utils.h +39 -0
- nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/utils.h +132 -0
- nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/vec.h +159 -0
- nvmath/internal/ndbuffer/copy_kernel/elementwise.h +53 -0
- nvmath/internal/ndbuffer/copy_kernel/transposed.h +58 -0
- nvmath/internal/package_ifc.py +168 -0
- nvmath/internal/package_ifc_cuda.py +57 -0
- nvmath/internal/package_ifc_cupy.py +67 -0
- nvmath/internal/package_ifc_torch.py +69 -0
- nvmath/internal/package_wrapper.py +14 -0
- nvmath/internal/tensor_ifc.py +179 -0
- nvmath/internal/tensor_ifc_cupy.py +234 -0
- nvmath/internal/tensor_ifc_ndbuffer.py +147 -0
- nvmath/internal/tensor_ifc_numpy.py +184 -0
- nvmath/internal/tensor_ifc_torch.py +178 -0
- nvmath/internal/tensor_wrapper.py +160 -0
- nvmath/internal/typemaps.py +113 -0
- nvmath/internal/utils.py +805 -0
- nvmath/linalg/__init__.py +56 -0
- nvmath/linalg/_internal/__init__.py +3 -0
- nvmath/linalg/_internal/algo_cap_ifc.py +82 -0
- nvmath/linalg/_internal/algo_config_ifc.py +43 -0
- nvmath/linalg/_internal/batch.py +234 -0
- nvmath/linalg/_internal/enum_to_tuples.py +64 -0
- nvmath/linalg/_internal/epilog_protocol.py +766 -0
- nvmath/linalg/_internal/layout.py +624 -0
- nvmath/linalg/_internal/matmul_desc_ifc.py +28 -0
- nvmath/linalg/_internal/matmul_pref_ifc.py +27 -0
- nvmath/linalg/_internal/matrix_layout_ifc.py +26 -0
- nvmath/linalg/_internal/solver_utils.py +432 -0
- nvmath/linalg/_internal/typemaps.py +144 -0
- nvmath/linalg/_internal/utils.py +157 -0
- nvmath/linalg/advanced/__init__.py +8 -0
- nvmath/linalg/advanced/_algorithmmod.py +170 -0
- nvmath/linalg/advanced/_configuration.py +351 -0
- nvmath/linalg/advanced/helpers/__init__.py +5 -0
- nvmath/linalg/advanced/helpers/matmul.py +1316 -0
- nvmath/linalg/advanced/matmulmod.py +3734 -0
- nvmath/linalg/generic/__init__.py +53 -0
- nvmath/linalg/generic/_configuration/__init__.py +39 -0
- nvmath/linalg/generic/_configuration/layout.py +263 -0
- nvmath/linalg/generic/_configuration/match.py +734 -0
- nvmath/linalg/generic/_configuration/qualifiers.py +493 -0
- nvmath/linalg/generic/_configuration/solver_configuration.py +59 -0
- nvmath/linalg/generic/_configuration/wrap.py +217 -0
- nvmath/linalg/generic/_dtype.py +15 -0
- nvmath/linalg/generic/matmulmod.py +2094 -0
- nvmath/linalg/generic/solvermod.py +1301 -0
- nvmath/memory.py +279 -0
- nvmath/sparse/__init__.py +38 -0
- nvmath/sparse/_internal/__init__.py +21 -0
- nvmath/sparse/_internal/common_utils.py +147 -0
- nvmath/sparse/_internal/cudss_config_ifc.py +702 -0
- nvmath/sparse/_internal/cudss_data_ifc.py +399 -0
- nvmath/sparse/_internal/cudss_utils.py +506 -0
- nvmath/sparse/_internal/cusparse_utils.py +382 -0
- nvmath/sparse/_internal/sparse_bsc_ifc.py +303 -0
- nvmath/sparse/_internal/sparse_bsr_ifc.py +305 -0
- nvmath/sparse/_internal/sparse_coo_ifc.py +256 -0
- nvmath/sparse/_internal/sparse_csc_ifc.py +268 -0
- nvmath/sparse/_internal/sparse_csr_ifc.py +288 -0
- nvmath/sparse/_internal/sparse_dia_ifc.py +242 -0
- nvmath/sparse/_internal/sparse_format_helpers.py +601 -0
- nvmath/sparse/_internal/sparse_tensor_ifc.py +133 -0
- nvmath/sparse/_internal/sparse_ust_ifc.py +141 -0
- nvmath/sparse/_internal/utils.py +56 -0
- nvmath/sparse/advanced/__init__.py +7 -0
- nvmath/sparse/advanced/_configuration.py +227 -0
- nvmath/sparse/advanced/direct_solver.py +2069 -0
- nvmath/sparse/generic/__init__.py +7 -0
- nvmath/sparse/generic/_configuration.py +129 -0
- nvmath/sparse/generic/_helpers.py +137 -0
- nvmath/sparse/generic/_thunks.py +21 -0
- nvmath/sparse/generic/matmulmod.py +2353 -0
- nvmath/sparse/ust/__init__.py +7 -0
- nvmath/sparse/ust/_converters.py +422 -0
- nvmath/sparse/ust/_cpp.py +28 -0
- nvmath/sparse/ust/_drawer.py +565 -0
- nvmath/sparse/ust/_emitter.py +1033 -0
- nvmath/sparse/ust/_jit.py +188 -0
- nvmath/sparse/ust/_kernel.py +282 -0
- nvmath/sparse/ust/_utils.py +149 -0
- nvmath/sparse/ust/interfaces/__init__.py +0 -0
- nvmath/sparse/ust/interfaces/torch_interface.py +476 -0
- nvmath/sparse/ust/tensor.py +1016 -0
- nvmath/sparse/ust/tensor_format.py +957 -0
- nvmath/tensor/__init__.py +6 -0
- nvmath/tensor/_configuration.py +120 -0
- nvmath/tensor/_internal/__init__.py +3 -0
- nvmath/tensor/_internal/cutensor_config_ifc.py +279 -0
- nvmath/tensor/_internal/cutensor_utils.py +230 -0
- nvmath/tensor/_internal/data.py +43 -0
- nvmath/tensor/_internal/einsum_parser.py +444 -0
- nvmath/tensor/_internal/typemaps.py +96 -0
- nvmath/tensor/contract.py +1900 -0
- nvmath_python-1.0.0.dist-info/METADATA +134 -0
- nvmath_python-1.0.0.dist-info/RECORD +332 -0
- nvmath_python-1.0.0.dist-info/WHEEL +5 -0
- nvmath_python-1.0.0.dist-info/licenses/LICENSE +177 -0
- nvmath_python-1.0.0.dist-info/top_level.txt +2 -0
|
@@ -0,0 +1,2069 @@
|
|
|
1
|
+
# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
|
|
2
|
+
#
|
|
3
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
__all__ = [
|
|
7
|
+
"direct_solver",
|
|
8
|
+
"DirectSolver",
|
|
9
|
+
"DirectSolverReorderingAlg",
|
|
10
|
+
"DirectSolverFactorizationAlg",
|
|
11
|
+
"DirectSolverSolveAlg",
|
|
12
|
+
"DirectSolverPivotEpsilonAlg",
|
|
13
|
+
"DirectSolverMatchingAlg",
|
|
14
|
+
"DirectSolverFactorizationConfig",
|
|
15
|
+
"DirectSolverFactorizationInfo",
|
|
16
|
+
"DirectSolverPlanConfig",
|
|
17
|
+
"DirectSolverPlanInfo",
|
|
18
|
+
"DirectSolverSolutionConfig",
|
|
19
|
+
"DirectSolverPlanPreferences",
|
|
20
|
+
"DirectSolverFactorizationPreferences",
|
|
21
|
+
"DirectSolverSolutionPreferences",
|
|
22
|
+
]
|
|
23
|
+
|
|
24
|
+
import itertools
|
|
25
|
+
import logging
|
|
26
|
+
import math
|
|
27
|
+
import operator
|
|
28
|
+
import os
|
|
29
|
+
from collections.abc import Sequence
|
|
30
|
+
from dataclasses import fields, replace
|
|
31
|
+
from typing import Any, TypeAlias
|
|
32
|
+
|
|
33
|
+
from nvmath.bindings import cudss
|
|
34
|
+
from nvmath.internal import formatters, tensor_wrapper, utils
|
|
35
|
+
from nvmath.internal.memory import get_device_memory_resource
|
|
36
|
+
from nvmath.internal.package_wrapper import StreamHolder
|
|
37
|
+
from nvmath.internal.typemaps import NAME_TO_DATA_TYPE
|
|
38
|
+
from nvmath.sparse._internal import common_utils as sp_utils
|
|
39
|
+
from nvmath.sparse._internal import cudss_config_ifc, cudss_data_ifc, cudss_utils
|
|
40
|
+
from nvmath.sparse.advanced._configuration import (
|
|
41
|
+
DirectSolverFactorizationPreferences,
|
|
42
|
+
DirectSolverOptions,
|
|
43
|
+
DirectSolverPlanPreferences,
|
|
44
|
+
DirectSolverSolutionPreferences,
|
|
45
|
+
ExecutionCUDA,
|
|
46
|
+
ExecutionHybrid,
|
|
47
|
+
HybridMemoryModeOptions,
|
|
48
|
+
)
|
|
49
|
+
|
|
50
|
+
VALID_INDEX_TYPES = ("int32", "int64")
|
|
51
|
+
|
|
52
|
+
VALID_DTYPES = ("float32", "float64", "complex64", "complex128")
|
|
53
|
+
|
|
54
|
+
# Qualified names for public export.
|
|
55
|
+
DirectSolverPlanConfig: TypeAlias = cudss_config_ifc.PlanConfig
|
|
56
|
+
DirectSolverFactorizationConfig: TypeAlias = cudss_config_ifc.FactorizationConfig
|
|
57
|
+
DirectSolverSolutionConfig: TypeAlias = cudss_config_ifc.SolutionConfig
|
|
58
|
+
DirectSolverPlanInfo: TypeAlias = cudss_data_ifc.PlanInfo
|
|
59
|
+
DirectSolverFactorizationInfo: TypeAlias = cudss_data_ifc.FactorizationInfo
|
|
60
|
+
DirectSolverReorderingAlg: TypeAlias = cudss.ReorderingAlg
|
|
61
|
+
DirectSolverFactorizationAlg: TypeAlias = cudss.FactorizationAlg
|
|
62
|
+
DirectSolverSolveAlg: TypeAlias = cudss.SolveAlg
|
|
63
|
+
DirectSolverPivotEpsilonAlg: TypeAlias = cudss.PivotEpsilonAlg
|
|
64
|
+
DirectSolverMatchingAlg: TypeAlias = cudss.MatchingAlg
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
def get_threading_lib(library=None):
|
|
68
|
+
"""
|
|
69
|
+
Return the name of the threading library, if defined using an environment variable and
|
|
70
|
+
the path is valid, or None.
|
|
71
|
+
"""
|
|
72
|
+
if library is None:
|
|
73
|
+
library = os.getenv("CUDSS_THREADING_LIB")
|
|
74
|
+
if library is None or not os.path.isfile(library):
|
|
75
|
+
return
|
|
76
|
+
return library
|
|
77
|
+
|
|
78
|
+
|
|
79
|
+
def set_async_workspace_allocator(handle: int, device_id: int, logger: logging.Logger):
|
|
80
|
+
"""
|
|
81
|
+
By default, cuDSS allocates workspace memory using synchronous cudaMalloc/cudaFree.
|
|
82
|
+
This call sets a custom allocator that forwards the alloc/dealloc calls to cuda
|
|
83
|
+
stream-aware device memory pool.
|
|
84
|
+
"""
|
|
85
|
+
with utils.device_ctx(device_id) as device:
|
|
86
|
+
if device.properties.memory_pools_supported:
|
|
87
|
+
pool = get_device_memory_resource(device_id)
|
|
88
|
+
if cudss.set_async_workspace_allocator(handle, int(pool.handle)):
|
|
89
|
+
logger.info("An async workspace allocator has been set for the provided handle.")
|
|
90
|
+
return True
|
|
91
|
+
else:
|
|
92
|
+
logger.info(
|
|
93
|
+
f"Could not set async workspace allocator for handle {handle} "
|
|
94
|
+
f"on device {device_id}. The provided handle already has "
|
|
95
|
+
f"a custom allocator set."
|
|
96
|
+
)
|
|
97
|
+
else:
|
|
98
|
+
logger.info(
|
|
99
|
+
f"The device {device_id} does not support memory pools. "
|
|
100
|
+
f"Workspace memory will be allocated using synchronous cudaMalloc/cudaFree."
|
|
101
|
+
)
|
|
102
|
+
return False
|
|
103
|
+
|
|
104
|
+
|
|
105
|
+
def check_dense_tensor_layout(shape: Sequence[int], strides: Sequence[int], *, explicitly_batched=None):
|
|
106
|
+
"""
|
|
107
|
+
Check that the dense vector or matrix (each sample matrix in a N-D tensor) is
|
|
108
|
+
in col-major format.
|
|
109
|
+
"""
|
|
110
|
+
|
|
111
|
+
num_dimensions = len(shape)
|
|
112
|
+
assert explicitly_batched is not None and len(strides) == num_dimensions, "Internal Error."
|
|
113
|
+
|
|
114
|
+
# For explicitly batched matrices, each sample must be a matrix or vector.
|
|
115
|
+
if explicitly_batched and num_dimensions > 2:
|
|
116
|
+
return False
|
|
117
|
+
|
|
118
|
+
first = int(num_dimensions > 1)
|
|
119
|
+
|
|
120
|
+
# Only column-major matrices are currently supported.
|
|
121
|
+
is_col_major = strides[-1 - first] == 1
|
|
122
|
+
|
|
123
|
+
# Check that the LD doesn't lead to overlapping memory.
|
|
124
|
+
# For batched matrices, LD has to be equal to the first matrix dimension.
|
|
125
|
+
comparison_op_for_ld = operator.eq if num_dimensions > 2 else operator.ge
|
|
126
|
+
|
|
127
|
+
def compare_ld(strides, shape):
|
|
128
|
+
if shape[-1] == 1:
|
|
129
|
+
# NOTE: a special case for F-order matrix with shape (..., m, 1)
|
|
130
|
+
# The stride of the last dimension is moot as we never need
|
|
131
|
+
# to move along that dimension. Meanwhile array frameworks
|
|
132
|
+
# may specify the stride of this dummy dimension to be 1 or m,
|
|
133
|
+
# or possibly some other value depending on how the array is created.
|
|
134
|
+
# Therefore we here skip the comparison.
|
|
135
|
+
return True
|
|
136
|
+
return comparison_op_for_ld(strides[-1], shape[-2])
|
|
137
|
+
|
|
138
|
+
if num_dimensions > 1:
|
|
139
|
+
return is_col_major and compare_ld(strides, shape)
|
|
140
|
+
|
|
141
|
+
return is_col_major
|
|
142
|
+
|
|
143
|
+
|
|
144
|
+
def check_rhs_sequence_layout(
|
|
145
|
+
shape: Sequence[Sequence[int]],
|
|
146
|
+
strides: Sequence[Sequence[int]],
|
|
147
|
+
):
|
|
148
|
+
return all(check_dense_tensor_layout(s, d, explicitly_batched=True) for s, d in zip(shape, strides, strict=True))
|
|
149
|
+
|
|
150
|
+
|
|
151
|
+
# TODO: unify to tensor_wrapper.to() and the axis utilities below.
|
|
152
|
+
def copy_single_or_sequence(operands, device_id, stream_holder):
|
|
153
|
+
if isinstance(operands, Sequence):
|
|
154
|
+
return tuple(o.to(device_id, stream_holder) for o in operands)
|
|
155
|
+
|
|
156
|
+
return operands.to(device_id, stream_holder)
|
|
157
|
+
|
|
158
|
+
|
|
159
|
+
def axis_order_in_memory(shape, strides):
|
|
160
|
+
"""
|
|
161
|
+
Compute the order in which the axes appear in memory.
|
|
162
|
+
"""
|
|
163
|
+
# The shape is used to resolve cases like (1, 2, 1) : (2, 1, 1) in CuTe notation.
|
|
164
|
+
_, _, axis_order = zip(*sorted(zip(strides, shape, range(len(strides)), strict=True)), strict=True)
|
|
165
|
+
|
|
166
|
+
return axis_order
|
|
167
|
+
|
|
168
|
+
|
|
169
|
+
def calculate_strides(shape, axis_order):
|
|
170
|
+
"""
|
|
171
|
+
Calculate the strides for the provided shape and axis order.
|
|
172
|
+
"""
|
|
173
|
+
strides = [None] * len(shape)
|
|
174
|
+
|
|
175
|
+
stride = 1
|
|
176
|
+
for axis in axis_order:
|
|
177
|
+
strides[axis] = stride
|
|
178
|
+
stride *= shape[axis]
|
|
179
|
+
|
|
180
|
+
return strides
|
|
181
|
+
|
|
182
|
+
|
|
183
|
+
def update_config_with_preferences(config, preferences):
|
|
184
|
+
"""
|
|
185
|
+
Update the configuration with the provided preferences
|
|
186
|
+
"""
|
|
187
|
+
if isinstance(config, DirectSolverPlanConfig):
|
|
188
|
+
preferences = utils.check_or_create_options(DirectSolverPlanPreferences, preferences, "plan preferences")
|
|
189
|
+
elif isinstance(config, DirectSolverFactorizationConfig):
|
|
190
|
+
preferences = utils.check_or_create_options(
|
|
191
|
+
DirectSolverFactorizationPreferences, preferences, "factorization preferences"
|
|
192
|
+
)
|
|
193
|
+
elif isinstance(config, DirectSolverSolutionConfig):
|
|
194
|
+
preferences = utils.check_or_create_options(DirectSolverSolutionPreferences, preferences, "solution preferences")
|
|
195
|
+
else:
|
|
196
|
+
raise ValueError(f"Invalid config type: {type(config)}")
|
|
197
|
+
for field in fields(preferences):
|
|
198
|
+
value = getattr(preferences, field.name)
|
|
199
|
+
if value is not None:
|
|
200
|
+
setattr(config, field.name, value)
|
|
201
|
+
return
|
|
202
|
+
|
|
203
|
+
|
|
204
|
+
def wrap_cudss_supported_lhs(lhs):
|
|
205
|
+
try:
|
|
206
|
+
lhs = sp_utils.wrap_sparse_operands(lhs)
|
|
207
|
+
except Exception as e:
|
|
208
|
+
raise TypeError(
|
|
209
|
+
"The LHS must be an N-D sparse CSR array/tensor or a sequence of 2D "
|
|
210
|
+
"sparse CSR array/tensor from one of the supported packages: CuPy, PyTorch, or SciPy."
|
|
211
|
+
) from e
|
|
212
|
+
if isinstance(lhs, Sequence):
|
|
213
|
+
# the sp_utils enforces all operands to have the same sparse format
|
|
214
|
+
# (and there's at least one), so we use just the first one to verify the format
|
|
215
|
+
format_name = lhs[0].format_name
|
|
216
|
+
else:
|
|
217
|
+
format_name = lhs.format_name
|
|
218
|
+
if format_name != "CSR":
|
|
219
|
+
raise TypeError(
|
|
220
|
+
f"DirectSolver does not support {format_name} sparse format. "
|
|
221
|
+
f"The LHS must be an N-D sparse CSR array/tensor or a sequence of 2D sparse CSR "
|
|
222
|
+
f"array/tensor from one of the supported packages: CuPy, PyTorch, or SciPy."
|
|
223
|
+
)
|
|
224
|
+
return lhs
|
|
225
|
+
|
|
226
|
+
|
|
227
|
+
SHARED_DSS_DOCUMENTATION = utils.COMMON_SHARED_DOC_MAP.copy()
|
|
228
|
+
SHARED_DSS_DOCUMENTATION.update(
|
|
229
|
+
{
|
|
230
|
+
"a": """\
|
|
231
|
+
The sparse operand (or sequence of operands) representing the left-hand side (LHS) of the system of equations. The LHS
|
|
232
|
+
operand may be a (sequence of) :class:`scipy.sparse.csr_matrix`, :class:`scipy.sparse.csr_array`,
|
|
233
|
+
:class:`cupyx.scipy.sparse.csr_matrix`, or :py:func:`torch.sparse_csr_tensor`. That is, the LHS is a sparse matrix or
|
|
234
|
+
tensor in Compressed Sparse Row (CSR) format from one of the supported packages: SciPy, CuPy, PyTorch. Refer to the
|
|
235
|
+
:ref:`semantics <Semantics>` section for details.
|
|
236
|
+
""".replace("\n", " "),
|
|
237
|
+
#
|
|
238
|
+
"b": """\
|
|
239
|
+
The ndarray/tensor or (sequence of ndarray/tensors) representing the dense right-hand side (RHS) of the system of equations.
|
|
240
|
+
The RHS operand may be a (sequence of) :class:`numpy.ndarray`, :class:`cupy.ndarray`, and :class:`torch.Tensor`. Refer to the
|
|
241
|
+
:ref:`semantics <Semantics>` section for details.
|
|
242
|
+
""".replace("\n", " "),
|
|
243
|
+
#
|
|
244
|
+
"options": """\
|
|
245
|
+
Specify options for the direct sparse solver as a :class:`DirectSolverOptions` object. Alternatively, a `dict` containing
|
|
246
|
+
the parameters for the ``DirectSolverOptions`` constructor can also be provided. If not specified, the value will be set
|
|
247
|
+
to the default-constructed ``DirectSolverOptions`` object.""".replace("\n", " "),
|
|
248
|
+
#
|
|
249
|
+
"execution": """\
|
|
250
|
+
Specify execution space options for the direct solver as a :class:`ExecutionCUDA` or :class:`ExecutionHybrid` object.
|
|
251
|
+
Alternatively, a string ('cuda' or 'hybrid'), or a `dict` with the 'name' key set to 'cuda' or 'hybrid' and optional
|
|
252
|
+
parameters relevant to the given execution space. The default execution space is 'cuda' and the corresponding
|
|
253
|
+
:class:`ExecutionCUDA` object will be default-constructed.""".replace("\n", " "),
|
|
254
|
+
#
|
|
255
|
+
"result": """\
|
|
256
|
+
The result of the specified sparse direct solve, which has the same shape, remains on the same device, and belongs to the
|
|
257
|
+
same package as the RHS ``b``. If ``b`` is a sequence, the result ``x`` is also a sequence of ndarray/tensor, each of which
|
|
258
|
+
has the same shape as the corresponding tensor in ``b``.
|
|
259
|
+
""".replace("\n", " "),
|
|
260
|
+
#
|
|
261
|
+
"semantics": r"""\
|
|
262
|
+
The sparse direct solver solves :math:`a @ x = b` for ``x`` given the left-hand side (LHS) ``a`` and the right-hand side
|
|
263
|
+
(RHS) ``b``.
|
|
264
|
+
|
|
265
|
+
* In the simplest version with no batching, ``a`` is sparse (square) matrix of size ``n`` in Compressed Sparse Row (CSR)
|
|
266
|
+
format, and ``b`` is a dense vector or matrix. A matrix (2D ndarray or tensor) in the RHS is treated as multiple
|
|
267
|
+
column vectors corresponding to multiple solution vectors (i.e. ``x`` is the same shape as the RHS) to solve for.
|
|
268
|
+
|
|
269
|
+
.. important:: Currently, only column-major (Fortran) layout is supported for the RHS.
|
|
270
|
+
|
|
271
|
+
* Batching can be specified either **explicitly** or **implicitly**.
|
|
272
|
+
|
|
273
|
+
* An **explicitly-batched** LHS is provided as a Python sequence of sparse CSR matrices :math:`[a_0, a_1, ..., a_n]`.
|
|
274
|
+
Likewise an explicitly-batched RHS is provided as a sequence of vectors or matrices :math:`[b_0, b_1, ..., b_n]`.
|
|
275
|
+
The solver will solve all :math:`n` systems :math:`a_i @ x_i = b_i` for the solution sequence :math:`x_i`.
|
|
276
|
+
Each sample in explicit batching can be of a different size, with the only constraint being that a given
|
|
277
|
+
LHS size is consistent with that of its corresponding RHS.
|
|
278
|
+
|
|
279
|
+
* An **implicitly-batched** LHS is provided as a higher-dimensional :math:`N \geq 3D` sparse tensor in CSR format,
|
|
280
|
+
where the leading :math:`N - 2` dimensions are the batch dimensions, and the last two dimensions correspond to that
|
|
281
|
+
of the :math:`n \times n` sparse system. Currently, only PyTorch supports higher-dimensional sparse CSR tensors.
|
|
282
|
+
Likewise, an implicitly-batched RHS is provided as a :math:`N \geq 3D` dense ndarray/tensor, where the leading
|
|
283
|
+
:math:`N - 2` dimensions are the batch dimensions, and the last two dimensions correspond to the :math:`n \times 1`
|
|
284
|
+
vector or :math:`n \times m` matrix for each sample. The solver solves :math:`a_i @ x_i = b_i` for each sample
|
|
285
|
+
:math:`i` in the batch, and the solution ``x`` has the same shape as the RHS ``b``.
|
|
286
|
+
|
|
287
|
+
* Each sample :math:`a_i` and :math:`b_i` in the (explicitly- or implicitly-specified) batch are essentially
|
|
288
|
+
treated as, and subject to, the same rules as the case with no batching.
|
|
289
|
+
|
|
290
|
+
* The LHS and RHS batch specification is independent: for example, the LHS can be explicitly-batched while the
|
|
291
|
+
RHS is implicitly-batched (or vice-versa). The same batch specification can be used for both as well.
|
|
292
|
+
|
|
293
|
+
* The solution ``x`` always has the same form as the RHS ``b``. It is a sequence of matrices or vectors if
|
|
294
|
+
``b`` is explicitly-batched, or a higher-dimensional ndarray/tensor if ``b`` is implicitly-batched.
|
|
295
|
+
|
|
296
|
+
.. tip:: For a description of the CSR sparse format, see
|
|
297
|
+
`here <https://docs.nvidia.com/nvpl/latest/sparse/storage_format/sparse_matrix.html#compressed-sparse-row-csr>`_
|
|
298
|
+
for example.
|
|
299
|
+
""".strip(),
|
|
300
|
+
#
|
|
301
|
+
"plan_preferences": """\
|
|
302
|
+
Specify plan preferences as a :class:`DirectSolverPlanPreferences` object. Alternatively, a `dict` containing
|
|
303
|
+
the parameters for the ``DirectSolverPlanPreferences`` constructor can also be provided. If not specified,
|
|
304
|
+
the default cuDSS plan configuration will be used.""".replace("\n", " "),
|
|
305
|
+
#
|
|
306
|
+
"factorization_preferences": """\
|
|
307
|
+
Specify factorization preferences as a :class:`DirectSolverFactorizationPreferences` object.
|
|
308
|
+
Alternatively, a `dict` containing the parameters for the ``DirectSolverFactorizationPreferences`` constructor can
|
|
309
|
+
also be provided. If not specified, the default cuDSS factorization configuration will be used.""".replace("\n", " "),
|
|
310
|
+
#
|
|
311
|
+
"solution_preferences": """\
|
|
312
|
+
Specify solution preferences as a :class:`DirectSolverSolutionPreferences` object.
|
|
313
|
+
Alternatively, a `dict` containing the parameters for the ``DirectSolverSolutionPreferences`` constructor
|
|
314
|
+
can also be provided. If not specified, the default cuDSS solution configuration will be used.""".replace("\n", " "),
|
|
315
|
+
#
|
|
316
|
+
"reset_operands_unchecked": utils._reset_operand_unchecked_docstring(True, experimental=False),
|
|
317
|
+
#
|
|
318
|
+
"release_operands": utils._release_operand_docstring(
|
|
319
|
+
True, version_added="0.9.0", execution_methods=("factorize", "solve")
|
|
320
|
+
),
|
|
321
|
+
#
|
|
322
|
+
"stream": """\
|
|
323
|
+
Provide the CUDA stream to use for executing the operation. Acceptable inputs include
|
|
324
|
+
``cudaStream_t`` (as Python :class:`int`), :class:`cupy.cuda.Stream`, and
|
|
325
|
+
:class:`torch.cuda.Stream`. If a stream is not provided, the current stream for the
|
|
326
|
+
operand device will be queried from the dense RHS operand ``b`` package.""".replace("\n", " "),
|
|
327
|
+
}
|
|
328
|
+
)
|
|
329
|
+
|
|
330
|
+
|
|
331
|
+
class InvalidDirectSolverState(Exception):
|
|
332
|
+
pass
|
|
333
|
+
|
|
334
|
+
|
|
335
|
+
def _check_cudss_version():
|
|
336
|
+
required = (0, 8)
|
|
337
|
+
available = cudss.get_property(0), cudss.get_property(1)
|
|
338
|
+
if available != required:
|
|
339
|
+
raise RuntimeError(
|
|
340
|
+
f"nvmath-python requires cuDSS version {'.'.join(str(i) for i in required)}, while you "
|
|
341
|
+
f"have cuDSS version {'.'.join(str(i) for i in available)}."
|
|
342
|
+
)
|
|
343
|
+
|
|
344
|
+
|
|
345
|
+
@utils.docstring_decorator(SHARED_DSS_DOCUMENTATION, skip_missing=False)
|
|
346
|
+
class DirectSolver:
|
|
347
|
+
"""
|
|
348
|
+
Create a stateful object that encapsulates the specified sparse direct solver
|
|
349
|
+
computations and required resources. This object ensures the validity of resources
|
|
350
|
+
during use and releases them when they are no longer needed to prevent misuse.
|
|
351
|
+
|
|
352
|
+
This object encompasses all functionalities of function-form API :func:`direct_solver`,
|
|
353
|
+
which is a convenience wrapper around it. The stateful object also allows for the
|
|
354
|
+
amortization of preparatory costs when the same solve operation is to be performed
|
|
355
|
+
with different left-hand side (LHS) and right-hand side (RHS) with the same problem
|
|
356
|
+
specification (see :meth:`reset_operands` for more details).
|
|
357
|
+
|
|
358
|
+
Using the stateful object typically involves the following steps:
|
|
359
|
+
|
|
360
|
+
1. **Problem Specification**: Initialize the object with the defined operation and
|
|
361
|
+
options.
|
|
362
|
+
2. **Preparation**: Use :meth:`plan` for reordering to minimize fill-in and
|
|
363
|
+
symbolic factorization for this specific direct sparse solver operation.
|
|
364
|
+
3. **Execution**: Factorize and solve the system.
|
|
365
|
+
4. **Resource Management**: Ensure all resources are released either by explicitly
|
|
366
|
+
calling :meth:`free` or by managing the stateful object within a context manager.
|
|
367
|
+
|
|
368
|
+
Detailed information on each step described above can be obtained by passing in a
|
|
369
|
+
:class:`logging.Logger` object to :class:`DirectSolverOptions` or by setting the
|
|
370
|
+
appropriate options in the root logger object, which is used by default:
|
|
371
|
+
|
|
372
|
+
>>> import logging
|
|
373
|
+
>>> logging.basicConfig(
|
|
374
|
+
... level=logging.INFO,
|
|
375
|
+
... format="%(asctime)s %(levelname)-8s %(message)s",
|
|
376
|
+
... datefmt="%m-%d %H:%M:%S",
|
|
377
|
+
... )
|
|
378
|
+
|
|
379
|
+
Args:
|
|
380
|
+
a: {a}
|
|
381
|
+
|
|
382
|
+
b: {b}
|
|
383
|
+
|
|
384
|
+
options: {options}
|
|
385
|
+
|
|
386
|
+
execution: {execution}
|
|
387
|
+
|
|
388
|
+
stream: {stream}
|
|
389
|
+
|
|
390
|
+
Semantics:
|
|
391
|
+
{semantics}
|
|
392
|
+
|
|
393
|
+
.. seealso::
|
|
394
|
+
:class:`DirectSolverPlanConfig`, :class:`DirectSolverFactorizationConfig`,
|
|
395
|
+
:class:`DirectSolverSolutionConfig`, :class:`DirectSolverPlanInfo`,
|
|
396
|
+
:class:`DirectSolverFactorizationInfo`, :class:`DirectSolverOptions`,
|
|
397
|
+
:class:`ExecutionCUDA`, :class:`ExecutionHybrid`, :meth:`plan`,
|
|
398
|
+
:meth:`release_operands`, :meth:`reset_operands`,
|
|
399
|
+
:meth:`factorize`, :meth:`solve`.
|
|
400
|
+
|
|
401
|
+
Examples:
|
|
402
|
+
|
|
403
|
+
>>> import numpy as np
|
|
404
|
+
>>> import scipy.sparse as sp
|
|
405
|
+
>>> import nvmath
|
|
406
|
+
|
|
407
|
+
Create a sparse float64 ndarray in CSR format on the CPU for the LHS.
|
|
408
|
+
|
|
409
|
+
>>> n = 16
|
|
410
|
+
>>> a = sp.random_array((n, n), density=0.5, format="csr", dtype="float64")
|
|
411
|
+
|
|
412
|
+
Ensure that the randomly-generated LHS is not singular.
|
|
413
|
+
|
|
414
|
+
>>> a += sp.diags_array([2.0] * n, format="csr", dtype="float64")
|
|
415
|
+
|
|
416
|
+
The RHS can be a vector or matrix. Here we create a random vector.
|
|
417
|
+
|
|
418
|
+
>>> b = np.random.rand(n).astype(dtype="float64")
|
|
419
|
+
|
|
420
|
+
We will define a sparse direct solver operation for solving the system a @ x = b
|
|
421
|
+
using the specialized sparse direct solver interface.
|
|
422
|
+
|
|
423
|
+
>>> solver = nvmath.sparse.advanced.DirectSolver(a, b)
|
|
424
|
+
|
|
425
|
+
Options can be provided above to control the behavior of the operation using the
|
|
426
|
+
`options` argument (see :class:`DirectSolverOptions`).
|
|
427
|
+
|
|
428
|
+
Next, plan the operation. The planning operation can be configured through
|
|
429
|
+
the :class:`DirectSolverPlanConfig` interface, which is accessed through
|
|
430
|
+
:attr:`plan_config`.
|
|
431
|
+
|
|
432
|
+
>>> plan_config = solver.plan_config
|
|
433
|
+
|
|
434
|
+
Here we set the reordering algorithm to choice 1.
|
|
435
|
+
|
|
436
|
+
>>> ReorderingAlg = nvmath.sparse.advanced.DirectSolverReorderingAlg
|
|
437
|
+
>>> plan_config.reordering_algorithm = ReorderingAlg.NESTED_DISSECTION
|
|
438
|
+
|
|
439
|
+
Plan the operation, which reorders the system to minimize fill-in and performs the
|
|
440
|
+
symbolic factorization. Planning returns a :class:`DirectSolverPlanInfo` object,
|
|
441
|
+
whose attributes (such as row or column permutation) can be queried.
|
|
442
|
+
|
|
443
|
+
>>> plan_info = solver.plan()
|
|
444
|
+
>>> plan_info.col_permutation
|
|
445
|
+
array([ 5, 4, 15, 13, 8, 12, 11, 9, 0, 14, 6, 3, 10, 2, 1, 7],
|
|
446
|
+
dtype=int32)
|
|
447
|
+
|
|
448
|
+
The next step is to perform the numerical factorization of the system using
|
|
449
|
+
:meth:`factorize`. Similar to planning, the numerical factorization step can
|
|
450
|
+
also be configured if desired.
|
|
451
|
+
|
|
452
|
+
>>> fac_config = solver.factorization_config
|
|
453
|
+
|
|
454
|
+
Here we set the pivot epsilon value to ``1e-14``, instead of the default ``1e-13``.
|
|
455
|
+
|
|
456
|
+
>>> fac_config.pivot_eps = 1e-14
|
|
457
|
+
|
|
458
|
+
Factorize the system, which returns a :class:`DirectSolverFactorizationInfo` object,
|
|
459
|
+
whose attributes can be inspected. We print the Sylvester inertia here.
|
|
460
|
+
|
|
461
|
+
>>> fac_info = solver.factorize()
|
|
462
|
+
>>> fac_info.inertia
|
|
463
|
+
array([0, 0], dtype=int32)
|
|
464
|
+
|
|
465
|
+
Now solve the factorized system, and obtain the result `x` as a NumPy ndarray on
|
|
466
|
+
the CPU.
|
|
467
|
+
|
|
468
|
+
>>> x = solver.solve()
|
|
469
|
+
|
|
470
|
+
Finally, free the object's resources. To avoid having to explicitly make this
|
|
471
|
+
call, it's recommended to use the DirectSolver object as a context manager as
|
|
472
|
+
shown below, if possible.
|
|
473
|
+
|
|
474
|
+
>>> solver.free()
|
|
475
|
+
|
|
476
|
+
.. note:: All :class:`DirectSolver` methods execute on the package current stream
|
|
477
|
+
by default. Alternatively, the `stream` argument can be used to run a method on
|
|
478
|
+
a specified stream.
|
|
479
|
+
|
|
480
|
+
Let's now look at a batched solve with PyTorch operands on the GPU.
|
|
481
|
+
|
|
482
|
+
Create a 3D complex128 PyTorch sparse tensor on the GPU representing the LHS, along
|
|
483
|
+
with the corresponding RHS:
|
|
484
|
+
|
|
485
|
+
>>> import torch
|
|
486
|
+
>>> n = 8
|
|
487
|
+
>>> batch = 2
|
|
488
|
+
>>> device_id = 0
|
|
489
|
+
|
|
490
|
+
Prepare sample input data. Create a diagonally-dominant random CSR matrix.
|
|
491
|
+
|
|
492
|
+
>>> a = torch.rand(n, n, dtype=torch.complex128) + torch.diag(torch.tensor([2.0] * n))
|
|
493
|
+
>>> a = torch.stack([a] * batch, dim=0)
|
|
494
|
+
>>> a = a.to_sparse_csr()
|
|
495
|
+
|
|
496
|
+
.. important:: PyTorch uses int64 for index buffers, whereas cuDSS currently
|
|
497
|
+
requires int32. So we'll have to convert the indices.
|
|
498
|
+
|
|
499
|
+
>>> a = torch.sparse_csr_tensor(
|
|
500
|
+
... a.crow_indices().to(dtype=torch.int32),
|
|
501
|
+
... a.col_indices().to(dtype=torch.int32),
|
|
502
|
+
... a.values(),
|
|
503
|
+
... size=a.size(),
|
|
504
|
+
... device=device_id,
|
|
505
|
+
... )
|
|
506
|
+
|
|
507
|
+
Create the RHS, which can be a matrix or vector in column-major layout.
|
|
508
|
+
|
|
509
|
+
>>> b = torch.ones(batch, 3, n, dtype=torch.complex128, device=device_id)
|
|
510
|
+
>>> b = b.permute(0, 2, 1)
|
|
511
|
+
|
|
512
|
+
Create a :class:`DirectSolver` object encapsulating the problem specification
|
|
513
|
+
described earlier and use it as a context manager.
|
|
514
|
+
|
|
515
|
+
>>> with nvmath.sparse.advanced.DirectSolver(a, b) as solver:
|
|
516
|
+
... plan_info = solver.plan()
|
|
517
|
+
...
|
|
518
|
+
... # Factorize the system.
|
|
519
|
+
... fac_info = solver.factorize()
|
|
520
|
+
...
|
|
521
|
+
... # Solve the factorized system.
|
|
522
|
+
... x1 = solver.solve()
|
|
523
|
+
...
|
|
524
|
+
... # Update the RHS b in-place (see reset_operands() for an alternative).
|
|
525
|
+
... b[...] = torch.rand(*b.shape, dtype=torch.complex128, device=device_id)
|
|
526
|
+
...
|
|
527
|
+
... # Solve again to get the new result.
|
|
528
|
+
... x2 = solver.solve()
|
|
529
|
+
|
|
530
|
+
All the resources used by the object are released at the end of the block.
|
|
531
|
+
|
|
532
|
+
Batching can be implicitly-specified as shown above, or explicitly-specified as
|
|
533
|
+
a sequence for both the LHS and the RHS. This, as well as other options and
|
|
534
|
+
usage patterns, are illustrated in the examples found in the
|
|
535
|
+
`nvmath/examples/sparse/advanced/direct_solver
|
|
536
|
+
<https://github.com/NVIDIA/nvmath-python/tree/main/examples/sparse/advanced/direct_solver>`_
|
|
537
|
+
directory.
|
|
538
|
+
""" # noqa: W505
|
|
539
|
+
|
|
540
|
+
def __init__(
|
|
541
|
+
self,
|
|
542
|
+
a,
|
|
543
|
+
b,
|
|
544
|
+
*,
|
|
545
|
+
options: DirectSolverOptions | dict[str, Any] | None = None,
|
|
546
|
+
execution: ExecutionCUDA | ExecutionHybrid | str | dict[str, Any] | None = None,
|
|
547
|
+
stream: utils.AnyStream | int | None = None,
|
|
548
|
+
):
|
|
549
|
+
# Check if the required cuDSS version is available.
|
|
550
|
+
_check_cudss_version()
|
|
551
|
+
|
|
552
|
+
# Process options.
|
|
553
|
+
self.options: DirectSolverOptions = utils.check_or_create_options(
|
|
554
|
+
DirectSolverOptions, options, "sparse direct solver options"
|
|
555
|
+
) # type: ignore[assignment]
|
|
556
|
+
|
|
557
|
+
# Process execution options. The default execution space is CUDA.
|
|
558
|
+
self.execution_options = utils.check_or_create_one_of_options(
|
|
559
|
+
(ExecutionCUDA, ExecutionHybrid), execution, "execution options", default_name="cuda"
|
|
560
|
+
)
|
|
561
|
+
if self.execution_options.name == "cuda":
|
|
562
|
+
self.execution_options = replace(
|
|
563
|
+
self.execution_options,
|
|
564
|
+
hybrid_memory_mode_options=utils.check_or_create_options(
|
|
565
|
+
HybridMemoryModeOptions,
|
|
566
|
+
self.execution_options.hybrid_memory_mode_options,
|
|
567
|
+
"hybrid memory mode options",
|
|
568
|
+
),
|
|
569
|
+
)
|
|
570
|
+
|
|
571
|
+
self.logger = self.options.logger if self.options.logger is not None else logging.getLogger()
|
|
572
|
+
self.logger.info("= SPECIFICATION PHASE =")
|
|
573
|
+
|
|
574
|
+
# Wrap the LHS.
|
|
575
|
+
self.a = wrap_cudss_supported_lhs(a)
|
|
576
|
+
|
|
577
|
+
# The LHS can be implicitly (N-D CSR tensor) or explicitly batched (provided as a
|
|
578
|
+
# sequence of CSR matrices).
|
|
579
|
+
self.implicitly_batched_lhs = False
|
|
580
|
+
self.explicitly_batched_lhs = isinstance(self.a, Sequence)
|
|
581
|
+
if not self.explicitly_batched_lhs:
|
|
582
|
+
self.implicitly_batched_lhs = self.a.num_dimensions > 2
|
|
583
|
+
|
|
584
|
+
# For explicitly batched LHS, check that all operands are matrices.
|
|
585
|
+
if self.explicitly_batched_lhs and any(a.num_dimensions > 2 for a in self.a):
|
|
586
|
+
raise TypeError(
|
|
587
|
+
f"""Every operator in the batched LHS provided as a sequence (explicit batching) must be a CSR matrix.
|
|
588
|
+
The specified LHS sequence = {a}."""
|
|
589
|
+
)
|
|
590
|
+
|
|
591
|
+
# The determination of a batched solve is purely based on the LHS. The RHS for each
|
|
592
|
+
# sample in the batch for explicit batching can be a vector or matrix, with the
|
|
593
|
+
# latter being considered multiple-RHS vectors as opposed to a batch. An
|
|
594
|
+
# implicitly-batched RHS is always 3D or higher dimension (implicitly-batched
|
|
595
|
+
# vectors are batched matrices with the last dimension of unit extent).
|
|
596
|
+
self.batched = self.explicitly_batched_lhs or self.implicitly_batched_lhs
|
|
597
|
+
|
|
598
|
+
# The LHS batch shape should be empty for explicit batching.
|
|
599
|
+
self.lhs_batch_shape = None if self.explicitly_batched_lhs else tuple(self.a.shape[:-2])
|
|
600
|
+
|
|
601
|
+
# Set the LHS package.
|
|
602
|
+
if self.explicitly_batched_lhs:
|
|
603
|
+
# The package is the same for explicitly batched LHS since we've have
|
|
604
|
+
# successfully wrapped them.
|
|
605
|
+
self.lhs_package = utils.infer_object_package(self.a[0].tensor)
|
|
606
|
+
else: # Single or implicitly-batched LHS.
|
|
607
|
+
self.lhs_package = utils.infer_object_package(self.a.tensor)
|
|
608
|
+
|
|
609
|
+
# Determine the batch count and batch indices.
|
|
610
|
+
self.batch_count = 0
|
|
611
|
+
self.batch_indices = None # Needed only for implicit batching.
|
|
612
|
+
if self.explicitly_batched_lhs:
|
|
613
|
+
self.batch_count = len(self.a)
|
|
614
|
+
elif self.implicitly_batched_lhs:
|
|
615
|
+
self.batch_count = math.prod(self.lhs_batch_shape) # type: ignore[arg-type]
|
|
616
|
+
# Create the sequence of batch coordinates to use for creating batched CSR
|
|
617
|
+
# matrix type.
|
|
618
|
+
self.batch_indices = tuple(itertools.product(*list(map(range, self.lhs_batch_shape)))) # type: ignore
|
|
619
|
+
|
|
620
|
+
self.logger.info(f"The LHS package is {self.lhs_package}.")
|
|
621
|
+
if self.explicitly_batched_lhs:
|
|
622
|
+
self.logger.info(f"The LHS is explicitly batched, with a batch count of {self.batch_count}.")
|
|
623
|
+
elif self.implicitly_batched_lhs:
|
|
624
|
+
self.logger.info(
|
|
625
|
+
f"The LHS is implicitly batched with shape = {self.a.shape}, dtype = {self.a.dtype}, \
|
|
626
|
+
and index type = {self.a.index_type}."
|
|
627
|
+
)
|
|
628
|
+
self.logger.info(f"The LHS batch shape is {self.lhs_batch_shape}, with a batch count of {self.batch_count}.")
|
|
629
|
+
self.logger.debug(f"The batch indices (generated from the LHS) are: {self.batch_indices}.")
|
|
630
|
+
|
|
631
|
+
# Wrap the RHS. It can be a N-D tensor or a sequence of matrices or vectors.
|
|
632
|
+
self.rhs_batch_shape = ()
|
|
633
|
+
self.implicitly_batched_rhs = self.explicitly_batched_rhs = False
|
|
634
|
+
if isinstance(b, Sequence):
|
|
635
|
+
self.explicitly_batched_rhs = True
|
|
636
|
+
self.b: Any = tensor_wrapper.wrap_operands(b)
|
|
637
|
+
# For explicitly batched RHS, check that all operands are vectors or matrices.
|
|
638
|
+
if any(len(r.shape) > 2 for r in self.b):
|
|
639
|
+
raise TypeError(
|
|
640
|
+
f"""Every RHS object in the batched RHS provided as a sequence (explicit batching) must be a dense matrix
|
|
641
|
+
or vector. The specified RHS sequence = {b}."""
|
|
642
|
+
)
|
|
643
|
+
rhs_batch_count = len(self.b)
|
|
644
|
+
else:
|
|
645
|
+
self.b = tensor_wrapper.wrap_operand(b) # type:ignore
|
|
646
|
+
self.implicitly_batched_rhs = len(self.b.shape) > 2
|
|
647
|
+
self.rhs_batch_shape = tuple(self.b.shape[:-2]) # type: ignore
|
|
648
|
+
if self.implicitly_batched_lhs and self.lhs_batch_shape != self.rhs_batch_shape:
|
|
649
|
+
raise TypeError(
|
|
650
|
+
f"The batch shapes for the LHS {self.lhs_batch_shape} and RHS {self.rhs_batch_shape} must match."
|
|
651
|
+
)
|
|
652
|
+
rhs_batch_count = math.prod(self.rhs_batch_shape) if self.rhs_batch_shape else 0
|
|
653
|
+
if self.batch_indices is None:
|
|
654
|
+
# Create the sequence of batch coordinates to use for creating batched dense
|
|
655
|
+
# matrix type.
|
|
656
|
+
self.batch_indices = tuple(itertools.product(*list(map(range, self.rhs_batch_shape))))
|
|
657
|
+
self.logger.debug(f"The batch indices (generated from the RHS) are: {self.batch_indices}.")
|
|
658
|
+
|
|
659
|
+
# The LHS can be implicitly or explicitly batched, independent from how the RHS is
|
|
660
|
+
# batched. So we need to check that the batch counts match.
|
|
661
|
+
if rhs_batch_count != self.batch_count:
|
|
662
|
+
raise TypeError(f"The batch count for the LHS {self.batch_count} and RHS {rhs_batch_count} must match.")
|
|
663
|
+
|
|
664
|
+
# Set the RHS package.
|
|
665
|
+
if self.explicitly_batched_rhs:
|
|
666
|
+
# The package is the same for explicitly batched RHS since we've have
|
|
667
|
+
# successfully wrapped them.
|
|
668
|
+
self.rhs_package = utils.infer_object_package(self.b[0].tensor)
|
|
669
|
+
else: # Single or implicitly-batched RHS
|
|
670
|
+
self.rhs_package = utils.infer_object_package(self.b.tensor)
|
|
671
|
+
|
|
672
|
+
self.logger.info(f"The RHS package is {self.rhs_package}.")
|
|
673
|
+
if self.explicitly_batched_rhs:
|
|
674
|
+
self.logger.info(f"The RHS is explicitly batched, with a batch count of {rhs_batch_count}.")
|
|
675
|
+
elif self.implicitly_batched_rhs:
|
|
676
|
+
self.logger.info(f"The RHS is implicitly batched with shape = {self.b.shape} and dtype = {self.b.dtype}")
|
|
677
|
+
self.logger.info(f"The RHS batch shape is {self.rhs_batch_shape}, with a batch count of {rhs_batch_count}.")
|
|
678
|
+
|
|
679
|
+
# Note that while the LHS and RHS packages can be different they must be compatible
|
|
680
|
+
# (such as scipy-numpy, cupyx-cupy).
|
|
681
|
+
if (self.lhs_package, self.rhs_package) not in cudss_utils.COMPATIBLE_LHS_RHS_PACKAGES:
|
|
682
|
+
raise TypeError(
|
|
683
|
+
f"""The LHS package {self.lhs_package} and RHS package {self.rhs_package} are not part of the
|
|
684
|
+
compatible choices: {cudss_utils.COMPATIBLE_LHS_RHS_PACKAGES}."""
|
|
685
|
+
)
|
|
686
|
+
|
|
687
|
+
# Get key LHS attributes and perform more basic checks.
|
|
688
|
+
if self.explicitly_batched_lhs:
|
|
689
|
+
# For sequence, the get_[attribute]() functions check for consistency within
|
|
690
|
+
# the sequence.
|
|
691
|
+
self.device_id = utils.get_operands_device_id(self.a)
|
|
692
|
+
self.value_type = utils.get_operands_dtype(self.a)
|
|
693
|
+
self.index_type = sp_utils.get_operands_index_type(self.a)
|
|
694
|
+
self.lhs_shape = shapes = tuple(o.shape for o in self.a)
|
|
695
|
+
if any(len(s) != 2 or s[0] != s[1] for s in shapes):
|
|
696
|
+
raise TypeError("Each object in an explicitly-batched LHS must be a CSR matrix of shape (N, N).")
|
|
697
|
+
self._N = tuple(s[0] for s in shapes)
|
|
698
|
+
self.lhs_nnz = tuple(o.values.size for o in self.a)
|
|
699
|
+
else: # Single or implicitly-batched LHS
|
|
700
|
+
self.device_id = self.a.device_id
|
|
701
|
+
self.value_type = self.a.dtype
|
|
702
|
+
self.index_type = self.a.index_type
|
|
703
|
+
rows, cols = self.a.shape[-2:]
|
|
704
|
+
if self.a.num_dimensions < 2 or rows != cols:
|
|
705
|
+
raise TypeError(
|
|
706
|
+
f"The LHS of type {type(self.a.tensor)} with shape {self.a.shape} must be a CSR matrix "
|
|
707
|
+
f"of shape (N, N) or CSR tensor with shape (..., N, N)."
|
|
708
|
+
)
|
|
709
|
+
self.lhs_shape = self.a.shape
|
|
710
|
+
self._N = rows
|
|
711
|
+
self.lhs_nnz = self.a.values.size
|
|
712
|
+
|
|
713
|
+
# Note that torch by default uses int64 which doesn't seem to give correct results
|
|
714
|
+
# for batched solves. SciPy and CuPy adapt the index type based on the dimension.
|
|
715
|
+
if self.index_type not in VALID_INDEX_TYPES:
|
|
716
|
+
raise TypeError(
|
|
717
|
+
f"The index type {self.index_type} is not supported. The supported index types are {VALID_INDEX_TYPES}."
|
|
718
|
+
)
|
|
719
|
+
|
|
720
|
+
if self.index_type == "int64" and self.batch_count > 1:
|
|
721
|
+
raise RuntimeError(
|
|
722
|
+
"The index type 'int64' is not supported for batched solve. The supported index types are: "
|
|
723
|
+
f"{', '.join(set(VALID_INDEX_TYPES) - {'int64'})}"
|
|
724
|
+
)
|
|
725
|
+
|
|
726
|
+
if self.value_type not in VALID_DTYPES:
|
|
727
|
+
raise TypeError(
|
|
728
|
+
f"The dtype (value type) {self.value_type} is not supported. The supported dtypes are {VALID_DTYPES}."
|
|
729
|
+
)
|
|
730
|
+
|
|
731
|
+
# Get key RHS attributes and perform more basic checks.
|
|
732
|
+
if self.explicitly_batched_rhs:
|
|
733
|
+
# For a sequence, the get_attribute functions check for consistency within
|
|
734
|
+
# the sequence.
|
|
735
|
+
rhs_device_id = utils.get_operands_device_id(self.b)
|
|
736
|
+
rhs_value_type = utils.get_operands_dtype(self.b)
|
|
737
|
+
self.rhs_shape = tuple(o.shape for o in self.b)
|
|
738
|
+
self.rhs_strides = tuple(o.strides for o in self.b)
|
|
739
|
+
if not check_rhs_sequence_layout(self.rhs_shape, self.rhs_strides):
|
|
740
|
+
raise TypeError(
|
|
741
|
+
f"Each object in an explicitly-batched RHS {[o.tensor for o in self.b]} must be a dense matrix or vector \
|
|
742
|
+
with compact col-major layout."
|
|
743
|
+
)
|
|
744
|
+
# For explicitly-batched RHS, the matrix or vector is compact so we can use
|
|
745
|
+
# the RHS strides for the result.
|
|
746
|
+
self.result_strides = self.rhs_strides
|
|
747
|
+
# Capture the RHS solution extent for compatibility check with the LHS.
|
|
748
|
+
self.rhs_n = tuple(s[0] for s in self.rhs_shape) # sample is not batched.
|
|
749
|
+
else: # Single or implicitly-batched RHS
|
|
750
|
+
rhs_device_id = self.b.device_id
|
|
751
|
+
rhs_value_type = self.b.dtype
|
|
752
|
+
self.rhs_shape = self.b.shape
|
|
753
|
+
self.rhs_strides = self.b.strides
|
|
754
|
+
if not check_dense_tensor_layout(self.rhs_shape, self.rhs_strides, explicitly_batched=False):
|
|
755
|
+
raise TypeError(
|
|
756
|
+
f"The RHS of type {type(self.b.tensor)} with shape {self.b.shape} must be a matrix or vector "
|
|
757
|
+
f"with col-major layout, and for implicitly-batched RHS (N-D >= 3), each matrix sample must have "
|
|
758
|
+
f"col-major layout (the second dimension from the end must have unit stride."
|
|
759
|
+
)
|
|
760
|
+
# For single or implicitly-batched RHS, the matrix may not be compact so we use
|
|
761
|
+
# the axis ordering to determine the strides.
|
|
762
|
+
axis_order = axis_order_in_memory(self.rhs_shape, self.rhs_strides)
|
|
763
|
+
self.result_strides = calculate_strides(self.rhs_shape, axis_order)
|
|
764
|
+
self.rhs_n = self.rhs_shape[-2] if len(self.rhs_shape) > 1 else self.rhs_shape[-1]
|
|
765
|
+
|
|
766
|
+
# Consistency within LHS and RHS sequence has been checked at this point.
|
|
767
|
+
# Now check that the extents match between a and b in Ax = b.
|
|
768
|
+
message = "The extent N corresponding to the number of equations is not consistent between the LHS and RHS."
|
|
769
|
+
if self.explicitly_batched_lhs:
|
|
770
|
+
if self.explicitly_batched_rhs:
|
|
771
|
+
if not all(x == y for (x, y) in zip(self._N, self.rhs_n, strict=True)):
|
|
772
|
+
raise TypeError(message)
|
|
773
|
+
elif self.implicitly_batched_rhs:
|
|
774
|
+
if not all(x == self.rhs_n for x in self._N):
|
|
775
|
+
raise TypeError(message)
|
|
776
|
+
else:
|
|
777
|
+
raise TypeError("The RHS is not batched, but the LHS is.")
|
|
778
|
+
elif self.implicitly_batched_lhs:
|
|
779
|
+
if self.explicitly_batched_rhs:
|
|
780
|
+
if not all(x == self._N for x in self.rhs_n):
|
|
781
|
+
raise TypeError(message)
|
|
782
|
+
elif self.implicitly_batched_rhs:
|
|
783
|
+
if self.rhs_n != self._N:
|
|
784
|
+
raise TypeError(message)
|
|
785
|
+
else:
|
|
786
|
+
raise TypeError("The RHS is not batched, but the LHS is.")
|
|
787
|
+
else:
|
|
788
|
+
if self.rhs_n != self._N:
|
|
789
|
+
raise TypeError(message)
|
|
790
|
+
|
|
791
|
+
if self.device_id != rhs_device_id:
|
|
792
|
+
raise TypeError(f"The LHS device ID {self.device_id} does not match the RHS device ID {rhs_device_id}.")
|
|
793
|
+
|
|
794
|
+
if self.value_type != rhs_value_type:
|
|
795
|
+
raise TypeError(f"The LHS dtype {self.value_type} does not match the RHS dtype {rhs_value_type}.")
|
|
796
|
+
|
|
797
|
+
# The RHS index type is currently always set to int32.
|
|
798
|
+
|
|
799
|
+
# The value types must currently be the same between a, x, and b.
|
|
800
|
+
self.result_data_type = self.value_type
|
|
801
|
+
|
|
802
|
+
self.logger.info(f"The device_id={self.device_id}, dtype = {self.value_type}, index type = {self.index_type}.")
|
|
803
|
+
self.logger.info(f"The number of equations = {self._N}.")
|
|
804
|
+
self.logger.debug(f"The LHS shape = {self.lhs_shape}.")
|
|
805
|
+
self.logger.debug(f"The RHS shape = {self.rhs_shape}, strides = {self.rhs_strides}.")
|
|
806
|
+
|
|
807
|
+
# Currently the value and index types must match between the LHS and RHS.
|
|
808
|
+
self.cuda_value_type = NAME_TO_DATA_TYPE[self.value_type]
|
|
809
|
+
self.cuda_index_type = NAME_TO_DATA_TYPE[self.index_type]
|
|
810
|
+
|
|
811
|
+
# For batched LHS and RHS check package, device_id, dtype, index_type is the same.
|
|
812
|
+
# Check consistency between LHS and RHS (same attributes, numpy -> scipy)
|
|
813
|
+
# Check that the size (_N) of each sample is consistent between LHS and RHS.
|
|
814
|
+
|
|
815
|
+
# Set the memory space "cuda" or "cpu".
|
|
816
|
+
self.memory_space = self.a[0].device if self.explicitly_batched_lhs else self.a.device
|
|
817
|
+
|
|
818
|
+
# Set the execution space "cuda" or "hybrid".
|
|
819
|
+
self.execution_space = self.execution_options.name
|
|
820
|
+
|
|
821
|
+
# Note #1: The device ID of the result is the same as that of the operands for
|
|
822
|
+
# hybrid execution since we don't copy them. Capture the original device here
|
|
823
|
+
# before it's potentially changed below. Also see note #3.
|
|
824
|
+
self.result_device_id = self.device_id
|
|
825
|
+
|
|
826
|
+
# Track potential issue with copying NumPy ndarray to GPU.
|
|
827
|
+
# TODO: This should be fixed when we use cuda.core for copying across memspace.
|
|
828
|
+
rhs_layout_flag = self.rhs_package == "numpy" and len(self.rhs_shape) > 2 and self.implicitly_batched_rhs
|
|
829
|
+
|
|
830
|
+
# NumPy has no stream support, use cuda.core stream interface if needed.
|
|
831
|
+
self.stream_package = "cuda" if self.rhs_package == "numpy" else self.rhs_package
|
|
832
|
+
if self.device_id == "cpu":
|
|
833
|
+
# For CPU operands, set the device ID based on the execution options.
|
|
834
|
+
self.device_id = self.execution_options.device_id
|
|
835
|
+
|
|
836
|
+
self.logger.info(
|
|
837
|
+
f"The operands' memory space is {self.memory_space}, and the execution space is on device {self.device_id}."
|
|
838
|
+
)
|
|
839
|
+
|
|
840
|
+
self.copy_across_memspace = False
|
|
841
|
+
if self.execution_space == "hybrid":
|
|
842
|
+
# No need to copy CPU operands.
|
|
843
|
+
if self.batched:
|
|
844
|
+
raise TypeError(f"Batching is not supported for hybrid execution: {self.execution_options}.")
|
|
845
|
+
|
|
846
|
+
# For non-batched b, matrix (multiple-RHS) is not supported for CPU memory
|
|
847
|
+
# space (seems to be ignored for CUDA as well).
|
|
848
|
+
if len(self.b.shape) > 1:
|
|
849
|
+
raise TypeError(f"Matrix RHS (multiple RHS) is not supported for hybrid execution: {self.execution_options}.")
|
|
850
|
+
else: # execute in the CUDA space.
|
|
851
|
+
# The operands must be on the GPU even for *hybrid memory* mode for CUDA
|
|
852
|
+
# execution.
|
|
853
|
+
self.copy_across_memspace = self.memory_space != "cuda"
|
|
854
|
+
# Note #3: For CUDA execution, the result's device ID should be that of the
|
|
855
|
+
# execution device (self.device_id), where the operands also reside (may or
|
|
856
|
+
# may not be copied, depending on their original memspace). Also see note #1.
|
|
857
|
+
self.result_device_id = self.device_id
|
|
858
|
+
|
|
859
|
+
# Flag whether to copy CPU operands to GPU or not, based on execution option.
|
|
860
|
+
# - if hybrid execution, check not batched and nrhs==1. Accept CPU or GPU
|
|
861
|
+
# operands.
|
|
862
|
+
# - if hybrid memory and CUDA execution, need to copy to GPU.
|
|
863
|
+
|
|
864
|
+
# Create the stream holder using the appropriate package for stream operations.
|
|
865
|
+
# NOTE: the resolved ctor stream holder is stored in ``_init_stream_holder``;
|
|
866
|
+
# currently it is only intended for use in the stateless function, which reuses that
|
|
867
|
+
# stream for all required method calls instead of resolving the stream repeatedly.
|
|
868
|
+
self._init_stream_holder = stream_holder = utils.get_or_create_stream(self.device_id, stream, self.stream_package)
|
|
869
|
+
self.logger.info(f"The specified stream for the DirectSolver ctor is {stream_holder.obj}.")
|
|
870
|
+
|
|
871
|
+
# cupy.asarray() doesn't preserve layout for > 2D arrays when copying from CPU
|
|
872
|
+
# to GPU.
|
|
873
|
+
if self.copy_across_memspace and rhs_layout_flag:
|
|
874
|
+
raise TypeError(
|
|
875
|
+
f"Implicit RHS batching for NumPy ndarrays is currently not supported with CUDA execution \
|
|
876
|
+
since the layout cannot be preserved when copying to the GPU (shape = {self.rhs_shape}, strides = {self.rhs_strides})."
|
|
877
|
+
)
|
|
878
|
+
|
|
879
|
+
# Copy operands to device if needed.
|
|
880
|
+
if self.copy_across_memspace:
|
|
881
|
+
self.a = copy_single_or_sequence(self.a, self.device_id, stream_holder)
|
|
882
|
+
self.b = copy_single_or_sequence(self.b, self.device_id, stream_holder)
|
|
883
|
+
|
|
884
|
+
# Create (batched or not, CSR or dense) matrix pointers for the LHS and RHS.
|
|
885
|
+
# The create wrappers return (dims, ptrs, matrix_ptr) where:
|
|
886
|
+
# dims: metadata arrays (shape/size) that don't change on update
|
|
887
|
+
# ptrs: device pointer arrays whose contents are updated in place
|
|
888
|
+
# inside reset_operands() / solve()
|
|
889
|
+
# For non-batched cases, both dims and ptrs are empty lists.
|
|
890
|
+
self.resources_a_dims, self.resources_a_ptrs, self.a_ptr = cudss_utils.create_cudss_csr_wrapper(
|
|
891
|
+
self.cuda_index_type,
|
|
892
|
+
self.cuda_value_type,
|
|
893
|
+
self.index_type,
|
|
894
|
+
self.options.sparse_system_type,
|
|
895
|
+
self.options.sparse_system_view,
|
|
896
|
+
self.batch_indices,
|
|
897
|
+
self.a,
|
|
898
|
+
stream_holder,
|
|
899
|
+
)
|
|
900
|
+
self.resources_b_dims, self.resources_b_ptrs, self.b_ptr = cudss_utils.create_cudss_dense_wrapper(
|
|
901
|
+
self.cuda_index_type, self.cuda_value_type, self.index_type, self.batch_indices, self.b, stream_holder
|
|
902
|
+
)
|
|
903
|
+
|
|
904
|
+
# Use `b` for creating the (potentially explicitly or implicitly batched) solution
|
|
905
|
+
# matrix or vector. The pointers will be updated later in execute.
|
|
906
|
+
self.resources_x_dims, self.resources_x_ptrs, self.x_ptr = cudss_utils.create_cudss_dense_wrapper(
|
|
907
|
+
self.cuda_index_type, self.cuda_value_type, self.index_type, self.batch_indices, self.b, stream_holder
|
|
908
|
+
)
|
|
909
|
+
|
|
910
|
+
# Track the allocation stream for pointer arrays so we can ensure proper
|
|
911
|
+
# ordering before releasing them in free() (dropping references triggers
|
|
912
|
+
# stream-ordered deallocation).
|
|
913
|
+
self.resources_ptrs_alloc_stream = (
|
|
914
|
+
stream_holder.obj if (self.resources_a_ptrs or self.resources_b_ptrs or self.resources_x_ptrs) else None
|
|
915
|
+
)
|
|
916
|
+
|
|
917
|
+
# Create or set handle, and create config and data pointers.
|
|
918
|
+
with utils.device_ctx(self.device_id):
|
|
919
|
+
if self.options.handle is not None:
|
|
920
|
+
self.own_handle = False
|
|
921
|
+
self.handle = self.options.handle
|
|
922
|
+
self.logger.info(f"The library handle has been set to the specified value: {self.handle}.")
|
|
923
|
+
else:
|
|
924
|
+
self.own_handle = True
|
|
925
|
+
self.handle = cudss.create()
|
|
926
|
+
self.logger.info(f"The library handle has been created: {self.handle}.")
|
|
927
|
+
|
|
928
|
+
self.config_ptr = cudss.config_create()
|
|
929
|
+
self.data_ptr = cudss.data_create(self.handle)
|
|
930
|
+
|
|
931
|
+
# Create the config interfaces for the various phases.
|
|
932
|
+
self._plan_config = cudss_config_ifc.PlanConfig(self)
|
|
933
|
+
self._factorization_config = cudss_config_ifc.FactorizationConfig(self)
|
|
934
|
+
self._solution_config = cudss_config_ifc.SolutionConfig(self)
|
|
935
|
+
|
|
936
|
+
# Create the data interfaces for the various phases.
|
|
937
|
+
self._plan_info = cudss_data_ifc.PlanInfo(self)
|
|
938
|
+
self._factorization_info = cudss_data_ifc.FactorizationInfo(self)
|
|
939
|
+
|
|
940
|
+
# Set the threading layer, if available.
|
|
941
|
+
threading_lib = get_threading_lib(self.options.multithreading_lib)
|
|
942
|
+
if threading_lib is not None:
|
|
943
|
+
cudss.set_threading_layer(self.handle, threading_lib)
|
|
944
|
+
else:
|
|
945
|
+
self.logger.warning(
|
|
946
|
+
"No multithreading interface library was specified using the \
|
|
947
|
+
DirectSolverOptions. The performance of CPU operations like planning will \
|
|
948
|
+
be significantly lower than if you provide a multithreading library."
|
|
949
|
+
)
|
|
950
|
+
|
|
951
|
+
# This doesn't guarantee that the library is usable, just that it is present.
|
|
952
|
+
self.multithreading = threading_lib is not None
|
|
953
|
+
|
|
954
|
+
# Set private attributes based on the options.
|
|
955
|
+
self._internal_config = cudss_config_ifc.InternalConfig(self)
|
|
956
|
+
|
|
957
|
+
# Set hybrid execution options. We have already checked that for batching and
|
|
958
|
+
# matrix RHS, which are not currently supported.
|
|
959
|
+
if self.execution_space == "hybrid":
|
|
960
|
+
self._internal_config.hybrid_execute_mode = 1
|
|
961
|
+
num_threads = self.execution_options.num_threads
|
|
962
|
+
if num_threads is not None:
|
|
963
|
+
if num_threads > 1 and threading_lib is None:
|
|
964
|
+
raise ValueError(
|
|
965
|
+
"The threading library must be specified if the number of threads is more than 1 "
|
|
966
|
+
f"(num_threads = {num_threads})."
|
|
967
|
+
)
|
|
968
|
+
self._internal_config.host_nthreads = num_threads
|
|
969
|
+
|
|
970
|
+
# Set CUDA execution options (including hybrid memory mode).
|
|
971
|
+
if self.execution_space == "cuda":
|
|
972
|
+
hmo = self.execution_options.hybrid_memory_mode_options
|
|
973
|
+
self._internal_config.hybrid_memory_mode = hmo.hybrid_memory_mode
|
|
974
|
+
# Set device memory limit for hybrid memory.
|
|
975
|
+
if hmo.hybrid_device_memory_limit is not None:
|
|
976
|
+
memory_limit = utils.get_memory_limit_from_device_id(hmo.hybrid_device_memory_limit, self.device_id)
|
|
977
|
+
self.logger.info(f"The hybrid memory limit is {formatters.MemoryStr(memory_limit)}.")
|
|
978
|
+
self._internal_config.hybrid_device_memory_limit = memory_limit
|
|
979
|
+
self._internal_config.use_cuda_register_memory = hmo.register_cuda_memory
|
|
980
|
+
|
|
981
|
+
# State tracking attributes.
|
|
982
|
+
self.solver_planned = False
|
|
983
|
+
self.solver_factorized = False
|
|
984
|
+
|
|
985
|
+
# Attribute to track the last compute event needed
|
|
986
|
+
# inside free() for stream ordering.
|
|
987
|
+
self.last_compute_event = None
|
|
988
|
+
|
|
989
|
+
# Track whether the user has called release_operands().
|
|
990
|
+
self._operands_released = False
|
|
991
|
+
|
|
992
|
+
# Set blocking or non-blocking behavior.
|
|
993
|
+
self.blocking = self.options.blocking is True or self.memory_space == "cpu"
|
|
994
|
+
if self.blocking:
|
|
995
|
+
self.call_prologue = "This call is blocking and will return only after the operation is complete."
|
|
996
|
+
else:
|
|
997
|
+
self.call_prologue = (
|
|
998
|
+
"This call is non-blocking and will return immediately after the operation is launched on the device."
|
|
999
|
+
)
|
|
1000
|
+
|
|
1001
|
+
if not self.blocking:
|
|
1002
|
+
set_async_workspace_allocator(self.handle, self.device_id, self.logger)
|
|
1003
|
+
|
|
1004
|
+
# The result (solution) class is that of the wrapped RHS.
|
|
1005
|
+
self.result_class = self.b[0].__class__ if self.explicitly_batched_rhs else self.b.__class__
|
|
1006
|
+
|
|
1007
|
+
# The result shape is a single value or a sequence, depending on whether the RHS
|
|
1008
|
+
# is explicitly or implicitly batched as set above.
|
|
1009
|
+
self.result_shape = self.rhs_shape
|
|
1010
|
+
|
|
1011
|
+
self.valid_state = True
|
|
1012
|
+
self.logger.info("The sparse direct solver operation has been created.")
|
|
1013
|
+
|
|
1014
|
+
def __enter__(self):
|
|
1015
|
+
return self
|
|
1016
|
+
|
|
1017
|
+
def __exit__(self, exc_type, exc_value, traceback):
|
|
1018
|
+
self.free()
|
|
1019
|
+
|
|
1020
|
+
def _check_valid_solver(self, *args, **kwargs):
|
|
1021
|
+
"""
|
|
1022
|
+
Check if the DirectSolver object is alive and well.
|
|
1023
|
+
"""
|
|
1024
|
+
if not self.valid_state:
|
|
1025
|
+
raise InvalidDirectSolverState("The DirectSolver object cannot be used after resources are free'd")
|
|
1026
|
+
|
|
1027
|
+
def _check_valid_operands(self, *args, **kwargs):
|
|
1028
|
+
"""
|
|
1029
|
+
Check if the operands are available for the operation.
|
|
1030
|
+
"""
|
|
1031
|
+
what = kwargs["what"]
|
|
1032
|
+
if self._operands_released:
|
|
1033
|
+
raise RuntimeError(
|
|
1034
|
+
f"{what} cannot be performed after the operands have been released. "
|
|
1035
|
+
f"Use reset_operands() to provide new operands before performing the {what.lower()}."
|
|
1036
|
+
)
|
|
1037
|
+
|
|
1038
|
+
def _check_planned(self, *args, **kwargs):
|
|
1039
|
+
what = kwargs["what"]
|
|
1040
|
+
if not self.solver_planned:
|
|
1041
|
+
raise RuntimeError(f"{what} cannot be performed before plan() has been called.")
|
|
1042
|
+
|
|
1043
|
+
def _check_factorized(self, *args, **kwargs):
|
|
1044
|
+
what = kwargs["what"]
|
|
1045
|
+
if not self.solver_factorized:
|
|
1046
|
+
raise RuntimeError(f"{what} cannot be performed before factorize() has been called.")
|
|
1047
|
+
|
|
1048
|
+
def _allocate_single_result(self, stream_holder: StreamHolder | None, log_debug):
|
|
1049
|
+
if log_debug:
|
|
1050
|
+
self.logger.debug("Beginning output (empty) tensor creation...")
|
|
1051
|
+
self.logger.debug(
|
|
1052
|
+
f"""The output tensor shape = {self.result_shape}, with strides = {self.result_strides}
|
|
1053
|
+
and data type '{self.result_data_type}'."""
|
|
1054
|
+
)
|
|
1055
|
+
result = utils.create_empty_tensor(
|
|
1056
|
+
self.result_class,
|
|
1057
|
+
self.result_shape,
|
|
1058
|
+
self.result_data_type,
|
|
1059
|
+
self.result_device_id, # see notes #1.
|
|
1060
|
+
stream_holder=None if self.result_device_id == "cpu" else stream_holder,
|
|
1061
|
+
verify_strides=False, # the strides are computed so that they are contiguous
|
|
1062
|
+
strides=self.result_strides,
|
|
1063
|
+
)
|
|
1064
|
+
if log_debug:
|
|
1065
|
+
self.logger.debug("The output (empty) tensor has been created.")
|
|
1066
|
+
return result
|
|
1067
|
+
|
|
1068
|
+
def _allocate_batched_result(self, stream_holder: StreamHolder | None, log_debug):
|
|
1069
|
+
if log_debug:
|
|
1070
|
+
self.logger.debug("Beginning output tensor sequence creation...")
|
|
1071
|
+
self.logger.debug(
|
|
1072
|
+
f"""The output tensor sequence shape = {self.result_shape}, with strides = {self.result_strides}
|
|
1073
|
+
and data type '{self.result_data_type}'."""
|
|
1074
|
+
)
|
|
1075
|
+
result = tuple(
|
|
1076
|
+
utils.create_empty_tensor(
|
|
1077
|
+
self.result_class,
|
|
1078
|
+
shape,
|
|
1079
|
+
self.result_data_type,
|
|
1080
|
+
self.result_device_id, # see notes #1.
|
|
1081
|
+
stream_holder=None if self.result_device_id == "cpu" else stream_holder,
|
|
1082
|
+
verify_strides=False, # the strides are computed so that they are contiguous
|
|
1083
|
+
strides=strides,
|
|
1084
|
+
)
|
|
1085
|
+
for (shape, strides) in zip(self.result_shape, self.result_strides, strict=True)
|
|
1086
|
+
)
|
|
1087
|
+
|
|
1088
|
+
if log_debug:
|
|
1089
|
+
self.logger.debug("The output tensor sequence has been created.")
|
|
1090
|
+
return result
|
|
1091
|
+
|
|
1092
|
+
@property
|
|
1093
|
+
def plan_config(self):
|
|
1094
|
+
"""
|
|
1095
|
+
An accessor to configure or query the solver planning phase attributes.
|
|
1096
|
+
|
|
1097
|
+
Returns:
|
|
1098
|
+
A :class:`DirectSolverPlanConfig` object, whose attributes can be set (or
|
|
1099
|
+
queried) to configure the planning phase.
|
|
1100
|
+
|
|
1101
|
+
.. seealso::
|
|
1102
|
+
:class:`DirectSolverPlanConfig`, :meth:`plan`.
|
|
1103
|
+
"""
|
|
1104
|
+
return self._plan_config
|
|
1105
|
+
|
|
1106
|
+
@property
|
|
1107
|
+
def factorization_config(self):
|
|
1108
|
+
"""
|
|
1109
|
+
An accessor to configure or query the solver factorization phase attributes.
|
|
1110
|
+
|
|
1111
|
+
Returns:
|
|
1112
|
+
A :class:`DirectSolverFactorizationConfig` object, whose attributes can be set
|
|
1113
|
+
(or queried) to configure the factorization phase.
|
|
1114
|
+
|
|
1115
|
+
.. seealso::
|
|
1116
|
+
:class:`DirectSolverFactorizationConfig`, :meth:`factorize`.
|
|
1117
|
+
"""
|
|
1118
|
+
return self._factorization_config
|
|
1119
|
+
|
|
1120
|
+
@property
|
|
1121
|
+
def solution_config(self):
|
|
1122
|
+
"""
|
|
1123
|
+
An accessor to configure or query the solver solution phase attributes.
|
|
1124
|
+
|
|
1125
|
+
Returns:
|
|
1126
|
+
A :class:`DirectSolverSolutionConfig` object, whose attributes can be set
|
|
1127
|
+
(or queried) to configure the factorization phase.
|
|
1128
|
+
|
|
1129
|
+
.. seealso::
|
|
1130
|
+
:class:`DirectSolverSolutionConfig`, :meth:`solve`.
|
|
1131
|
+
"""
|
|
1132
|
+
return self._solution_config
|
|
1133
|
+
|
|
1134
|
+
@property
|
|
1135
|
+
def plan_info(self):
|
|
1136
|
+
"""
|
|
1137
|
+
An accessor to get information about the solver planning phase.
|
|
1138
|
+
|
|
1139
|
+
Returns:
|
|
1140
|
+
A :class:`DirectSolverPlanInfo` object, whose attributes can be queried for
|
|
1141
|
+
information regarding the planning phase.
|
|
1142
|
+
|
|
1143
|
+
.. seealso::
|
|
1144
|
+
:class:`DirectSolverPlanInfo`, :meth:`plan`.
|
|
1145
|
+
"""
|
|
1146
|
+
return self._plan_info
|
|
1147
|
+
|
|
1148
|
+
@property
|
|
1149
|
+
def factorization_info(self):
|
|
1150
|
+
"""
|
|
1151
|
+
Query solver factorization information
|
|
1152
|
+
(see :class:`nvmath.sparse.advanced.DirectSolverFactorizationInfo`).
|
|
1153
|
+
An accessor to get information about the solver factorization phase.
|
|
1154
|
+
|
|
1155
|
+
Returns:
|
|
1156
|
+
A :class:`DirectSolverFactorizationInfo` object, whose attributes can be
|
|
1157
|
+
queried for information regarding the factorization phase.
|
|
1158
|
+
|
|
1159
|
+
.. seealso::
|
|
1160
|
+
:class:`DirectSolverFactorizationInfo`, :meth:`factorize`.
|
|
1161
|
+
"""
|
|
1162
|
+
return self._factorization_info
|
|
1163
|
+
|
|
1164
|
+
@utils.precondition(_check_valid_solver)
|
|
1165
|
+
def reset_operands(
|
|
1166
|
+
self,
|
|
1167
|
+
*,
|
|
1168
|
+
a=None,
|
|
1169
|
+
b=None,
|
|
1170
|
+
stream: utils.AnyStream | int | None = None,
|
|
1171
|
+
):
|
|
1172
|
+
"""
|
|
1173
|
+
Reset one or both operands held by this :class:`DirectSolver` instance.
|
|
1174
|
+
|
|
1175
|
+
.. versionchanged:: 0.9
|
|
1176
|
+
All parameters are now keyword-only.
|
|
1177
|
+
|
|
1178
|
+
Args:
|
|
1179
|
+
a: {a}
|
|
1180
|
+
|
|
1181
|
+
b: {b}
|
|
1182
|
+
|
|
1183
|
+
stream: {stream}
|
|
1184
|
+
|
|
1185
|
+
Semantics:
|
|
1186
|
+
- Only the operands explicitly passed are updated. At least one operand
|
|
1187
|
+
is required (all of them after :meth:`release_operands`), otherwise
|
|
1188
|
+
a :class:`ValueError` is raised.
|
|
1189
|
+
|
|
1190
|
+
- This method will perform various checks on the new operands to make sure:
|
|
1191
|
+
|
|
1192
|
+
- The shapes, index and data types match those of the old ones.
|
|
1193
|
+
- The packages that the operands belong to match those of the old ones.
|
|
1194
|
+
- If input tensors are on GPU, the device must match.
|
|
1195
|
+
|
|
1196
|
+
Examples:
|
|
1197
|
+
|
|
1198
|
+
>>> import cupy as cp
|
|
1199
|
+
>>> import cupyx.scipy.sparse as sp
|
|
1200
|
+
>>> import nvmath
|
|
1201
|
+
|
|
1202
|
+
Prepare sample input data.
|
|
1203
|
+
|
|
1204
|
+
>>> n = 8
|
|
1205
|
+
>>> a = sp.random(n, n, density=0.15, format="csr", dtype="float64")
|
|
1206
|
+
>>> a += sp.diags([2.0] * n, format="csr", dtype="float64")
|
|
1207
|
+
|
|
1208
|
+
Create the RHS, which can be a matrix or vector in column-major layout.
|
|
1209
|
+
|
|
1210
|
+
>>> b = cp.ones((n,), dtype="float64")
|
|
1211
|
+
|
|
1212
|
+
Specify, plan, factorize and solve a @ x = b. Use the stateful object as a
|
|
1213
|
+
context manager to automatically release resources.
|
|
1214
|
+
|
|
1215
|
+
>>> with nvmath.sparse.advanced.DirectSolver(a, b) as solver:
|
|
1216
|
+
... # Plan the operation.
|
|
1217
|
+
... plan_info = solver.plan()
|
|
1218
|
+
...
|
|
1219
|
+
... # Factorize the system.
|
|
1220
|
+
... fac_info = solver.factorize()
|
|
1221
|
+
...
|
|
1222
|
+
... # Solve the factorized system for the first result.
|
|
1223
|
+
... x1 = solver.solve()
|
|
1224
|
+
...
|
|
1225
|
+
... # Reset the RHS to a new CuPy ndarray.
|
|
1226
|
+
... c = cp.random.rand(n, dtype="float64")
|
|
1227
|
+
... solver.reset_operands(b=c)
|
|
1228
|
+
...
|
|
1229
|
+
... # Solve for the second result corresponding to the updated operands.
|
|
1230
|
+
... x2 = solver.solve()
|
|
1231
|
+
|
|
1232
|
+
.. tip:: If only a subset of operands are reset, the operands that are not
|
|
1233
|
+
reset hold their original values.
|
|
1234
|
+
|
|
1235
|
+
With :meth:`reset_operands`, minimal overhead is achieved as problem
|
|
1236
|
+
specification and planning are only performed once.
|
|
1237
|
+
|
|
1238
|
+
For the particular example above, the operands are on the GPU, so
|
|
1239
|
+
calling :meth:`reset_operands` only updates internal references
|
|
1240
|
+
and is efficient. An alternative for that case would be to modify the
|
|
1241
|
+
the operand `b` in-place, i.e, replacing ``solver.reset_operands(b=c)``
|
|
1242
|
+
with ``b[:]=c``, but doing so triggers a copy of the data and
|
|
1243
|
+
has performance implications.
|
|
1244
|
+
|
|
1245
|
+
.. danger:: Updating the operand in-place can only yield the expected result
|
|
1246
|
+
under the additional constraints below:
|
|
1247
|
+
|
|
1248
|
+
- The operand is on the GPU (more precisely, the operand memory space should
|
|
1249
|
+
be accessible from the execution space).
|
|
1250
|
+
|
|
1251
|
+
- The user has called :meth:`factorize` if needed (they are not relying on
|
|
1252
|
+
iterative refinement for example)
|
|
1253
|
+
|
|
1254
|
+
For more details, please refer to `the reset operand example
|
|
1255
|
+
<https://github.com/NVIDIA/nvmath-python/tree/main/examples/sparse/advanced/direct_solver/example05_reset_operands.py>`_.
|
|
1256
|
+
"""
|
|
1257
|
+
|
|
1258
|
+
# If operands have been released, all required operands must be provided.
|
|
1259
|
+
if self._operands_released:
|
|
1260
|
+
if a is None or b is None:
|
|
1261
|
+
raise ValueError("After release_operands(), both 'a' and 'b' must be provided to reset_operands().")
|
|
1262
|
+
elif a is None and b is None:
|
|
1263
|
+
# All arguments are None: there is nothing to update, so reject the call.
|
|
1264
|
+
raise ValueError("reset_operands() requires at least one operand to be provided.")
|
|
1265
|
+
|
|
1266
|
+
stream_holder = utils.get_or_create_stream(self.device_id, stream, self.stream_package)
|
|
1267
|
+
|
|
1268
|
+
# Update LHS.
|
|
1269
|
+
if a is not None:
|
|
1270
|
+
if isinstance(a, Sequence) and not self.explicitly_batched_lhs:
|
|
1271
|
+
raise TypeError(f"The specified type for 'a' is a sequence while the original type is {type(self.a.tensor)}.")
|
|
1272
|
+
|
|
1273
|
+
# Wrap A.
|
|
1274
|
+
a = wrap_cudss_supported_lhs(a)
|
|
1275
|
+
|
|
1276
|
+
explicitly_batched = isinstance(a, Sequence)
|
|
1277
|
+
if explicitly_batched:
|
|
1278
|
+
# Successfully wrapping A means that these package is the same for all
|
|
1279
|
+
# sparse operands.
|
|
1280
|
+
lhs_package = utils.infer_object_package(a[0].tensor)
|
|
1281
|
+
# The get_() helpers ensure that the attributes are the same for all
|
|
1282
|
+
# operands.
|
|
1283
|
+
device_id = utils.get_operands_device_id(a)
|
|
1284
|
+
memory_space = a[0].device
|
|
1285
|
+
value_type = utils.get_operands_dtype(a)
|
|
1286
|
+
index_type = sp_utils.get_operands_index_type(a)
|
|
1287
|
+
|
|
1288
|
+
shape = tuple(o.shape for o in a)
|
|
1289
|
+
nnz = tuple(o.values.size for o in a)
|
|
1290
|
+
|
|
1291
|
+
else: # Single or implicitly-batched LHS
|
|
1292
|
+
lhs_package = utils.infer_object_package(a.tensor)
|
|
1293
|
+
device_id = a.device_id
|
|
1294
|
+
memory_space = a.device
|
|
1295
|
+
value_type = a.dtype
|
|
1296
|
+
index_type = a.index_type
|
|
1297
|
+
|
|
1298
|
+
shape = a.shape
|
|
1299
|
+
nnz = a.values.size
|
|
1300
|
+
|
|
1301
|
+
# Check package, device ID, dtype, and index type.
|
|
1302
|
+
if lhs_package != self.lhs_package:
|
|
1303
|
+
raise TypeError(f"The package for 'a' ({lhs_package}) doesn't match the original one ({self.lhs_package}).")
|
|
1304
|
+
|
|
1305
|
+
if memory_space != self.memory_space:
|
|
1306
|
+
raise TypeError(
|
|
1307
|
+
f"The memory space for 'a' ({memory_space}) doesn't match the original one ({self.memory_space})."
|
|
1308
|
+
)
|
|
1309
|
+
|
|
1310
|
+
if device_id != "cpu" and device_id != self.device_id:
|
|
1311
|
+
raise TypeError(f"The device id for 'a' ({device_id}) doesn't match the original one ({self.device_id}).")
|
|
1312
|
+
|
|
1313
|
+
if value_type != self.value_type:
|
|
1314
|
+
raise TypeError(f"The dtype for 'a' ({value_type}) doesn't match the original one ({self.value_type}).")
|
|
1315
|
+
|
|
1316
|
+
if index_type != self.index_type:
|
|
1317
|
+
raise TypeError(f"The index type for 'a' ({index_type}) doesn't match the original one ({self.index_type}).")
|
|
1318
|
+
|
|
1319
|
+
# Checking that the shape is consistent also checks the batch count for both
|
|
1320
|
+
# implicit and explicit batching.
|
|
1321
|
+
if shape != self.lhs_shape:
|
|
1322
|
+
raise TypeError(f"The shape of 'a' ({shape}) doesn't match the original one ({self.lhs_shape}).")
|
|
1323
|
+
|
|
1324
|
+
if nnz != self.lhs_nnz:
|
|
1325
|
+
raise TypeError(f"The number of non-zeros of 'a' ({nnz}) doesn't match the original one ({self.lhs_nnz}).")
|
|
1326
|
+
|
|
1327
|
+
# Copy operand if needed, and replace object reference.
|
|
1328
|
+
if self.copy_across_memspace:
|
|
1329
|
+
# Copy operand into original buffer if it exists or create new ones.
|
|
1330
|
+
log_warning = False
|
|
1331
|
+
if explicitly_batched:
|
|
1332
|
+
if not self._operands_released:
|
|
1333
|
+
for x, y in zip(self.a, a, strict=True):
|
|
1334
|
+
x.copy_(y, stream_holder)
|
|
1335
|
+
else:
|
|
1336
|
+
self.a = [x.to(self.device_id, stream_holder) for x in a]
|
|
1337
|
+
log_warning = True
|
|
1338
|
+
else:
|
|
1339
|
+
if not self._operands_released:
|
|
1340
|
+
self.a.copy_(a, stream_holder)
|
|
1341
|
+
else:
|
|
1342
|
+
self.a = a.to(self.device_id, stream_holder)
|
|
1343
|
+
log_warning = True
|
|
1344
|
+
if log_warning:
|
|
1345
|
+
self.logger.warning(
|
|
1346
|
+
"The LHS buffer pointers have changed when copying between CPU-GPU, which requires calling \
|
|
1347
|
+
plan() and factorize() again even if the sparsity structure is identical."
|
|
1348
|
+
)
|
|
1349
|
+
else:
|
|
1350
|
+
# Invalidate the plan, since the buffer pointers could have changed.
|
|
1351
|
+
self.solver_planned = self.solver_factorized = False
|
|
1352
|
+
self.logger.warning(
|
|
1353
|
+
"The specified LHS may have different buffers for the compressed row or column indices \
|
|
1354
|
+
or values, which requires calling plan() and factorize() again even if the sparsity structure is identical. To avoid this, it \
|
|
1355
|
+
is recommended to update the values in place and refactorize if needed."
|
|
1356
|
+
)
|
|
1357
|
+
self.a = a
|
|
1358
|
+
|
|
1359
|
+
# Update the pointer values in the existing device buffers.
|
|
1360
|
+
cudss_utils.update_cudss_csr_ptr_wrapper(
|
|
1361
|
+
existing_ptrs=self.resources_a_ptrs,
|
|
1362
|
+
lhs_ptr=self.a_ptr,
|
|
1363
|
+
batch_indices=self.batch_indices,
|
|
1364
|
+
new_lhs=self.a,
|
|
1365
|
+
stream_holder=stream_holder,
|
|
1366
|
+
)
|
|
1367
|
+
|
|
1368
|
+
self.logger.warning(
|
|
1369
|
+
"Resetting the LHS 'a' typically requires calling factorize() again. An exception is the use of \
|
|
1370
|
+
iterative refinement during solve(), but it's the user's responsibility to check that the solution has converged."
|
|
1371
|
+
)
|
|
1372
|
+
|
|
1373
|
+
# Update RHS.
|
|
1374
|
+
if b is not None:
|
|
1375
|
+
# Wrap b.
|
|
1376
|
+
explicitly_batched = isinstance(b, Sequence)
|
|
1377
|
+
if explicitly_batched:
|
|
1378
|
+
b = tensor_wrapper.wrap_operands(b)
|
|
1379
|
+
rhs_package = utils.get_operands_package(b)
|
|
1380
|
+
|
|
1381
|
+
# The get_() helpers ensure that the attributes are the same for all
|
|
1382
|
+
# operands.
|
|
1383
|
+
device_id = utils.get_operands_device_id(b)
|
|
1384
|
+
memory_space = b[0].device
|
|
1385
|
+
value_type = utils.get_operands_dtype(b)
|
|
1386
|
+
shape = tuple(o.shape for o in b)
|
|
1387
|
+
strides = tuple(o.strides for o in b)
|
|
1388
|
+
else: # Single or implicitly-batched RHS
|
|
1389
|
+
b = tensor_wrapper.wrap_operand(b)
|
|
1390
|
+
rhs_package = utils.infer_object_package(b.tensor)
|
|
1391
|
+
|
|
1392
|
+
device_id = b.device_id
|
|
1393
|
+
memory_space = b.device
|
|
1394
|
+
value_type = b.dtype
|
|
1395
|
+
shape = b.shape
|
|
1396
|
+
strides = b.strides
|
|
1397
|
+
|
|
1398
|
+
# Check package, device ID, shape, strides, and dtype.
|
|
1399
|
+
if rhs_package != self.rhs_package:
|
|
1400
|
+
raise TypeError(f"The package for 'b' ({rhs_package}) doesn't match the original one ({self.rhs_package}).")
|
|
1401
|
+
|
|
1402
|
+
if memory_space != self.memory_space:
|
|
1403
|
+
raise TypeError(
|
|
1404
|
+
f"The memory space for 'b' ({memory_space}) doesn't match the original one ({self.memory_space})."
|
|
1405
|
+
)
|
|
1406
|
+
|
|
1407
|
+
if device_id != "cpu" and device_id != self.device_id:
|
|
1408
|
+
raise TypeError(f"The device id for 'b' ({device_id}) doesn't match the original one ({self.device_id}).")
|
|
1409
|
+
|
|
1410
|
+
if value_type != self.value_type:
|
|
1411
|
+
raise TypeError(f"The dtype for 'b' ({value_type}) doesn't match the original one ({self.value_type}).")
|
|
1412
|
+
|
|
1413
|
+
# Checking that the shape is consistent also checks the batch count for both
|
|
1414
|
+
# implicit and explicit batching.
|
|
1415
|
+
if shape != self.rhs_shape:
|
|
1416
|
+
raise TypeError(f"The shape of 'b' ({shape}) doesn't match the original one ({self.rhs_shape}).")
|
|
1417
|
+
|
|
1418
|
+
if strides != self.rhs_strides:
|
|
1419
|
+
raise TypeError(f"The strides of 'b' ({strides}) don't match the original one ({self.rhs_strides}).")
|
|
1420
|
+
|
|
1421
|
+
# Copy operand if needed, and replace object reference.
|
|
1422
|
+
if self.copy_across_memspace:
|
|
1423
|
+
# Copy operand into original buffer if it exists or create new ones.
|
|
1424
|
+
if explicitly_batched:
|
|
1425
|
+
if not self._operands_released:
|
|
1426
|
+
for x, y in zip(self.b, b, strict=True):
|
|
1427
|
+
x.copy_(y, stream_holder)
|
|
1428
|
+
else:
|
|
1429
|
+
self.b = [x.to(self.device_id, stream_holder) for x in b]
|
|
1430
|
+
else:
|
|
1431
|
+
if not self._operands_released:
|
|
1432
|
+
self.b.copy_(b, stream_holder)
|
|
1433
|
+
else:
|
|
1434
|
+
self.b = b.to(self.device_id, stream_holder)
|
|
1435
|
+
else:
|
|
1436
|
+
self.b = b
|
|
1437
|
+
|
|
1438
|
+
# Update the pointer values in the existing device buffers.
|
|
1439
|
+
cudss_utils.update_cudss_dense_ptr_wrapper(
|
|
1440
|
+
existing_ptrs=self.resources_b_ptrs,
|
|
1441
|
+
rhs_ptr=self.b_ptr,
|
|
1442
|
+
batch_indices=self.batch_indices,
|
|
1443
|
+
new_rhs=self.b,
|
|
1444
|
+
stream_holder=stream_holder,
|
|
1445
|
+
)
|
|
1446
|
+
|
|
1447
|
+
# Clear the operands released flag
|
|
1448
|
+
self._operands_released = False
|
|
1449
|
+
|
|
1450
|
+
def _reset_operands_unchecked_cross_memspace(self, a, b, stream):
|
|
1451
|
+
"""
|
|
1452
|
+
(private) Reset operands unchecked for cross-memspace scenario,
|
|
1453
|
+
i.e., when a copy across memory spaces is required.
|
|
1454
|
+
"""
|
|
1455
|
+
# We need to create a stream holder here because for this
|
|
1456
|
+
# cross-memspace scenario, we need to copy the operands across memory spaces.
|
|
1457
|
+
stream_holder = utils.get_or_create_stream(self.device_id, stream, self.stream_package)
|
|
1458
|
+
|
|
1459
|
+
if a is not None:
|
|
1460
|
+
a = sp_utils.wrap_sparse_operands(a)
|
|
1461
|
+
if self.explicitly_batched_lhs:
|
|
1462
|
+
if not self._operands_released:
|
|
1463
|
+
for x, y in zip(self.a, a, strict=True):
|
|
1464
|
+
x.copy_(y, stream_holder)
|
|
1465
|
+
else:
|
|
1466
|
+
self.a = [x.to(self.device_id, stream_holder) for x in a]
|
|
1467
|
+
else:
|
|
1468
|
+
if not self._operands_released:
|
|
1469
|
+
self.a.copy_(a, stream_holder)
|
|
1470
|
+
else:
|
|
1471
|
+
self.a = a.to(self.device_id, stream_holder)
|
|
1472
|
+
|
|
1473
|
+
cudss_utils.update_cudss_csr_ptr_wrapper(
|
|
1474
|
+
existing_ptrs=self.resources_a_ptrs,
|
|
1475
|
+
lhs_ptr=self.a_ptr,
|
|
1476
|
+
batch_indices=self.batch_indices,
|
|
1477
|
+
new_lhs=self.a,
|
|
1478
|
+
stream_holder=stream_holder,
|
|
1479
|
+
)
|
|
1480
|
+
|
|
1481
|
+
if b is not None:
|
|
1482
|
+
b = tensor_wrapper.wrap_operands(b) if self.explicitly_batched_rhs else tensor_wrapper.wrap_operand(b)
|
|
1483
|
+
if self.explicitly_batched_rhs:
|
|
1484
|
+
if not self._operands_released:
|
|
1485
|
+
for x, y in zip(self.b, b, strict=True):
|
|
1486
|
+
x.copy_(y, stream_holder)
|
|
1487
|
+
else:
|
|
1488
|
+
self.b = [x.to(self.device_id, stream_holder) for x in b]
|
|
1489
|
+
else:
|
|
1490
|
+
if not self._operands_released:
|
|
1491
|
+
self.b.copy_(b, stream_holder)
|
|
1492
|
+
else:
|
|
1493
|
+
self.b = b.to(self.device_id, stream_holder)
|
|
1494
|
+
|
|
1495
|
+
cudss_utils.update_cudss_dense_ptr_wrapper(
|
|
1496
|
+
existing_ptrs=self.resources_b_ptrs,
|
|
1497
|
+
rhs_ptr=self.b_ptr,
|
|
1498
|
+
batch_indices=self.batch_indices,
|
|
1499
|
+
new_rhs=self.b,
|
|
1500
|
+
stream_holder=stream_holder,
|
|
1501
|
+
)
|
|
1502
|
+
|
|
1503
|
+
def _reset_operands_unchecked_same_memspace(self, a, b, stream):
|
|
1504
|
+
"""
|
|
1505
|
+
(private) Reset operands unchecked for not cross-memspace scenario,
|
|
1506
|
+
i.e., when no copy across memory spaces is required and we can
|
|
1507
|
+
directly reference the user-provided operands.
|
|
1508
|
+
"""
|
|
1509
|
+
# The stream holder is created lazily because the non-batched paths
|
|
1510
|
+
# (update_cudss_*_matrix_ptr inside cudss_utils.update_cudss_*_ptr_wrapper)
|
|
1511
|
+
# do not need one, so avoid it if we can, since it has non-trivial overhead.
|
|
1512
|
+
stream_holder = None
|
|
1513
|
+
if self.batched:
|
|
1514
|
+
stream_holder = utils.get_or_create_stream(self.device_id, stream, self.stream_package)
|
|
1515
|
+
|
|
1516
|
+
if a is not None:
|
|
1517
|
+
if self.explicitly_batched_lhs:
|
|
1518
|
+
for ai_wrapper, ai_raw in zip(self.a, a, strict=True):
|
|
1519
|
+
ai_wrapper.reset_unchecked_keeping_shell(ai_raw)
|
|
1520
|
+
else:
|
|
1521
|
+
self.a.reset_unchecked_keeping_shell(a)
|
|
1522
|
+
|
|
1523
|
+
cudss_utils.update_cudss_csr_ptr_wrapper(
|
|
1524
|
+
existing_ptrs=self.resources_a_ptrs,
|
|
1525
|
+
lhs_ptr=self.a_ptr,
|
|
1526
|
+
batch_indices=self.batch_indices,
|
|
1527
|
+
new_lhs=self.a,
|
|
1528
|
+
stream_holder=stream_holder,
|
|
1529
|
+
)
|
|
1530
|
+
self.solver_planned = self.solver_factorized = False
|
|
1531
|
+
|
|
1532
|
+
if b is not None:
|
|
1533
|
+
if self.explicitly_batched_rhs:
|
|
1534
|
+
for bi_wrapper, bi_raw in zip(self.b, b, strict=True):
|
|
1535
|
+
bi_wrapper.tensor = bi_raw
|
|
1536
|
+
else:
|
|
1537
|
+
self.b.tensor = b
|
|
1538
|
+
|
|
1539
|
+
cudss_utils.update_cudss_dense_ptr_wrapper(
|
|
1540
|
+
existing_ptrs=self.resources_b_ptrs,
|
|
1541
|
+
rhs_ptr=self.b_ptr,
|
|
1542
|
+
batch_indices=self.batch_indices,
|
|
1543
|
+
new_rhs=self.b,
|
|
1544
|
+
stream_holder=stream_holder,
|
|
1545
|
+
)
|
|
1546
|
+
|
|
1547
|
+
def reset_operands_unchecked(
|
|
1548
|
+
self,
|
|
1549
|
+
*,
|
|
1550
|
+
a=None,
|
|
1551
|
+
b=None,
|
|
1552
|
+
stream: utils.AnyStream | int | None = None,
|
|
1553
|
+
):
|
|
1554
|
+
"""
|
|
1555
|
+
{reset_operands_unchecked}
|
|
1556
|
+
"""
|
|
1557
|
+
|
|
1558
|
+
if self.copy_across_memspace:
|
|
1559
|
+
self._reset_operands_unchecked_cross_memspace(a, b, stream)
|
|
1560
|
+
else:
|
|
1561
|
+
self._reset_operands_unchecked_same_memspace(a, b, stream)
|
|
1562
|
+
|
|
1563
|
+
self._operands_released = False
|
|
1564
|
+
|
|
1565
|
+
@utils.precondition(_check_valid_solver)
|
|
1566
|
+
def release_operands(self):
|
|
1567
|
+
"""
|
|
1568
|
+
{release_operands}
|
|
1569
|
+
"""
|
|
1570
|
+
if self._operands_released:
|
|
1571
|
+
self.logger.info("Operands have already been released; nothing to do.")
|
|
1572
|
+
return
|
|
1573
|
+
|
|
1574
|
+
# When copy_across_memspace is False
|
|
1575
|
+
# (CUDA execution with GPU operands, or hybrid execution), self.a
|
|
1576
|
+
# and self.b hold direct references to user-provided tensors.
|
|
1577
|
+
# When copy_across_memspace is True (CUDA execution with CPU operands),
|
|
1578
|
+
# self.a and self.b hold internal device mirrors.
|
|
1579
|
+
# In both cases, for both a and b, we release the tensor references
|
|
1580
|
+
# held by the wrapper objects while keeping the wrapper alive.
|
|
1581
|
+
# This achieves the desired effect while preserving any metadata
|
|
1582
|
+
# (shape, device, etc.) so that reset_operands_unchecked can reuse.
|
|
1583
|
+
|
|
1584
|
+
if self.explicitly_batched_lhs:
|
|
1585
|
+
for ai in self.a:
|
|
1586
|
+
ai.release_keeping_shell()
|
|
1587
|
+
else:
|
|
1588
|
+
self.a.release_keeping_shell()
|
|
1589
|
+
|
|
1590
|
+
if self.explicitly_batched_rhs:
|
|
1591
|
+
for bi in self.b:
|
|
1592
|
+
bi.tensor = None
|
|
1593
|
+
else:
|
|
1594
|
+
self.b.tensor = None
|
|
1595
|
+
|
|
1596
|
+
self._operands_released = True
|
|
1597
|
+
self.logger.info("User-provided operands have been released.")
|
|
1598
|
+
|
|
1599
|
+
@utils.precondition(_check_valid_solver)
|
|
1600
|
+
@utils.precondition(_check_valid_operands, "Planning")
|
|
1601
|
+
def plan(self, *, stream: utils.AnyStream | None = None):
|
|
1602
|
+
"""
|
|
1603
|
+
Plan the sparse direct solve (reordering to minimize fill-in, and symbolic
|
|
1604
|
+
factorization). The planning phase can be optionally configured through
|
|
1605
|
+
the property :attr:`plan_config` (an object of type
|
|
1606
|
+
:class:`DirectSolverPlanConfig`). Planning returns a :class:`DirectSolverPlanInfo`
|
|
1607
|
+
object, which can also be accessed through the property :attr:`plan_info`.
|
|
1608
|
+
|
|
1609
|
+
Args:
|
|
1610
|
+
stream: {stream}
|
|
1611
|
+
|
|
1612
|
+
Returns:
|
|
1613
|
+
A :class:`DirectSolverPlanInfo` object, whose attributes can be queried for
|
|
1614
|
+
information regarding the plan.
|
|
1615
|
+
|
|
1616
|
+
.. seealso::
|
|
1617
|
+
:attr:`plan_config`, :class:`DirectSolverPlanConfig`,
|
|
1618
|
+
:class:`DirectSolverPlanInfo`.
|
|
1619
|
+
|
|
1620
|
+
Examples:
|
|
1621
|
+
|
|
1622
|
+
>>> import cupy as cp
|
|
1623
|
+
>>> import cupyx.scipy.sparse as sp
|
|
1624
|
+
>>> import nvmath
|
|
1625
|
+
|
|
1626
|
+
Prepare sample input data.
|
|
1627
|
+
|
|
1628
|
+
>>> n = 8
|
|
1629
|
+
>>> a = sp.random(n, n, density=0.5, format="csr", dtype="float64")
|
|
1630
|
+
>>> a += sp.diags([2.0] * n, format="csr", dtype="float64")
|
|
1631
|
+
|
|
1632
|
+
Create the RHS, which can be a matrix or vector in column-major layout.
|
|
1633
|
+
|
|
1634
|
+
>>> b = cp.ones((n,), dtype="float64")
|
|
1635
|
+
|
|
1636
|
+
Specify, configure, and plan a @ x = b. Use the stateful object as a context
|
|
1637
|
+
manager to automatically release resources.
|
|
1638
|
+
|
|
1639
|
+
>>> with nvmath.sparse.advanced.DirectSolver(a, b) as solver:
|
|
1640
|
+
... # Configure the reordering algorithm for the plan.
|
|
1641
|
+
... plan_config = solver.plan_config
|
|
1642
|
+
... plan_config.reordering_algorithm = (
|
|
1643
|
+
... nvmath.sparse.advanced.DirectSolverReorderingAlg.NESTED_DISSECTION
|
|
1644
|
+
... )
|
|
1645
|
+
... # Plan the operation using the specified plan configuration, which
|
|
1646
|
+
... # returns a DirectSolverPlanInfo object.
|
|
1647
|
+
... plan_info = solver.plan()
|
|
1648
|
+
... # Query the column permutation, memory estimates, ...
|
|
1649
|
+
... col_perm = plan_info.col_permutation
|
|
1650
|
+
|
|
1651
|
+
Further examples can be found in the `nvmath/examples/sparse/advanced/direct_solver
|
|
1652
|
+
<https://github.com/NVIDIA/nvmath-python/tree/main/examples/sparse/advanced/direct_solver>`_
|
|
1653
|
+
directory.
|
|
1654
|
+
""" # noqa: W505
|
|
1655
|
+
r = self._execute(phase=cudss.Phase.ANALYSIS, stream=stream)
|
|
1656
|
+
self.solver_planned = True
|
|
1657
|
+
return r
|
|
1658
|
+
|
|
1659
|
+
@utils.precondition(_check_valid_solver)
|
|
1660
|
+
@utils.precondition(_check_planned, "Factorization")
|
|
1661
|
+
@utils.precondition(_check_valid_operands, "Factorization")
|
|
1662
|
+
def factorize(self, *, stream: utils.AnyStream | None = None):
|
|
1663
|
+
"""
|
|
1664
|
+
Factorize the system of equations. Numerical factorization is required each time
|
|
1665
|
+
the values in the LHS change, unless (for example) iterative refinement is used
|
|
1666
|
+
during the solve and it converges (see :attr:`solution_config`).
|
|
1667
|
+
|
|
1668
|
+
This phase can be optionally configured through the property
|
|
1669
|
+
:attr:`factorization_config` (an object of type
|
|
1670
|
+
:class:`DirectSolverFactorizationConfig`). Factorization returns a
|
|
1671
|
+
:class:`DirectSolverFactorizationInfo` object, which can also be accessed through
|
|
1672
|
+
the property :attr:`factorization_info`.
|
|
1673
|
+
|
|
1674
|
+
Args:
|
|
1675
|
+
stream: {stream}
|
|
1676
|
+
|
|
1677
|
+
Returns:
|
|
1678
|
+
A :class:`DirectSolverFactorizationInfo` object, whose attributes can be
|
|
1679
|
+
queried for information regarding the numerical factorization.
|
|
1680
|
+
|
|
1681
|
+
.. seealso::
|
|
1682
|
+
:attr:`factorization_config`, :class:`DirectSolverFactorizationConfig`,
|
|
1683
|
+
:class:`DirectSolverFactorizationInfo`.
|
|
1684
|
+
|
|
1685
|
+
Examples:
|
|
1686
|
+
|
|
1687
|
+
>>> import cupy as cp
|
|
1688
|
+
>>> import cupyx.scipy.sparse as sp
|
|
1689
|
+
>>> import nvmath
|
|
1690
|
+
|
|
1691
|
+
Prepare sample input data.
|
|
1692
|
+
|
|
1693
|
+
>>> n = 8
|
|
1694
|
+
>>> a = sp.random(n, n, density=0.25, format="csr", dtype="float64")
|
|
1695
|
+
>>> a += sp.diags([2.0] * n, format="csr", dtype="float64")
|
|
1696
|
+
|
|
1697
|
+
Create the RHS, which can be a matrix or vector in column-major layout.
|
|
1698
|
+
|
|
1699
|
+
>>> b = cp.ones((n, 2), dtype="float64", order="F")
|
|
1700
|
+
|
|
1701
|
+
Specify, plan, and factorize a @ x = b. Use the stateful object as a context
|
|
1702
|
+
manager to automatically release resources.
|
|
1703
|
+
|
|
1704
|
+
>>> with nvmath.sparse.advanced.DirectSolver(a, b) as solver:
|
|
1705
|
+
... # Plan the operation (configure if desired).
|
|
1706
|
+
... plan_info = solver.plan()
|
|
1707
|
+
... # (Optionally) configure the factorization.
|
|
1708
|
+
... fac_config = solver.factorization_config
|
|
1709
|
+
... fac_config.pivot_eps = 1e-14
|
|
1710
|
+
... # Factorize using the specified factorization configuration, which
|
|
1711
|
+
... # returns a DirectSolverFactorizationInfo object.
|
|
1712
|
+
... fac_info = solver.factorize()
|
|
1713
|
+
... # Query the number of non-zeros, inertia, ...
|
|
1714
|
+
... fac_info.lu_nnz
|
|
1715
|
+
40
|
|
1716
|
+
|
|
1717
|
+
Further examples can be found in the `nvmath/examples/sparse/advanced/direct_solver
|
|
1718
|
+
<https://github.com/NVIDIA/nvmath-python/tree/main/examples/sparse/advanced/direct_solver>`_
|
|
1719
|
+
directory.
|
|
1720
|
+
"""
|
|
1721
|
+
r = self._execute(phase=cudss.Phase.FACTORIZATION, stream=stream)
|
|
1722
|
+
self.solver_factorized = True
|
|
1723
|
+
return r
|
|
1724
|
+
|
|
1725
|
+
@utils.precondition(_check_valid_solver)
|
|
1726
|
+
@utils.precondition(_check_planned, "Solver")
|
|
1727
|
+
@utils.precondition(_check_factorized, "Solver")
|
|
1728
|
+
@utils.precondition(_check_valid_operands, "Solver")
|
|
1729
|
+
def solve(self, *, stream: utils.AnyStream | None = None):
|
|
1730
|
+
"""
|
|
1731
|
+
Solve the factorized system of equations.
|
|
1732
|
+
|
|
1733
|
+
This phase can be optionally configured through the property
|
|
1734
|
+
:attr:`solution_config` (an object of type
|
|
1735
|
+
:class:`DirectSolverSolutionConfig`).
|
|
1736
|
+
|
|
1737
|
+
Args:
|
|
1738
|
+
stream: {stream}
|
|
1739
|
+
|
|
1740
|
+
Returns:
|
|
1741
|
+
{result}
|
|
1742
|
+
|
|
1743
|
+
.. seealso::
|
|
1744
|
+
:attr:`solution_config`, :class:`DirectSolverSolutionConfig`.
|
|
1745
|
+
|
|
1746
|
+
Examples:
|
|
1747
|
+
|
|
1748
|
+
>>> import cupy as cp
|
|
1749
|
+
>>> import cupyx.scipy.sparse as sp
|
|
1750
|
+
>>> import nvmath
|
|
1751
|
+
|
|
1752
|
+
Prepare sample input data.
|
|
1753
|
+
|
|
1754
|
+
>>> n = 8
|
|
1755
|
+
>>> a = sp.random(n, n, density=0.25, format="csr", dtype="float64")
|
|
1756
|
+
>>> a += sp.diags([2.0] * n, format="csr", dtype="float64")
|
|
1757
|
+
|
|
1758
|
+
Create the RHS, which can be a matrix or vector in column-major layout.
|
|
1759
|
+
|
|
1760
|
+
>>> b = cp.ones((n, 2), dtype="float64", order="F")
|
|
1761
|
+
|
|
1762
|
+
Specify, plan, factorize, and solve a @ x = b for x. Use the stateful
|
|
1763
|
+
object as a context manager to automatically release resources.
|
|
1764
|
+
|
|
1765
|
+
>>> with nvmath.sparse.advanced.DirectSolver(a, b) as solver:
|
|
1766
|
+
... # Plan the operation (configure if desired).
|
|
1767
|
+
... plan_info = solver.plan()
|
|
1768
|
+
... # Factorize the system (configure if desired).
|
|
1769
|
+
... fac_info = solver.factorize()
|
|
1770
|
+
... # (Optionally) configure the solve.
|
|
1771
|
+
... solution_config = solver.solution_config
|
|
1772
|
+
... solution_config.ir_num_steps = 10
|
|
1773
|
+
... # Solve the system based on the solution configuration set above.
|
|
1774
|
+
... x = solver.solve()
|
|
1775
|
+
|
|
1776
|
+
Further examples can be found in the `nvmath/examples/sparse/advanced/direct_solver
|
|
1777
|
+
<https://github.com/NVIDIA/nvmath-python/tree/main/examples/sparse/advanced/direct_solver>`_
|
|
1778
|
+
directory.
|
|
1779
|
+
"""
|
|
1780
|
+
return self._execute(phase=cudss.Phase.SOLVE, stream=stream)
|
|
1781
|
+
|
|
1782
|
+
def _execute(self, *, phase: cudss.Phase, stream=None):
|
|
1783
|
+
"""
|
|
1784
|
+
For internal use only. Execute the specified operation (reordering, factorization,
|
|
1785
|
+
and solve).
|
|
1786
|
+
|
|
1787
|
+
Args:
|
|
1788
|
+
phase: The operation phase (as Python `int` or enumeration of type Phase).
|
|
1789
|
+
|
|
1790
|
+
stream: {stream}
|
|
1791
|
+
|
|
1792
|
+
Returns:
|
|
1793
|
+
{plan|factorization]_info object for reordering and factorization, and the
|
|
1794
|
+
solution in the expected memory space for solve.
|
|
1795
|
+
"""
|
|
1796
|
+
|
|
1797
|
+
assert phase in (cudss.Phase.ANALYSIS, cudss.Phase.FACTORIZATION, cudss.Phase.SOLVE), "Internal error."
|
|
1798
|
+
|
|
1799
|
+
log_info = self.logger.isEnabledFor(logging.INFO)
|
|
1800
|
+
log_debug = self.logger.isEnabledFor(logging.DEBUG)
|
|
1801
|
+
|
|
1802
|
+
stream_holder = utils.get_or_create_stream(self.device_id, stream, self.stream_package)
|
|
1803
|
+
|
|
1804
|
+
cudss.set_stream(self.handle, stream_holder.ptr)
|
|
1805
|
+
|
|
1806
|
+
# Allocate result for only the solve phase and update the pointer.
|
|
1807
|
+
result = None
|
|
1808
|
+
if phase == cudss.Phase.SOLVE:
|
|
1809
|
+
# The result can be either a single tensor or a sequence of tensors, depending
|
|
1810
|
+
# on whether the RHS is explicitly batched.
|
|
1811
|
+
result_allocator = self._allocate_batched_result if self.explicitly_batched_rhs else self._allocate_single_result
|
|
1812
|
+
result = result_allocator(stream_holder, log_debug)
|
|
1813
|
+
|
|
1814
|
+
# Update the pointer values in the existing device buffers.
|
|
1815
|
+
cudss_utils.update_cudss_dense_ptr_wrapper(
|
|
1816
|
+
existing_ptrs=self.resources_x_ptrs,
|
|
1817
|
+
rhs_ptr=self.x_ptr,
|
|
1818
|
+
batch_indices=self.batch_indices,
|
|
1819
|
+
new_rhs=result,
|
|
1820
|
+
stream_holder=stream_holder,
|
|
1821
|
+
)
|
|
1822
|
+
|
|
1823
|
+
if log_info:
|
|
1824
|
+
self.logger.info(f"Starting solver phase {phase.name}...")
|
|
1825
|
+
self.logger.info(f"{self.call_prologue}")
|
|
1826
|
+
|
|
1827
|
+
with utils.cuda_call_ctx(stream_holder, self.blocking, timing=log_info) as (
|
|
1828
|
+
self.last_compute_event,
|
|
1829
|
+
elapsed,
|
|
1830
|
+
):
|
|
1831
|
+
cudss.execute(self.handle, phase, self.config_ptr, self.data_ptr, self.a_ptr, self.x_ptr, self.b_ptr)
|
|
1832
|
+
|
|
1833
|
+
if log_info and elapsed.data is not None:
|
|
1834
|
+
self.logger.info(f"The solver phase {phase.name} took {elapsed.data:.3f} ms to complete.")
|
|
1835
|
+
|
|
1836
|
+
if phase == cudss.Phase.ANALYSIS:
|
|
1837
|
+
return self.plan_info
|
|
1838
|
+
|
|
1839
|
+
if phase == cudss.Phase.FACTORIZATION:
|
|
1840
|
+
return self.factorization_info
|
|
1841
|
+
|
|
1842
|
+
assert result is not None, "Internal Error."
|
|
1843
|
+
|
|
1844
|
+
# Ideally, we should set the x_ptr to 0, but this adds overhead for the batched
|
|
1845
|
+
# case.
|
|
1846
|
+
|
|
1847
|
+
# Copy result to required memory spaces if needed.
|
|
1848
|
+
if self.copy_across_memspace:
|
|
1849
|
+
result = copy_single_or_sequence(result, "cpu", stream_holder)
|
|
1850
|
+
|
|
1851
|
+
# Extract the result tensor from the wrapped result for the solution phase.
|
|
1852
|
+
if isinstance(result, Sequence):
|
|
1853
|
+
out = tuple(r.tensor for r in result)
|
|
1854
|
+
else:
|
|
1855
|
+
out = result.tensor
|
|
1856
|
+
|
|
1857
|
+
return out
|
|
1858
|
+
|
|
1859
|
+
def free(self):
|
|
1860
|
+
"""Free DirectSolver resources.
|
|
1861
|
+
|
|
1862
|
+
It is recommended that the :class:`DirectSolver` object be used within a context,
|
|
1863
|
+
but if it is not possible then this method must be called explicitly to ensure
|
|
1864
|
+
that the sparse direct solver resources (especially internal library objects) are
|
|
1865
|
+
properly cleaned up.
|
|
1866
|
+
"""
|
|
1867
|
+
|
|
1868
|
+
if not self.valid_state:
|
|
1869
|
+
return
|
|
1870
|
+
|
|
1871
|
+
try:
|
|
1872
|
+
# Ensure ordering with respect to the last computation
|
|
1873
|
+
# to avoid race conditions when releasing internal resources.
|
|
1874
|
+
if self.last_compute_event is not None:
|
|
1875
|
+
if self.resources_ptrs_alloc_stream is not None:
|
|
1876
|
+
self.resources_ptrs_alloc_stream.wait(self.last_compute_event)
|
|
1877
|
+
self.last_compute_event = None
|
|
1878
|
+
|
|
1879
|
+
# Currently, workspace is allocated internally by the library.
|
|
1880
|
+
|
|
1881
|
+
# Release internal resource references.
|
|
1882
|
+
self.resources_a_dims = self.resources_a_ptrs = None
|
|
1883
|
+
self.resources_b_dims = self.resources_b_ptrs = None
|
|
1884
|
+
self.resources_x_dims = self.resources_x_ptrs = None
|
|
1885
|
+
self.resources_ptrs_alloc_stream = None
|
|
1886
|
+
self.a = self.b = None
|
|
1887
|
+
|
|
1888
|
+
# Free matrix pointers.
|
|
1889
|
+
cudss.matrix_destroy(self.x_ptr)
|
|
1890
|
+
cudss.matrix_destroy(self.b_ptr)
|
|
1891
|
+
cudss.matrix_destroy(self.a_ptr)
|
|
1892
|
+
self.x_ptr = self.b_ptr = self.a_ptr = None
|
|
1893
|
+
|
|
1894
|
+
# Free config and data pointers.
|
|
1895
|
+
cudss.data_destroy(self.handle, self.data_ptr)
|
|
1896
|
+
cudss.config_destroy(self.config_ptr)
|
|
1897
|
+
self.data_ptr = self.config_ptr = None
|
|
1898
|
+
|
|
1899
|
+
# Free handle if we own it.
|
|
1900
|
+
if self.handle is not None and self.own_handle:
|
|
1901
|
+
cudss.destroy(self.handle)
|
|
1902
|
+
self.handle, self.own_handle = None, False
|
|
1903
|
+
|
|
1904
|
+
# Set all attributes to None except for logger and valid_state.
|
|
1905
|
+
for attr in list(vars(self)):
|
|
1906
|
+
if attr not in {"logger", "valid_state"}:
|
|
1907
|
+
setattr(self, attr, None)
|
|
1908
|
+
|
|
1909
|
+
except Exception as e:
|
|
1910
|
+
self.logger.critical("Internal error: only part of the DirectSolver object's resources have been released.")
|
|
1911
|
+
self.logger.critical(str(e))
|
|
1912
|
+
raise e
|
|
1913
|
+
finally:
|
|
1914
|
+
self.valid_state = False
|
|
1915
|
+
|
|
1916
|
+
self.logger.info("The DirectSolver object's resources have been released.")
|
|
1917
|
+
|
|
1918
|
+
|
|
1919
|
+
@utils.docstring_decorator(SHARED_DSS_DOCUMENTATION, skip_missing=False)
|
|
1920
|
+
def direct_solver(
|
|
1921
|
+
a,
|
|
1922
|
+
b,
|
|
1923
|
+
/,
|
|
1924
|
+
*,
|
|
1925
|
+
options: DirectSolverOptions | dict[str, Any] | None = None,
|
|
1926
|
+
plan_preferences: DirectSolverPlanPreferences | None = None,
|
|
1927
|
+
factorization_preferences: DirectSolverFactorizationPreferences | None = None,
|
|
1928
|
+
solution_preferences: DirectSolverSolutionPreferences | None = None,
|
|
1929
|
+
execution: ExecutionCUDA | ExecutionHybrid | str | dict[str, Any] | None = None,
|
|
1930
|
+
stream: utils.AnyStream | int | None = None,
|
|
1931
|
+
):
|
|
1932
|
+
"""
|
|
1933
|
+
Solve :math:`a @ x = b` for :math:`x`. This function-form API is a wrapper around the
|
|
1934
|
+
stateful :class:`DirectSolver` object APIs and is meant for *single* use (the user needs
|
|
1935
|
+
to perform just one sparse direct solve, for example), in which case there is no
|
|
1936
|
+
possibility of amortizing preparatory costs.
|
|
1937
|
+
|
|
1938
|
+
Detailed information on what's happening within this function can be obtained by passing
|
|
1939
|
+
in a :class:`logging.Logger` object to :class:`DirectSolverOptions` or by setting the
|
|
1940
|
+
appropriate options in the root logger object, which is used by default:
|
|
1941
|
+
|
|
1942
|
+
>>> import logging
|
|
1943
|
+
>>> logging.basicConfig(
|
|
1944
|
+
... level=logging.INFO,
|
|
1945
|
+
... format="%(asctime)s %(levelname)-8s %(message)s",
|
|
1946
|
+
... datefmt="%m-%d %H:%M:%S",
|
|
1947
|
+
... )
|
|
1948
|
+
|
|
1949
|
+
A user can select the desired logging level and, in general, take advantage of all of
|
|
1950
|
+
the functionality offered by the Python :mod:`logging` module.
|
|
1951
|
+
|
|
1952
|
+
Args:
|
|
1953
|
+
a: {a}
|
|
1954
|
+
|
|
1955
|
+
b: {b}
|
|
1956
|
+
|
|
1957
|
+
options: {options}
|
|
1958
|
+
|
|
1959
|
+
plan_preferences: {plan_preferences}
|
|
1960
|
+
|
|
1961
|
+
factorization_preferences: {factorization_preferences}
|
|
1962
|
+
|
|
1963
|
+
solution_preferences: {solution_preferences}
|
|
1964
|
+
|
|
1965
|
+
execution: {execution}
|
|
1966
|
+
|
|
1967
|
+
stream: {stream}
|
|
1968
|
+
|
|
1969
|
+
Returns:
|
|
1970
|
+
{result}
|
|
1971
|
+
|
|
1972
|
+
.. _Semantics:
|
|
1973
|
+
|
|
1974
|
+
Semantics:
|
|
1975
|
+
{semantics}
|
|
1976
|
+
|
|
1977
|
+
.. seealso::
|
|
1978
|
+
:class:`DirectSolver`, :class:`DirectSolverOptions`, :class:`ExecutionCUDA`,
|
|
1979
|
+
:class:`ExecutionHybrid`.
|
|
1980
|
+
|
|
1981
|
+
Examples:
|
|
1982
|
+
|
|
1983
|
+
>>> import cupy as cp
|
|
1984
|
+
>>> import cupyx.scipy.sparse as sp
|
|
1985
|
+
>>> import nvmath
|
|
1986
|
+
|
|
1987
|
+
Create a sparse float32 ndarray in CSR format on the CPU for the LHS.
|
|
1988
|
+
|
|
1989
|
+
>>> n = 16
|
|
1990
|
+
>>> a = sp.random(n, n, density=0.5, format="csr", dtype="float32")
|
|
1991
|
+
|
|
1992
|
+
Ensure that the randomly-generated LHS is not singular.
|
|
1993
|
+
|
|
1994
|
+
>>> a += sp.diags([2.0] * n, format="csr", dtype="float32")
|
|
1995
|
+
|
|
1996
|
+
The RHS can be a vector or matrix. Here we create a random matrix with 4 columns
|
|
1997
|
+
(indicating 4 vectors to be solved for) in column-major format.
|
|
1998
|
+
|
|
1999
|
+
>>> b = cp.random.rand(4, n, dtype="float32").T
|
|
2000
|
+
|
|
2001
|
+
Solve a @ x = b for x.
|
|
2002
|
+
|
|
2003
|
+
>>> x = nvmath.sparse.advanced.direct_solver(a, b)
|
|
2004
|
+
|
|
2005
|
+
Batching can be specified, explicitly or implicitly, following the semantics
|
|
2006
|
+
described above. Here we explicitly batch the LHS since CuPy doesn't support 3D CSR,
|
|
2007
|
+
while we implicitly batch the RHS.
|
|
2008
|
+
|
|
2009
|
+
Create an explicit batch of two CSR matrices:
|
|
2010
|
+
|
|
2011
|
+
>>> batch = 2
|
|
2012
|
+
>>> a = [a] * batch
|
|
2013
|
+
>>> a[1] *= 10.0
|
|
2014
|
+
|
|
2015
|
+
Create a 3D ndarray, with each sample in the batch having column-major layout.
|
|
2016
|
+
|
|
2017
|
+
>>> b = cp.random.rand(batch, 4, n, dtype="float32").transpose(0, 2, 1)
|
|
2018
|
+
|
|
2019
|
+
Solve the batched system a @ x = b for x, where x has the same shape as b.
|
|
2020
|
+
|
|
2021
|
+
>>> x = nvmath.sparse.advanced.direct_solver(a, b)
|
|
2022
|
+
|
|
2023
|
+
Options can be provided to the sparse direct solver using
|
|
2024
|
+
:class:`DirectSolverOptions`, and the execution space can be specified using the
|
|
2025
|
+
``execution`` option. Refer to :class:`DirectSolver` and the GitHub link below for
|
|
2026
|
+
examples.
|
|
2027
|
+
|
|
2028
|
+
Notes:
|
|
2029
|
+
|
|
2030
|
+
- This function is a convenience wrapper around :class:`DirectSolver` and is
|
|
2031
|
+
specifically meant for *single* use.
|
|
2032
|
+
|
|
2033
|
+
Further examples can be found in the `nvmath/examples/sparse/advanced/direct_solver
|
|
2034
|
+
<https://github.com/NVIDIA/nvmath-python/tree/main/examples/sparse/advanced/direct_solver>`_
|
|
2035
|
+
directory.
|
|
2036
|
+
"""
|
|
2037
|
+
|
|
2038
|
+
with DirectSolver(
|
|
2039
|
+
a,
|
|
2040
|
+
b,
|
|
2041
|
+
options=options,
|
|
2042
|
+
execution=execution,
|
|
2043
|
+
stream=stream,
|
|
2044
|
+
) as solver:
|
|
2045
|
+
# Update the configurations with the provided preferences if any.
|
|
2046
|
+
if plan_preferences is not None:
|
|
2047
|
+
update_config_with_preferences(solver.plan_config, plan_preferences)
|
|
2048
|
+
if factorization_preferences is not None:
|
|
2049
|
+
update_config_with_preferences(solver.factorization_config, factorization_preferences)
|
|
2050
|
+
if solution_preferences is not None:
|
|
2051
|
+
update_config_with_preferences(solver.solution_config, solution_preferences)
|
|
2052
|
+
|
|
2053
|
+
# Stateless API calls are all executed on the same stream, so we can reuse the
|
|
2054
|
+
# stream resolved in DirectSolver.__init__ (self._init_stream_holder) for plan,
|
|
2055
|
+
# factorize, and solve. Otherwise, in the most common use case when stream = None
|
|
2056
|
+
# is passed, package.cuda.get_current_stream() is repeatedly called inside each
|
|
2057
|
+
# method with non-negligible overhead.
|
|
2058
|
+
resolved_stream = solver._init_stream_holder.external if solver._init_stream_holder is not None else None
|
|
2059
|
+
|
|
2060
|
+
# Planning and factorization information cannot be returned without making copies
|
|
2061
|
+
# because of scope (the interfaces need the solver object, which is released when
|
|
2062
|
+
# the function returns).
|
|
2063
|
+
solver.plan(stream=resolved_stream)
|
|
2064
|
+
|
|
2065
|
+
solver.factorize(stream=resolved_stream)
|
|
2066
|
+
|
|
2067
|
+
r = solver.solve(stream=resolved_stream)
|
|
2068
|
+
|
|
2069
|
+
return r
|