nvmath-python 0.9.0__cp314-cp314-win_amd64.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (309) hide show
  1. nvmath/__init__.pxd +0 -0
  2. nvmath/__init__.py +45 -0
  3. nvmath/_internal/__init__.py +0 -0
  4. nvmath/_internal/attribute_ifc_factory.py +330 -0
  5. nvmath/_internal/layout.py +70 -0
  6. nvmath/_internal/templates.py +383 -0
  7. nvmath/_utils.py +150 -0
  8. nvmath/bindings/__init__.py +51 -0
  9. nvmath/bindings/_internal/__init__.pxd +0 -0
  10. nvmath/bindings/_internal/__init__.py +0 -0
  11. nvmath/bindings/_internal/cublas.cp314-win_amd64.pyd +0 -0
  12. nvmath/bindings/_internal/cublas.pxd +530 -0
  13. nvmath/bindings/_internal/cublasLt.cp314-win_amd64.pyd +0 -0
  14. nvmath/bindings/_internal/cublasLt.pxd +59 -0
  15. nvmath/bindings/_internal/cublasMp.pxd +52 -0
  16. nvmath/bindings/_internal/cudss.cp314-win_amd64.pyd +0 -0
  17. nvmath/bindings/_internal/cudss.pxd +47 -0
  18. nvmath/bindings/_internal/cufft.cp314-win_amd64.pyd +0 -0
  19. nvmath/bindings/_internal/cufft.pxd +70 -0
  20. nvmath/bindings/_internal/cufftMp.pxd +77 -0
  21. nvmath/bindings/_internal/curand.cp314-win_amd64.pyd +0 -0
  22. nvmath/bindings/_internal/curand.pxd +42 -0
  23. nvmath/bindings/_internal/cusolver.cp314-win_amd64.pyd +0 -0
  24. nvmath/bindings/_internal/cusolver.pxd +15 -0
  25. nvmath/bindings/_internal/cusolverDn.cp314-win_amd64.pyd +0 -0
  26. nvmath/bindings/_internal/cusolverDn.pxd +406 -0
  27. nvmath/bindings/_internal/cusparse.cp314-win_amd64.pyd +0 -0
  28. nvmath/bindings/_internal/cusparse.pxd +277 -0
  29. nvmath/bindings/_internal/cusparseLt.cp314-win_amd64.pyd +0 -0
  30. nvmath/bindings/_internal/cusparseLt.pxd +48 -0
  31. nvmath/bindings/_internal/cutensor.cp314-win_amd64.pyd +0 -0
  32. nvmath/bindings/_internal/cutensor.pxd +58 -0
  33. nvmath/bindings/_internal/mathdx.cp314-win_amd64.pyd +0 -0
  34. nvmath/bindings/_internal/mathdx.pxd +116 -0
  35. nvmath/bindings/_internal/nvshmem.pxd +29 -0
  36. nvmath/bindings/_internal/utils.cp314-win_amd64.pyd +0 -0
  37. nvmath/bindings/_internal/utils.pxd +179 -0
  38. nvmath/bindings/_internal/utils.pyi +10 -0
  39. nvmath/bindings/cublas.cp314-win_amd64.pyd +0 -0
  40. nvmath/bindings/cublas.pxd +558 -0
  41. nvmath/bindings/cublas.pyi +746 -0
  42. nvmath/bindings/cublasLt.cp314-win_amd64.pyd +0 -0
  43. nvmath/bindings/cublasLt.pxd +109 -0
  44. nvmath/bindings/cublasLt.pyi +1337 -0
  45. nvmath/bindings/cublasMp.pxd +85 -0
  46. nvmath/bindings/cublasMp.pyi +219 -0
  47. nvmath/bindings/cudss.cp314-win_amd64.pyd +0 -0
  48. nvmath/bindings/cudss.pxd +86 -0
  49. nvmath/bindings/cudss.pyi +282 -0
  50. nvmath/bindings/cufft.cp314-win_amd64.pyd +0 -0
  51. nvmath/bindings/cufft.pxd +118 -0
  52. nvmath/bindings/cufft.pyi +241 -0
  53. nvmath/bindings/cufftMp.pxd +124 -0
  54. nvmath/bindings/cufftMp.pyi +260 -0
  55. nvmath/bindings/curand.cp314-win_amd64.pyd +0 -0
  56. nvmath/bindings/curand.pxd +71 -0
  57. nvmath/bindings/curand.pyi +159 -0
  58. nvmath/bindings/cusolver.cp314-win_amd64.pyd +0 -0
  59. nvmath/bindings/cusolver.pxd +62 -0
  60. nvmath/bindings/cusolver.pyi +242 -0
  61. nvmath/bindings/cusolverDn.cp314-win_amd64.pyd +0 -0
  62. nvmath/bindings/cusolverDn.pxd +430 -0
  63. nvmath/bindings/cusolverDn.pyi +416 -0
  64. nvmath/bindings/cusparse.cp314-win_amd64.pyd +0 -0
  65. nvmath/bindings/cusparse.pxd +338 -0
  66. nvmath/bindings/cusparse.pyi +659 -0
  67. nvmath/bindings/cusparseLt.cp314-win_amd64.pyd +0 -0
  68. nvmath/bindings/cusparseLt.pxd +99 -0
  69. nvmath/bindings/cusparseLt.pyi +198 -0
  70. nvmath/bindings/cutensor.cp314-win_amd64.pyd +0 -0
  71. nvmath/bindings/cutensor.pxd +98 -0
  72. nvmath/bindings/cutensor.pyi +294 -0
  73. nvmath/bindings/cycublas.cp314-win_amd64.pyd +0 -0
  74. nvmath/bindings/cycublas.pxd +684 -0
  75. nvmath/bindings/cycublasLt.cp314-win_amd64.pyd +0 -0
  76. nvmath/bindings/cycublasLt.pxd +1055 -0
  77. nvmath/bindings/cycublasMp.pxd +183 -0
  78. nvmath/bindings/cycudss.cp314-win_amd64.pyd +0 -0
  79. nvmath/bindings/cycudss.pxd +224 -0
  80. nvmath/bindings/cycufft.cp314-win_amd64.pyd +0 -0
  81. nvmath/bindings/cycufft.pxd +326 -0
  82. nvmath/bindings/cycufftMp.pxd +334 -0
  83. nvmath/bindings/cycurand.cp314-win_amd64.pyd +0 -0
  84. nvmath/bindings/cycurand.pxd +146 -0
  85. nvmath/bindings/cycusolver.cp314-win_amd64.pyd +0 -0
  86. nvmath/bindings/cycusolver.pxd +154 -0
  87. nvmath/bindings/cycusolverDn.cp314-win_amd64.pyd +0 -0
  88. nvmath/bindings/cycusolverDn.pxd +446 -0
  89. nvmath/bindings/cycusparse.cp314-win_amd64.pyd +0 -0
  90. nvmath/bindings/cycusparse.pxd +470 -0
  91. nvmath/bindings/cycusparseLt.cp314-win_amd64.pyd +0 -0
  92. nvmath/bindings/cycusparseLt.pxd +150 -0
  93. nvmath/bindings/cycutensor.cp314-win_amd64.pyd +0 -0
  94. nvmath/bindings/cycutensor.pxd +192 -0
  95. nvmath/bindings/cymathdx.cp314-win_amd64.pyd +0 -0
  96. nvmath/bindings/cymathdx.pxd +552 -0
  97. nvmath/bindings/cynvshmem.pxd +126 -0
  98. nvmath/bindings/mathdx.cp314-win_amd64.pyd +0 -0
  99. nvmath/bindings/mathdx.pxd +182 -0
  100. nvmath/bindings/mathdx.pyi +1244 -0
  101. nvmath/bindings/nvpl/__init__.pxd +0 -0
  102. nvmath/bindings/nvpl/__init__.py +13 -0
  103. nvmath/bindings/nvpl/_internal/__init__.pxd +0 -0
  104. nvmath/bindings/nvpl/_internal/__init__.py +0 -0
  105. nvmath/bindings/nvpl/_internal/blas.cp314-win_amd64.pyd +0 -0
  106. nvmath/bindings/nvpl/_internal/blas.pxd +125 -0
  107. nvmath/bindings/nvpl/_internal/fft.cp314-win_amd64.pyd +0 -0
  108. nvmath/bindings/nvpl/_internal/fft.pxd +36 -0
  109. nvmath/bindings/nvpl/blas.cp314-win_amd64.pyd +0 -0
  110. nvmath/bindings/nvpl/blas.pxd +85 -0
  111. nvmath/bindings/nvpl/blas.pyi +122 -0
  112. nvmath/bindings/nvpl/cyblas.cp314-win_amd64.pyd +0 -0
  113. nvmath/bindings/nvpl/cyblas.pxd +166 -0
  114. nvmath/bindings/nvpl/cyfft.cp314-win_amd64.pyd +0 -0
  115. nvmath/bindings/nvpl/cyfft.pxd +92 -0
  116. nvmath/bindings/nvpl/fft.cp314-win_amd64.pyd +0 -0
  117. nvmath/bindings/nvpl/fft.pxd +100 -0
  118. nvmath/bindings/nvpl/fft.pyi +100 -0
  119. nvmath/bindings/nvshmem.pxd +54 -0
  120. nvmath/bindings/nvshmem.pyi +179 -0
  121. nvmath/device/__init__.py +21 -0
  122. nvmath/device/_deprecated.py +33 -0
  123. nvmath/device/common.py +313 -0
  124. nvmath/device/common_backend.py +131 -0
  125. nvmath/device/common_cuda.py +201 -0
  126. nvmath/device/common_numba.py +310 -0
  127. nvmath/device/common_opaque_tensor.py +201 -0
  128. nvmath/device/cublasdx.py +1805 -0
  129. nvmath/device/cublasdx_backend.py +807 -0
  130. nvmath/device/cublasdx_numba.py +1612 -0
  131. nvmath/device/cufftdx.py +510 -0
  132. nvmath/device/cufftdx_backend.py +196 -0
  133. nvmath/device/cufftdx_numba.py +196 -0
  134. nvmath/device/curand_kernel.py +9147 -0
  135. nvmath/device/cusolverdx.py +2690 -0
  136. nvmath/device/cusolverdx_backend.py +440 -0
  137. nvmath/device/cusolverdx_numba.py +1624 -0
  138. nvmath/device/llvm_array.py +29 -0
  139. nvmath/device/random.py +445 -0
  140. nvmath/device/random_helpers.py +23 -0
  141. nvmath/device/random_states.py +187 -0
  142. nvmath/device/types.py +147 -0
  143. nvmath/device/vector_types_numba.py +203 -0
  144. nvmath/distributed/__init__.py +205 -0
  145. nvmath/distributed/_internal/__init__.py +0 -0
  146. nvmath/distributed/_internal/nvshmem.py +302 -0
  147. nvmath/distributed/_internal/tensor_ifc.py +67 -0
  148. nvmath/distributed/_internal/tensor_ifc_cupy.py +62 -0
  149. nvmath/distributed/_internal/tensor_ifc_host_device.py +165 -0
  150. nvmath/distributed/_internal/tensor_ifc_numpy.py +41 -0
  151. nvmath/distributed/_internal/tensor_ifc_torch.py +141 -0
  152. nvmath/distributed/_internal/tensor_wrapper.py +78 -0
  153. nvmath/distributed/_utils.py +167 -0
  154. nvmath/distributed/distribution.py +770 -0
  155. nvmath/distributed/fft/__init__.py +7 -0
  156. nvmath/distributed/fft/_configuration.py +79 -0
  157. nvmath/distributed/fft/fft.py +2801 -0
  158. nvmath/distributed/linalg/__init__.py +12 -0
  159. nvmath/distributed/linalg/_internal/__init__.py +3 -0
  160. nvmath/distributed/linalg/_internal/epilog_protocol.py +496 -0
  161. nvmath/distributed/linalg/_internal/matmul_desc_ifc.py +28 -0
  162. nvmath/distributed/linalg/advanced/__init__.py +8 -0
  163. nvmath/distributed/linalg/advanced/_configuration.py +166 -0
  164. nvmath/distributed/linalg/advanced/matmulmod.py +2908 -0
  165. nvmath/distributed/process_group.py +408 -0
  166. nvmath/distributed/reshape/__init__.py +6 -0
  167. nvmath/distributed/reshape/_configuration.py +39 -0
  168. nvmath/distributed/reshape/reshape.py +1256 -0
  169. nvmath/fft/__init__.py +7 -0
  170. nvmath/fft/_configuration.py +208 -0
  171. nvmath/fft/_exec_utils.py +109 -0
  172. nvmath/fft/_helpers.py +237 -0
  173. nvmath/fft/fft.py +2903 -0
  174. nvmath/internal/__init__.pxd +3 -0
  175. nvmath/internal/__init__.py +10 -0
  176. nvmath/internal/_bindings.cp314-win_amd64.pyd +0 -0
  177. nvmath/internal/_bindings.pxd +12 -0
  178. nvmath/internal/_device_utils.py +54 -0
  179. nvmath/internal/enum_utils.py +142 -0
  180. nvmath/internal/formatters.py +87 -0
  181. nvmath/internal/mem_limit.py +51 -0
  182. nvmath/internal/memory.cp314-win_amd64.pyd +0 -0
  183. nvmath/internal/memory.pxd +10 -0
  184. nvmath/internal/memory.pyi +46 -0
  185. nvmath/internal/ndbuffer/__init__.pxd +3 -0
  186. nvmath/internal/ndbuffer/__init__.py +7 -0
  187. nvmath/internal/ndbuffer/copy_kernel/args.h +34 -0
  188. nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/array_view.h +52 -0
  189. nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/elementwise.h +68 -0
  190. nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/grid_indexer.h +69 -0
  191. nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/transposed.h +240 -0
  192. nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/type_utils.h +39 -0
  193. nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/utils.h +132 -0
  194. nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/vec.h +159 -0
  195. nvmath/internal/ndbuffer/copy_kernel/elementwise.h +53 -0
  196. nvmath/internal/ndbuffer/copy_kernel/transposed.h +58 -0
  197. nvmath/internal/ndbuffer/copy_kernel.cp314-win_amd64.pyd +0 -0
  198. nvmath/internal/ndbuffer/copy_kernel.pxd +9 -0
  199. nvmath/internal/ndbuffer/copy_kernel.pyi +9 -0
  200. nvmath/internal/ndbuffer/data_layout.cp314-win_amd64.pyd +0 -0
  201. nvmath/internal/ndbuffer/data_layout.pxd +72 -0
  202. nvmath/internal/ndbuffer/data_layout.pyi +23 -0
  203. nvmath/internal/ndbuffer/jit.cp314-win_amd64.pyd +0 -0
  204. nvmath/internal/ndbuffer/jit.pxd +7 -0
  205. nvmath/internal/ndbuffer/jit.pyi +6 -0
  206. nvmath/internal/ndbuffer/ndbuffer.cp314-win_amd64.pyd +0 -0
  207. nvmath/internal/ndbuffer/ndbuffer.pxd +36 -0
  208. nvmath/internal/ndbuffer/ndbuffer.pyi +41 -0
  209. nvmath/internal/ndbuffer/package_utils.cp314-win_amd64.pyd +0 -0
  210. nvmath/internal/ndbuffer/package_utils.pxd +11 -0
  211. nvmath/internal/ndbuffer/package_utils.pyi +12 -0
  212. nvmath/internal/package_ifc.py +147 -0
  213. nvmath/internal/package_ifc_cuda.py +49 -0
  214. nvmath/internal/package_ifc_cupy.py +49 -0
  215. nvmath/internal/package_ifc_torch.py +31 -0
  216. nvmath/internal/package_wrapper.py +14 -0
  217. nvmath/internal/tensor_ifc.py +194 -0
  218. nvmath/internal/tensor_ifc_cupy.py +241 -0
  219. nvmath/internal/tensor_ifc_ndbuffer.py +153 -0
  220. nvmath/internal/tensor_ifc_numpy.py +183 -0
  221. nvmath/internal/tensor_ifc_torch.py +176 -0
  222. nvmath/internal/tensor_wrapper.py +157 -0
  223. nvmath/internal/typemaps.py +162 -0
  224. nvmath/internal/utils.py +771 -0
  225. nvmath/linalg/__init__.py +46 -0
  226. nvmath/linalg/_internal/__init__.py +3 -0
  227. nvmath/linalg/_internal/algo_cap_ifc.py +82 -0
  228. nvmath/linalg/_internal/algo_config_ifc.py +43 -0
  229. nvmath/linalg/_internal/batch.py +217 -0
  230. nvmath/linalg/_internal/enum_to_tuples.py +64 -0
  231. nvmath/linalg/_internal/epilog_protocol.py +766 -0
  232. nvmath/linalg/_internal/layout.py +624 -0
  233. nvmath/linalg/_internal/matmul_desc_ifc.py +28 -0
  234. nvmath/linalg/_internal/matmul_pref_ifc.py +27 -0
  235. nvmath/linalg/_internal/matrix_layout_ifc.py +26 -0
  236. nvmath/linalg/_internal/typemaps.py +144 -0
  237. nvmath/linalg/_internal/utils.py +134 -0
  238. nvmath/linalg/advanced/__init__.py +8 -0
  239. nvmath/linalg/advanced/_algorithmmod.py +149 -0
  240. nvmath/linalg/advanced/_configuration.py +349 -0
  241. nvmath/linalg/advanced/helpers/__init__.py +5 -0
  242. nvmath/linalg/advanced/helpers/matmul.py +1348 -0
  243. nvmath/linalg/advanced/matmulmod.py +3937 -0
  244. nvmath/linalg/generic/__init__.py +43 -0
  245. nvmath/linalg/generic/_configuration/__init__.py +37 -0
  246. nvmath/linalg/generic/_configuration/layout.py +263 -0
  247. nvmath/linalg/generic/_configuration/match.py +609 -0
  248. nvmath/linalg/generic/_configuration/qualifiers.py +494 -0
  249. nvmath/linalg/generic/_configuration/wrap.py +217 -0
  250. nvmath/linalg/generic/_dtype.py +15 -0
  251. nvmath/linalg/generic/matmulmod.py +1060 -0
  252. nvmath/memory.py +282 -0
  253. nvmath/sparse/__init__.py +38 -0
  254. nvmath/sparse/_internal/__init__.py +21 -0
  255. nvmath/sparse/_internal/common_utils.py +144 -0
  256. nvmath/sparse/_internal/cudss_config_ifc.py +741 -0
  257. nvmath/sparse/_internal/cudss_data_ifc.py +399 -0
  258. nvmath/sparse/_internal/cudss_utils.py +445 -0
  259. nvmath/sparse/_internal/cusparse_utils.py +382 -0
  260. nvmath/sparse/_internal/sparse_bsc_ifc.py +300 -0
  261. nvmath/sparse/_internal/sparse_bsr_ifc.py +303 -0
  262. nvmath/sparse/_internal/sparse_coo_ifc.py +254 -0
  263. nvmath/sparse/_internal/sparse_csc_ifc.py +266 -0
  264. nvmath/sparse/_internal/sparse_csr_ifc.py +266 -0
  265. nvmath/sparse/_internal/sparse_dia_ifc.py +239 -0
  266. nvmath/sparse/_internal/sparse_format_helpers.py +601 -0
  267. nvmath/sparse/_internal/sparse_tensor_ifc.py +121 -0
  268. nvmath/sparse/_internal/sparse_ust_ifc.py +146 -0
  269. nvmath/sparse/_internal/utils.py +56 -0
  270. nvmath/sparse/advanced/__init__.py +7 -0
  271. nvmath/sparse/advanced/_configuration.py +230 -0
  272. nvmath/sparse/advanced/direct_solver.py +1884 -0
  273. nvmath/sparse/generic/__init__.py +7 -0
  274. nvmath/sparse/generic/_configuration.py +129 -0
  275. nvmath/sparse/generic/_helpers.py +137 -0
  276. nvmath/sparse/generic/_thunks.py +21 -0
  277. nvmath/sparse/generic/matmulmod.py +2426 -0
  278. nvmath/sparse/ust/__init__.py +7 -0
  279. nvmath/sparse/ust/_converters.py +378 -0
  280. nvmath/sparse/ust/_drawer.py +552 -0
  281. nvmath/sparse/ust/_emitter.py +553 -0
  282. nvmath/sparse/ust/_jit.py +212 -0
  283. nvmath/sparse/ust/_kernel.py +346 -0
  284. nvmath/sparse/ust/_semiring.py +71 -0
  285. nvmath/sparse/ust/_utils.py +102 -0
  286. nvmath/sparse/ust/interfaces/__init__.py +0 -0
  287. nvmath/sparse/ust/interfaces/torch_interface.py +476 -0
  288. nvmath/sparse/ust/semiring_ops.py +71 -0
  289. nvmath/sparse/ust/tensor.py +1027 -0
  290. nvmath/sparse/ust/tensor_drawer.py +552 -0
  291. nvmath/sparse/ust/tensor_emitter.py +558 -0
  292. nvmath/sparse/ust/tensor_format.py +935 -0
  293. nvmath/sparse/ust/tensor_jit.py +212 -0
  294. nvmath/sparse/ust/tensor_kernel.py +348 -0
  295. nvmath/sparse/ust/tensor_utils.py +380 -0
  296. nvmath/tensor/__init__.py +6 -0
  297. nvmath/tensor/_configuration.py +95 -0
  298. nvmath/tensor/_internal/__init__.py +3 -0
  299. nvmath/tensor/_internal/cutensor_config_ifc.py +247 -0
  300. nvmath/tensor/_internal/cutensor_utils.py +162 -0
  301. nvmath/tensor/_internal/data.py +43 -0
  302. nvmath/tensor/_internal/einsum_parser.py +444 -0
  303. nvmath/tensor/_internal/typemaps.py +96 -0
  304. nvmath/tensor/contract.py +1861 -0
  305. nvmath_python-0.9.0.dist-info/METADATA +127 -0
  306. nvmath_python-0.9.0.dist-info/RECORD +309 -0
  307. nvmath_python-0.9.0.dist-info/WHEEL +5 -0
  308. nvmath_python-0.9.0.dist-info/licenses/LICENSE +177 -0
  309. nvmath_python-0.9.0.dist-info/top_level.txt +2 -0
@@ -0,0 +1,2690 @@
1
+ # Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ __all__ = [
6
+ "Solver",
7
+ "CholeskySolver",
8
+ "LUSolver",
9
+ "LUPivotSolver",
10
+ "TriangularSolver",
11
+ "QRFactorize",
12
+ "LQFactorize",
13
+ "QRMultiply",
14
+ "LQMultiply",
15
+ "LeastSquaresSolver",
16
+ "compile_solver_execute",
17
+ ]
18
+
19
+ from collections.abc import Sequence
20
+ from functools import cached_property
21
+ from typing import Any, Literal
22
+
23
+ import numpy as np
24
+
25
+ from nvmath._utils import get_nvrtc_version
26
+ from nvmath.bindings import mathdx
27
+ from nvmath.internal.utils import docstring_decorator
28
+
29
+ from .common import SHARED_DEVICE_DOCSTRINGS, check_code_type, parse_code_type, parse_sm
30
+ from .common_backend import get_isa_version, get_lto
31
+ from .common_cuda import (
32
+ Code,
33
+ ComputeCapability,
34
+ )
35
+ from .cusolverdx_backend import (
36
+ _ENABLE_CUSOLVERDX_0_3,
37
+ ALLOWED_ARRANGEMENT,
38
+ ALLOWED_CUSOLVERDX_FUNCTIONS,
39
+ ALLOWED_DATA_TYPE,
40
+ ALLOWED_DIAG,
41
+ ALLOWED_EXECUTION,
42
+ ALLOWED_FILL_MODE,
43
+ ALLOWED_REAL_NP_TYPES,
44
+ ALLOWED_SIDE,
45
+ ALLOWED_TRANSPOSE_MODE,
46
+ CUSOLVERDX_0_3_ALLOWED_FUNCTIONS,
47
+ DIAG_SUPPORTED_FUNCTIONS,
48
+ FILL_MODE_SUPPORTED_FUNCTIONS,
49
+ JOB_SUPPORT_MAP,
50
+ JOB_SUPPORTED_FUNCTIONS,
51
+ SIDE_SUPPORTED_FUNCTIONS,
52
+ generate_code,
53
+ generate_SOLVER,
54
+ get_int_traits,
55
+ get_str_trait,
56
+ validate,
57
+ )
58
+ from .types import Complex, complex64, complex128
59
+
60
+ # ==========================
61
+ # docs
62
+ # ==========================
63
+
64
+ SOLVER_DOCSTRING_SIZE_BASE = """Problem size specified as a sequence of 1 to 3 elements:
65
+ ``(M,)`` (treated as ``(M, M, 1)``), ``(M, N)`` (treated as ``(M, N, 1)``), or ``(M, N, K)``.""".replace("\n", " ")
66
+
67
+ SOLVER_DOCSTRING = {
68
+ "function": f"""Solver function to be executed on execute() method. List of available options:
69
+ {", ".join(f"``'{k}'`` ({v})" for k, v in ALLOWED_CUSOLVERDX_FUNCTIONS.items())}.
70
+ Functions {", ".join(f"``'{k}'``" for k in CUSOLVERDX_0_3_ALLOWED_FUNCTIONS)} require libmathdx 0.3.2 or later.""".replace(
71
+ "\n", " "
72
+ ),
73
+ #
74
+ "size": f"""{SOLVER_DOCSTRING_SIZE_BASE} Please refer to cuSOLVERDx functionalities for detailed meaning.
75
+ """.replace("\n", " "),
76
+ #
77
+ "arrangement": f"""Storage layout for matrices A and B, specified as a sequence of 2 elements ``(arr_A, arr_B)``.
78
+ Each element can be one of: {", ".join(f"``'{v}'``" for v in ALLOWED_ARRANGEMENT)}.
79
+ Defaults to ``("col_major", "col_major")``.""".replace("\n", " "),
80
+ #
81
+ "transpose_mode": f"""Transpose mode of matrix A. Refer to cuSOLVERDx documentation,
82
+ as some functions do not support this option.
83
+ Can be one of: {", ".join(f"``'{v}'``" for v in ALLOWED_TRANSPOSE_MODE)}.
84
+ Defaults to ``'non_transposed'``.""".replace("\n", " "),
85
+ #
86
+ "side": f"""Side of matrix A in a multiplication operation.
87
+ Required and supported only by functions: {", ".join(f"``'{v}'``" for v in SIDE_SUPPORTED_FUNCTIONS)}.
88
+ Can be one of: {", ".join(f"``'{v}'``" for v in ALLOWED_SIDE)}.""".replace("\n", " "),
89
+ #
90
+ "diag": f"""Indicates whether the diagonal elements of matrix A are ones or not.
91
+ Required and supported only by {", ".join(f"``'{v}'``" for v in DIAG_SUPPORTED_FUNCTIONS)} functions.
92
+ Can be one of: {", ".join(f"``'{v}'``" for v in ALLOWED_DIAG)}.""".replace("\n", " "),
93
+ #
94
+ "fill_mode": f"""Indicates which part of matrix A is filled and should be used by function.
95
+ Required and supported only by functions: {", ".join(f"``'{v}'``" for v in FILL_MODE_SUPPORTED_FUNCTIONS)}.
96
+ Can be one of: {", ".join(f"``'{v}'``" for v in ALLOWED_FILL_MODE)}.""".replace("\n", " "),
97
+ #
98
+ "job": f"""Job type for eigenvalue computation.
99
+ Required and supported only by functions: {", ".join(f"``'{v}'``" for v in JOB_SUPPORTED_FUNCTIONS)}.
100
+ For ``'htev'``: {", ".join(f"``'{v}'``" for v in JOB_SUPPORT_MAP["htev"])}.
101
+ For ``'heev'``: {", ".join(f"``'{v}'``" for v in JOB_SUPPORT_MAP["heev"])}.
102
+ Requires libmathdx 0.3.2 or later.""".replace("\n", " "),
103
+ #
104
+ "batches_per_block": """Number of batches to compute in parallel in a single CUDA block.
105
+ Can be a non-zero integer or the string ``'suggested'`` for automatic selection of an optimal value.
106
+ We recommend using 1 for matrix A size larger than or equal to 16 x 16,
107
+ and using ``'suggested'`` for smaller sizes to achieve optimal performance.
108
+ Defaults to 1.""".replace("\n", " "),
109
+ #
110
+ "data_type": f"""The data type of the input matrices,
111
+ can be one of: {", ".join(f"``'{v}'``" for v in ALLOWED_DATA_TYPE)}.
112
+ Defaults to ``'real'``.""".replace("\n", " "),
113
+ #
114
+ "leading_dimensions": """The leading dimensions for input matrices A and B,
115
+ specified as a sequence of 2 elements (``lda``, ``ldb``) or ``None``.
116
+ If not provided, it will be automatically deduced from ``size`` and ``arrangement``.
117
+ """.replace("\n", " "),
118
+ #
119
+ "block_dim": """The block dimension for launching the CUDA kernel,
120
+ specified as a 1 to 3 integer sequence (x, y, z) where missing dimensions are assumed to be 1.
121
+ Can be a sequence of 1 to 3 positive integers, the string ``'suggested'``
122
+ for optimal value selection, or ``None`` for the default value.""".replace("\n", " "),
123
+ #
124
+ "execution": f"""A string specifying the execution method.
125
+ Supported values: {", ".join(f"``'{v}'``" for v in ALLOWED_EXECUTION)}.""".replace("\n", " "),
126
+ #
127
+ "precision": f"""The computation precision specified as a numpy float dtype.
128
+ Currently supports: {", ".join(f"``numpy.{v.__name__}``" for v in ALLOWED_REAL_NP_TYPES)}.""".replace("\n", " "),
129
+ #
130
+ "sm": SHARED_DEVICE_DOCSTRINGS["sm"],
131
+ }
132
+
133
+ # ==========================
134
+ # Solver class
135
+ # ==========================
136
+
137
+
138
+ @docstring_decorator(SOLVER_DOCSTRING, skip_missing=False)
139
+ class Solver:
140
+ """
141
+ A class that encapsulates a partial dense matrix factorization and solve
142
+ device function.
143
+
144
+ **Memory Layout Requirements:**
145
+
146
+ Matrices must be stored in shared memory according
147
+ to their arrangement and leading dimension (ld):
148
+
149
+ * **Column-major arrangement**: Matrix shape ``(batches_per_block, rows, cols)``
150
+ with strides ``(ld * cols, 1, ld)``
151
+ * **Row-major arrangement**: Matrix shape ``(batches_per_block, rows, cols)``
152
+ with strides ``(ld * rows, ld, 1)``
153
+
154
+ Args:
155
+ function (str): {function}
156
+
157
+ size (Sequence[int]): {size}
158
+
159
+ precision (type[np.floating]): {precision}
160
+
161
+ execution (str): {execution}
162
+
163
+ sm (ComputeCapability): {sm}
164
+
165
+ arrangement (Sequence[str], optional): {arrangement}
166
+
167
+ transpose_mode (str, optional): {transpose_mode}
168
+
169
+ side (str, optional): {side}
170
+
171
+ diag (str, optional): {diag}
172
+
173
+ fill_mode (str, optional): {fill_mode}
174
+
175
+ batches_per_block (int | Literal["suggested"], optional): {batches_per_block}
176
+
177
+ data_type (str, optional): {data_type}
178
+
179
+ leading_dimensions (Sequence[int], optional): {leading_dimensions}
180
+
181
+ block_dim (Sequence[int] | Literal["suggested"], optional): {block_dim}
182
+
183
+ job (str, optional): {job}
184
+
185
+ See Also:
186
+ The attributes of this class provide a 1:1
187
+ mapping with the CUDA C++ cuSOLVERDx APIs.
188
+ For further details, please refer to :cusolverdx_doc:`cuSOLVERDx documentation
189
+ <index.html>`.
190
+ """
191
+
192
+ def __init__(
193
+ self,
194
+ function: str,
195
+ size: Sequence[int],
196
+ precision: type[np.floating],
197
+ execution: str,
198
+ *,
199
+ sm=None,
200
+ arrangement: Sequence[str] | None = None,
201
+ transpose_mode: str | None = None,
202
+ side: str | None = None,
203
+ diag: str | None = None,
204
+ fill_mode: str | None = None,
205
+ batches_per_block: int | Literal["suggested"] | None = None,
206
+ data_type: str | None = None,
207
+ leading_dimensions: Sequence[int] | None = None,
208
+ block_dim: Sequence[int] | Literal["suggested"] | None = None,
209
+ job: str | None = None,
210
+ ):
211
+ if get_nvrtc_version() < (12, 6, 85):
212
+ raise RuntimeError("cuSOLVERDx requires CUDA Toolkit 12.6 Update 3 or later.")
213
+
214
+ sm = parse_sm(sm)
215
+
216
+ validate(
217
+ function=function,
218
+ size=size,
219
+ precision=precision,
220
+ execution=execution,
221
+ sm=sm,
222
+ arrangement=arrangement,
223
+ transpose_mode=transpose_mode,
224
+ side=side,
225
+ diag=diag,
226
+ fill_mode=fill_mode,
227
+ batches_per_block=batches_per_block,
228
+ data_type=data_type,
229
+ leading_dimensions=leading_dimensions,
230
+ block_dim=block_dim,
231
+ job=job,
232
+ )
233
+
234
+ if len(size) == 1:
235
+ size = (size[0], size[0], 1)
236
+ elif len(size) == 2:
237
+ size = (size[0], size[1], 1)
238
+ else:
239
+ size = (size[0], size[1], size[2])
240
+
241
+ if block_dim is not None and block_dim != "suggested":
242
+ if len(block_dim) == 1:
243
+ block_dim = (block_dim[0], 1, 1)
244
+ elif len(block_dim) == 2:
245
+ block_dim = (block_dim[0], block_dim[1], 1)
246
+ else:
247
+ block_dim = (block_dim[0], block_dim[1], block_dim[2])
248
+
249
+ self._function = function
250
+ self._size = size
251
+ self._precision = precision
252
+ self._execution = execution
253
+ self._sm = sm
254
+ self._arrangement = tuple(arrangement) if arrangement is not None else None
255
+ self._transpose_mode = transpose_mode
256
+ self._side = side
257
+ self._diag = diag
258
+ self._fill_mode = fill_mode
259
+ self._data_type = data_type
260
+ self._job = job
261
+
262
+ self._batches_per_block = batches_per_block if batches_per_block != "suggested" else None
263
+ self._leading_dimensions = tuple(leading_dimensions) if leading_dimensions is not None else None
264
+ self._block_dim = block_dim if block_dim != "suggested" and block_dim is not None else None
265
+
266
+ if batches_per_block == "suggested" or block_dim == "suggested":
267
+ # Update suggested fields
268
+ traits = _SolverTraits(self)
269
+
270
+ if batches_per_block == "suggested":
271
+ self._batches_per_block = traits.suggested_batches_per_block
272
+
273
+ if block_dim == "suggested":
274
+ self._block_dim = traits.suggested_block_dim
275
+
276
+ # ==========================
277
+ # Operator properties
278
+ # ==========================
279
+
280
+ @property
281
+ def function(self) -> str:
282
+ return self._function
283
+
284
+ @property
285
+ def size(self) -> tuple[int, int, int]:
286
+ return self._size
287
+
288
+ @property
289
+ def m(self) -> int:
290
+ m, _, _ = self.size
291
+ return m
292
+
293
+ @property
294
+ def n(self) -> int:
295
+ _, n, _ = self.size
296
+ return n
297
+
298
+ @property
299
+ def k(self) -> int:
300
+ _, _, k = self.size
301
+ return k
302
+
303
+ @property
304
+ def precision(self) -> type[np.floating]:
305
+ return self._precision
306
+
307
+ @property
308
+ def execution(self) -> str:
309
+ return self._execution
310
+
311
+ @property
312
+ def sm(self) -> ComputeCapability:
313
+ return self._sm
314
+
315
+ @property
316
+ def arrangement(self) -> Sequence[str]:
317
+ if self._arrangement is None:
318
+ # TODO: replace with libmathdx trait when supported
319
+ return ("col_major", "col_major")
320
+ if len(self._arrangement) == 1:
321
+ return (self._arrangement[0], "col_major")
322
+ return self._arrangement
323
+
324
+ @property
325
+ def a_arrangement(self) -> str:
326
+ arr_a, _ = self.arrangement
327
+ return arr_a
328
+
329
+ @property
330
+ def b_arrangement(self) -> str:
331
+ _, arr_b = self.arrangement
332
+ return arr_b
333
+
334
+ @property
335
+ def transpose_mode(self) -> str | None:
336
+ return self._transpose_mode
337
+
338
+ @property
339
+ def side(self) -> str | None:
340
+ return self._side
341
+
342
+ @property
343
+ def diag(self) -> str | None:
344
+ return self._diag
345
+
346
+ @property
347
+ def fill_mode(self) -> str | None:
348
+ return self._fill_mode
349
+
350
+ @property
351
+ def batches_per_block(self) -> int:
352
+ if self._batches_per_block is None:
353
+ # TODO: replace with libmathdx trait when supported
354
+ return 1
355
+ return self._batches_per_block
356
+
357
+ @property
358
+ def data_type(self) -> str:
359
+ if self._data_type is None:
360
+ # TODO: replace with libmathdx trait when supported
361
+ return "real"
362
+ return self._data_type
363
+
364
+ @property
365
+ def leading_dimensions(self) -> tuple:
366
+ if self._leading_dimensions is None:
367
+ # TODO: replace with libmathdx trait when supported
368
+ return self._calculate_default_leading_dimensions()
369
+
370
+ if len(self._leading_dimensions) == 1:
371
+ return (self._leading_dimensions[0], self._calculate_default_leading_dimensions()[1])
372
+ return self._leading_dimensions
373
+
374
+ @property
375
+ def lda(self) -> int:
376
+ lda, _ = self.leading_dimensions
377
+ return lda
378
+
379
+ @property
380
+ def ldb(self) -> int:
381
+ _, ldb = self.leading_dimensions
382
+ return ldb
383
+
384
+ @property
385
+ def block_dim(self) -> tuple:
386
+ if self._block_dim is None:
387
+ return self._traits.block_dim
388
+ return self._block_dim
389
+
390
+ @property
391
+ def block_size(self) -> int:
392
+ return self.block_dim[0] * self.block_dim[1] * self.block_dim[2]
393
+
394
+ @property
395
+ def job(self) -> str | None:
396
+ return self._job
397
+
398
+ # ==========================
399
+ # Trait properties
400
+ # ==========================
401
+
402
+ @cached_property
403
+ def _traits(self):
404
+ return _SolverTraits(self)
405
+
406
+ @property
407
+ def workspace_size(self) -> int:
408
+ if not _ENABLE_CUSOLVERDX_0_3:
409
+ raise RuntimeError("workspace size requires libmathdx 0.3.2")
410
+
411
+ return self._traits.workspace_size
412
+
413
+ @property
414
+ def value_type(self) -> type[np.floating] | Complex:
415
+ if self.data_type == "complex":
416
+ return complex64 if self.precision == np.float32 else complex128
417
+ return self.precision
418
+
419
+ @property
420
+ def info_type(self) -> type[np.signedinteger]:
421
+ return np.int32
422
+
423
+ @property
424
+ def ipiv_type(self) -> type[np.signedinteger]:
425
+ return np.int32
426
+
427
+ @property
428
+ def tau_type(self) -> type[np.floating] | Complex:
429
+ return self.value_type
430
+
431
+ # ==========================
432
+ # execute()
433
+ # ==========================
434
+
435
+ def execute(*args):
436
+ raise RuntimeError("execute is a device function and cannot be called on the host.")
437
+
438
+ # ==========================
439
+ # Private methods
440
+ # ==========================
441
+
442
+ def _generate_SOLVER(self, execution_api):
443
+ return generate_SOLVER(
444
+ function=self._function,
445
+ size=self._size,
446
+ precision=self._precision,
447
+ execution=self._execution,
448
+ sm=self._sm,
449
+ arrangement=self.arrangement, # TODO: remove property when libmathdx supports traits
450
+ transpose_mode=self._transpose_mode,
451
+ side=self._side,
452
+ diag=self._diag,
453
+ fill_mode=self._fill_mode,
454
+ batches_per_block=self.batches_per_block, # TODO: remove property when libmathdx supports traits
455
+ data_type=self.data_type, # TODO: remove property when libmathdx supports traits
456
+ leading_dimensions=self.leading_dimensions, # TODO: remove property when libmathdx supports traits
457
+ block_dim=self._block_dim,
458
+ execution_api=execution_api,
459
+ job=self._job,
460
+ )
461
+
462
+ def _calculate_default_leading_dimensions(self) -> tuple:
463
+ """
464
+ Calculates default leading dimensions based on problem
465
+ size (m, n, k) for different solver functions.
466
+ """
467
+
468
+ def calculate_ld(rows, cols, arr):
469
+ return rows if arr == "col_major" else cols
470
+
471
+ arr_a, arr_b = self.arrangement
472
+ m, n, k = self.size
473
+
474
+ match self.function:
475
+ case (
476
+ "potrf"
477
+ | "potrs"
478
+ | "posv"
479
+ | "getrf_no_pivot"
480
+ | "getrs_no_pivot"
481
+ | "gesv_no_pivot"
482
+ | "getrf_partial_pivot"
483
+ | "getrs_partial_pivot"
484
+ | "gesv_partial_pivot"
485
+ ):
486
+ return (calculate_ld(m, n, arr_a), calculate_ld(n, k, arr_b))
487
+ case "trsm":
488
+ assert self.side is not None
489
+ return (m if self.side == "left" else n, calculate_ld(m, n, arr_b))
490
+ case "geqrf" | "gelqf" | "htev" | "unglq" | "ungqr" | "heev":
491
+ return (calculate_ld(m, n, arr_a), 0)
492
+ case "unmqr":
493
+ return (calculate_ld(m if self.side == "left" else n, k, arr_a), calculate_ld(m, n, arr_b))
494
+ case "unmlq":
495
+ return (calculate_ld(k, m if self.side == "left" else n, arr_a), calculate_ld(m, n, arr_b))
496
+ case "gels":
497
+ return (calculate_ld(m, n, arr_a), calculate_ld(max(m, n), k, arr_b)) # mirrors cusolverdx behaviour
498
+ case "gtsv_no_pivot":
499
+ return (0, calculate_ld(n, k, arr_b))
500
+ case _:
501
+ raise NotImplementedError(f"Function {self.function} is not supported for leading dimension calculation")
502
+
503
+
504
+ # ==========================
505
+ # Internal traits class
506
+ # ==========================
507
+
508
+
509
+ class _SolverTraits:
510
+ def __init__(self, solver: Solver):
511
+ h = solver._generate_SOLVER("compiled_leading_dim")
512
+
513
+ try:
514
+ self.suggested_batches_per_block = int(
515
+ mathdx.cusolverdx_get_trait_int64(h.descriptor, mathdx.CusolverdxTraitType.SUGGESTED_BATCHES_PER_BLOCK)
516
+ )
517
+
518
+ self.suggested_block_dim = tuple(get_int_traits(h.descriptor, mathdx.CusolverdxTraitType.SUGGESTED_BLOCK_DIM, 3))
519
+ self.block_dim = tuple(get_int_traits(h.descriptor, mathdx.CusolverdxTraitType.BLOCK_DIM, 3))
520
+
521
+ if _ENABLE_CUSOLVERDX_0_3:
522
+ self.workspace_size = int(
523
+ mathdx.cusolverdx_get_trait_int64(h.descriptor, mathdx.CusolverdxTraitType.WORKSPACE_SIZE)
524
+ )
525
+
526
+ except mathdx.LibMathDxError as e:
527
+ raise RuntimeError(
528
+ "Failed to compile the solver. This may indicate incompatible "
529
+ f"parameter combination. Please refer to the cuSOLVERDx documentation. Details: {e}"
530
+ ) from e
531
+
532
+
533
+ # ==========================
534
+ # Compile function
535
+ # ==========================
536
+
537
+
538
+ def compile_solver_execute(
539
+ solver: Solver,
540
+ code_type: Any,
541
+ execution_api: str,
542
+ ) -> tuple[Code, str]:
543
+ code_type = parse_code_type(code_type)
544
+ check_code_type(code_type, "cuSOLVERdx")
545
+
546
+ h = solver._generate_SOLVER(execution_api).descriptor
547
+
548
+ try:
549
+ code = generate_code(h, code_type.cc)
550
+ except mathdx.LibMathDxError as e:
551
+ raise RuntimeError(
552
+ "Failed to compile the solver. This may indicate incompatible "
553
+ f"parameter combination. Please refer to the cuSOLVERDx documentation. Details: {e}"
554
+ ) from e
555
+
556
+ lto = get_lto(code.descriptor)
557
+ isa_version = get_isa_version(code.descriptor)
558
+ symbol = get_str_trait(h, mathdx.CusolverdxTraitType.SYMBOL_NAME)
559
+
560
+ return Code(code_type, isa_version, lto), symbol
561
+
562
+
563
+ # ==========================
564
+ # Pythonic Adapters
565
+ # ==========================
566
+
567
+
568
+ def _calculate_strides(shape: tuple[int, int], ld: int, arrangement: str) -> tuple[int, int, int]:
569
+ return (ld * shape[1], 1, ld) if arrangement == "col_major" else (ld * shape[0], ld, 1)
570
+
571
+
572
+ class _SolverProperties:
573
+ def __init__(self, source: Solver):
574
+ self._properties_source = source
575
+
576
+ @property
577
+ def size(self) -> tuple[int, int, int]:
578
+ return self._properties_source.size
579
+
580
+ @property
581
+ def m(self) -> int:
582
+ return self._properties_source.m
583
+
584
+ @property
585
+ def n(self) -> int:
586
+ return self._properties_source.n
587
+
588
+ @property
589
+ def k(self) -> int:
590
+ return self._properties_source.k
591
+
592
+ @property
593
+ def precision(self) -> type[np.floating]:
594
+ return self._properties_source.precision
595
+
596
+ @property
597
+ def execution(self) -> str:
598
+ return self._properties_source.execution
599
+
600
+ @property
601
+ def sm(self) -> ComputeCapability:
602
+ return self._properties_source.sm
603
+
604
+ @property
605
+ def arrangement(self) -> Sequence[str]:
606
+ return self._properties_source.arrangement
607
+
608
+ @property
609
+ def a_arrangement(self) -> str:
610
+ return self._properties_source.a_arrangement
611
+
612
+ @property
613
+ def b_arrangement(self) -> str:
614
+ return self._properties_source.b_arrangement
615
+
616
+ @property
617
+ def batches_per_block(self) -> int:
618
+ return self._properties_source.batches_per_block
619
+
620
+ @property
621
+ def data_type(self) -> str:
622
+ return self._properties_source.data_type
623
+
624
+ @property
625
+ def leading_dimensions(self) -> tuple:
626
+ return self._properties_source.leading_dimensions
627
+
628
+ @property
629
+ def lda(self) -> int:
630
+ return self._properties_source.lda
631
+
632
+ @property
633
+ def ldb(self) -> int:
634
+ return self._properties_source.ldb
635
+
636
+ @property
637
+ def block_dim(self) -> tuple:
638
+ return self._properties_source.block_dim
639
+
640
+ @property
641
+ def block_size(self) -> int:
642
+ return self._properties_source.block_size
643
+
644
+ @property
645
+ def value_type(self) -> type[np.floating] | Complex:
646
+ return self._properties_source.value_type
647
+
648
+
649
+ class _LinearSolverProperties(_SolverProperties):
650
+ def __init__(self, source: Solver):
651
+ super().__init__(source)
652
+
653
+ @property
654
+ def info_type(self) -> type[np.signedinteger]:
655
+ return self._properties_source.info_type
656
+
657
+ @property
658
+ def info_shape(self) -> tuple[int]:
659
+ return (self.batches_per_block,)
660
+
661
+ @property
662
+ def info_strides(self) -> tuple[int]:
663
+ return (1,)
664
+
665
+ @property
666
+ def a_shape(self) -> tuple[int, int, int]:
667
+ return (self.batches_per_block, self.m, self.n)
668
+
669
+ @property
670
+ def b_shape(self) -> tuple[int, int, int]:
671
+ return (self.batches_per_block, self.n, self.k)
672
+
673
+ def a_strides(self, *, lda: int | None = None) -> tuple[int, int, int]:
674
+ lda = self.lda if lda is None else lda
675
+ return _calculate_strides(self.a_shape[1:], lda, self.a_arrangement)
676
+
677
+ def b_strides(self, *, ldb: int | None = None) -> tuple[int, int, int]:
678
+ ldb = self.ldb if ldb is None else ldb
679
+ return _calculate_strides(self.b_shape[1:], ldb, self.b_arrangement)
680
+
681
+ def a_size(self, *, lda: int | None = None) -> int:
682
+ return self.a_strides(lda=lda)[0] * self.a_shape[0]
683
+
684
+ def b_size(self, *, ldb: int | None = None) -> int:
685
+ return self.b_strides(ldb=ldb)[0] * self.b_shape[0]
686
+
687
+
688
+ # ==========================
689
+ # Cholesky Solver
690
+ # ==========================
691
+
692
+ ADAPTERS_API_LD_DOCSTRING = """
693
+ Note: When provided in the constructor, leading dimensions are set at compile-time.
694
+ To use runtime leading dimensions (avoiding recompilation for different leading dimensions),
695
+ provide the leading dimension parameters directly to the device methods instead.
696
+ """.replace("\n", " ")
697
+
698
+ CHOLESKY_SOLVER_DOCSTRING = SOLVER_DOCSTRING.copy()
699
+ del CHOLESKY_SOLVER_DOCSTRING["function"]
700
+ del CHOLESKY_SOLVER_DOCSTRING["transpose_mode"]
701
+ del CHOLESKY_SOLVER_DOCSTRING["side"]
702
+ del CHOLESKY_SOLVER_DOCSTRING["diag"]
703
+ del CHOLESKY_SOLVER_DOCSTRING["job"]
704
+
705
+ CHOLESKY_SOLVER_DOCSTRING["size"] = f"""{SOLVER_DOCSTRING_SIZE_BASE}
706
+ ``M`` represents the dimension of the square matrix A (``M`` x ``M``) used in factorization, ``N`` must be equal to ``M``.
707
+ ``K`` represents the number of columns in the right-hand side matrix B
708
+ (dimensions ``M`` x ``K``) for the solve operation.""".replace("\n", " ")
709
+ CHOLESKY_SOLVER_DOCSTRING["leading_dimensions"] = SOLVER_DOCSTRING["leading_dimensions"] + ADAPTERS_API_LD_DOCSTRING
710
+ CHOLESKY_SOLVER_DOCSTRING["fill_mode"] = f"""Indicates which part of matrix A is filled and should be used by function.
711
+ Can be one of: {", ".join(f"``'{v}'``" for v in ALLOWED_FILL_MODE)}.""".replace("\n", " ")
712
+
713
+
714
+ @docstring_decorator(CHOLESKY_SOLVER_DOCSTRING, skip_missing=False)
715
+ class CholeskySolver(_LinearSolverProperties):
716
+ """
717
+ A class that encapsulates Cholesky factorization and solve device functions
718
+ for symmetric positive definite matrices.
719
+
720
+ **Available operations:**
721
+
722
+ * factorize: Computes the Cholesky factorization
723
+ A = L @ L^H (lower) or A = U^H @ U (upper),
724
+ where L is a lower triangular matrix and U is an upper triangular matrix.
725
+ The choice depends on the fill_mode parameter.
726
+ * solve: Solves the system Ax = B using a previously computed Cholesky factorization
727
+
728
+ **Memory Layout Requirements:**
729
+
730
+ Matrices must be stored in shared memory according
731
+ to their arrangement and leading dimension (ld):
732
+
733
+ **For matrix A (M x N):**
734
+
735
+ * **Column-major arrangement**: Matrix shape ``(batches_per_block, M, N)``
736
+ with strides ``(lda * N, 1, lda)``
737
+ * **Row-major arrangement**: Matrix shape ``(batches_per_block, M, N)``
738
+ with strides ``(lda * M, lda, 1)``
739
+
740
+ **For matrix B (N x K):**
741
+
742
+ * **Column-major arrangement**: Matrix shape ``(batches_per_block, N, K)``
743
+ with strides ``(ldb * K, 1, ldb)``
744
+ * **Row-major arrangement**: Matrix shape ``(batches_per_block, N, K)``
745
+ with strides ``(ldb * N, ldb, 1)``
746
+
747
+ Args:
748
+ size (Sequence[int]): {size}
749
+
750
+ precision (type[np.floating]): {precision}
751
+
752
+ execution (str): {execution}
753
+
754
+ sm (ComputeCapability): {sm}
755
+
756
+ fill_mode (str): {fill_mode}
757
+
758
+ arrangement (Sequence[str], optional): {arrangement}
759
+
760
+ batches_per_block (int | Literal["suggested"], optional): {batches_per_block}
761
+
762
+ data_type (str, optional): {data_type}
763
+
764
+ leading_dimensions (Sequence[int], optional): {leading_dimensions}
765
+
766
+ block_dim (Sequence[int] | Literal["suggested"], optional): {block_dim}
767
+
768
+ See Also:
769
+ For further details, please refer to the cuSOLVERDx documentation:
770
+
771
+ * :cusolverdx_doc:`factorize <get_started/functions/potrf.html>`
772
+ * :cusolverdx_doc:`solve <get_started/functions/potrs.html>`
773
+ """
774
+
775
+ # ==========================
776
+ # Constructor
777
+ # ==========================
778
+
779
+ def __init__(
780
+ self,
781
+ size: Sequence[int],
782
+ precision: type[np.floating],
783
+ execution: str,
784
+ fill_mode: str,
785
+ *,
786
+ sm=None,
787
+ arrangement: Sequence[str] | None = None,
788
+ batches_per_block: int | Literal["suggested"] | None = None,
789
+ data_type: str | None = None,
790
+ leading_dimensions: Sequence[int] | None = None,
791
+ block_dim: Sequence[int] | Literal["suggested"] | None = None,
792
+ ):
793
+ if fill_mode is None:
794
+ raise ValueError("fill_mode must be provided for CholeskySolver")
795
+
796
+ def construct(function):
797
+ return Solver(
798
+ function=function,
799
+ size=size,
800
+ precision=precision,
801
+ execution=execution,
802
+ sm=sm,
803
+ arrangement=arrangement,
804
+ fill_mode=fill_mode,
805
+ batches_per_block=batches_per_block,
806
+ data_type=data_type,
807
+ leading_dimensions=leading_dimensions,
808
+ block_dim=block_dim,
809
+ )
810
+
811
+ self._factorize = construct("potrf")
812
+ self._solve = construct("potrs")
813
+
814
+ super().__init__(self._solve)
815
+
816
+ if self.m != self.n:
817
+ raise ValueError("A must be a square matrix for CholeskySolver.")
818
+
819
+ # ==========================
820
+ # Property methods
821
+ # ==========================
822
+
823
+ @property
824
+ def fill_mode(self) -> str:
825
+ return self._solve.fill_mode
826
+
827
+ # ==========================
828
+ # Device function methods
829
+ # ==========================
830
+
831
+ def factorize(self, a, info, lda=None) -> None:
832
+ """
833
+ Computes the Cholesky factorization of a symmetric positive definite matrix A.
834
+
835
+ This device function computes A = L @ L^H (if fill_mode = ``'lower'``)
836
+ or A = U^H @ U (if fill_mode = ``'upper'``). Uses cuSOLVERDx ``'potrf'``.
837
+
838
+ If ``lda`` is provided, uses runtime version with the specified
839
+ leading dimension. If ``lda`` is not provided (``None``),
840
+ uses compile-time version with
841
+ default or constructor-provided leading dimensions.
842
+
843
+ For more details, see: :cusolverdx_doc:`get_started/functions/potrf.html`
844
+
845
+ Args:
846
+ a: Pointer to an array in shared memory, storing
847
+ the matrix according to the specified
848
+ arrangement and leading dimension (see :meth:`__init__`).
849
+ On entry, contains the symmetric positive definite matrix.
850
+ On exit, contains the triangular factor L (lower) or U (upper).
851
+ info: Pointer to a 1D array of ``int32``. On exit, ``info[batch_id] = 0``
852
+ indicates success for that batch, ``info[batch_id] != 0``
853
+ indicates the matrix is not positive definite.
854
+ lda: Optional runtime leading dimension of matrix A.
855
+ If not specified, the compile-time ``lda`` is used.
856
+ """
857
+ raise RuntimeError("factorize is a device function and cannot be called on the host.")
858
+
859
+ def solve(self, a, b, lda=None, ldb=None) -> None:
860
+ """
861
+ Solves a system of linear equations Ax = B using the Cholesky factorization.
862
+
863
+ This device function uses the previously computed
864
+ factorization A = L @ L^H (lower) or A = U^H @ U (upper)
865
+ to solve the system. Uses cuSOLVERDx ``'potrs'``.
866
+
867
+ If ``lda`` and ``ldb`` are provided, uses
868
+ runtime version with the specified leading dimensions.
869
+ If not provided (``None``), uses
870
+ compile-time version with default or constructor-provided leading dimensions.
871
+
872
+ For more details, see: :cusolverdx_doc:`get_started/functions/potrs.html`
873
+
874
+ Args:
875
+ a: Pointer to an array in shared memory, storing
876
+ the triangular factor L (lower) or U (upper)
877
+ from the Cholesky factorization, according
878
+ to the specified arrangement and leading dimension (see :meth:`__init__`).
879
+ b: Pointer to an array in shared memory, storing the
880
+ matrix according to the specified
881
+ arrangement and leading dimension (see :meth:`__init__`).
882
+ The matrix is overwritten in place with the solution matrix x.
883
+ lda: Optional runtime leading dimension of matrix A.
884
+ The ``lda`` and ``ldb`` must be specified together.
885
+ If not specified, the compile-time ``lda`` is used.
886
+ ldb: Optional runtime leading dimension of matrix B.
887
+ The ``lda`` and ``ldb`` must be specified together.
888
+ If not specified, the compile-time ``ldb`` is used.
889
+ """
890
+ raise RuntimeError("solve is a device function and cannot be called on the host.")
891
+
892
+
893
+ # ==========================
894
+ # LU Solver
895
+ # ==========================
896
+
897
+ LU_SOLVER_DOCSTRING = SOLVER_DOCSTRING.copy()
898
+ del LU_SOLVER_DOCSTRING["function"]
899
+ del LU_SOLVER_DOCSTRING["fill_mode"]
900
+ del LU_SOLVER_DOCSTRING["side"]
901
+ del LU_SOLVER_DOCSTRING["diag"]
902
+ del LU_SOLVER_DOCSTRING["job"]
903
+
904
+ LU_SOLVER_DOCSTRING["size"] = f"""{SOLVER_DOCSTRING_SIZE_BASE}
905
+ ``M`` and ``N`` represent the dimensions of the matrix A used in factorization.
906
+ ``K`` represents the number of columns in the right-hand side matrix B (dimensions ``N`` x ``K``) for the ``solve`` operation.
907
+ To use :meth:`solve`, ``N`` must be equal to ``M``,
908
+ otherwise an exception will be thrown when ``solver.solve()`` is used.""".replace("\n", " ")
909
+ LU_SOLVER_DOCSTRING["leading_dimensions"] = SOLVER_DOCSTRING["leading_dimensions"] + ADAPTERS_API_LD_DOCSTRING
910
+ LU_SOLVER_DOCSTRING["transpose_mode"] = f"""Transpose mode of matrix A for the solve operation.
911
+ Can be one of: {", ".join(f"``'{v}'``" for v in ALLOWED_TRANSPOSE_MODE)}.
912
+ Defaults to ``'non_transposed'``.""".replace("\n", " ")
913
+
914
+
915
+ @docstring_decorator(LU_SOLVER_DOCSTRING, skip_missing=False)
916
+ class LUSolver(_LinearSolverProperties):
917
+ """
918
+ A class that encapsulates cuSOLVERDx LU factorization without pivoting
919
+ and linear solver for general matrices.
920
+
921
+ **Available operations:**
922
+
923
+ * factorize: Computes the LU factorization A = L @ U,
924
+ where L is a unit lower triangular matrix and U is an upper triangular matrix.
925
+ * solve: Solves the system Ax = B using a previously computed LU factorization.
926
+
927
+ **Memory Layout Requirements:**
928
+
929
+ Matrices must be stored in shared memory according
930
+ to their arrangement and leading dimension (ld):
931
+
932
+ **For matrix A (M x N):**
933
+
934
+ * **Column-major arrangement**: Matrix shape ``(batches_per_block, M, N)``
935
+ with strides ``(lda * N, 1, lda)``
936
+ * **Row-major arrangement**: Matrix shape ``(batches_per_block, M, N)``
937
+ with strides ``(lda * M, lda, 1)``
938
+
939
+ **For matrix B (N x K):**
940
+
941
+ * **Column-major arrangement**: Matrix shape ``(batches_per_block, N, K)``
942
+ with strides ``(ldb * K, 1, ldb)``
943
+ * **Row-major arrangement**: Matrix shape ``(batches_per_block, N, K)``
944
+ with strides ``(ldb * N, ldb, 1)``
945
+
946
+ .. note::
947
+ If a nonsingular matrix A is diagonal dominant, then it is safe to factorize
948
+ without pivoting. If a matrix is not diagonal dominant, then pivoting is usually
949
+ required to ensure numerical stability (see :class:`LUPivotSolver`).
950
+
951
+ Args:
952
+ size (Sequence[int]): {size}
953
+
954
+ precision (type[np.floating]): {precision}
955
+
956
+ execution (str): {execution}
957
+
958
+ sm (ComputeCapability): {sm}
959
+
960
+ transpose_mode (str, optional): {transpose_mode}
961
+
962
+ arrangement (Sequence[str], optional): {arrangement}
963
+
964
+ batches_per_block (int | Literal["suggested"], optional): {batches_per_block}
965
+
966
+ data_type (str, optional): {data_type}
967
+
968
+ leading_dimensions (Sequence[int], optional): {leading_dimensions}
969
+
970
+ block_dim (Sequence[int] | Literal["suggested"], optional): {block_dim}
971
+
972
+ See Also:
973
+ For further details, please refer to the cuSOLVERDx documentation:
974
+
975
+ * :cusolverdx_doc:`factorize <get_started/functions/getrf.html>`
976
+ * :cusolverdx_doc:`solve <get_started/functions/getrs.html>`
977
+
978
+ """
979
+
980
+ # ==========================
981
+ # Constructor
982
+ # ==========================
983
+
984
+ def __init__(
985
+ self,
986
+ size: Sequence[int],
987
+ precision: type[np.floating],
988
+ execution: str,
989
+ *,
990
+ sm=None,
991
+ transpose_mode: str = "non_transposed",
992
+ arrangement: Sequence[str] | None = None,
993
+ batches_per_block: int | Literal["suggested"] | None = None,
994
+ data_type: str | None = None,
995
+ leading_dimensions: Sequence[int] | None = None,
996
+ block_dim: Sequence[int] | Literal["suggested"] | None = None,
997
+ ):
998
+ def construct(function, transpose_mode):
999
+ return Solver(
1000
+ function=function,
1001
+ size=size,
1002
+ precision=precision,
1003
+ execution=execution,
1004
+ sm=sm,
1005
+ arrangement=arrangement,
1006
+ batches_per_block=batches_per_block,
1007
+ data_type=data_type,
1008
+ leading_dimensions=leading_dimensions,
1009
+ block_dim=block_dim,
1010
+ transpose_mode=transpose_mode,
1011
+ )
1012
+
1013
+ if transpose_mode is None:
1014
+ raise ValueError("transpose_mode must be provided for LUSolver")
1015
+
1016
+ self._factorize = construct("getrf_no_pivot", "non_transposed")
1017
+ self._transpose_mode = transpose_mode
1018
+ self._solve = construct("getrs_no_pivot", transpose_mode) if self._factorize.m == self._factorize.n else None
1019
+
1020
+ super().__init__(self._factorize)
1021
+
1022
+ # ==========================
1023
+ # Property methods
1024
+ # ==========================
1025
+
1026
+ @property
1027
+ def transpose_mode(self) -> str:
1028
+ return self._transpose_mode
1029
+
1030
+ # ==========================
1031
+ # Device function methods
1032
+ # ==========================
1033
+
1034
+ def factorize(self, a, info, lda=None) -> None:
1035
+ """
1036
+ Computes the LU factorization of a general matrix A without pivoting.
1037
+
1038
+ This device function computes A = L @ U, where L is a unit lower triangular matrix
1039
+ and U is an upper triangular matrix. This variant is suitable for diagonally
1040
+ dominant matrices or when pivoting is not required.
1041
+ Uses cuSOLVERDx ``'getrf_no_pivot'``.
1042
+
1043
+ If ``lda`` is provided, uses runtime version with the specified leading dimension.
1044
+ If ``lda`` is not provided (``None``), uses compile-time version with default
1045
+ or constructor-provided leading dimensions.
1046
+
1047
+ .. note::
1048
+ The ``transpose_mode`` parameter does not affect factorization.
1049
+ This operation always treats the input matrix as-is (non-transposed).
1050
+
1051
+ For more details, see: :cusolverdx_doc:`get_started/functions/getrf.html`
1052
+
1053
+ Args:
1054
+ a: Pointer to an array in shared memory, storing the batched matrix according
1055
+ to the specified arrangement and leading dimension (see :meth:`__init__`).
1056
+ The matrix is overwritten in place.
1057
+ On exit, contains the factors L and U from the factorization A = L @ U.
1058
+ The unit diagonal elements of L are not stored.
1059
+ info: Pointer to a 1D array of ``int32``.
1060
+ On exit, ``info[batch_id] = 0`` indicates success for that batch,
1061
+ ``info[batch_id] = i > 0`` indicates ``U(i,i)`` is exactly zero,
1062
+ meaning the factorization has been completed but the factor U
1063
+ is singular and division by zero will occur if it is used to solve
1064
+ a system of equations.
1065
+ lda: Optional runtime leading dimension of matrix A.
1066
+ If not specified, the compile-time ``lda`` is used.
1067
+ """
1068
+ raise RuntimeError("factorize is a device function and cannot be called on the host.")
1069
+
1070
+ def solve(self, a, b, lda=None, ldb=None) -> None:
1071
+ """
1072
+ Solves a system of linear equations Ax = B
1073
+ using the LU factorization without pivoting.
1074
+ The ``a`` operand must be a square matrix (``M == N``),
1075
+ otherwise this function will throw an exception.
1076
+
1077
+ This device function uses the previously computed factorization A = L @ U
1078
+ to solve the system. Uses cuSOLVERDx ``'getrs_no_pivot'``.
1079
+
1080
+ If ``lda`` and ``ldb`` are provided, uses runtime version with
1081
+ the specified leading dimensions. If not provided (``None``),
1082
+ uses compile-time version with default or constructor-provided
1083
+ leading dimensions.
1084
+
1085
+ .. note::
1086
+ The ``transpose_mode`` parameter (set in constructor) determines which
1087
+ system is solved: A*x=B (``'non_transposed'``), A^T*x=B (``'transposed'``),
1088
+ or A^H*x=B (``'conj_transposed'`` for complex matrices).
1089
+
1090
+ For more details, see: :cusolverdx_doc:`get_started/functions/getrs.html`
1091
+
1092
+ Args:
1093
+ a: Pointer to an array in shared memory, storing the batched factors L and U
1094
+ from the LU factorization, according to the specified arrangement
1095
+ and leading dimension (see :meth:`__init__`).
1096
+ The unit diagonal elements of L are not stored.
1097
+ See the :meth:`factorize` documentation for details.
1098
+ b: Pointer to an array in shared memory, storing the batched matrix according
1099
+ to the specified arrangement and leading dimension (see :meth:`__init__`).
1100
+ The matrix is overwritten in place with the solution matrix x.
1101
+ lda: Optional runtime leading dimension of matrix A.
1102
+ The ``lda`` and ``ldb`` must be specified together.
1103
+ If not specified, the compile-time ``lda`` is used.
1104
+ ldb: Optional runtime leading dimension of matrix B.
1105
+ The ``lda`` and ``ldb`` must be specified together.
1106
+ If not specified, the compile-time ``ldb`` is used.
1107
+ """
1108
+ raise RuntimeError("solve is a device function and cannot be called on the host.")
1109
+
1110
+
1111
+ # ==========================
1112
+ # Triangular Solver
1113
+ # ==========================
1114
+
1115
+ TRIANGULAR_SOLVER_DOCSTRING = SOLVER_DOCSTRING.copy()
1116
+ del TRIANGULAR_SOLVER_DOCSTRING["function"]
1117
+ del TRIANGULAR_SOLVER_DOCSTRING["job"]
1118
+
1119
+ TRIANGULAR_SOLVER_DOCSTRING["size"] = f"""{SOLVER_DOCSTRING_SIZE_BASE}
1120
+ ``M`` and ``N`` represent the dimensions of matrices A and B.
1121
+ When ``side='left'``, A is ``M`` x ``M``, otherwise when ``side='right'``, A is ``N`` x ``N``.
1122
+ B is always ``M`` x ``N``.""".replace("\n", " ")
1123
+
1124
+ TRIANGULAR_SOLVER_DOCSTRING["side"] = f"""Side of matrix A in the triangular solve operation (required for TRSM).
1125
+ Can be one of: {", ".join(f"``'{v}'``" for v in ALLOWED_SIDE)}.
1126
+ If ``side='left'``, solves op(A) * X = B where A is ``M`` x ``M``.
1127
+ If ``side='right'``, solves X * op(A) = B where A is ``N`` x ``N``.""".replace("\n", " ")
1128
+
1129
+ TRIANGULAR_SOLVER_DOCSTRING["fill_mode"] = f"""Indicates which part of triangular matrix A is filled and should be used.
1130
+ Can be one of: {", ".join(f"``'{v}'``" for v in ALLOWED_FILL_MODE)}.
1131
+ For lower fill mode, only the diagonal and lower triangular part of A is processed, the upper part is untouched.
1132
+ For upper fill mode, only the diagonal and upper triangular part of A is processed, the lower part is untouched.""".replace(
1133
+ "\n", " "
1134
+ )
1135
+
1136
+ TRIANGULAR_SOLVER_DOCSTRING["diag"] = f"""Indicates whether the diagonal elements of matrix A are unity or not.
1137
+ Can be one of: {", ".join(f"``'{v}'``" for v in ALLOWED_DIAG)}.
1138
+ For unit diagonal mode, the diagonal elements of A are unity and are not accessed.
1139
+ For non-unit diagonal mode, the diagonal elements of A are used in the computation.""".replace("\n", " ")
1140
+
1141
+ TRIANGULAR_SOLVER_DOCSTRING["transpose_mode"] = f"""Transpose mode for operation op(A) applied to matrix A.
1142
+ Can be one of: {", ".join(f"``'{v}'``" for v in ALLOWED_TRANSPOSE_MODE)}.
1143
+ Defaults to ``'non_transposed'``.""".replace("\n", " ")
1144
+
1145
+ TRIANGULAR_SOLVER_DOCSTRING["leading_dimensions"] = SOLVER_DOCSTRING["leading_dimensions"] + ADAPTERS_API_LD_DOCSTRING
1146
+
1147
+
1148
+ @docstring_decorator(TRIANGULAR_SOLVER_DOCSTRING, skip_missing=False)
1149
+ class TriangularSolver(_SolverProperties):
1150
+ """
1151
+ A class that encapsulates triangular matrix-matrix solve device function (``'trsm'``).
1152
+
1153
+ TRSM (TRiangular Solve for Matrix) solves a triangular linear system
1154
+ with multiple right-hand sides:
1155
+
1156
+ * op(A) * X = B (if ``side='left'``)
1157
+ * X * op(A) = B (if ``side='right'``)
1158
+
1159
+ where:
1160
+
1161
+ * A is the input batched triangular matrix stored in lower or upper mode
1162
+ * B is the batched right-hand side matrix, overwritten by the result X on exit
1163
+ * Operation op(A) indicates if matrix A is ``'non_transposed'``,
1164
+ ``'transposed'`` (for real data type),
1165
+ or ``'conj_transposed'`` (for complex data type)
1166
+
1167
+ **Memory Layout Requirements:**
1168
+
1169
+ Matrices must be stored in shared memory according
1170
+ to their arrangement and leading dimension (ld):
1171
+
1172
+ **For matrix A (M x M) with ``side='left'``:**
1173
+
1174
+ * **Column-major arrangement**: Matrix shape ``(batches_per_block, M, M)``
1175
+ with strides ``(lda * M, 1, lda)``
1176
+ * **Row-major arrangement**: Matrix shape ``(batches_per_block, M, M)``
1177
+ with strides ``(lda * M, lda, 1)``
1178
+
1179
+ **For matrix A (N x N) with ``side='right'``:**
1180
+
1181
+ * **Column-major arrangement**: Matrix shape ``(batches_per_block, N, N)``
1182
+ with strides ``(lda * N, 1, lda)``
1183
+ * **Row-major arrangement**: Matrix shape ``(batches_per_block, N, N)``
1184
+ with strides ``(lda * N, lda, 1)``
1185
+
1186
+ **For matrix B (M x N):**
1187
+
1188
+ * **Column-major arrangement**: Matrix shape ``(batches_per_block, M, N)``
1189
+ with strides ``(ldb * N, 1, ldb)``
1190
+ * **Row-major arrangement**: Matrix shape ``(batches_per_block, M, N)``
1191
+ with strides ``(ldb * M, ldb, 1)``
1192
+
1193
+ .. note::
1194
+ The TRSM function is temporarily exposed in cuSolverDx library
1195
+ and will be moved to cuBLASDx library in a future release.
1196
+
1197
+ Args:
1198
+ size (Sequence[int]): {size}
1199
+
1200
+ precision (type[np.floating]): {precision}
1201
+
1202
+ execution (str): {execution}
1203
+
1204
+ sm (ComputeCapability): {sm}
1205
+
1206
+ side (str): {side}
1207
+
1208
+ fill_mode (str): {fill_mode}
1209
+
1210
+ diag (str): {diag}
1211
+
1212
+ transpose_mode (str, optional): {transpose_mode}
1213
+
1214
+ arrangement (Sequence[str], optional): {arrangement}
1215
+
1216
+ batches_per_block (int | Literal["suggested"], optional): {batches_per_block}
1217
+
1218
+ data_type (str, optional): {data_type}
1219
+
1220
+ leading_dimensions (Sequence[int], optional): {leading_dimensions}
1221
+
1222
+ block_dim (Sequence[int] | Literal["suggested"], optional): {block_dim}
1223
+
1224
+ See Also:
1225
+ For further details, please refer to the cuSOLVERDx documentation:
1226
+
1227
+ * :cusolverdx_doc:`trsm <get_started/functions/trsm.html>`
1228
+ """
1229
+
1230
+ # ==========================
1231
+ # Constructor
1232
+ # ==========================
1233
+
1234
+ def __init__(
1235
+ self,
1236
+ size: Sequence[int],
1237
+ precision: type[np.floating],
1238
+ execution: str,
1239
+ side: str,
1240
+ fill_mode: str,
1241
+ diag: str,
1242
+ transpose_mode: str = "non_transposed",
1243
+ *,
1244
+ sm=None,
1245
+ arrangement: Sequence[str] | None = None,
1246
+ batches_per_block: int | Literal["suggested"] | None = None,
1247
+ data_type: str | None = None,
1248
+ leading_dimensions: Sequence[int] | None = None,
1249
+ block_dim: Sequence[int] | Literal["suggested"] | None = None,
1250
+ ):
1251
+ self._solve: Solver = Solver(
1252
+ function="trsm",
1253
+ size=size,
1254
+ precision=precision,
1255
+ execution=execution,
1256
+ sm=sm,
1257
+ side=side,
1258
+ fill_mode=fill_mode,
1259
+ diag=diag,
1260
+ transpose_mode=transpose_mode,
1261
+ arrangement=arrangement,
1262
+ batches_per_block=batches_per_block,
1263
+ data_type=data_type,
1264
+ leading_dimensions=leading_dimensions,
1265
+ block_dim=block_dim,
1266
+ )
1267
+
1268
+ super().__init__(self._solve)
1269
+
1270
+ if not isinstance(self._solve.side, str):
1271
+ raise ValueError("side must be provided for TriangularSolver")
1272
+
1273
+ if not isinstance(self._solve.fill_mode, str):
1274
+ raise ValueError("fill_mode must be provided for TriangularSolver")
1275
+
1276
+ if not isinstance(self._solve.diag, str):
1277
+ raise ValueError("diag must be provided for TriangularSolver")
1278
+
1279
+ # ==========================
1280
+ # Property methods
1281
+ # ==========================
1282
+
1283
+ @property
1284
+ def side(self) -> str:
1285
+ assert self._solve.side is not None
1286
+ return self._solve.side
1287
+
1288
+ @property
1289
+ def fill_mode(self) -> str:
1290
+ assert self._solve.fill_mode is not None
1291
+ return self._solve.fill_mode
1292
+
1293
+ @property
1294
+ def diag(self) -> str:
1295
+ assert self._solve.diag is not None
1296
+ return self._solve.diag
1297
+
1298
+ @property
1299
+ def transpose_mode(self) -> str:
1300
+ assert self._solve.transpose_mode is not None
1301
+ return self._solve.transpose_mode
1302
+
1303
+ @property
1304
+ def a_shape(self) -> tuple[int, int, int]:
1305
+ dim = self.m if self.side == "left" else self.n
1306
+ return (self.batches_per_block, dim, dim)
1307
+
1308
+ @property
1309
+ def b_shape(self) -> tuple[int, int, int]:
1310
+ return (self.batches_per_block, self.m, self.n)
1311
+
1312
+ def a_strides(self, *, lda: int | None = None) -> tuple[int, int, int]:
1313
+ lda = self.lda if lda is None else lda
1314
+ return _calculate_strides(self.a_shape[1:], lda, self.a_arrangement)
1315
+
1316
+ def b_strides(self, *, ldb: int | None = None) -> tuple[int, int, int]:
1317
+ ldb = self.ldb if ldb is None else ldb
1318
+ return _calculate_strides(self.b_shape[1:], ldb, self.b_arrangement)
1319
+
1320
+ def a_size(self, *, lda: int | None = None) -> int:
1321
+ return self.a_strides(lda=lda)[0] * self.a_shape[0]
1322
+
1323
+ def b_size(self, *, ldb: int | None = None) -> int:
1324
+ return self.b_strides(ldb=ldb)[0] * self.b_shape[0]
1325
+
1326
+ # ==========================
1327
+ # Device function methods
1328
+ # ==========================
1329
+
1330
+ def solve(self, a, b, lda=None, ldb=None) -> None:
1331
+ """
1332
+ Solves a triangular linear system with multiple right-hand sides:
1333
+ ``op(A) * X = B`` (if ``side='left'``)
1334
+ ``X * op(A) = B`` (if ``side='right'``)
1335
+
1336
+ This device function solves a triangular system where A is a triangular matrix.
1337
+ Uses cuSOLVERDx ``'trsm'``. The operation is in-place: result X overwrites B.
1338
+
1339
+ If ``lda`` and ``ldb`` are provided, uses runtime version with the
1340
+ specified leading dimensions. If not provided (``None``), uses compile-time
1341
+ version with default or constructor-provided leading dimensions.
1342
+
1343
+ For more details, see: :cusolverdx_doc:`get_started/functions/trsm.html`
1344
+
1345
+ Args:
1346
+ a: Pointer to an array in shared memory, storing the batched triangular matrix
1347
+ according to the specified arrangement
1348
+ and leading dimension (see :meth:`__init__`).
1349
+ The ``fill_mode`` parameter denotes which
1350
+ part of the matrix is used (the other part is ignored).
1351
+ For unit diagonal mode (``diag='unit'``),
1352
+ diagonal elements are unity and not accessed.
1353
+ b: Pointer to an array in shared memory,
1354
+ storing the batched ``M`` x ``N`` right-hand side
1355
+ matrix according to the specified arrangement
1356
+ and leading dimension (see :meth:`__init__`).
1357
+ The operation is in-place: result X overwrites B.
1358
+ lda: Optional runtime leading dimension for matrix A.
1359
+ The ``lda`` and ``ldb`` must be specified together.
1360
+ If not specified, the compile-time ``lda`` is used.
1361
+ ldb: Optional runtime leading dimension for matrix B.
1362
+ The ``lda`` and ``ldb`` must be specified together.
1363
+ If not specified, the compile-time ``ldb`` is used.
1364
+ """
1365
+ raise RuntimeError("solve is a device function and cannot be called on the host.")
1366
+
1367
+
1368
+ # ==========================
1369
+ # LU Pivot Solver
1370
+ # ==========================
1371
+
1372
+
1373
+ @docstring_decorator(LU_SOLVER_DOCSTRING, skip_missing=False)
1374
+ class LUPivotSolver(_LinearSolverProperties):
1375
+ """
1376
+ A class that encapsulates cuSOLVERDx LU factorization with partial pivoting
1377
+ and linear solver for general matrices.
1378
+
1379
+ **Available operations:**
1380
+
1381
+ * factorize: Computes the LU factorization P @ A = L @ U with partial pivoting,
1382
+ where P is a permutation matrix, L is a
1383
+ lower triangular matrix and U is an upper triangular matrix.
1384
+ * solve: Solves the system Ax = B using a previously
1385
+ computed LU factorization with partial pivoting
1386
+
1387
+ **Memory Layout Requirements:**
1388
+
1389
+ Matrices must be stored in shared memory according
1390
+ to their arrangement and leading dimension (ld):
1391
+
1392
+ **For matrix A (M x N):**
1393
+
1394
+ * **Column-major arrangement**: Matrix shape ``(batches_per_block, M, N)``
1395
+ with strides ``(lda * N, 1, lda)``
1396
+ * **Row-major arrangement**: Matrix shape ``(batches_per_block, M, N)``
1397
+ with strides ``(lda * M, lda, 1)``
1398
+
1399
+ **For matrix B (N x K):**
1400
+
1401
+ * **Column-major arrangement**: Matrix shape ``(batches_per_block, N, K)``
1402
+ with strides ``(ldb * K, 1, ldb)``
1403
+ * **Row-major arrangement**: Matrix shape ``(batches_per_block, N, K)``
1404
+ with strides ``(ldb * N, ldb, 1)``
1405
+
1406
+ .. note::
1407
+ This solver uses partial pivoting for improved numerical stability and is
1408
+ suitable for general matrices. If your matrix is diagonally dominant,
1409
+ you may consider using :class:`LUSolver` which does not use pivoting
1410
+ and may be faster.
1411
+
1412
+ Args:
1413
+ size (Sequence[int]): {size}
1414
+
1415
+ precision (type[np.floating]): {precision}
1416
+
1417
+ execution (str): {execution}
1418
+
1419
+ sm (ComputeCapability): {sm}
1420
+
1421
+ transpose_mode (str, optional): {transpose_mode}
1422
+
1423
+ arrangement (Sequence[str], optional): {arrangement}
1424
+
1425
+ batches_per_block (int | Literal["suggested"], optional): {batches_per_block}
1426
+
1427
+ data_type (str, optional): {data_type}
1428
+
1429
+ leading_dimensions (Sequence[int], optional): {leading_dimensions}
1430
+
1431
+ block_dim (Sequence[int] | Literal["suggested"], optional): {block_dim}
1432
+
1433
+ See Also:
1434
+ For further details, please refer to the cuSOLVERDx documentation:
1435
+
1436
+ * :cusolverdx_doc:`factorize <get_started/functions/getrf.html>`
1437
+ * :cusolverdx_doc:`solve <get_started/functions/getrs.html>`
1438
+
1439
+ """
1440
+
1441
+ # ==========================
1442
+ # Constructor
1443
+ # ==========================
1444
+
1445
+ def __init__(
1446
+ self,
1447
+ size: Sequence[int],
1448
+ precision: type[np.floating],
1449
+ execution: str,
1450
+ *,
1451
+ sm=None,
1452
+ transpose_mode: str = "non_transposed",
1453
+ arrangement: Sequence[str] | None = None,
1454
+ batches_per_block: int | Literal["suggested"] | None = None,
1455
+ data_type: str | None = None,
1456
+ leading_dimensions: Sequence[int] | None = None,
1457
+ block_dim: Sequence[int] | Literal["suggested"] | None = None,
1458
+ ):
1459
+ def construct(function, transpose_mode):
1460
+ return Solver(
1461
+ function=function,
1462
+ size=size,
1463
+ precision=precision,
1464
+ execution=execution,
1465
+ sm=sm,
1466
+ arrangement=arrangement,
1467
+ batches_per_block=batches_per_block,
1468
+ data_type=data_type,
1469
+ leading_dimensions=leading_dimensions,
1470
+ block_dim=block_dim,
1471
+ transpose_mode=transpose_mode,
1472
+ )
1473
+
1474
+ if transpose_mode is None:
1475
+ raise ValueError("transpose_mode must be provided for LUPivotSolver")
1476
+
1477
+ self._factorize = construct("getrf_partial_pivot", "non_transposed")
1478
+ self._solve = construct("getrs_partial_pivot", transpose_mode) if self._factorize.m == self._factorize.n else None
1479
+ self._transpose_mode = transpose_mode
1480
+
1481
+ super().__init__(self._factorize)
1482
+
1483
+ # ==========================
1484
+ # Property methods
1485
+ # ==========================
1486
+
1487
+ @property
1488
+ def transpose_mode(self) -> str:
1489
+ return self._transpose_mode
1490
+
1491
+ @property
1492
+ def ipiv_type(self) -> type[np.signedinteger]:
1493
+ return self._factorize.ipiv_type
1494
+
1495
+ @property
1496
+ def ipiv_shape(self) -> tuple[int, int]:
1497
+ return (self.batches_per_block, min(self.m, self.n))
1498
+
1499
+ @property
1500
+ def ipiv_strides(self) -> tuple[int, int]:
1501
+ return (self.ipiv_shape[1], 1)
1502
+
1503
+ @property
1504
+ def ipiv_size(self) -> int:
1505
+ return self.ipiv_shape[0] * self.ipiv_shape[1]
1506
+
1507
+ # ==========================
1508
+ # Device function methods
1509
+ # ==========================
1510
+
1511
+ def factorize(self, a, ipiv, info, lda=None) -> None:
1512
+ """
1513
+ Computes the LU factorization of a general matrix A with partial pivoting.
1514
+
1515
+ This device function computes P @ A = L @ U, where P is a permutation matrix,
1516
+ L is a unit lower triangular matrix and U is an upper triangular matrix.
1517
+ This variant uses partial pivoting for improved numerical stability
1518
+ and is suitable for general matrices. Uses cuSOLVERDx ``'getrf_partial_pivot'``.
1519
+
1520
+ If ``lda`` is provided, uses runtime version with the specified leading dimension.
1521
+ If ``lda`` is not provided (``None``), uses compile-time version with default
1522
+ or constructor-provided leading dimensions.
1523
+
1524
+ .. note::
1525
+ The ``transpose_mode`` parameter does not affect factorization.
1526
+ This operation always treats the input matrix as-is (non-transposed).
1527
+
1528
+ For more details, see: :cusolverdx_doc:`get_started/functions/getrf.html`
1529
+
1530
+ Args:
1531
+ a: Pointer to an array in shared memory, storing the batched matrix according
1532
+ to the specified arrangement and leading dimension (see :meth:`__init__`).
1533
+ The matrix is overwritten in place.
1534
+ On exit, contains the factors L and U from the factorization P @ A = L @ U.
1535
+ The unit diagonal elements of L are not stored.
1536
+ ipiv: Pointer to a 1D array of ``int32``, storing pivot indices.
1537
+ The array has size min(M, N) for each batch.
1538
+ On exit, ``ipiv[batch_id * min(M, N) + i]`` indicates that row i
1539
+ was interchanged with row ``ipiv[batch_id * min(M, N) + i] - 1``
1540
+ in the batch_id-th batch of A.
1541
+ info: Pointer to a 1D array of ``int32``. On exit, ``info[batch_id] = 0``
1542
+ indicates success for that batch, ``info[batch_id] = i > 0``
1543
+ indicates ``U(i,i)`` is exactly zero,
1544
+ meaning the factorization has been completed but the factor U
1545
+ is singular and division by zero will occur if it is used to solve
1546
+ a system of equations.
1547
+ lda: Optional runtime leading dimension of matrix A.
1548
+ If not specified, the compile-time ``lda`` is used.
1549
+ """
1550
+ raise RuntimeError("factorize is a device function and cannot be called on the host.")
1551
+
1552
+ def solve(self, a, ipiv, b, lda=None, ldb=None) -> None:
1553
+ """
1554
+ Solves a system of linear equations Ax = B
1555
+ using the LU factorization with partial pivoting.
1556
+ The ``a`` operand must be a square matrix (``M == N``),
1557
+ otherwise this function will throw an exception.
1558
+
1559
+ This device function uses the previously computed factorization P @ A = L @ U
1560
+ to solve the system. Uses cuSOLVERDx ``'getrs_partial_pivot'``.
1561
+
1562
+ If ``lda`` and ``ldb`` are provided, uses runtime version with
1563
+ the specified leading dimensions. If not provided (``None``),
1564
+ uses compile-time version with default or constructor-provided
1565
+ leading dimensions.
1566
+
1567
+ .. note::
1568
+ The ``transpose_mode`` parameter (set in constructor) determines which
1569
+ system is solved: A*x=B (``'non_transposed'``), A^T*x=B (``'transposed'``),
1570
+ or A^H*x=B (``'conj_transposed'`` for complex matrices).
1571
+
1572
+ For more details, see: :cusolverdx_doc:`get_started/functions/getrs.html`
1573
+
1574
+ Args:
1575
+ a: Pointer to an array in shared memory, storing the batched factors L and U
1576
+ from the LU factorization with partial pivoting, according
1577
+ to the specified arrangement and leading dimension (see :meth:`__init__`).
1578
+ The unit diagonal elements of L are not stored.
1579
+ See the :meth:`factorize` documentation for details.
1580
+ ipiv: Pointer to a 1D array of ``int32`` in shared or global memory
1581
+ storing pivot indices.
1582
+ The array has size min(M, N) for each batch. The ipiv array should contain
1583
+ the pivot information from the :meth:`factorize` call.
1584
+ ``ipiv[batch_id * min(M, N) + i]`` indicates that row i was interchanged
1585
+ with row ``ipiv[batch_id * min(M, N) + i]`` in the batch_id-th batch of A.
1586
+ b: Pointer to an array in shared memory, storing the batched matrix according
1587
+ to the specified arrangement and leading dimension (see :meth:`__init__`).
1588
+ The matrix is overwritten in place with the solution matrix x.
1589
+ lda: Optional runtime leading dimension of matrix A.
1590
+ The ``lda`` and ``ldb`` must be specified together.
1591
+ If not specified, the compile-time ``lda`` is used.
1592
+ ldb: Optional runtime leading dimension of matrix B.
1593
+ The ``lda`` and ``ldb`` must be specified together.
1594
+ If not specified, the compile-time ``ldb`` is used.
1595
+ """
1596
+ raise RuntimeError("solve is a device function and cannot be called on the host.")
1597
+
1598
+
1599
+ # ==========================
1600
+ # QR/LQ Factorizers Base
1601
+ # ==========================
1602
+
1603
+ ORTOGHONAL_FACTORIZER_DOCSTRING = SOLVER_DOCSTRING.copy()
1604
+ del ORTOGHONAL_FACTORIZER_DOCSTRING["function"]
1605
+ del ORTOGHONAL_FACTORIZER_DOCSTRING["fill_mode"]
1606
+ del ORTOGHONAL_FACTORIZER_DOCSTRING["side"]
1607
+ del ORTOGHONAL_FACTORIZER_DOCSTRING["diag"]
1608
+ del ORTOGHONAL_FACTORIZER_DOCSTRING["transpose_mode"]
1609
+ del ORTOGHONAL_FACTORIZER_DOCSTRING["job"]
1610
+
1611
+ ORTOGHONAL_FACTORIZER_DOCSTRING["size"] = f"""{SOLVER_DOCSTRING_SIZE_BASE}
1612
+ ``M`` and ``N`` represent the dimensions of the matrix A used in factorization.
1613
+ ``K`` is ignored if specified.""".replace("\n", " ")
1614
+
1615
+ ORTOGHONAL_FACTORIZER_DOCSTRING["arrangement"] = (
1616
+ """Storage layout for matrix A.
1617
+ Can be one of: ``'col_major'``, ``'row_major'``. Defaults to ``'col_major'``.""".replace("\n", " ")
1618
+ + " "
1619
+ + ADAPTERS_API_LD_DOCSTRING
1620
+ )
1621
+
1622
+ ORTOGHONAL_FACTORIZER_DOCSTRING["leading_dimension"] = (
1623
+ """The leading dimension for input matrix A, or ``None``.
1624
+ If not provided, it will be automatically deduced from ``size`` and ``arrangement``.""".replace("\n", " ")
1625
+ + " "
1626
+ + ADAPTERS_API_LD_DOCSTRING
1627
+ )
1628
+
1629
+
1630
+ class _OrthogonalFactorizerProperties:
1631
+ def __init__(self, source: Solver):
1632
+ self._properties_source = source
1633
+
1634
+ @property
1635
+ def size(self) -> tuple[int, int, int]:
1636
+ return self._properties_source.size
1637
+
1638
+ @property
1639
+ def m(self) -> int:
1640
+ return self._properties_source.m
1641
+
1642
+ @property
1643
+ def n(self) -> int:
1644
+ return self._properties_source.n
1645
+
1646
+ @property
1647
+ def precision(self) -> type[np.floating]:
1648
+ return self._properties_source.precision
1649
+
1650
+ @property
1651
+ def execution(self) -> str:
1652
+ return self._properties_source.execution
1653
+
1654
+ @property
1655
+ def sm(self) -> ComputeCapability:
1656
+ return self._properties_source.sm
1657
+
1658
+ @property
1659
+ def a_arrangement(self) -> str:
1660
+ return self._properties_source.a_arrangement
1661
+
1662
+ @property
1663
+ def batches_per_block(self) -> int:
1664
+ return self._properties_source.batches_per_block
1665
+
1666
+ @property
1667
+ def data_type(self) -> str:
1668
+ return self._properties_source.data_type
1669
+
1670
+ @property
1671
+ def lda(self) -> int:
1672
+ return self._properties_source.lda
1673
+
1674
+ @property
1675
+ def block_dim(self) -> tuple:
1676
+ return self._properties_source.block_dim
1677
+
1678
+ @property
1679
+ def block_size(self) -> int:
1680
+ return self._properties_source.block_size
1681
+
1682
+ @property
1683
+ def value_type(self) -> type[np.floating] | Complex:
1684
+ return self._properties_source.value_type
1685
+
1686
+ @property
1687
+ def tau_type(self) -> type[np.floating] | Complex:
1688
+ return self._properties_source.tau_type
1689
+
1690
+ @property
1691
+ def tau_shape(self) -> tuple[int, int]:
1692
+ return (self.batches_per_block, min(self.m, self.n))
1693
+
1694
+ @property
1695
+ def tau_strides(self) -> tuple[int, int]:
1696
+ return (self.tau_shape[1], 1)
1697
+
1698
+ @property
1699
+ def a_shape(self) -> tuple[int, int, int]:
1700
+ return (self.batches_per_block, self.m, self.n)
1701
+
1702
+ def a_strides(self, *, lda: int | None = None) -> tuple[int, int, int]:
1703
+ lda = self.lda if lda is None else lda
1704
+ return _calculate_strides(self.a_shape[1:], lda, self.a_arrangement)
1705
+
1706
+ def a_size(self, *, lda: int | None = None) -> int:
1707
+ return self.a_strides(lda=lda)[0] * self.a_shape[0]
1708
+
1709
+ @property
1710
+ def tau_size(self) -> int:
1711
+ return self.tau_shape[0] * self.tau_shape[1]
1712
+
1713
+
1714
+ # ==========================
1715
+ # QR Factorize
1716
+ # ==========================
1717
+
1718
+
1719
+ @docstring_decorator(ORTOGHONAL_FACTORIZER_DOCSTRING, skip_missing=False)
1720
+ class QRFactorize(_OrthogonalFactorizerProperties):
1721
+ """
1722
+ A class that encapsulates QR orthogonal factorization device function
1723
+ for general matrices using Householder reflections.
1724
+
1725
+ **Available operation:**
1726
+
1727
+ * factorize: Computes the QR factorization A = Q @ R,
1728
+ where Q is a unitary M x M matrix
1729
+ and R is an upper triangular matrix (if M >= N)
1730
+ or upper trapezoidal matrix (if M < N).
1731
+
1732
+ The factorization uses Householder reflection transformations and does not explicitly
1733
+ form the unitary matrix Q. Instead, Q is represented as a product of Householder vectors
1734
+ stored in the input matrix A along with the tau array.
1735
+
1736
+ **Memory Layout Requirements:**
1737
+
1738
+ Matrices must be stored in shared memory according
1739
+ to their arrangement and leading dimension (ld):
1740
+
1741
+ **For matrix A (M x N):**
1742
+
1743
+ * **Column-major arrangement**: Matrix shape ``(batches_per_block, M, N)``
1744
+ with strides ``(lda * N, 1, lda)``
1745
+ * **Row-major arrangement**: Matrix shape ``(batches_per_block, M, N)``
1746
+ with strides ``(lda * M, lda, 1)``
1747
+
1748
+ Args:
1749
+ size (Sequence[int]): {size}
1750
+
1751
+ precision (type[np.floating]): {precision}
1752
+
1753
+ execution (str): {execution}
1754
+
1755
+ sm (ComputeCapability): {sm}
1756
+
1757
+ arrangement (str, optional): {arrangement}
1758
+
1759
+ batches_per_block (int | Literal["suggested"], optional): {batches_per_block}
1760
+
1761
+ data_type (str, optional): {data_type}
1762
+
1763
+ leading_dimension (int, optional): {leading_dimension}
1764
+
1765
+ block_dim (Sequence[int] | Literal["suggested"], optional): {block_dim}
1766
+
1767
+ See Also:
1768
+ For further details, please refer to the cuSOLVERDx documentation:
1769
+
1770
+ * :cusolverdx_doc:`factorize (geqrf) <get_started/functions/geqrf.html>`
1771
+ """
1772
+
1773
+ # ==========================
1774
+ # Constructor
1775
+ # ==========================
1776
+
1777
+ def __init__(
1778
+ self,
1779
+ size: Sequence[int],
1780
+ precision: type[np.floating],
1781
+ execution: str,
1782
+ *,
1783
+ sm=None,
1784
+ arrangement: str | None = None,
1785
+ batches_per_block: int | Literal["suggested"] | None = None,
1786
+ data_type: str | None = None,
1787
+ leading_dimension: int | None = None,
1788
+ block_dim: Sequence[int] | Literal["suggested"] | None = None,
1789
+ ):
1790
+ self._factorize = Solver(
1791
+ function="geqrf",
1792
+ size=size,
1793
+ precision=precision,
1794
+ execution=execution,
1795
+ sm=sm,
1796
+ arrangement=(arrangement,) if arrangement is not None else None,
1797
+ batches_per_block=batches_per_block,
1798
+ data_type=data_type,
1799
+ leading_dimensions=(leading_dimension,) if leading_dimension is not None else None,
1800
+ block_dim=block_dim,
1801
+ )
1802
+ super().__init__(self._factorize)
1803
+
1804
+ # ==========================
1805
+ # Device function methods
1806
+ # ==========================
1807
+
1808
+ def factorize(self, a, tau, lda=None) -> None:
1809
+ """
1810
+ Computes the QR factorization of a general matrix A using Householder reflections.
1811
+
1812
+ This device function computes A = Q @ R, where Q is a unitary M x M matrix
1813
+ and R is an upper triangular matrix (if M >= N)
1814
+ or upper trapezoidal matrix (if M < N).
1815
+ Uses cuSOLVERDx ``'geqrf'``.
1816
+
1817
+ If ``lda`` is provided, uses runtime version with the
1818
+ specified leading dimension. If ``lda`` is not provided (``None``),
1819
+ uses compile-time version with
1820
+ default or constructor-provided leading dimensions.
1821
+
1822
+ Matrix Q is not explicitly formed. Instead, Q is represented as a product of
1823
+ min(M, N) Householder vectors: Q = H(0) * H(1) * ... * H(min(M, N) - 1).
1824
+
1825
+ Each Householder vector has the form H(i) = I - tau[i] * v * v^H, where:
1826
+
1827
+ * v is a vector of size M for each batch
1828
+ * v[0:i-1] = 0, v[i] = 1
1829
+ * v[i+1:M] is stored on exit in A[i+1:M, i]
1830
+
1831
+ For more details, see: :cusolverdx_doc:`get_started/functions/geqrf.html`
1832
+
1833
+ Args:
1834
+ a: Pointer to an array in shared memory, storing
1835
+ the batched matrix according
1836
+ to the specified arrangement and leading dimension (see :meth:`__init__`).
1837
+ The matrix is overwritten in place.
1838
+ On exit, the upper triangular or upper trapezoidal part (including diagonal)
1839
+ contains the matrix R. The elements below the diagonal, with the array tau,
1840
+ represent the unitary matrix Q as a product of Householder vectors.
1841
+ tau: Pointer to a 1D array of size min(M, N) for each batch.
1842
+ Contains the scalar factors of the Householder reflections.
1843
+ The tau array, together with the Householder vectors stored in A,
1844
+ defines the unitary matrix Q.
1845
+ lda: Optional runtime leading dimension of matrix A.
1846
+ If not specified, the compile-time ``lda`` is used.
1847
+ """
1848
+ raise RuntimeError("factorize is a device function and cannot be called on the host.")
1849
+
1850
+
1851
+ # ==========================
1852
+ # LQ Factorize
1853
+ # ==========================
1854
+
1855
+
1856
+ @docstring_decorator(ORTOGHONAL_FACTORIZER_DOCSTRING, skip_missing=False)
1857
+ class LQFactorize(_OrthogonalFactorizerProperties):
1858
+ """
1859
+ A class that encapsulates LQ orthogonal factorization device function
1860
+ for general matrices using Householder reflections.
1861
+
1862
+ **Available operation:**
1863
+
1864
+ * factorize: Computes the LQ factorization A = L @ Q,
1865
+ where L is a lower triangular matrix (if M <= N)
1866
+ or lower trapezoidal matrix (if M > N)
1867
+ and Q is a unitary N x N matrix.
1868
+
1869
+ The factorization uses Householder reflection transformations and does not explicitly
1870
+ form the unitary matrix Q. Instead, Q is represented as a product of Householder vectors
1871
+ stored in the input matrix A along with the tau array.
1872
+
1873
+ **Memory Layout Requirements:**
1874
+
1875
+ Matrices must be stored in shared memory according
1876
+ to their arrangement and leading dimension (ld):
1877
+
1878
+ **For matrix A (M x N):**
1879
+
1880
+ * **Column-major arrangement**: Matrix shape ``(batches_per_block, M, N)``
1881
+ with strides ``(lda * N, 1, lda)``
1882
+ * **Row-major arrangement**: Matrix shape ``(batches_per_block, M, N)``
1883
+ with strides ``(lda * M, lda, 1)``
1884
+
1885
+ .. note::
1886
+ The LQ factorization is essentially equivalent to the
1887
+ QR factorization of A^T (or A^H for complex types).
1888
+
1889
+ Args:
1890
+ size (Sequence[int]): {size}
1891
+
1892
+ precision (type[np.floating]): {precision}
1893
+
1894
+ execution (str): {execution}
1895
+
1896
+ sm (ComputeCapability): {sm}
1897
+
1898
+ arrangement (str, optional): {arrangement}
1899
+
1900
+ batches_per_block (int | Literal["suggested"], optional): {batches_per_block}
1901
+
1902
+ data_type (str, optional): {data_type}
1903
+
1904
+ leading_dimension (int, optional): {leading_dimension}
1905
+
1906
+ block_dim (Sequence[int] | Literal["suggested"], optional): {block_dim}
1907
+
1908
+ See Also:
1909
+ For further details, please refer to the cuSOLVERDx documentation:
1910
+
1911
+ * :cusolverdx_doc:`factorize (gelqf) <get_started/functions/gelqf.html>`
1912
+ """
1913
+
1914
+ # ==========================
1915
+ # Constructor
1916
+ # ==========================
1917
+
1918
+ def __init__(
1919
+ self,
1920
+ size: Sequence[int],
1921
+ precision: type[np.floating],
1922
+ execution: str,
1923
+ *,
1924
+ sm=None,
1925
+ arrangement: str | None = None,
1926
+ batches_per_block: int | Literal["suggested"] | None = None,
1927
+ data_type: str | None = None,
1928
+ leading_dimension: int | None = None,
1929
+ block_dim: Sequence[int] | Literal["suggested"] | None = None,
1930
+ ):
1931
+ self._factorize = Solver(
1932
+ function="gelqf",
1933
+ size=size,
1934
+ precision=precision,
1935
+ execution=execution,
1936
+ sm=sm,
1937
+ arrangement=(arrangement,) if arrangement is not None else None,
1938
+ batches_per_block=batches_per_block,
1939
+ data_type=data_type,
1940
+ leading_dimensions=(leading_dimension,) if leading_dimension is not None else None,
1941
+ block_dim=block_dim,
1942
+ )
1943
+ super().__init__(self._factorize)
1944
+
1945
+ # ==========================
1946
+ # Device function methods
1947
+ # ==========================
1948
+
1949
+ def factorize(self, a, tau, lda=None) -> None:
1950
+ """
1951
+ Computes the LQ factorization of a general matrix A using Householder reflections.
1952
+
1953
+ This device function computes A = L @ Q, where L is a lower triangular matrix
1954
+ (if M <= N) or lower trapezoidal matrix (if M > N),
1955
+ and Q is a unitary N x N matrix.
1956
+ Uses cuSOLVERDx ``'gelqf'``.
1957
+
1958
+ If ``lda`` is provided, uses runtime version with the
1959
+ specified leading dimension. If ``lda`` is not provided (``None``),
1960
+ uses compile-time version with
1961
+ default or constructor-provided leading dimensions.
1962
+
1963
+ The LQ factorization is essentially the same as the QR factorization of A^T
1964
+ (or A^H for complex data types).
1965
+
1966
+ Matrix Q is not explicitly formed. Instead, Q is represented as a product of
1967
+ min(M, N) Householder vectors: Q = H(min(M, N) - 1)^H * ... * H(1)^H * H(0)^H.
1968
+
1969
+ Each Householder vector has the form H(i) = I - tau[i] * v * v^H, where:
1970
+
1971
+ * v is a vector of size N for each batch
1972
+ * v[0:i-1] = 0, v[i] = 1
1973
+ * conjugate(v[i+1:N]) is stored on exit in A[i, i+1:N]
1974
+
1975
+ For more details, see: :cusolverdx_doc:`get_started/functions/gelqf.html`
1976
+
1977
+ Args:
1978
+ a: Pointer to an array in shared memory, storing
1979
+ the batched M x N matrix according
1980
+ to the specified arrangement and leading dimension (see :meth:`__init__`).
1981
+ The matrix is overwritten in place.
1982
+ On exit, the lower triangular or lower trapezoidal part (including diagonal)
1983
+ contains the matrix L. The elements above the diagonal, with the array tau,
1984
+ represent the unitary matrix Q as a product of Householder vectors.
1985
+ tau: Pointer to a 1D array of size min(M, N) for each batch.
1986
+ Contains the scalar factors of the Householder reflections.
1987
+ The tau array, together with the Householder vectors stored in A,
1988
+ defines the unitary matrix Q.
1989
+ lda: Optional runtime leading dimension of matrix A.
1990
+ If not specified, the compile-time ``lda`` is used.
1991
+ """
1992
+ raise RuntimeError("factorize is a device function and cannot be called on the host.")
1993
+
1994
+
1995
+ # ==========================
1996
+ # QR Multiplier (UNMQR)
1997
+ # ==========================
1998
+
1999
+ QR_MULTIPLIER_DOCSTRING = SOLVER_DOCSTRING.copy()
2000
+ del QR_MULTIPLIER_DOCSTRING["function"]
2001
+ del QR_MULTIPLIER_DOCSTRING["fill_mode"]
2002
+ del QR_MULTIPLIER_DOCSTRING["diag"]
2003
+ del QR_MULTIPLIER_DOCSTRING["job"]
2004
+
2005
+ QR_MULTIPLIER_DOCSTRING["size"] = f"""{SOLVER_DOCSTRING_SIZE_BASE}
2006
+ ``M`` and ``N`` represent the dimensions of matrix C.
2007
+ ``K`` represents the number of Householder reflections from the QR factorization.
2008
+ If ``side='left'``, then ``K <= M`` and A is ``M`` x ``K``.
2009
+ If ``side='right'``, then ``K <= N`` and A is ``N`` x ``K``.""".replace("\n", " ")
2010
+
2011
+ QR_MULTIPLIER_DOCSTRING["side"] = f"""Side of matrix Q in the multiplication operation.
2012
+ Can be one of: {", ".join(f"``'{v}'``" for v in ALLOWED_SIDE)}.
2013
+ If ``side='left'``, computes op(Q) * C where Q is ``M`` x ``M``.
2014
+ If ``side='right'``, computes C * op(Q) where Q is ``N`` x ``N``.""".replace("\n", " ")
2015
+
2016
+ QR_MULTIPLIER_DOCSTRING["transpose_mode"] = f"""Transpose mode for operation op(Q) applied to matrix Q.
2017
+ Can be one of: {", ".join(f"``'{v}'``" for v in ALLOWED_TRANSPOSE_MODE)}.
2018
+ Defaults to ``'non_transposed'``.""".replace("\n", " ")
2019
+
2020
+ QR_MULTIPLIER_DOCSTRING["leading_dimensions"] = SOLVER_DOCSTRING["leading_dimensions"] + ADAPTERS_API_LD_DOCSTRING
2021
+
2022
+
2023
+ @docstring_decorator(QR_MULTIPLIER_DOCSTRING, skip_missing=False)
2024
+ class QRMultiply(_SolverProperties):
2025
+ """
2026
+ A class that encapsulates the multiplication of a matrix C by the unitary matrix Q
2027
+ from a QR factorization (UNMQR operation).
2028
+
2029
+ **Memory Layout Requirements:**
2030
+
2031
+ Matrices must be stored in shared memory according
2032
+ to their arrangement and leading dimension (ld):
2033
+
2034
+ **For matrix A (containing Householder vectors):**
2035
+
2036
+ * If ``side='left'``: A is ``M`` x ``K``
2037
+ * If ``side='right'``: A is ``N`` x ``K``
2038
+ * **Column-major arrangement**: Matrix shape ``(batches_per_block, rows, K)``
2039
+ with strides ``(lda * K, 1, lda)``
2040
+ * **Row-major arrangement**: Matrix shape ``(batches_per_block, rows, K)``
2041
+ with strides ``(lda * rows, lda, 1)``
2042
+
2043
+ **For matrix C (M x N):**
2044
+
2045
+ * **Column-major arrangement**: Matrix shape ``(batches_per_block, M, N)``
2046
+ with strides ``(ldb * N, 1, ldb)``
2047
+ * **Row-major arrangement**: Matrix shape ``(batches_per_block, M, N)``
2048
+ with strides ``(ldb * M, ldb, 1)``
2049
+
2050
+ Args:
2051
+ size (Sequence[int]): {size}
2052
+
2053
+ precision (type[np.floating]): {precision}
2054
+
2055
+ execution (str): {execution}
2056
+
2057
+ sm (ComputeCapability): {sm}
2058
+
2059
+ side (str): {side}
2060
+
2061
+ transpose_mode (str, optional): {transpose_mode}
2062
+
2063
+ arrangement (Sequence[str], optional): {arrangement}
2064
+
2065
+ batches_per_block (int | Literal["suggested"], optional): {batches_per_block}
2066
+
2067
+ data_type (str, optional): {data_type}
2068
+
2069
+ leading_dimensions (Sequence[int], optional): {leading_dimensions}
2070
+
2071
+ block_dim (Sequence[int] | Literal["suggested"], optional): {block_dim}
2072
+
2073
+ See Also:
2074
+ For further details, please refer to the cuSOLVERDx documentation:
2075
+
2076
+ * :cusolverdx_doc:`unmqr <get_started/functions/unmqr.html>`
2077
+ """
2078
+
2079
+ # ==========================
2080
+ # Constructor
2081
+ # ==========================
2082
+
2083
+ def __init__(
2084
+ self,
2085
+ size: Sequence[int],
2086
+ precision: type[np.floating],
2087
+ execution: str,
2088
+ side: str,
2089
+ *,
2090
+ sm=None,
2091
+ transpose_mode: str = "non_transposed",
2092
+ arrangement: Sequence[str] | None = None,
2093
+ batches_per_block: int | Literal["suggested"] | None = None,
2094
+ data_type: str | None = None,
2095
+ leading_dimensions: Sequence[int] | None = None,
2096
+ block_dim: Sequence[int] | Literal["suggested"] | None = None,
2097
+ ):
2098
+ self._multiply: Solver = Solver(
2099
+ function="unmqr",
2100
+ size=size,
2101
+ precision=precision,
2102
+ execution=execution,
2103
+ sm=sm,
2104
+ side=side,
2105
+ transpose_mode=transpose_mode,
2106
+ arrangement=arrangement,
2107
+ batches_per_block=batches_per_block,
2108
+ data_type=data_type,
2109
+ leading_dimensions=leading_dimensions,
2110
+ block_dim=block_dim,
2111
+ )
2112
+
2113
+ super().__init__(self._multiply)
2114
+
2115
+ if not isinstance(self._multiply.transpose_mode, str):
2116
+ raise ValueError("transpose_mode must be provided for QRMultiply")
2117
+
2118
+ if not isinstance(self._multiply.side, str):
2119
+ raise ValueError("side must be provided for QRMultiply")
2120
+
2121
+ # ==========================
2122
+ # Property methods
2123
+ # ==========================
2124
+
2125
+ @property
2126
+ def side(self) -> str:
2127
+ assert self._multiply.side is not None
2128
+ return self._multiply.side
2129
+
2130
+ @property
2131
+ def transpose_mode(self) -> str:
2132
+ assert self._multiply.transpose_mode is not None
2133
+ return self._multiply.transpose_mode
2134
+
2135
+ @property
2136
+ def tau_type(self) -> type[np.floating] | Complex:
2137
+ return self.value_type
2138
+
2139
+ @property
2140
+ def tau_shape(self) -> tuple[int, int]:
2141
+ return (self.batches_per_block, self.k)
2142
+
2143
+ @property
2144
+ def tau_strides(self) -> tuple[int, int]:
2145
+ return (self.tau_shape[1], 1)
2146
+
2147
+ @property
2148
+ def a_shape(self) -> tuple[int, int, int]:
2149
+ rows = self.m if self.side == "left" else self.n
2150
+ return (self.batches_per_block, rows, self.k)
2151
+
2152
+ @property
2153
+ def c_shape(self) -> tuple[int, int, int]:
2154
+ return (self.batches_per_block, self.m, self.n)
2155
+
2156
+ def a_strides(self, *, lda: int | None = None) -> tuple[int, int, int]:
2157
+ lda = self.lda if lda is None else lda
2158
+ return _calculate_strides(self.a_shape[1:], lda, self.a_arrangement)
2159
+
2160
+ def c_strides(self, *, ldc: int | None = None) -> tuple[int, int, int]:
2161
+ ldc = self.ldb if ldc is None else ldc
2162
+ return _calculate_strides(self.c_shape[1:], ldc, self.b_arrangement)
2163
+
2164
+ def a_size(self, *, lda: int | None = None) -> int:
2165
+ return self.a_strides(lda=lda)[0] * self.a_shape[0]
2166
+
2167
+ def c_size(self, *, ldc: int | None = None) -> int:
2168
+ return self.c_strides(ldc=ldc)[0] * self.c_shape[0]
2169
+
2170
+ @property
2171
+ def tau_size(self) -> int:
2172
+ return self.tau_shape[0] * self.tau_shape[1]
2173
+
2174
+ # ==========================
2175
+ # Device function methods
2176
+ # ==========================
2177
+
2178
+ def multiply(self, a, tau, c, lda=None, ldc=None) -> None:
2179
+ """
2180
+ Multiplies matrix C by the unitary matrix Q from a QR factorization.
2181
+
2182
+ This device function computes:
2183
+ ``op(Q) * C`` (if ``side='left'``)
2184
+ ``C * op(Q)`` (if ``side='right'``)
2185
+
2186
+ where Q is the unitary matrix from the QR factorization, represented
2187
+ by Householder vectors stored in A and the tau array.
2188
+ Uses cuSOLVERDx ``'unmqr'``. The result overwrites matrix C.
2189
+
2190
+ If ``lda`` and ``ldc`` are provided, uses runtime version with the
2191
+ specified leading dimensions. If not provided (``None``), uses compile-time
2192
+ version with default or constructor-provided leading dimensions.
2193
+
2194
+ For more details, see: :cusolverdx_doc:`get_started/functions/unmqr.html`
2195
+
2196
+ Args:
2197
+ a: Pointer to an array in shared memory, storing the batched matrix containing
2198
+ Householder vectors from the QR factorization, according to the specified
2199
+ arrangement and leading dimension (see :meth:`__init__`).
2200
+ The elements below the diagonal of A, with the array tau, represent the
2201
+ unitary matrix Q as a product of Householder reflections.
2202
+ If ``side='left'``, A is ``M`` x ``K``.
2203
+ If ``side='right'``, A is ``N`` x ``K``.
2204
+ tau: Pointer to a 1D array of size K for each batch, containing the scalar
2205
+ factors of the Householder reflections from the QR factorization.
2206
+ The tau array, together with the Householder vectors in A, defines Q.
2207
+ c: Pointer to an array in shared memory,
2208
+ storing the batched ``M`` x ``N`` matrix
2209
+ according to the specified arrangement
2210
+ and leading dimension (see :meth:`__init__`).
2211
+ The operation is in-place: result overwrites C.
2212
+ lda: Optional runtime leading dimension for matrix A.
2213
+ The ``lda`` and ``ldc`` must be specified together.
2214
+ If not specified, the compile-time ``lda`` is used.
2215
+ ldc: Optional runtime leading dimension for matrix C.
2216
+ The ``lda`` and ``ldc`` must be specified together.
2217
+ If not specified, the compile-time ``ldc`` is used.
2218
+ """
2219
+ raise RuntimeError("multiply is a device function and cannot be called on the host.")
2220
+
2221
+
2222
+ # ==========================
2223
+ # LQ Multiplier (UNMLQ)
2224
+ # ==========================
2225
+
2226
+ LQ_MULTIPLIER_DOCSTRING = SOLVER_DOCSTRING.copy()
2227
+ del LQ_MULTIPLIER_DOCSTRING["function"]
2228
+ del LQ_MULTIPLIER_DOCSTRING["fill_mode"]
2229
+ del LQ_MULTIPLIER_DOCSTRING["diag"]
2230
+ del LQ_MULTIPLIER_DOCSTRING["job"]
2231
+
2232
+ LQ_MULTIPLIER_DOCSTRING["size"] = f"""{SOLVER_DOCSTRING_SIZE_BASE}
2233
+ ``M`` and ``N`` represent the dimensions of matrix C.
2234
+ ``K`` represents the number of Householder reflections from the LQ factorization.
2235
+ If ``side='left'``, then ``K <= M`` and A is ``K`` x ``M``.
2236
+ If ``side='right'``, then ``K <= N`` and A is ``K`` x ``N``.""".replace("\n", " ")
2237
+
2238
+ LQ_MULTIPLIER_DOCSTRING["side"] = f"""Side of matrix Q in the multiplication operation.
2239
+ Can be one of: {", ".join(f"``'{v}'``" for v in ALLOWED_SIDE)}.
2240
+ If ``side='left'``, computes op(Q) * C where Q is ``M`` x ``M``.
2241
+ If ``side='right'``, computes C * op(Q) where Q is ``N`` x ``N``.""".replace("\n", " ")
2242
+
2243
+ LQ_MULTIPLIER_DOCSTRING["transpose_mode"] = f"""Transpose mode for operation op(Q) applied to matrix Q.
2244
+ Can be one of: {", ".join(f"``'{v}'``" for v in ALLOWED_TRANSPOSE_MODE)}.
2245
+ Defaults to ``'non_transposed'``.""".replace("\n", " ")
2246
+
2247
+ LQ_MULTIPLIER_DOCSTRING["leading_dimensions"] = SOLVER_DOCSTRING["leading_dimensions"] + ADAPTERS_API_LD_DOCSTRING
2248
+
2249
+
2250
+ @docstring_decorator(LQ_MULTIPLIER_DOCSTRING, skip_missing=False)
2251
+ class LQMultiply(_SolverProperties):
2252
+ """
2253
+ A class that encapsulates the multiplication of a matrix C by the unitary matrix Q
2254
+ from an LQ factorization (UNMLQ operation).
2255
+
2256
+ **Memory Layout Requirements:**
2257
+
2258
+ Matrices must be stored in shared memory according
2259
+ to their arrangement and leading dimension (ld):
2260
+
2261
+ **For matrix A (containing Householder vectors):**
2262
+
2263
+ * If ``side='left'``: A is ``K`` x ``M``
2264
+ * If ``side='right'``: A is ``K`` x ``N``
2265
+ * **Column-major arrangement**: Matrix shape ``(batches_per_block, K, cols)``
2266
+ with strides ``(lda * cols, 1, lda)``
2267
+ * **Row-major arrangement**: Matrix shape ``(batches_per_block, K, cols)``
2268
+ with strides ``(lda * K, lda, 1)``
2269
+
2270
+ **For matrix C (M x N):**
2271
+
2272
+ * **Column-major arrangement**: Matrix shape ``(batches_per_block, M, N)``
2273
+ with strides ``(ldc * N, 1, ldc)``
2274
+ * **Row-major arrangement**: Matrix shape ``(batches_per_block, M, N)``
2275
+ with strides ``(ldc * M, ldc, 1)``
2276
+
2277
+ Args:
2278
+ size (Sequence[int]): {size}
2279
+
2280
+ precision (type[np.floating]): {precision}
2281
+
2282
+ execution (str): {execution}
2283
+
2284
+ sm (ComputeCapability): {sm}
2285
+
2286
+ side (str): {side}
2287
+
2288
+ transpose_mode (str, optional): {transpose_mode}
2289
+
2290
+ arrangement (Sequence[str], optional): {arrangement}
2291
+
2292
+ batches_per_block (int | Literal["suggested"], optional): {batches_per_block}
2293
+
2294
+ data_type (str, optional): {data_type}
2295
+
2296
+ leading_dimensions (Sequence[int], optional): {leading_dimensions}
2297
+
2298
+ block_dim (Sequence[int] | Literal["suggested"], optional): {block_dim}
2299
+
2300
+ See Also:
2301
+ For further details, please refer to the cuSOLVERDx documentation:
2302
+
2303
+ * :cusolverdx_doc:`unmlq <get_started/functions/unmlq.html>`
2304
+ """
2305
+
2306
+ # ==========================
2307
+ # Constructor
2308
+ # ==========================
2309
+
2310
+ def __init__(
2311
+ self,
2312
+ size: Sequence[int],
2313
+ precision: type[np.floating],
2314
+ execution: str,
2315
+ side: str,
2316
+ *,
2317
+ sm=None,
2318
+ transpose_mode: str = "non_transposed",
2319
+ arrangement: Sequence[str] | None = None,
2320
+ batches_per_block: int | Literal["suggested"] | None = None,
2321
+ data_type: str | None = None,
2322
+ leading_dimensions: Sequence[int] | None = None,
2323
+ block_dim: Sequence[int] | Literal["suggested"] | None = None,
2324
+ ):
2325
+ self._multiply: Solver = Solver(
2326
+ function="unmlq",
2327
+ size=size,
2328
+ precision=precision,
2329
+ execution=execution,
2330
+ sm=sm,
2331
+ side=side,
2332
+ transpose_mode=transpose_mode,
2333
+ arrangement=arrangement,
2334
+ batches_per_block=batches_per_block,
2335
+ data_type=data_type,
2336
+ leading_dimensions=leading_dimensions,
2337
+ block_dim=block_dim,
2338
+ )
2339
+
2340
+ super().__init__(self._multiply)
2341
+
2342
+ if not isinstance(self._multiply.transpose_mode, str):
2343
+ raise ValueError("transpose_mode must be provided for LQMultiply")
2344
+
2345
+ if not isinstance(self._multiply.side, str):
2346
+ raise ValueError("side must be provided for LQMultiply")
2347
+
2348
+ # ==========================
2349
+ # Property methods
2350
+ # ==========================
2351
+
2352
+ @property
2353
+ def side(self) -> str:
2354
+ assert self._multiply.side is not None
2355
+ return self._multiply.side
2356
+
2357
+ @property
2358
+ def transpose_mode(self) -> str:
2359
+ assert self._multiply.transpose_mode is not None
2360
+ return self._multiply.transpose_mode
2361
+
2362
+ @property
2363
+ def tau_type(self) -> type[np.floating] | Complex:
2364
+ return self.value_type
2365
+
2366
+ @property
2367
+ def tau_shape(self) -> tuple[int, int]:
2368
+ return (self.batches_per_block, self.k)
2369
+
2370
+ @property
2371
+ def tau_strides(self) -> tuple[int, int]:
2372
+ return (self.tau_shape[1], 1)
2373
+
2374
+ @property
2375
+ def a_shape(self) -> tuple[int, int, int]:
2376
+ cols = self.m if self.side == "left" else self.n
2377
+ return (self.batches_per_block, self.k, cols)
2378
+
2379
+ @property
2380
+ def c_shape(self) -> tuple[int, int, int]:
2381
+ return (self.batches_per_block, self.m, self.n)
2382
+
2383
+ def a_strides(self, *, lda: int | None = None) -> tuple[int, int, int]:
2384
+ lda = self.lda if lda is None else lda
2385
+ return _calculate_strides(self.a_shape[1:], lda, self.a_arrangement)
2386
+
2387
+ def c_strides(self, *, ldc: int | None = None) -> tuple[int, int, int]:
2388
+ ldc = self.ldb if ldc is None else ldc
2389
+ return _calculate_strides(self.c_shape[1:], ldc, self.b_arrangement)
2390
+
2391
+ def a_size(self, *, lda: int | None = None) -> int:
2392
+ return self.a_strides(lda=lda)[0] * self.a_shape[0]
2393
+
2394
+ def c_size(self, *, ldc: int | None = None) -> int:
2395
+ return self.c_strides(ldc=ldc)[0] * self.c_shape[0]
2396
+
2397
+ @property
2398
+ def tau_size(self) -> int:
2399
+ return self.tau_shape[0] * self.tau_shape[1]
2400
+
2401
+ # ==========================
2402
+ # Device function methods
2403
+ # ==========================
2404
+
2405
+ def multiply(self, a, tau, c, lda=None, ldc=None) -> None:
2406
+ """
2407
+ Multiplies matrix C by the unitary matrix Q from an LQ factorization.
2408
+
2409
+ This device function computes:
2410
+ ``op(Q) * C`` (if ``side='left'``)
2411
+ ``C * op(Q)`` (if ``side='right'``)
2412
+
2413
+ where Q is the unitary matrix from the LQ factorization, represented
2414
+ by Householder vectors stored in A and the tau array.
2415
+ Uses cuSOLVERDx ``'unmlq'``. The result overwrites matrix C.
2416
+
2417
+ If ``lda`` and ``ldc`` are provided, uses runtime version with the
2418
+ specified leading dimensions. If not provided (``None``), uses compile-time
2419
+ version with default or constructor-provided leading dimensions.
2420
+
2421
+ For more details, see: :cusolverdx_doc:`get_started/functions/unmlq.html`
2422
+
2423
+ Args:
2424
+ a: Pointer to an array in shared memory, storing the batched matrix containing
2425
+ Householder vectors from the LQ factorization, according to the specified
2426
+ arrangement and leading dimension (see :meth:`__init__`).
2427
+ The elements above the diagonal of A, with the array tau, represent the
2428
+ unitary matrix Q as a product of Householder reflections.
2429
+ If ``side='left'``, A is ``K`` x ``M``.
2430
+ If ``side='right'``, A is ``K`` x ``N``.
2431
+ tau: Pointer to a 1D array of size K for each batch, containing the scalar
2432
+ factors of the Householder reflections from the LQ factorization.
2433
+ The tau array, together with the Householder vectors in A, defines Q.
2434
+ c: Pointer to an array in shared memory,
2435
+ storing the batched ``M`` x ``N`` matrix
2436
+ according to the specified arrangement
2437
+ and leading dimension (see :meth:`__init__`).
2438
+ The operation is in-place: result overwrites C.
2439
+ lda: Optional runtime leading dimension for matrix A.
2440
+ The ``lda`` and ``ldc`` must be specified together.
2441
+ If not specified, the compile-time ``lda`` is used.
2442
+ ldc: Optional runtime leading dimension for matrix C.
2443
+ The ``lda`` and ``ldc`` must be specified together.
2444
+ If not specified, the compile-time ``ldc`` is used.
2445
+ """
2446
+ raise RuntimeError("multiply is a device function and cannot be called on the host.")
2447
+
2448
+
2449
+ # ==========================
2450
+ # Least Squares Solver
2451
+ # ==========================
2452
+
2453
+ LEAST_SQUARES_SOLVER_DOCSTRING = SOLVER_DOCSTRING.copy()
2454
+ del LEAST_SQUARES_SOLVER_DOCSTRING["function"]
2455
+ del LEAST_SQUARES_SOLVER_DOCSTRING["fill_mode"]
2456
+ del LEAST_SQUARES_SOLVER_DOCSTRING["side"]
2457
+ del LEAST_SQUARES_SOLVER_DOCSTRING["diag"]
2458
+
2459
+ LEAST_SQUARES_SOLVER_DOCSTRING["size"] = f"""{SOLVER_DOCSTRING_SIZE_BASE}
2460
+ ``M`` and ``N`` represent the dimensions of matrix A (``M`` x ``N``).
2461
+ ``K`` represents the number of columns in the right-hand side matrix B.""".replace("\n", " ")
2462
+
2463
+ LEAST_SQUARES_SOLVER_DOCSTRING["transpose_mode"] = f"""Transpose mode for operation op(A) applied to matrix A.
2464
+ Can be one of: {", ".join(f"``'{v}'``" for v in ALLOWED_TRANSPOSE_MODE)}.
2465
+ Defaults to ``'non_transposed'``.""".replace("\n", " ")
2466
+
2467
+ LEAST_SQUARES_SOLVER_DOCSTRING["leading_dimensions"] = SOLVER_DOCSTRING["leading_dimensions"] + ADAPTERS_API_LD_DOCSTRING
2468
+
2469
+
2470
+ @docstring_decorator(LEAST_SQUARES_SOLVER_DOCSTRING, skip_missing=False)
2471
+ class LeastSquaresSolver(_SolverProperties):
2472
+ """
2473
+ A class that encapsulates least squares solver device function (``'gels'``).
2474
+ GELS (GEneral Least Square) solves overdetermined
2475
+ or underdetermined least squares problems:
2476
+
2477
+ .. math::
2478
+ \\min \\| op(A) * X - B \\|_2
2479
+
2480
+ using the QR or LQ factorization of A, and overwriting B with the solution X.
2481
+
2482
+ The configurations supported by GELS are:
2483
+
2484
+ **Memory layout requirements:**
2485
+
2486
+ Matrices must be stored in shared memory according
2487
+ to their arrangement and leading dimension (ld):
2488
+
2489
+ **For matrix A (M x N):**
2490
+
2491
+ * **Column-major arrangement**: Matrix shape ``(batches_per_block, M, N)``
2492
+ with strides ``(lda * N, 1, lda)``
2493
+ * **Row-major arrangement**: Matrix shape ``(batches_per_block, M, N)``
2494
+ with strides ``(lda * M, lda, 1)``
2495
+
2496
+ **Matrix B and X** are stored in the same buffer (B is overwritten by X in place).
2497
+ The buffer has shape ``(batches_per_block, max(M, N), K)``. The second leading
2498
+ dimension ``ldb`` refers to this shared B/X buffer
2499
+ and must satisfy ``ldb >= max(M, N)``.
2500
+ Logical shapes differ by :attr:`transpose_mode`.
2501
+ Use :attr:`b_shape` and :attr:`x_shape` for B and X.
2502
+
2503
+ **Matrix B** (right-hand side):
2504
+
2505
+ * **Logical shape**: :attr:`b_shape` - ``(batches_per_block, M, K)`` if non-transposed,
2506
+ ``(batches_per_block, N, K)`` if transposed.
2507
+ * **Column-major arrangement**: strides ``(ldb * K, 1, ldb)``.
2508
+ * **Row-major arrangement**: strides ``(ldb * max(M, N), ldb, 1)``.
2509
+
2510
+ **Matrix X** (solution):
2511
+
2512
+ * **Logical shape**: :attr:`x_shape` - ``(batches_per_block, N, K)`` if non-transposed,
2513
+ ``(batches_per_block, M, K)`` if transposed.
2514
+ * **Column-major arrangement**: strides ``(ldb * K, 1, ldb)``.
2515
+ * **Row-major arrangement**: strides ``(ldb * max(M, N), ldb, 1)``.
2516
+
2517
+ .. note::
2518
+ GELS is an in-place function. Matrix A is overwritten by the QR or LQ factorization,
2519
+ and matrix B is overwritten by the solution X. Both B and X use the single buffer
2520
+ of shape ``(max(M, N), K)`` per batch.
2521
+
2522
+ Args:
2523
+ size (Sequence[int]): {size}
2524
+
2525
+ precision (type[np.floating]): {precision}
2526
+
2527
+ execution (str): {execution}
2528
+
2529
+ sm (ComputeCapability): {sm}
2530
+
2531
+ transpose_mode (str, optional): {transpose_mode}
2532
+
2533
+ arrangement (Sequence[str], optional): {arrangement}
2534
+
2535
+ batches_per_block (int | Literal["suggested"], optional): {batches_per_block}
2536
+
2537
+ data_type (str, optional): {data_type}
2538
+
2539
+ leading_dimensions (Sequence[int], optional): {leading_dimensions}
2540
+
2541
+ block_dim (Sequence[int] | Literal["suggested"], optional): {block_dim}
2542
+
2543
+ See Also:
2544
+ For further details, please refer to the cuSOLVERDx documentation:
2545
+
2546
+ * :cusolverdx_doc:`gels <get_started/functions/gels.html>`
2547
+ """
2548
+
2549
+ # ==========================
2550
+ # Constructor
2551
+ # ==========================
2552
+
2553
+ def __init__(
2554
+ self,
2555
+ size: Sequence[int],
2556
+ precision: type[np.floating],
2557
+ execution: str,
2558
+ *,
2559
+ sm=None,
2560
+ transpose_mode: str = "non_transposed",
2561
+ arrangement: Sequence[str] | None = None,
2562
+ batches_per_block: int | Literal["suggested"] | None = None,
2563
+ data_type: str | None = None,
2564
+ leading_dimensions: Sequence[int] | None = None,
2565
+ block_dim: Sequence[int] | Literal["suggested"] | None = None,
2566
+ ):
2567
+ self._solve: Solver = Solver(
2568
+ function="gels",
2569
+ size=size,
2570
+ precision=precision,
2571
+ execution=execution,
2572
+ sm=sm,
2573
+ transpose_mode=transpose_mode,
2574
+ arrangement=arrangement,
2575
+ batches_per_block=batches_per_block,
2576
+ data_type=data_type,
2577
+ leading_dimensions=leading_dimensions,
2578
+ block_dim=block_dim,
2579
+ )
2580
+
2581
+ super().__init__(self._solve)
2582
+
2583
+ # ==========================
2584
+ # Property methods
2585
+ # ==========================
2586
+
2587
+ @property
2588
+ def transpose_mode(self) -> str:
2589
+ assert self._solve.transpose_mode is not None
2590
+ return self._solve.transpose_mode
2591
+
2592
+ @property
2593
+ def tau_type(self) -> type[np.floating] | Complex:
2594
+ return self._solve.tau_type
2595
+
2596
+ @property
2597
+ def tau_shape(self) -> tuple[int, int]:
2598
+ return (self.batches_per_block, min(self.m, self.n))
2599
+
2600
+ @property
2601
+ def tau_strides(self) -> tuple[int, int]:
2602
+ return (self.tau_shape[1], 1)
2603
+
2604
+ @property
2605
+ def a_shape(self) -> tuple[int, int, int]:
2606
+ return (self.batches_per_block, self.m, self.n)
2607
+
2608
+ @property
2609
+ def b_shape(self) -> tuple[int, int, int]:
2610
+ return (self.batches_per_block, self.m if self.transpose_mode == "non_transposed" else self.n, self.k)
2611
+
2612
+ @property
2613
+ def x_shape(self) -> tuple[int, int, int]:
2614
+ return (self.batches_per_block, self.n if self.transpose_mode == "non_transposed" else self.m, self.k)
2615
+
2616
+ def a_strides(self, *, lda: int | None = None) -> tuple[int, int, int]:
2617
+ lda = self.lda if lda is None else lda
2618
+ return _calculate_strides(self.a_shape[1:], lda, self.a_arrangement)
2619
+
2620
+ def bx_strides(self, *, ldb: int | None = None) -> tuple[int, int, int]:
2621
+ ldb = self.ldb if ldb is None else ldb
2622
+ return _calculate_strides((max(self.m, self.n), self.k), ldb, self.b_arrangement)
2623
+
2624
+ def a_size(self, *, lda: int | None = None) -> int:
2625
+ return self.a_strides(lda=lda)[0] * self.a_shape[0]
2626
+
2627
+ def bx_size(self, *, ldb: int | None = None) -> int:
2628
+ return self.bx_strides(ldb=ldb)[0] * self.batches_per_block
2629
+
2630
+ @property
2631
+ def tau_size(self) -> int:
2632
+ return self.tau_shape[0] * self.tau_shape[1]
2633
+
2634
+ # ==========================
2635
+ # Device function methods
2636
+ # ==========================
2637
+
2638
+ def solve(self, a, tau, b, lda=None, ldb=None) -> None:
2639
+ """
2640
+ Solves a least squares problem using QR or LQ factorization.
2641
+
2642
+ This device function solves:
2643
+
2644
+ .. math::
2645
+ \\min \\| op(A) * X - B \\|_2
2646
+
2647
+ using the QR or LQ factorization of A, and overwrites B with the solution X.
2648
+ Uses cuSOLVERDx ``'gels'``. The operation is in-place: matrix A is overwritten
2649
+ by the factorization, and matrix B is overwritten by the solution X.
2650
+
2651
+ If ``lda`` and ``ldb`` are provided, uses runtime version with the
2652
+ specified leading dimensions. If not provided (``None``), uses compile-time
2653
+ version with default or constructor-provided leading dimensions.
2654
+
2655
+ .. note::
2656
+ The choice between QR and LQ factorization depends on the problem dimensions
2657
+ and transpose mode:
2658
+
2659
+ * If ``op(A)`` is ``'non_transposed'`` and ``M >= N``: uses QR factorization
2660
+ * If ``op(A)`` is ``'non_transposed'`` and ``M < N``: uses LQ factorization
2661
+ * If ``op(A)`` is ``'transposed'`` or ``'conj_transposed'``
2662
+ and ``M >= N``: uses LQ factorization
2663
+ * If ``op(A)`` is ``'transposed'`` or ``'conj_transposed'``
2664
+ and ``M < N``: uses QR factorization
2665
+
2666
+ For more details, see: :cusolverdx_doc:`get_started/functions/gels.html`
2667
+
2668
+ Args:
2669
+ a: Pointer to an array in shared memory, storing the batched matrix
2670
+ according to the specified arrangement
2671
+ and leading dimension (see :meth:`__init__`).
2672
+ The matrix is overwritten in place by the QR or LQ factorization.
2673
+ tau: Pointer to a 1D array of size min(M, N) for each batch.
2674
+ Contains the scalar factors of the Householder reflections.
2675
+ The tau array, together with the Householder vectors stored in A,
2676
+ defines the unitary matrix Q.
2677
+ b: Pointer to an array in shared memory,
2678
+ storing the batched right-hand side matrix
2679
+ according to the specified arrangement
2680
+ and leading dimension (see :meth:`__init__`).
2681
+ The storage size is ``max(M, N) x K`` per batch.
2682
+ The operation is in-place: result X overwrites B.
2683
+ lda: Optional runtime leading dimension for matrix A.
2684
+ The ``lda`` and ``ldb`` must be specified together.
2685
+ If not specified, the compile-time ``lda`` is used.
2686
+ ldb: Optional runtime leading dimension for matrix B.
2687
+ The ``lda`` and ``ldb`` must be specified together.
2688
+ If not specified, the compile-time ``ldb`` is used.
2689
+ """
2690
+ raise RuntimeError("solve is a device function and cannot be called on the host.")