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