nvmath-python 1.0.0__cp314-cp314t-win_amd64.whl

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