nvmath-python 1.0.0__cp314-cp314t-win_amd64.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (332) hide show
  1. nvmath/__init__.pxd +0 -0
  2. nvmath/__init__.py +45 -0
  3. nvmath/_internal/__init__.py +0 -0
  4. nvmath/_internal/attribute_ifc_factory.py +330 -0
  5. nvmath/_internal/layout.py +70 -0
  6. nvmath/_internal/templates.py +130 -0
  7. nvmath/_internal/threadsafe.py +106 -0
  8. nvmath/_internal/utils.py +43 -0
  9. nvmath/_internal/workspace.py +490 -0
  10. nvmath/_utils.py +147 -0
  11. nvmath/bindings/__init__.py +60 -0
  12. nvmath/bindings/_internal/__init__.pxd +0 -0
  13. nvmath/bindings/_internal/__init__.py +0 -0
  14. nvmath/bindings/_internal/common_types.pxd +31 -0
  15. nvmath/bindings/_internal/cublas.cp314t-win_amd64.pyd +0 -0
  16. nvmath/bindings/_internal/cublas.pxd +530 -0
  17. nvmath/bindings/_internal/cublasLt.cp314t-win_amd64.pyd +0 -0
  18. nvmath/bindings/_internal/cublasLt.pxd +59 -0
  19. nvmath/bindings/_internal/cublasMp.pxd +52 -0
  20. nvmath/bindings/_internal/cudss.cp314t-win_amd64.pyd +0 -0
  21. nvmath/bindings/_internal/cudss.pxd +54 -0
  22. nvmath/bindings/_internal/cufft.cp314t-win_amd64.pyd +0 -0
  23. nvmath/bindings/_internal/cufft.pxd +70 -0
  24. nvmath/bindings/_internal/cufftMp.pxd +77 -0
  25. nvmath/bindings/_internal/curand.cp314t-win_amd64.pyd +0 -0
  26. nvmath/bindings/_internal/curand.pxd +42 -0
  27. nvmath/bindings/_internal/cusolver.cp314t-win_amd64.pyd +0 -0
  28. nvmath/bindings/_internal/cusolver.pxd +15 -0
  29. nvmath/bindings/_internal/cusolverDn.cp314t-win_amd64.pyd +0 -0
  30. nvmath/bindings/_internal/cusolverDn.pxd +406 -0
  31. nvmath/bindings/_internal/cusolverMp.pxd +71 -0
  32. nvmath/bindings/_internal/cusolverSp.cp314t-win_amd64.pyd +0 -0
  33. nvmath/bindings/_internal/cusolverSp.pxd +75 -0
  34. nvmath/bindings/_internal/cusparse.cp314t-win_amd64.pyd +0 -0
  35. nvmath/bindings/_internal/cusparse.pxd +471 -0
  36. nvmath/bindings/_internal/cusparseLt.cp314t-win_amd64.pyd +0 -0
  37. nvmath/bindings/_internal/cusparseLt.pxd +48 -0
  38. nvmath/bindings/_internal/cutensor.cp314t-win_amd64.pyd +0 -0
  39. nvmath/bindings/_internal/cutensor.pxd +58 -0
  40. nvmath/bindings/_internal/mathdx.cp314t-win_amd64.pyd +0 -0
  41. nvmath/bindings/_internal/mathdx.pxd +116 -0
  42. nvmath/bindings/_internal/nvshmem.pxd +29 -0
  43. nvmath/bindings/_internal/utils.cp314t-win_amd64.pyd +0 -0
  44. nvmath/bindings/_internal/utils.pxd +174 -0
  45. nvmath/bindings/_internal/utils.pyi +10 -0
  46. nvmath/bindings/cublas.cp314t-win_amd64.pyd +0 -0
  47. nvmath/bindings/cublas.pxd +558 -0
  48. nvmath/bindings/cublas.pyi +812 -0
  49. nvmath/bindings/cublasLt.cp314t-win_amd64.pyd +0 -0
  50. nvmath/bindings/cublasLt.pxd +109 -0
  51. nvmath/bindings/cublasLt.pyi +1461 -0
  52. nvmath/bindings/cublasMp.pxd +85 -0
  53. nvmath/bindings/cublasMp.pyi +267 -0
  54. nvmath/bindings/cudss.cp314t-win_amd64.pyd +0 -0
  55. nvmath/bindings/cudss.pxd +98 -0
  56. nvmath/bindings/cudss.pyi +443 -0
  57. nvmath/bindings/cufft.cp314t-win_amd64.pyd +0 -0
  58. nvmath/bindings/cufft.pxd +118 -0
  59. nvmath/bindings/cufft.pyi +301 -0
  60. nvmath/bindings/cufftMp.pxd +124 -0
  61. nvmath/bindings/cufftMp.pyi +326 -0
  62. nvmath/bindings/curand.cp314t-win_amd64.pyd +0 -0
  63. nvmath/bindings/curand.pxd +71 -0
  64. nvmath/bindings/curand.pyi +189 -0
  65. nvmath/bindings/cusolver.cp314t-win_amd64.pyd +0 -0
  66. nvmath/bindings/cusolver.pxd +62 -0
  67. nvmath/bindings/cusolver.pyi +320 -0
  68. nvmath/bindings/cusolverDn.cp314t-win_amd64.pyd +0 -0
  69. nvmath/bindings/cusolverDn.pxd +430 -0
  70. nvmath/bindings/cusolverDn.pyi +422 -0
  71. nvmath/bindings/cusolverMp.pxd +98 -0
  72. nvmath/bindings/cusolverMp.pyi +114 -0
  73. nvmath/bindings/cusolverSp.cp314t-win_amd64.pyd +0 -0
  74. nvmath/bindings/cusolverSp.pxd +95 -0
  75. nvmath/bindings/cusolverSp.pyi +70 -0
  76. nvmath/bindings/cusparse.cp314t-win_amd64.pyd +0 -0
  77. nvmath/bindings/cusparse.pxd +546 -0
  78. nvmath/bindings/cusparse.pyi +1017 -0
  79. nvmath/bindings/cusparseLt.cp314t-win_amd64.pyd +0 -0
  80. nvmath/bindings/cusparseLt.pxd +99 -0
  81. nvmath/bindings/cusparseLt.pyi +252 -0
  82. nvmath/bindings/cutensor.cp314t-win_amd64.pyd +0 -0
  83. nvmath/bindings/cutensor.pxd +98 -0
  84. nvmath/bindings/cutensor.pyi +324 -0
  85. nvmath/bindings/cycublas.cp314t-win_amd64.pyd +0 -0
  86. nvmath/bindings/cycublas.pxd +664 -0
  87. nvmath/bindings/cycublasLt.cp314t-win_amd64.pyd +0 -0
  88. nvmath/bindings/cycublasLt.pxd +1045 -0
  89. nvmath/bindings/cycublasMp.pxd +171 -0
  90. nvmath/bindings/cycudss.cp314t-win_amd64.pyd +0 -0
  91. nvmath/bindings/cycudss.pxd +277 -0
  92. nvmath/bindings/cycufft.cp314t-win_amd64.pyd +0 -0
  93. nvmath/bindings/cycufft.pxd +333 -0
  94. nvmath/bindings/cycufftMp.pxd +342 -0
  95. nvmath/bindings/cycurand.cp314t-win_amd64.pyd +0 -0
  96. nvmath/bindings/cycurand.pxd +141 -0
  97. nvmath/bindings/cycusolver.cp314t-win_amd64.pyd +0 -0
  98. nvmath/bindings/cycusolver.pxd +137 -0
  99. nvmath/bindings/cycusolverDn.cp314t-win_amd64.pyd +0 -0
  100. nvmath/bindings/cycusolverDn.pxd +443 -0
  101. nvmath/bindings/cycusolverMp.pxd +107 -0
  102. nvmath/bindings/cycusolverSp.cp314t-win_amd64.pyd +0 -0
  103. nvmath/bindings/cycusolverSp.pxd +93 -0
  104. nvmath/bindings/cycusparse.cp314t-win_amd64.pyd +0 -0
  105. nvmath/bindings/cycusparse.pxd +679 -0
  106. nvmath/bindings/cycusparseLt.cp314t-win_amd64.pyd +0 -0
  107. nvmath/bindings/cycusparseLt.pxd +135 -0
  108. nvmath/bindings/cycutensor.cp314t-win_amd64.pyd +0 -0
  109. nvmath/bindings/cycutensor.pxd +189 -0
  110. nvmath/bindings/cymathdx.cp314t-win_amd64.pyd +0 -0
  111. nvmath/bindings/cymathdx.pxd +552 -0
  112. nvmath/bindings/cynvshmem.pxd +118 -0
  113. nvmath/bindings/mathdx.cp314t-win_amd64.pyd +0 -0
  114. nvmath/bindings/mathdx.pxd +182 -0
  115. nvmath/bindings/mathdx.pyi +1562 -0
  116. nvmath/bindings/nvpl/__init__.pxd +0 -0
  117. nvmath/bindings/nvpl/__init__.py +13 -0
  118. nvmath/bindings/nvpl/_internal/__init__.pxd +0 -0
  119. nvmath/bindings/nvpl/_internal/__init__.py +0 -0
  120. nvmath/bindings/nvpl/_internal/blas.cp314t-win_amd64.pyd +0 -0
  121. nvmath/bindings/nvpl/_internal/blas.pxd +237 -0
  122. nvmath/bindings/nvpl/_internal/fft.cp314t-win_amd64.pyd +0 -0
  123. nvmath/bindings/nvpl/_internal/fft.pxd +36 -0
  124. nvmath/bindings/nvpl/blas.cp314t-win_amd64.pyd +0 -0
  125. nvmath/bindings/nvpl/blas.pxd +131 -0
  126. nvmath/bindings/nvpl/blas.pyi +168 -0
  127. nvmath/bindings/nvpl/cyblas.cp314t-win_amd64.pyd +0 -0
  128. nvmath/bindings/nvpl/cyblas.pxd +280 -0
  129. nvmath/bindings/nvpl/cyfft.cp314t-win_amd64.pyd +0 -0
  130. nvmath/bindings/nvpl/cyfft.pxd +93 -0
  131. nvmath/bindings/nvpl/fft.cp314t-win_amd64.pyd +0 -0
  132. nvmath/bindings/nvpl/fft.pxd +100 -0
  133. nvmath/bindings/nvpl/fft.pyi +168 -0
  134. nvmath/bindings/nvshmem.pxd +54 -0
  135. nvmath/bindings/nvshmem.pyi +191 -0
  136. nvmath/device/__init__.py +38 -0
  137. nvmath/device/_deprecated.py +33 -0
  138. nvmath/device/common.py +315 -0
  139. nvmath/device/common_backend.py +131 -0
  140. nvmath/device/common_cuda.py +201 -0
  141. nvmath/device/common_numba.py +300 -0
  142. nvmath/device/common_numba_cuda_mlir.py +202 -0
  143. nvmath/device/common_opaque_tensor.py +201 -0
  144. nvmath/device/cublasdx.py +1606 -0
  145. nvmath/device/cublasdx_backend.py +860 -0
  146. nvmath/device/cublasdx_numba.py +1534 -0
  147. nvmath/device/cublasdx_numba_cuda_mlir.py +208 -0
  148. nvmath/device/cufftdx.py +373 -0
  149. nvmath/device/cufftdx_backend.py +220 -0
  150. nvmath/device/cufftdx_numba.py +140 -0
  151. nvmath/device/cufftdx_numba_cuda_mlir.py +79 -0
  152. nvmath/device/curand_kernel.py +9147 -0
  153. nvmath/device/cusolverdx.py +2708 -0
  154. nvmath/device/cusolverdx_backend.py +440 -0
  155. nvmath/device/cusolverdx_numba.py +567 -0
  156. nvmath/device/cusolverdx_numba_cuda_mlir.py +604 -0
  157. nvmath/device/cusolverdx_overload_backend.py +1029 -0
  158. nvmath/device/llvm_array.py +29 -0
  159. nvmath/device/random.py +441 -0
  160. nvmath/device/random_helpers.py +23 -0
  161. nvmath/device/random_states.py +187 -0
  162. nvmath/device/types.py +138 -0
  163. nvmath/device/vector_types_numba.py +259 -0
  164. nvmath/distributed/__init__.py +200 -0
  165. nvmath/distributed/_internal/__init__.py +0 -0
  166. nvmath/distributed/_internal/nccl.py +86 -0
  167. nvmath/distributed/_internal/nvshmem.py +307 -0
  168. nvmath/distributed/_internal/symmetric_memory.py +35 -0
  169. nvmath/distributed/_internal/tensor_ifc.py +70 -0
  170. nvmath/distributed/_internal/tensor_ifc_cupy.py +68 -0
  171. nvmath/distributed/_internal/tensor_ifc_host_device.py +172 -0
  172. nvmath/distributed/_internal/tensor_ifc_numpy.py +46 -0
  173. nvmath/distributed/_internal/tensor_ifc_torch.py +162 -0
  174. nvmath/distributed/_internal/tensor_wrapper.py +81 -0
  175. nvmath/distributed/_utils.py +167 -0
  176. nvmath/distributed/distribution/__init__.py +30 -0
  177. nvmath/distributed/distribution/_configuration.py +39 -0
  178. nvmath/distributed/distribution/distributions.py +1024 -0
  179. nvmath/distributed/distribution/redistribute.py +1284 -0
  180. nvmath/distributed/fft/__init__.py +7 -0
  181. nvmath/distributed/fft/_configuration.py +82 -0
  182. nvmath/distributed/fft/fft.py +2742 -0
  183. nvmath/distributed/linalg/__init__.py +22 -0
  184. nvmath/distributed/linalg/_internal/__init__.py +3 -0
  185. nvmath/distributed/linalg/_internal/epilog_protocol.py +586 -0
  186. nvmath/distributed/linalg/_internal/matmul_desc_ifc.py +28 -0
  187. nvmath/distributed/linalg/advanced/__init__.py +8 -0
  188. nvmath/distributed/linalg/advanced/_configuration.py +171 -0
  189. nvmath/distributed/linalg/advanced/matmulmod.py +3573 -0
  190. nvmath/distributed/linalg/generic/__init__.py +8 -0
  191. nvmath/distributed/linalg/generic/_caching.py +66 -0
  192. nvmath/distributed/linalg/generic/_configuration.py +61 -0
  193. nvmath/distributed/linalg/generic/_factorization.py +172 -0
  194. nvmath/distributed/linalg/generic/_initialization.py +966 -0
  195. nvmath/distributed/linalg/generic/_problem_spec.py +511 -0
  196. nvmath/distributed/linalg/generic/solvermod.py +1368 -0
  197. nvmath/distributed/process_group.py +408 -0
  198. nvmath/fft/__init__.py +7 -0
  199. nvmath/fft/_configuration.py +189 -0
  200. nvmath/fft/_exec_utils.py +82 -0
  201. nvmath/fft/_helpers.py +237 -0
  202. nvmath/fft/fft.py +3122 -0
  203. nvmath/internal/__init__.pxd +3 -0
  204. nvmath/internal/__init__.py +10 -0
  205. nvmath/internal/_bindings.cp314t-win_amd64.pyd +0 -0
  206. nvmath/internal/_bindings.pxd +18 -0
  207. nvmath/internal/_device_utils.py +45 -0
  208. nvmath/internal/_layout/__init__.pxd +3 -0
  209. nvmath/internal/_layout/__init__.py +7 -0
  210. nvmath/internal/_layout/_layout.cp314t-win_amd64.pyd +0 -0
  211. nvmath/internal/_layout/_layout.pxd +1303 -0
  212. nvmath/internal/_layout/_layout.pyi +1145 -0
  213. nvmath/internal/enum_utils.py +142 -0
  214. nvmath/internal/formatters.py +87 -0
  215. nvmath/internal/mem_limit.py +51 -0
  216. nvmath/internal/memory.cp314t-win_amd64.pyd +0 -0
  217. nvmath/internal/memory.pxd +13 -0
  218. nvmath/internal/memory.pyi +50 -0
  219. nvmath/internal/ndbuffer/__init__.pxd +3 -0
  220. nvmath/internal/ndbuffer/__init__.py +9 -0
  221. nvmath/internal/ndbuffer/_copy_kernel.cp314t-win_amd64.pyd +0 -0
  222. nvmath/internal/ndbuffer/_copy_kernel.pxd +10 -0
  223. nvmath/internal/ndbuffer/_jit.cp314t-win_amd64.pyd +0 -0
  224. nvmath/internal/ndbuffer/_jit.pxd +7 -0
  225. nvmath/internal/ndbuffer/_ndbuffer.cp314t-win_amd64.pyd +0 -0
  226. nvmath/internal/ndbuffer/_ndbuffer.pxd +40 -0
  227. nvmath/internal/ndbuffer/_ndbuffer.pyi +463 -0
  228. nvmath/internal/ndbuffer/copy_kernel/args.h +34 -0
  229. nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/array_view.h +52 -0
  230. nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/elementwise.h +68 -0
  231. nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/grid_indexer.h +69 -0
  232. nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/transposed.h +242 -0
  233. nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/type_utils.h +39 -0
  234. nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/utils.h +132 -0
  235. nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/vec.h +159 -0
  236. nvmath/internal/ndbuffer/copy_kernel/elementwise.h +53 -0
  237. nvmath/internal/ndbuffer/copy_kernel/transposed.h +58 -0
  238. nvmath/internal/package_ifc.py +168 -0
  239. nvmath/internal/package_ifc_cuda.py +57 -0
  240. nvmath/internal/package_ifc_cupy.py +67 -0
  241. nvmath/internal/package_ifc_torch.py +69 -0
  242. nvmath/internal/package_wrapper.py +14 -0
  243. nvmath/internal/tensor_ifc.py +179 -0
  244. nvmath/internal/tensor_ifc_cupy.py +234 -0
  245. nvmath/internal/tensor_ifc_ndbuffer.py +147 -0
  246. nvmath/internal/tensor_ifc_numpy.py +184 -0
  247. nvmath/internal/tensor_ifc_torch.py +178 -0
  248. nvmath/internal/tensor_wrapper.py +160 -0
  249. nvmath/internal/typemaps.py +113 -0
  250. nvmath/internal/utils.py +805 -0
  251. nvmath/linalg/__init__.py +56 -0
  252. nvmath/linalg/_internal/__init__.py +3 -0
  253. nvmath/linalg/_internal/algo_cap_ifc.py +82 -0
  254. nvmath/linalg/_internal/algo_config_ifc.py +43 -0
  255. nvmath/linalg/_internal/batch.py +234 -0
  256. nvmath/linalg/_internal/enum_to_tuples.py +64 -0
  257. nvmath/linalg/_internal/epilog_protocol.py +766 -0
  258. nvmath/linalg/_internal/layout.py +624 -0
  259. nvmath/linalg/_internal/matmul_desc_ifc.py +28 -0
  260. nvmath/linalg/_internal/matmul_pref_ifc.py +27 -0
  261. nvmath/linalg/_internal/matrix_layout_ifc.py +26 -0
  262. nvmath/linalg/_internal/solver_utils.py +432 -0
  263. nvmath/linalg/_internal/typemaps.py +144 -0
  264. nvmath/linalg/_internal/utils.py +157 -0
  265. nvmath/linalg/advanced/__init__.py +8 -0
  266. nvmath/linalg/advanced/_algorithmmod.py +170 -0
  267. nvmath/linalg/advanced/_configuration.py +351 -0
  268. nvmath/linalg/advanced/helpers/__init__.py +5 -0
  269. nvmath/linalg/advanced/helpers/matmul.py +1316 -0
  270. nvmath/linalg/advanced/matmulmod.py +3734 -0
  271. nvmath/linalg/generic/__init__.py +53 -0
  272. nvmath/linalg/generic/_configuration/__init__.py +39 -0
  273. nvmath/linalg/generic/_configuration/layout.py +263 -0
  274. nvmath/linalg/generic/_configuration/match.py +734 -0
  275. nvmath/linalg/generic/_configuration/qualifiers.py +493 -0
  276. nvmath/linalg/generic/_configuration/solver_configuration.py +59 -0
  277. nvmath/linalg/generic/_configuration/wrap.py +217 -0
  278. nvmath/linalg/generic/_dtype.py +15 -0
  279. nvmath/linalg/generic/matmulmod.py +2094 -0
  280. nvmath/linalg/generic/solvermod.py +1301 -0
  281. nvmath/memory.py +279 -0
  282. nvmath/sparse/__init__.py +38 -0
  283. nvmath/sparse/_internal/__init__.py +21 -0
  284. nvmath/sparse/_internal/common_utils.py +147 -0
  285. nvmath/sparse/_internal/cudss_config_ifc.py +702 -0
  286. nvmath/sparse/_internal/cudss_data_ifc.py +399 -0
  287. nvmath/sparse/_internal/cudss_utils.py +506 -0
  288. nvmath/sparse/_internal/cusparse_utils.py +382 -0
  289. nvmath/sparse/_internal/sparse_bsc_ifc.py +303 -0
  290. nvmath/sparse/_internal/sparse_bsr_ifc.py +305 -0
  291. nvmath/sparse/_internal/sparse_coo_ifc.py +256 -0
  292. nvmath/sparse/_internal/sparse_csc_ifc.py +268 -0
  293. nvmath/sparse/_internal/sparse_csr_ifc.py +288 -0
  294. nvmath/sparse/_internal/sparse_dia_ifc.py +242 -0
  295. nvmath/sparse/_internal/sparse_format_helpers.py +601 -0
  296. nvmath/sparse/_internal/sparse_tensor_ifc.py +133 -0
  297. nvmath/sparse/_internal/sparse_ust_ifc.py +141 -0
  298. nvmath/sparse/_internal/utils.py +56 -0
  299. nvmath/sparse/advanced/__init__.py +7 -0
  300. nvmath/sparse/advanced/_configuration.py +227 -0
  301. nvmath/sparse/advanced/direct_solver.py +2069 -0
  302. nvmath/sparse/generic/__init__.py +7 -0
  303. nvmath/sparse/generic/_configuration.py +129 -0
  304. nvmath/sparse/generic/_helpers.py +137 -0
  305. nvmath/sparse/generic/_thunks.py +21 -0
  306. nvmath/sparse/generic/matmulmod.py +2353 -0
  307. nvmath/sparse/ust/__init__.py +7 -0
  308. nvmath/sparse/ust/_converters.py +422 -0
  309. nvmath/sparse/ust/_cpp.py +28 -0
  310. nvmath/sparse/ust/_drawer.py +565 -0
  311. nvmath/sparse/ust/_emitter.py +1033 -0
  312. nvmath/sparse/ust/_jit.py +188 -0
  313. nvmath/sparse/ust/_kernel.py +282 -0
  314. nvmath/sparse/ust/_utils.py +149 -0
  315. nvmath/sparse/ust/interfaces/__init__.py +0 -0
  316. nvmath/sparse/ust/interfaces/torch_interface.py +476 -0
  317. nvmath/sparse/ust/tensor.py +1016 -0
  318. nvmath/sparse/ust/tensor_format.py +957 -0
  319. nvmath/tensor/__init__.py +6 -0
  320. nvmath/tensor/_configuration.py +120 -0
  321. nvmath/tensor/_internal/__init__.py +3 -0
  322. nvmath/tensor/_internal/cutensor_config_ifc.py +279 -0
  323. nvmath/tensor/_internal/cutensor_utils.py +230 -0
  324. nvmath/tensor/_internal/data.py +43 -0
  325. nvmath/tensor/_internal/einsum_parser.py +444 -0
  326. nvmath/tensor/_internal/typemaps.py +96 -0
  327. nvmath/tensor/contract.py +1900 -0
  328. nvmath_python-1.0.0.dist-info/METADATA +134 -0
  329. nvmath_python-1.0.0.dist-info/RECORD +332 -0
  330. nvmath_python-1.0.0.dist-info/WHEEL +5 -0
  331. nvmath_python-1.0.0.dist-info/licenses/LICENSE +177 -0
  332. nvmath_python-1.0.0.dist-info/top_level.txt +2 -0
@@ -0,0 +1,805 @@
1
+ # Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ """
6
+ A collection of (internal use) helper functions.
7
+ """
8
+
9
+ import contextlib
10
+ import functools
11
+ import time
12
+ import typing
13
+ from collections.abc import Callable, Sequence
14
+
15
+ import cuda.bindings.runtime as cbr
16
+ from cuda.core import Device, Event
17
+
18
+ from nvmath._internal.layout import is_contiguous_and_dense
19
+
20
+ from . import formatters, mem_limit, package_wrapper
21
+ from .package_ifc import AnyStream, Stream, StreamHolder
22
+ from .tensor_ifc import AnyTensor, Tensor, TensorHolder
23
+
24
+
25
+ def infer_object_package(obj: object) -> str:
26
+ """
27
+ Infer the package that defines this object.
28
+ """
29
+ module = obj.__class__.__module__
30
+ return module.split(".")[0]
31
+
32
+
33
+ """A generic type for Options classes"""
34
+ Options = typing.TypeVar("Options")
35
+
36
+
37
+ def check_or_create_options(
38
+ cls: type[Options], options: None | dict[str, typing.Any] | Options, options_description: str, *, keep_none=False
39
+ ) -> Options | None:
40
+ """
41
+ Create the specified options dataclass from a dictionary of options or None.
42
+ """
43
+
44
+ if options is None:
45
+ if keep_none:
46
+ return options
47
+ options = cls()
48
+ elif isinstance(options, dict):
49
+ options = cls(**options)
50
+
51
+ if not isinstance(options, cls):
52
+ raise TypeError(
53
+ f"The {options_description} must be provided as an object "
54
+ f"of type {cls.__name__} or as a dict with valid {options_description}. "
55
+ f"The provided object is '{options}'."
56
+ )
57
+
58
+ return options
59
+
60
+
61
+ def check_or_create_one_of_options(
62
+ clss: tuple[type, ...],
63
+ options: None | str | dict[str, typing.Any] | object,
64
+ options_description: str,
65
+ *,
66
+ cls_key: str = "name",
67
+ default_name: str | None = None,
68
+ ) -> typing.Any:
69
+ """
70
+ Create one of the specified options dataclasses by name or from a dictionary of options.
71
+ """
72
+
73
+ assert isinstance(clss, tuple)
74
+ assert all(hasattr(cls, "name") and isinstance(cls.name, str) for cls in clss)
75
+
76
+ if options is None:
77
+ assert default_name is not None, "Internal error in execution options"
78
+ options = _create_one_of_options_from_name(clss, default_name, options_description, cls_key=cls_key)
79
+ elif isinstance(options, str):
80
+ options = _create_one_of_options_from_name(clss, options, options_description, cls_key=cls_key)
81
+ elif isinstance(options, dict):
82
+ options = _create_one_of_options_from_dict(
83
+ clss,
84
+ options,
85
+ options_description,
86
+ cls_key=cls_key,
87
+ default_name=default_name,
88
+ )
89
+ else:
90
+ if not isinstance(options, clss):
91
+ _raise_invalid_one_of_options(clss, options, options_description, cls_key=cls_key)
92
+ return options
93
+
94
+
95
+ def _create_one_of_options_from_name(
96
+ clss: tuple[type[typing.Any], ...],
97
+ cls_name: str,
98
+ options_description: str,
99
+ *,
100
+ cls_key: str = "name",
101
+ ) -> typing.Any:
102
+ try:
103
+ _cls_name = cls_name.lower()
104
+ return next(cls for cls in clss if cls.name == _cls_name)()
105
+ except StopIteration:
106
+ raise _raise_invalid_one_of_options(clss, cls_name, options_description, cls_key=cls_key) from None
107
+
108
+
109
+ def _create_one_of_options_from_dict(
110
+ clss: tuple[type[typing.Any], ...],
111
+ options: dict[str, typing.Any],
112
+ options_description: str,
113
+ *,
114
+ cls_key: str = "name",
115
+ default_name: str | None = None,
116
+ ) -> typing.Any:
117
+ cls_name = options.get(cls_key)
118
+ if cls_name is None:
119
+ cls_name = default_name
120
+ if not isinstance(cls_name, str):
121
+ raise _raise_invalid_one_of_options(clss, cls_name, options_description, cls_key=cls_key)
122
+ try:
123
+ _cls_name = cls_name.lower()
124
+ cls = next(cls for cls in clss if cls.name == _cls_name)
125
+ except StopIteration:
126
+ raise _raise_invalid_one_of_options(clss, cls_name, options_description, cls_key=cls_key) from None
127
+ return cls(**{key: name for key, name in options.items() if key != cls_key})
128
+
129
+
130
+ def _raise_invalid_one_of_options(
131
+ clss: tuple[type[typing.Any], ...],
132
+ options: object,
133
+ options_description: str,
134
+ *,
135
+ cls_key: str = "name",
136
+ ) -> typing.NoReturn:
137
+ accepted_names = ", ".join(f"'{cls.name}'" for cls in clss)
138
+ accepted_types = ", ".join(str(cls) for cls in clss)
139
+ raise ValueError(
140
+ f"The {options_description} must be: \n"
141
+ f"1. an object of one of the following types {accepted_types} or\n"
142
+ f"2. a string with the type name ({accepted_names}) or\n"
143
+ f"3. a dict with a '{cls_key}' key and one of the type names {accepted_names}' as a value, "
144
+ f"and optional options valid for that type. \n"
145
+ f"The provided object is '{options}'."
146
+ )
147
+
148
+
149
+ class device_ctx:
150
+ """
151
+ NOTE, using classic class-based context manager here as it has lower overhead
152
+ than the generator-based context manager.
153
+
154
+ Semantics:
155
+
156
+ 1. The device context manager makes the specified device current from the point of entry
157
+ until the point of exit.
158
+
159
+ 2. When the context manager exits, the current device is reset to what it was when the
160
+ context manager was entered.
161
+
162
+ 3. Any explicit setting of the device within the context manager (using
163
+ cupy.cuda.Device().use(), torch.cuda.set_device(), etc) will overrule the device set
164
+ by the context manager from that point onwards till the context manager exits. In
165
+ other words, the context manager provides a local device scope and the current device
166
+ can be explicitly reset for the remainder of that scope.
167
+
168
+ 4. The context manager is single-use.
169
+
170
+ Corollary: if any library function resets the device globally and this is an undesired
171
+ side-effect, such functions must be called from within the device context manager.
172
+
173
+ Device context managers can be arbitrarily nested.
174
+ """
175
+
176
+ __slots__ = ("new_device_id", "_old_device")
177
+
178
+ def __init__(self, new_device_id):
179
+ self.new_device_id = new_device_id
180
+ self._old_device = None
181
+
182
+ def __enter__(self):
183
+ if self._old_device is not None:
184
+ raise RuntimeError("Reusing a device_ctx instance is not allowed.")
185
+ self._old_device = old_device = Device()
186
+ new_device_id = self.new_device_id
187
+ if old_device.device_id != new_device_id:
188
+ device = Device(new_device_id)
189
+ device.set_current()
190
+ return device
191
+ else:
192
+ return old_device
193
+
194
+ def __exit__(self, type, value, traceback):
195
+ # We should always restore the old device at exit.
196
+ self._old_device.set_current()
197
+
198
+
199
+ def is_hashable(obj: object) -> bool:
200
+ try:
201
+ hash(obj)
202
+ except TypeError:
203
+ return False
204
+ return True
205
+
206
+
207
+ @functools.lru_cache(maxsize=128)
208
+ def cached_get_or_create_stream(
209
+ device_id: int,
210
+ stream: Stream | int,
211
+ op_package: str,
212
+ ) -> StreamHolder[Stream]:
213
+ op_package_ifc = package_wrapper.PACKAGE[op_package]
214
+ if isinstance(stream, int):
215
+ ptr = stream
216
+ external = op_package_ifc.create_external_stream(device_id, ptr)
217
+ ctx = op_package_ifc.to_stream_context(external)
218
+ obj = op_package_ifc.create_stream(external, device_id)
219
+ return StreamHolder(ctx=ctx, device_id=device_id, external=external, obj=obj, ptr=ptr, package=op_package)
220
+
221
+ stream_package = infer_object_package(stream)
222
+ if stream_package != op_package:
223
+ message = (
224
+ f"The stream object must belong to the same package as the tensor network operands. "
225
+ f"Stream package: {stream_package}, Tensor package: {op_package}"
226
+ )
227
+ raise TypeError(message)
228
+
229
+ ctx = op_package_ifc.to_stream_context(stream)
230
+ ptr = op_package_ifc.to_stream_pointer(stream)
231
+ external = stream
232
+ obj = op_package_ifc.create_stream(stream, device_id)
233
+ return StreamHolder(ctx=ctx, device_id=device_id, external=external, obj=obj, ptr=ptr, package=op_package)
234
+
235
+
236
+ def get_or_create_stream(
237
+ device_id: int,
238
+ stream: Stream | int | None,
239
+ op_package: str,
240
+ ) -> StreamHolder[Stream]:
241
+ """
242
+ Create a stream object from a stream pointer or extract the stream pointer from a stream
243
+ object, or use the current stream.
244
+
245
+ Args:
246
+ device_id: The device ID.
247
+
248
+ stream: A stream object, stream pointer, or None.
249
+
250
+ op_package: The package the tensor network operands belong to.
251
+
252
+ Returns:
253
+ StreamHolder: Hold a CuPy stream object, package stream context, stream pointer, ...
254
+ """
255
+ op_package_ifc = package_wrapper.PACKAGE[op_package]
256
+ if stream is None:
257
+ stream = op_package_ifc.get_current_stream(device_id)
258
+ return cached_get_or_create_stream(device_id, stream, op_package)
259
+
260
+
261
+ @functools.lru_cache(maxsize=128)
262
+ def get_memory_limit_from_device_id(memory_limit: int | float | str, device_id: int) -> int:
263
+ with device_ctx(device_id):
264
+ status, _, total_memory = cbr.cudaMemGetInfo()
265
+ if status != 0 or total_memory is None:
266
+ raise RuntimeError(
267
+ f"cudaMemGetInfo failed with status {status}, total_memory={total_memory}. A possible cause is \
268
+ an inconsistent version of cuda-bindings for the CTK version used with nvmath-python."
269
+ )
270
+ return _get_memory_limit(memory_limit, total_memory)
271
+
272
+
273
+ def _get_memory_limit(memory_limit: int | float | str, total_memory: int) -> int:
274
+ """
275
+ Parse user provided memory limit and return the memory limit in bytes.
276
+ """
277
+ if isinstance(memory_limit, int):
278
+ if memory_limit < 0:
279
+ raise ValueError("The specified memory limit must be greater than or equal to 0.")
280
+ return memory_limit
281
+
282
+ if isinstance(memory_limit, float):
283
+ if memory_limit < 0:
284
+ raise ValueError("The specified memory limit must be greater than or equal to 0.")
285
+ if memory_limit <= 1.0:
286
+ memory_limit *= total_memory
287
+ return int(memory_limit)
288
+
289
+ m = mem_limit.MEM_LIMIT_RE_PCT.match(memory_limit)
290
+ if m:
291
+ factor = float(m.group(1))
292
+ if factor <= 0 or factor > 100:
293
+ raise ValueError("The memory limit percentage must be in the range [0, 100].")
294
+ return int(factor * total_memory / 100.0)
295
+
296
+ m = mem_limit.MEM_LIMIT_RE_VAL.match(memory_limit)
297
+ if not m:
298
+ raise ValueError(mem_limit.MEM_LIMIT_DOC.format(kind="memory limit", value=memory_limit))
299
+
300
+ base = 1000
301
+ if m.group("binary"):
302
+ base = 1024
303
+
304
+ powers = {"": 0, "k": 1, "m": 2, "g": 3}
305
+ unit = m.group("units").lower() if m.group("units") else ""
306
+ multiplier = base ** powers[unit]
307
+
308
+ value = float(m.group("value"))
309
+ memory_limit = int(value * multiplier)
310
+
311
+ return memory_limit
312
+
313
+
314
+ def create_empty_tensor(
315
+ cls: type[TensorHolder[Tensor]],
316
+ extents: Sequence[int],
317
+ dtype: str,
318
+ device_id: int | typing.Literal["cpu"],
319
+ stream_holder: StreamHolder | None,
320
+ verify_strides: bool,
321
+ strides: Sequence[int] | None = None,
322
+ **context,
323
+ ) -> TensorHolder[Tensor]:
324
+ """
325
+ Create a wrapped tensor of the same type as (the wrapped) cls on the specified device
326
+ having the specified extents and dtype.
327
+
328
+ The tensor is created within a stream context to allow for asynchronous memory
329
+ allocators like CuPy's MemoryAsyncPool.
330
+
331
+ Note, the function assumes the `strides` are dense (possibly permuted).
332
+ Otherwise, the behaviour is framework specific and tensor creation may fail
333
+ or created tensor may be corrupted. Set `verify_strides` to True to check
334
+ the layout and drop the strides if the layout is not dense.
335
+ """
336
+ assert isinstance(device_id, int) or device_id == "cpu", f"Internal Error: device_id must be int or 'cpu'; not {device_id}"
337
+ # if device id is none the stream holder must be too
338
+ assert device_id != "cpu" or stream_holder is None
339
+ if strides is not None and verify_strides and not is_contiguous_and_dense(extents, strides):
340
+ strides = None
341
+ tensor = cls.empty(extents, dtype=dtype, device_id=device_id, strides=strides, stream_holder=stream_holder, **context)
342
+ return tensor
343
+
344
+
345
+ def get_operands_device_id(operands: Sequence[TensorHolder[AnyTensor]]):
346
+ """
347
+ Return the id (ordinal) of the device the operands are on, or None if it is on the CPU.
348
+ """
349
+ device_id = operands[0].device_id
350
+ if not all(operand.device_id == device_id for operand in operands):
351
+ devices = {operand.device_id for operand in operands}
352
+ raise ValueError(f"All operands are not on the same device. Devices = {devices}.")
353
+
354
+ return device_id
355
+
356
+
357
+ def get_operands_dtype(operands: Sequence[TensorHolder[AnyTensor]]):
358
+ """
359
+ Return the data type name of the tensors.
360
+ """
361
+ dtype = operands[0].dtype
362
+ if not all(operand.dtype == dtype for operand in operands[1:]):
363
+ dtypes = {operand.dtype for operand in operands}
364
+ raise ValueError(f"All tensors must have the same data type. Data types found = {dtypes}.")
365
+ return dtype
366
+
367
+
368
+ def get_operands_package(operands: Sequence[TensorHolder[AnyTensor]]) -> str:
369
+ """
370
+ Return the package name of the tensors.
371
+ """
372
+ package = infer_object_package(operands[0].tensor)
373
+ if not all(infer_object_package(operand.tensor) == package for operand in operands[1:]):
374
+ packages = {infer_object_package(operand.tensor) for operand in operands}
375
+ raise TypeError(f"All tensors must be from the same library package. Packages found = {packages}.")
376
+ return package
377
+
378
+
379
+ def check_operands_match(
380
+ orig_operands: Sequence[object] | object, new_operands: Sequence[object] | object, attribute: str, description: str
381
+ ):
382
+ """
383
+ Check if the specified attribute matches between the corresponding new and old operands,
384
+ and raise an exception if it doesn't.
385
+ """
386
+ if isinstance(orig_operands, Sequence):
387
+ assert isinstance(new_operands, Sequence), "Internal Error. Cannot compare operand sequence with non-sequence."
388
+ checks: list[bool] = [
389
+ getattr(o, attribute) == getattr(n, attribute) for o, n in zip(orig_operands, new_operands, strict=True)
390
+ ]
391
+
392
+ if not all(checks):
393
+ mismatch = [
394
+ f"{location}: {getattr(orig_operands[location], attribute)} => {getattr(new_operands[location], attribute)}"
395
+ for location, predicate in enumerate(checks)
396
+ if predicate is False
397
+ ]
398
+ mismatch = formatters.array2string(mismatch)
399
+ message = f"""\
400
+ The {description} of each new operand must match the {description} of the corresponding original operand.
401
+
402
+ The mismatch in {description} as a sequence of "position: original {description} => new {description}" is: \n{mismatch}"""
403
+ raise ValueError(message)
404
+ else:
405
+ assert isinstance(new_operands, object)
406
+ check: bool = getattr(orig_operands, attribute) == getattr(new_operands, attribute)
407
+ if not check:
408
+ message = f"""The {description} of the new operand must match the {description} of the original operand."""
409
+ raise ValueError(message)
410
+
411
+
412
+ def check_attribute_match(orig_attribute: object, new_attribute: object, description: str):
413
+ """
414
+ Check if the specified attribute matches between the corresponding new and old operands,
415
+ and raise an exception if it doesn't.
416
+ """
417
+ check = orig_attribute == new_attribute
418
+ if not check:
419
+ message = f"""The {description} of the new operand must match the {description} of the original operand."""
420
+ raise ValueError(message)
421
+
422
+
423
+ def get_ptr_from_memory_pointer(mem_ptr):
424
+ """
425
+ Access the value associated with one of the attributes 'device_ptr', 'device_pointer',
426
+ 'ptr'.
427
+
428
+ This is a compatibility function so that external memory pointer implementations from
429
+ numba and cupy can be used instead of MemoryPointer from nvmath.memory.
430
+ """
431
+ attributes = (
432
+ "device_ptr", # nvmath.memory
433
+ "device_pointer",
434
+ "ptr", # cupy.cuda.MemoryPointer
435
+ "pointer", # numba.cuda.MemoryPointer
436
+ )
437
+ for attr in attributes:
438
+ if hasattr(mem_ptr, attr):
439
+ return getattr(mem_ptr, attr)
440
+
441
+ message = f"Memory pointer objects should have one of the following attributes specifying the device pointer: {attributes}"
442
+ raise AttributeError(message)
443
+
444
+
445
+ class Value:
446
+ """
447
+ A simple value wrapper holding a default value.
448
+ """
449
+
450
+ def __init__(self, default, *, validator: Callable[[object], bool]):
451
+ """
452
+ Args:
453
+ default: The default value to use.
454
+ validator: A callable that validates the provided value.
455
+ """
456
+ self.validator = validator
457
+ self._data = default
458
+
459
+ @property
460
+ def data(self):
461
+ return self._data
462
+
463
+ @data.setter
464
+ def data(self, value):
465
+ self._data = self._validate(value)
466
+
467
+ def _validate(self, value):
468
+ if self.validator(value):
469
+ return value
470
+ raise ValueError(f"Internal Error: value '{value}' is not valid.")
471
+
472
+
473
+ @contextlib.contextmanager
474
+ def cuda_call_ctx(stream_holder: StreamHolder[AnyStream], blocking=True, timing=True) -> typing.Iterator[tuple[Event, Value]]:
475
+ """
476
+ A simple context manager that provides (non-)blocking behavior depending on the
477
+ `blocking` parameter for CUDA calls. The call is timed only for blocking behavior when
478
+ timing is requested.
479
+
480
+ An `end` event is recorded after the CUDA call for use in establishing stream ordering
481
+ for non-blocking calls. This event is returned together with a `Value` object that
482
+ stores the elapsed time if the call is blocking and timing is requested, or None
483
+ otherwise.
484
+
485
+ Inside the context, the current device is the device which owns the provided stream.
486
+ """
487
+ device_id = stream_holder.device_id
488
+ stream = stream_holder.obj
489
+
490
+ with device_ctx(device_id):
491
+ if timing:
492
+ start = stream.record(options={timing_enabled_name(): blocking})
493
+
494
+ end = stream.device.create_event(options={timing_enabled_name(): (timing and blocking)})
495
+
496
+ time = Value(None, validator=lambda v: True)
497
+ yield end, time
498
+
499
+ stream.record(end)
500
+
501
+ if not blocking:
502
+ return
503
+
504
+ end.sync()
505
+
506
+ if timing:
507
+ time.data = end - start
508
+
509
+
510
+ @contextlib.contextmanager
511
+ def host_call_ctx(timing=False):
512
+ elapsed = Value(None, validator=lambda v: True)
513
+
514
+ if timing:
515
+ start_time = time.perf_counter_ns()
516
+
517
+ yield elapsed
518
+
519
+ if timing:
520
+ elapsed.data = (time.perf_counter_ns() - start_time) * 1e-6
521
+
522
+
523
+ # Decorator definitions
524
+
525
+
526
+ def atomic(
527
+ handler: Callable[[typing.Any, Exception | None], bool] | Callable[[Exception | None], bool], method: bool = False
528
+ ) -> Callable:
529
+ """
530
+ A decorator that provides "succeed or roll-back" semantics. A typical use for this is to
531
+ release partial resources if an exception occurs.
532
+
533
+ Args:
534
+ handler: A function to call when an exception occurs. The handler takes a single
535
+ argument, which is the exception object, and returns a boolean stating whether
536
+ the same exception should be reraised. We assume that this function does not
537
+ raise an exception.
538
+
539
+ method: Specify if the wrapped function as well as the exception handler are methods
540
+ bound to the same object (method = True) or they are free functions (method =
541
+ False).
542
+
543
+ Returns:
544
+ Callable: A decorator that creates the wrapping.
545
+ """
546
+
547
+ def outer(wrapped_function):
548
+ """
549
+ A decorator that actually wraps the function for exception handling.
550
+ """
551
+
552
+ @functools.wraps(wrapped_function)
553
+ def inner(*args, **kwargs):
554
+ """
555
+ Call the wrapped function and return the result. If an exception occurs, then
556
+ call the exception handler and reraise the exception.
557
+ """
558
+ try:
559
+ result = wrapped_function(*args, **kwargs)
560
+ except BaseException as e:
561
+ if method:
562
+ flag = handler(args[0], e)
563
+ else:
564
+ flag = handler(e)
565
+
566
+ if flag:
567
+ raise e
568
+
569
+ return result
570
+
571
+ return inner
572
+
573
+ return outer
574
+
575
+
576
+ def precondition(checker: Callable[..., None], what: str = "") -> Callable:
577
+ """
578
+ A decorator that adds checks to ensure any preconditions are met.
579
+
580
+ Args:
581
+ checker: The function to call to check whether the preconditions are met. It has the
582
+ same signature as the wrapped function with the addition of the keyword argument
583
+ `what`.
584
+ what: A string that is passed in to `checker` to provide context information.
585
+
586
+ Returns:
587
+ Callable: A decorator that creates the wrapping.
588
+ """
589
+
590
+ def outer(wrapped_function):
591
+ """
592
+ A decorator that actually wraps the function for checking preconditions.
593
+ """
594
+
595
+ @functools.wraps(wrapped_function)
596
+ def inner(*args, **kwargs):
597
+ """
598
+ Check preconditions and if they are met, call the wrapped function.
599
+ """
600
+ checker(*args, **kwargs, what=what)
601
+ result = wrapped_function(*args, **kwargs)
602
+
603
+ return result
604
+
605
+ return inner
606
+
607
+ return outer
608
+
609
+
610
+ COMMON_SHARED_DOC_MAP = {
611
+ "operand": """\
612
+ A tensor (ndarray-like object). The currently supported types are :class:`numpy.ndarray`,
613
+ :class:`cupy.ndarray`, and :class:`torch.Tensor`.""".replace("\n", " "),
614
+ #
615
+ "stream": """\
616
+ Provide the CUDA stream to use for executing the operation. Acceptable inputs include
617
+ ``cudaStream_t`` (as Python :class:`int`), :class:`cupy.cuda.Stream`, and
618
+ :class:`torch.cuda.Stream`. If a stream is not provided, the current stream from the operand
619
+ package will be used. See :ref:`stream-semantics-guide` for more details on stream
620
+ handling.""".replace("\n", " "),
621
+ #
622
+ "release_workspace": """\
623
+ A value of `True` specifies that the stateful object should release workspace memory back to
624
+ the package memory pool on function return, while a value of `False` specifies that the
625
+ object should retain the memory. This option may be set to `True` if the application
626
+ performs other operations that consume a lot of memory between successive calls to the (same
627
+ or different) :meth:`execute` API, but incurs a small overhead due to obtaining and
628
+ releasing workspace memory from and to the package memory pool on every call. The default is
629
+ `False`.""".replace("\n", " "),
630
+ }
631
+
632
+
633
+ def _release_operand_docstring(plural, version_added=None, execution_methods=("execute",)):
634
+ operands = "operands" if plural else "operand"
635
+ their = "their" if plural else "its"
636
+ counts = "counts" if plural else "count"
637
+ reside = "reside" if plural else "resides"
638
+ reset_method = "reset_operands" if plural else "reset_operand"
639
+ reset_method_unchecked = "reset_operands_unchecked" if plural else "reset_operand_unchecked"
640
+ new_operands = "new operands" if plural else "a new operand"
641
+ are = "are" if plural else "is"
642
+ # The method(s) that perform the computation. Matmul/FFT/tensor use :meth:`execute`;
643
+ # direct solvers use :meth:`factorize` / :meth:`solve` instead.
644
+ if not execution_methods:
645
+ raise ValueError("execution_methods must contain at least one method name")
646
+ execution_ref = " / ".join(f":meth:`{name}`" for name in execution_methods)
647
+ version_line = f"\n\n .. versionadded:: {version_added}" if version_added else ""
648
+ return f"""\
649
+ {version_line}
650
+
651
+ This method does two things:
652
+
653
+ - Releases internal references to the user-provided {operands},
654
+ so that this instance no longer contributes to {their} reference {counts}.
655
+
656
+ - Frees any internal copies (mirrors) that were created when the
657
+ user-provided {operands} {reside} in a different memory space than
658
+ the execution (i.e., copies made during construction or within
659
+ :meth:`{reset_method}` / :meth:`{reset_method_unchecked}`).
660
+
661
+ This functionality can be useful in memory-constrained scenarios, e.g. where
662
+ multiple stateful objects need to coexist. Leveraging this functionality,
663
+ the caller can reduce memory usage while retaining the planned state.
664
+
665
+ Args:
666
+ None
667
+
668
+ Returns:
669
+ None
670
+
671
+ Semantics:
672
+ - Preserves the planned state of the stateful object.
673
+
674
+ - Repeated calls are safe: if the {operands} {are} already released, the
675
+ call is a no-op and an informational message is logged.
676
+
677
+ - After calling this method, :meth:`{reset_method}` (or
678
+ :meth:`{reset_method_unchecked}`) must be called to
679
+ supply {new_operands} before calling {execution_ref} again.
680
+ Failure to do so will result in a runtime error.
681
+ Device-side copies will be re-allocated as needed.
682
+
683
+ - For cross-space scenarios (e.g. CPU {operands} with GPU execution,
684
+ or GPU {operands} with CPU execution): execution is guaranteed to be
685
+ always blocking. It is therefore always safe to call this
686
+ method after calling {execution_ref} without additional synchronization.
687
+
688
+ - When the {operands} {are} in the same memory space as the execution
689
+ (e.g. GPU {operands} with GPU execution): in such case, this method
690
+ drops this instance's internal reference to the user-provided {operands}.
691
+ If the reference count of the {operands} reaches zero, {their} memory
692
+ may be freed, so particular attention should be paid. The caller is
693
+ responsible to ensure that if such deallocation happens, it is
694
+ ordered after pending computation (e.g. by retaining a reference until
695
+ the computation is complete, or by synchronizing the stream).
696
+ Failure to do so is analogous to use-after-free.
697
+
698
+ See :ref:`nvmath overview`, :ref:`stateful_apis_guide` for operand
699
+ lifecycle and usage patterns, and :ref:`stream-semantics-guide`
700
+ for stream ordering rules.
701
+ """
702
+
703
+
704
+ COMMON_SHARED_DOC_MAP["release_operands"] = _release_operand_docstring(True, version_added="0.9.0")
705
+ COMMON_SHARED_DOC_MAP["release_operand"] = _release_operand_docstring(False, version_added="0.9.0")
706
+
707
+
708
+ def _reset_operand_unchecked_docstring(
709
+ plural, version_added=None, validation_examples="package match, data type match", experimental=True
710
+ ):
711
+ operands = "operands" if plural else "operand"
712
+ reset_method = "reset_operands" if plural else "reset_operand"
713
+ release_method = "release_operands" if plural else "release_operand"
714
+ experimental_line = ".. experimental:: method" if experimental else ""
715
+ version_line = f"\n\n .. versionadded:: {version_added}" if version_added else ""
716
+ return f"""\
717
+ {experimental_line}
718
+ {version_line}
719
+
720
+ This method is a performance-optimized alternative to :meth:`{reset_method}`
721
+ that eliminates validation and logging overhead, making it ideal for
722
+ performance-critical loops where {operands}
723
+ compatibility is guaranteed by the caller.
724
+
725
+ This method accepts the same parameters as :meth:`{reset_method}`.
726
+
727
+ Semantics:
728
+ The semantics are the same as in :meth:`{reset_method}`,
729
+ except that this method does not perform any validation
730
+ (e.g. {validation_examples}, etc.) or logging.
731
+
732
+ When to Use:
733
+ - Performance-critical loops with repeated executions on different {operands}
734
+
735
+ - After verifying correctness with :meth:`{reset_method}` during development
736
+
737
+ - When {operands} compatibility is guaranteed by construction or invariant
738
+
739
+ .. seealso::
740
+ :meth:`{reset_method}`, :meth:`{release_method}`
741
+ """
742
+
743
+
744
+ COMMON_SHARED_DOC_MAP["reset_operands_unchecked"] = _reset_operand_unchecked_docstring(True, version_added="0.9.0")
745
+ COMMON_SHARED_DOC_MAP["reset_operand_unchecked"] = _reset_operand_unchecked_docstring(False, version_added="0.9.0")
746
+
747
+
748
+ class DefaultDocstring(dict):
749
+ def __missing__(self, key):
750
+ return "{" + key + "}"
751
+
752
+
753
+ def docstring_decorator(doc_map, skip_missing=False):
754
+ assert isinstance(doc_map, dict)
755
+
756
+ def _format_doc(doc):
757
+ if doc is not None:
758
+ if skip_missing:
759
+ # Using a default class to handle missing keys in doc_map
760
+ doc = doc.format_map(DefaultDocstring(**doc_map))
761
+ else:
762
+ doc = doc.format(**doc_map)
763
+ return doc
764
+
765
+ def decorator(func_or_class):
766
+ if isinstance(func_or_class, type): # class decorator
767
+ # update the docstring of all public methods with docstrings
768
+ static_methods = [] # staticmethods appear to require special handling
769
+ for name, method in vars(func_or_class).items():
770
+ if isinstance(method, (staticmethod, classmethod)):
771
+ static_methods.append(name)
772
+ continue
773
+ if isinstance(method, property):
774
+ if method.fget and method.fget.__doc__:
775
+ # Create a new property with updated docstring
776
+ setattr(
777
+ func_or_class,
778
+ name,
779
+ property(method.fget, method.fset, method.fdel, _format_doc(method.fget.__doc__)),
780
+ )
781
+ continue
782
+ if callable(method) and (not name.startswith("_")) and method.__doc__:
783
+ method.__doc__ = _format_doc(method.__doc__)
784
+ # update the docstring of the constructor
785
+ func_or_class.__doc__ = _format_doc(func_or_class.__doc__)
786
+ for name in static_methods:
787
+ method = func_or_class.__dict__[name]
788
+ method.__func__.__doc__ = _format_doc(method.__func__.__doc__)
789
+ return func_or_class
790
+ else: # function decorator
791
+ func_or_class.__doc__ = _format_doc(func_or_class.__doc__)
792
+ return func_or_class
793
+
794
+ return decorator
795
+
796
+
797
+ @functools.cache
798
+ def timing_enabled_name() -> str:
799
+ """
800
+ The name of the `timing_enabled` attribute in cuda.core.EventOptions.
801
+ """
802
+ from cuda.core import __version__ as version
803
+
804
+ major = int(version.split(".")[0])
805
+ return "enable_timing" if major == 0 else "timing_enabled"