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,300 @@
1
+ # Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ from collections.abc import Callable
6
+ from typing import Any
7
+
8
+ import llvmlite.ir as llvmir
9
+ import numpy as np
10
+ from llvmlite import ir
11
+ from numba import types
12
+ from numba.core import cgutils, typing
13
+ from numba.core.base import BaseContext
14
+ from numba.core.errors import TypingError
15
+ from numba.cuda.cudaimpl import lower_constant
16
+ from numba.cuda.extending import intrinsic, overload, overload_attribute, register_model
17
+ from numba.extending import models, typeof_impl
18
+
19
+ from nvmath.bindings import mathdx
20
+
21
+ from .llvm_array import LLVMArray
22
+ from .types import Complex, Vector, complex32, complex64, complex128, half2, half4, np_float16x2, np_float16x4
23
+ from .vector_types_numba import VectorTypeClass, float16x2_type, float16x4_type, float32x2_type, float64x2_type
24
+
25
+ NP_TYPES_TO_NUMBA_FE_TYPES = {
26
+ np.float16: np.float16,
27
+ np.float32: np.float32,
28
+ np.float64: np.float64,
29
+ np_float16x2: float16x2_type,
30
+ np_float16x4: float16x4_type,
31
+ np.complex64: float32x2_type,
32
+ np.complex128: float64x2_type,
33
+ np.int8: np.int8,
34
+ np.int16: np.int16,
35
+ np.int32: np.int32,
36
+ np.int64: np.int64,
37
+ np.uint8: np.uint8,
38
+ np.uint16: np.uint16,
39
+ np.uint32: np.uint32,
40
+ np.uint64: np.uint64,
41
+ }
42
+
43
+ NUMBA_FE_TYPES_TO_NUMBA_IR = {
44
+ np.float16: types.float16,
45
+ np.float32: types.float32,
46
+ np.float64: types.float64,
47
+ complex32: float16x2_type,
48
+ complex64: float32x2_type,
49
+ complex128: float64x2_type,
50
+ half2: float16x2_type,
51
+ half4: float16x4_type,
52
+ float16x2_type: float16x2_type,
53
+ float16x4_type: float16x4_type,
54
+ float32x2_type: float32x2_type,
55
+ float64x2_type: float64x2_type,
56
+ np.int8: types.int8,
57
+ np.int16: types.int16,
58
+ np.int32: types.int32,
59
+ np.int64: types.int64,
60
+ np.uint8: types.uint8,
61
+ np.uint16: types.uint16,
62
+ np.uint32: types.uint32,
63
+ np.uint64: types.uint64,
64
+ }
65
+
66
+
67
+ class EmptyStructModel(models.StructModel):
68
+ """Data model that does not take space. Intended to be used with types that
69
+ are presented only at typing stage and not represented in memory."""
70
+
71
+ def __init__(self, dmm, fe_type):
72
+ members = []
73
+ super().__init__(dmm, fe_type, members)
74
+
75
+
76
+ def overload_type_attribute(numba_type, attribute_base, attribute):
77
+ """Make type attribute available inside jitted code."""
78
+ assert issubclass(numba_type, types.Type)
79
+
80
+ @overload_attribute(numba_type, attribute, inline="always")
81
+ def ol_blas_attribute(blas_numba):
82
+ tp = blas_numba
83
+ if attribute_base != "":
84
+ tp = getattr(tp, attribute_base)
85
+ val = getattr(tp, attribute)
86
+ return lambda blas_numba: val
87
+
88
+
89
+ @typeof_impl.register(Complex)
90
+ def typeof_complex(val: Complex, c: typing.Context) -> Any:
91
+ return VectorTypeClass(val._numba_type)
92
+
93
+
94
+ @typeof_impl.register(Vector)
95
+ def typeof_vector(val: Vector, c: typing.Context) -> Any:
96
+ return VectorTypeClass(val._numba_type)
97
+
98
+
99
+ @intrinsic
100
+ def get_array_ptr(typingctx: typing.Context, arr):
101
+ """Get raw pointer to the data of a Numba array."""
102
+ assert isinstance(arr, types.Array)
103
+
104
+ sig = typing.signature(types.CPointer(arr.dtype), arr)
105
+
106
+ def codegen(context: BaseContext, builder, sig, args):
107
+ arrTy = sig.args[0]
108
+ arr = args[0]
109
+
110
+ dtype = arrTy.dtype
111
+ valueTy = context.get_value_type(dtype)
112
+ ptrTy = ir.PointerType(valueTy)
113
+ if arr is None:
114
+ ptr = ptrTy(None)
115
+ else:
116
+ ptr = cgutils.create_struct_proxy(arrTy)(context, builder, arr).data
117
+
118
+ # Future release of numba-cuda may have support for address spaces.
119
+ # It is not supported to pass a non generic pointer to device function
120
+ # call.
121
+ if ptr.type.addrspace != 0:
122
+ ptr = builder.addrspacecast(ptr, ir.PointerType(ptr.type.pointee), "generic")
123
+
124
+ return ptr
125
+
126
+ return sig, codegen
127
+
128
+
129
+ def build_get_value_ptr(target_type=None):
130
+ @intrinsic
131
+ def get_value_ptr_impl(typingctx: typing.Context, value):
132
+ """Get raw pointer to the value."""
133
+ if value not in [float16x2_type, float16x4_type, float32x2_type, float64x2_type] and not isinstance( # noqa: UP038
134
+ value, (types.Float, types.Complex, types.Integer)
135
+ ):
136
+ raise TypingError(f"get_value_ptr does not support type {value}")
137
+
138
+ return_type = target_type if target_type is not None else value
139
+ sig = typing.signature(types.CPointer(return_type), value)
140
+
141
+ def codegen(context: BaseContext, builder, sig, args):
142
+ if target_type is not None and sig.args[0] != target_type:
143
+ value = context.cast(builder, args[0], sig.args[0], target_type)
144
+ else:
145
+ value = args[0]
146
+ return cgutils.alloca_once_value(builder, value)
147
+
148
+ return sig, codegen
149
+
150
+ return get_value_ptr_impl
151
+
152
+
153
+ get_value_ptr = build_get_value_ptr()
154
+ get_uint32_value_ptr = build_get_value_ptr(target_type=types.uint32)
155
+
156
+
157
+ class MathdxOpaqueTensorType(types.Type):
158
+ def __init__(self, dynamic_strides_size=0, dynamic_shape_size=0):
159
+ self.dynamic_strides_size = dynamic_strides_size
160
+ self.dynamic_shape_size = dynamic_shape_size
161
+ super().__init__(f"MathdxOpaqueTensor({dynamic_strides_size}, {dynamic_shape_size})")
162
+
163
+
164
+ @register_model(MathdxOpaqueTensorType)
165
+ class MathdxOpaqueTensorModel(models.StructModel):
166
+ def __init__(self, dmm, fe_type: MathdxOpaqueTensorType):
167
+ members = [("ptr", types.voidptr)]
168
+
169
+ if fe_type.dynamic_strides_size > 0:
170
+ array_type = LLVMArray(types.int64, fe_type.dynamic_strides_size)
171
+ members += [("strides", array_type)]
172
+
173
+ if fe_type.dynamic_shape_size > 0:
174
+ array_type = LLVMArray(types.int64, fe_type.dynamic_shape_size)
175
+ members += [("shape", array_type)]
176
+
177
+ models.StructModel.__init__(self, dmm, fe_type, members)
178
+
179
+ def get_value_type(self):
180
+ if self._value_type is None:
181
+ name = f"struct.libmathdx_tensor_{self._fe_type.dynamic_strides_size}s_{self._fe_type.dynamic_shape_size}s"
182
+ t = llvmir.global_context.get_identified_type(name)
183
+ if t.is_opaque:
184
+ t.set_body(*[m.get_value_type() for m in self._models])
185
+ self._value_type = t
186
+ return self._value_type
187
+
188
+
189
+ class MathdxPipelineType(types.Type):
190
+ def __init__(self):
191
+ super().__init__("MathdxPipeline()")
192
+
193
+
194
+ @register_model(MathdxPipelineType)
195
+ class MathdxPipelineModel(models.StructModel):
196
+ def __init__(self, dmm, fe_type: MathdxPipelineType):
197
+ members = [("ptr", types.voidptr)]
198
+ models.StructModel.__init__(self, dmm, fe_type, members)
199
+
200
+ def get_value_type(self):
201
+ if self._value_type is None:
202
+ name = "struct.libmathdx_pipeline"
203
+ t = llvmir.global_context.get_identified_type(name)
204
+ if t.is_opaque:
205
+ t.set_body(*[m.get_value_type() for m in self._models])
206
+ self._value_type = t
207
+ return self._value_type
208
+
209
+
210
+ @intrinsic
211
+ def cast_to_void_pointer(typingctx: typing.Context, ptr_ty):
212
+ if not isinstance(ptr_ty, types.CPointer):
213
+ raise TypingError("cast_to_void_pointer expects a CPointer type")
214
+ i8_ptr_ty = types.CPointer(types.uint8)
215
+ sig = i8_ptr_ty(ptr_ty)
216
+
217
+ def codegen(context: BaseContext, builder, sig, args):
218
+ return builder.bitcast(args[0], context.get_value_type(i8_ptr_ty))
219
+
220
+ return sig, codegen
221
+
222
+
223
+ _cabi_device_registry: dict[str | tuple[str, typing.Signature], Callable] = {}
224
+
225
+
226
+ def declare_cabi_device(symbol: str, sig: typing.Signature, link=None):
227
+ """declare_cabi_device is an analog of cuda.declare_device but uses C ABI
228
+ calling convention instead of Numba ABI calling convention. It means that
229
+ the first argument is not a return value pointer."""
230
+ key: str | tuple[str, typing.Signature] = symbol
231
+ if mathdx.get_version_ex() < (0, 3, 0):
232
+ key = (symbol, sig)
233
+ device_func = _cabi_device_registry.get(key)
234
+ if device_func is None:
235
+ device_func = _declare_cabi_device(symbol, sig, link)
236
+ _cabi_device_registry[key] = device_func
237
+
238
+ return device_func
239
+
240
+
241
+ def _declare_cabi_device(symbol: str, sig: typing.Signature, link=None):
242
+ intrinsic_sig = sig.return_type(types.Tuple(sig.args))
243
+
244
+ @intrinsic
245
+ def call_device(typingctx: typing.Context, args):
246
+ def codegen(context: BaseContext, builder, sig: typing.Signature, args):
247
+ if link is not None:
248
+ context.active_code_library.add_linking_file(link)
249
+
250
+ args = cgutils.unpack_tuple(builder, args[0])
251
+ argTypes = sig.args[0]
252
+
253
+ assert len(args) == len(argTypes)
254
+
255
+ # CUDA target maps types.void to i8*, but we need void for abi correctness.
256
+ if sig.return_type == types.void:
257
+ retTy = ir.VoidType()
258
+ else:
259
+ retTy = context.get_value_type(sig.return_type)
260
+ fnTy = ir.FunctionType(retTy, [context.get_value_type(argTy) for argTy in argTypes])
261
+ fn = cgutils.get_or_insert_function(builder.module, fnTy, symbol)
262
+ builder.call(fn, args)
263
+
264
+ return intrinsic_sig, codegen
265
+
266
+ def device_func():
267
+ pass
268
+
269
+ @overload(device_func, jit_options={"forceinline": True})
270
+ def ol_device_func(*args):
271
+ if len(args) != len(sig.args):
272
+ raise RuntimeError(
273
+ f"Invalid number of arguments for device function {symbol}: expected {len(sig.args)}, got {len(args)}"
274
+ )
275
+ for expected_t, provided_t in zip(sig.args, args, strict=True):
276
+ if expected_t == provided_t:
277
+ continue
278
+ raise RuntimeError(f"Invalid argument type for device function {symbol}: expected {expected_t}, got {provided_t}")
279
+ return lambda *args: call_device(tuple(args))
280
+
281
+ return device_func
282
+
283
+
284
+ def register_dummy_numba_type(numba_type, base_type, name, attributes):
285
+ """Register and lower a dummy Numba type that does not take space in memory.
286
+ Intended for types that are present only at typing stage and not represented in memory,
287
+ but are allowed to be accessed inside kernel."""
288
+ for attribute in attributes:
289
+ overload_type_attribute(numba_type, name, attribute)
290
+
291
+ @typeof_impl.register(base_type)
292
+ def typeof_numba(val, context):
293
+ return numba_type(val)
294
+
295
+ register_model(numba_type)(EmptyStructModel)
296
+
297
+ @lower_constant(numba_type)
298
+ def constant_dummy(context, builder, typ, pyval):
299
+ struct_ptr = cgutils.create_struct_proxy(typ)(context, builder)
300
+ return struct_ptr._getvalue()
@@ -0,0 +1,202 @@
1
+ # Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+
6
+ from collections.abc import Iterable
7
+ from typing import Any
8
+
9
+ import cffi
10
+ import numpy as np
11
+ from numba_cuda_mlir import cuda
12
+ from numba_cuda_mlir._mlir.dialects import llvm
13
+ from numba_cuda_mlir._mlir.extras import types as T
14
+ from numba_cuda_mlir.cuda import vector_types
15
+ from numba_cuda_mlir.extending import lowering_registry, overload, overload_attribute, typing_registry
16
+ from numba_cuda_mlir.lowering_utilities import to_mlir_type, unverified_convert
17
+ from numba_cuda_mlir.lowering_utilities.type_conversions import to_numba_type
18
+ from numba_cuda_mlir.models import PrimitiveModel
19
+ from numba_cuda_mlir.models import register_model as numba_cuda_mlir_register_model
20
+ from numba_cuda_mlir.numba_cuda import types
21
+ from numba_cuda_mlir.numba_cuda.extending import typeof_impl
22
+ from numba_cuda_mlir.numba_cuda.np import numpy_support
23
+ from numba_cuda_mlir.type_defs.vector_types import VectorType
24
+
25
+ from .types import Complex, Vector
26
+
27
+ __all__ = [
28
+ "HostDescriptorBase",
29
+ "overload_type_attribute",
30
+ "register_dummy_type",
31
+ "get_array_pointer",
32
+ ]
33
+
34
+ # ==========================
35
+ # Argument-conversion helpers for JIT bodies
36
+ # ==========================
37
+
38
+ _FFI = cffi.FFI()
39
+
40
+
41
+ @cuda.jit(device=True, forceinline=True)
42
+ def get_array_pointer(arr):
43
+ return _FFI.from_buffer(arr)
44
+
45
+
46
+ @cuda.jit(device=True, inline="always")
47
+ def int_to_uint32_ptr(val):
48
+ local = cuda.local.array(shape=(1,), dtype=np.uint32)
49
+ local[0] = val
50
+ return _FFI.from_buffer(local)
51
+
52
+
53
+ @cuda.jit(device=True, forceinline=True)
54
+ def passthrough(val):
55
+ return val
56
+
57
+
58
+ def get_value_ptr(value: Any) -> int:
59
+ raise NotImplementedError("get_value_ptr is callable only from jitted code")
60
+
61
+
62
+ @overload(get_value_ptr, strict=False, inline="always", typing_registry=typing_registry)
63
+ def ol_get_value_ptr(value):
64
+ # TODO(numba-cuda-mlir): Known complex types work correctly,
65
+ # but types.Complex("complex32",types.float16) currently fails.
66
+ # For now, handle all complex types here explicitly.
67
+ # Ideally, all of this should be replaced by a generic as_dtype path.
68
+ dtype = value if isinstance(value, (VectorType, types.Complex)) else numpy_support.as_dtype(value).type
69
+
70
+ def impl(value):
71
+ local = cuda.local.array(shape=(1,), dtype=dtype)
72
+ local[0] = value
73
+ return _FFI.from_buffer(local)
74
+
75
+ return impl
76
+
77
+
78
+ # ==========================
79
+ # Type registration helpers
80
+ # ==========================
81
+
82
+
83
+ class HostDescriptorBase(types.Type):
84
+ """
85
+ Base numba type wrapping a host descriptor,
86
+ subclass per API (cufftdx, cusolverdx, ...).
87
+ """
88
+
89
+ def __init__(
90
+ self,
91
+ host_descriptor,
92
+ definition_args: Iterable[str],
93
+ type_kind: str,
94
+ ) -> None:
95
+ self._host_descriptor = host_descriptor
96
+
97
+ attributes = [
98
+ f"{attr}={getattr(host_descriptor, attr)}" for attr in definition_args if getattr(host_descriptor, attr, None)
99
+ ]
100
+ attributes.sort()
101
+
102
+ super().__init__(f"{type_kind}({','.join(attributes)})")
103
+
104
+ @property
105
+ def host_descriptor(self):
106
+ return self._host_descriptor
107
+
108
+
109
+ def overload_type_attribute(
110
+ numba_type: type[types.Type],
111
+ attribute_base: str,
112
+ attribute: str,
113
+ ) -> None:
114
+ """Make type attribute available inside jitted code."""
115
+ assert issubclass(numba_type, types.Type)
116
+
117
+ @overload_attribute(
118
+ numba_type,
119
+ attribute,
120
+ inline="always",
121
+ typing_registry=typing_registry,
122
+ lowering_registry=lowering_registry,
123
+ )
124
+ def ol_attribute(api_numba):
125
+ tp = api_numba
126
+ if attribute_base != "":
127
+ tp = getattr(tp, attribute_base)
128
+ val = getattr(tp, attribute)
129
+ return lambda api_numba: val
130
+
131
+
132
+ class _DummyMlirModel(PrimitiveModel):
133
+ """Shared data model for numba types that have no runtime memory footprint."""
134
+
135
+ def __init__(self, dmm, fe_type) -> None:
136
+ super().__init__(dmm, fe_type, T.i8())
137
+
138
+
139
+ def register_dummy_type(
140
+ numba_type: type[types.Type],
141
+ base_type: type,
142
+ attributes: Iterable[str],
143
+ *,
144
+ attribute_base: str | None = None,
145
+ ) -> None:
146
+ """
147
+ Register a dummy numba-cuda-mlir type that exists only at typing time.
148
+
149
+ Bundles five effects:
150
+ * typeof_impl.register(base_type)
151
+ * register_model(numba_type) with the shared dummy PrimitiveModel
152
+ * to_mlir_type.register(numba_type) returning T.i8()
153
+ * unverified_convert.register(base_type) returning llvm.mlir_undef(T.i8())
154
+ * an overload_type_attribute(numba_type, attribute_base, attribute)
155
+ for each name in attributes.
156
+ """
157
+ if attribute_base is None and issubclass(numba_type, HostDescriptorBase):
158
+ attribute_base = "host_descriptor"
159
+ assert attribute_base is not None
160
+
161
+ @typeof_impl.register(base_type)
162
+ def _typeof(val, context):
163
+ return numba_type(val)
164
+
165
+ numba_cuda_mlir_register_model(numba_type)(_DummyMlirModel)
166
+
167
+ @to_mlir_type.register(numba_type)
168
+ def _to_mlir(val, context):
169
+ return T.i8()
170
+
171
+ @unverified_convert.register(base_type)
172
+ def _unverified_convert(val, target_type, signed: bool = False):
173
+ assert target_type == T.i8()
174
+ return llvm.mlir_undef(T.i8())
175
+
176
+ for attribute in attributes:
177
+ overload_type_attribute(numba_type, attribute_base, attribute)
178
+
179
+
180
+ # ==========================
181
+ # Vector/Complex typing
182
+ # ==========================
183
+
184
+
185
+ @to_numba_type.register(Vector)
186
+ def _(val: Vector) -> types.Type:
187
+ return getattr(vector_types, f"{np.dtype(val.real_dtype).name}x{val.size}")
188
+
189
+
190
+ @to_numba_type.register(Complex)
191
+ def _(val: Complex) -> types.Type:
192
+ # real_type = to_numba_type(val.real_dtype)
193
+ real_type = to_numba_type(np.dtype(val.real_dtype))
194
+ assert isinstance(real_type, types.Float)
195
+ return types.Complex(f"complex{real_type.bitwidth * 2}", real_type)
196
+
197
+
198
+ # TODO: do we need both to_numba_type and typeof_impl for Vector and Complex?
199
+ @typeof_impl.register(Vector)
200
+ @typeof_impl.register(Complex)
201
+ def _typeof_host_type(val, c):
202
+ return typeof_impl(to_numba_type(val), c)
@@ -0,0 +1,201 @@
1
+ # Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ from functools import cached_property
6
+
7
+ import numpy
8
+ from llvmlite import ir as llir
9
+ from numba.core import cgutils
10
+ from numba.core.base import BaseContext
11
+ from numba.cuda.extending import lower_builtin, make_attribute_wrapper, type_callable
12
+ from numba.cuda.models import register_model
13
+ from numba.extending import models, types
14
+
15
+ from nvmath.device.common import Layout, OpaqueTensor
16
+ from nvmath.device.common_numba import NUMBA_FE_TYPES_TO_NUMBA_IR, overload_type_attribute
17
+ from nvmath.device.llvm_array import LLVMArray
18
+
19
+ _LIBMATHDX_RUNTIME = -9223372036854775808
20
+
21
+
22
+ class OpaqueLayout(Layout):
23
+ """
24
+ A generic tensor layout that can represent arbitrary shapes and strides.
25
+ It is the same as libmathdx opaque tensor, but without the pointer
26
+ to data.
27
+ """
28
+
29
+ _dtype: numpy.number
30
+ _shape: tuple[int, ...]
31
+ _strides: tuple[int, ...]
32
+
33
+ def __init__(self, shape: tuple[int, ...], strides: tuple[int, ...], dtype: numpy.number):
34
+ assert len(shape) == len(strides)
35
+ self._shape = shape
36
+ self._strides = strides
37
+ self._dtype = dtype
38
+
39
+ @cached_property
40
+ def dynamic_strides_size(self) -> int:
41
+ return self._strides.count(_LIBMATHDX_RUNTIME)
42
+
43
+ @cached_property
44
+ def dynamic_shape_size(self) -> int:
45
+ return self._shape.count(_LIBMATHDX_RUNTIME)
46
+
47
+ @property
48
+ def dtype(self) -> numpy.number:
49
+ return self._dtype
50
+
51
+ @cached_property
52
+ def ndim(self) -> int:
53
+ return len(self._shape)
54
+
55
+ @property
56
+ def shape(self) -> tuple[int, ...]:
57
+ return self._shape
58
+
59
+ @property
60
+ def strides(self) -> tuple[int, ...]:
61
+ return self._strides
62
+
63
+
64
+ class LayoutType(types.Type):
65
+ """
66
+ Type class associated with opaque tensor layouts.
67
+ """
68
+
69
+ def __init__(
70
+ self,
71
+ layout: OpaqueLayout,
72
+ ):
73
+ self._layout = layout
74
+ super().__init__(
75
+ f"Layout(dynamic_shape_size={layout.dynamic_shape_size}, dynamic_strides_size={layout.dynamic_strides_size})"
76
+ )
77
+
78
+ @property
79
+ def layout(self) -> OpaqueLayout:
80
+ return self._layout
81
+
82
+ @cached_property
83
+ def dtype(self) -> types.Number:
84
+ return NUMBA_FE_TYPES_TO_NUMBA_IR[self._layout.dtype]
85
+
86
+
87
+ @register_model(LayoutType)
88
+ class LayoutModel(models.StructModel):
89
+ def __init__(self, dmm, fe_type: LayoutType):
90
+ members = []
91
+
92
+ strides_size = fe_type.layout.dynamic_strides_size
93
+ if strides_size > 0:
94
+ array_type = LLVMArray(types.int64, strides_size)
95
+ members += [("strides", array_type)]
96
+
97
+ shape_size = fe_type.layout.dynamic_shape_size
98
+ if shape_size > 0:
99
+ array_type = LLVMArray(types.int64, shape_size)
100
+ members += [("shape", array_type)]
101
+
102
+ models.StructModel.__init__(self, dmm, fe_type, members)
103
+
104
+
105
+ make_attribute_wrapper(LayoutType, "strides", "strides")
106
+ make_attribute_wrapper(LayoutType, "shape", "shape")
107
+
108
+ for attribute in ("uid", "alignment", "storage_bytes"):
109
+ overload_type_attribute(LayoutType, "", attribute)
110
+
111
+
112
+ class OpaqueTensorType(types.Type):
113
+ def __init__(self, layout: LayoutType):
114
+ super().__init__(f"OpaqueTensor(layout={layout})")
115
+ self._layout = layout
116
+
117
+ @property
118
+ def uid(self):
119
+ return self._layout.layout.uid
120
+
121
+ @property
122
+ def layout(self) -> LayoutType:
123
+ return self._layout
124
+
125
+ @cached_property
126
+ def ndim(self):
127
+ return 2
128
+
129
+ @cached_property
130
+ def dtype(self):
131
+ return self._layout.dtype
132
+
133
+
134
+ class OpaqueTensorModel(models.StructModel):
135
+ def __init__(self, dmm, fe_type: OpaqueTensorType):
136
+ members = [
137
+ ("ptr", types.voidptr),
138
+ ]
139
+ layout = fe_type.layout.layout
140
+
141
+ strides_size = layout.dynamic_strides_size
142
+ if strides_size > 0:
143
+ array_type = LLVMArray(types.int64, strides_size)
144
+ members += [("strides", array_type)]
145
+
146
+ shape_size = layout.dynamic_shape_size
147
+ if shape_size > 0:
148
+ array_type = LLVMArray(types.int64, shape_size)
149
+ members += [("shape", array_type)]
150
+
151
+ models.StructModel.__init__(self, dmm, fe_type, members)
152
+
153
+
154
+ register_model(OpaqueTensorType)(OpaqueTensorModel)
155
+
156
+
157
+ @type_callable(OpaqueTensor)
158
+ def type_callable_OpaqueTensor(context):
159
+ def typer(buffer_ty, layout_ty):
160
+ if not isinstance(buffer_ty, types.Array):
161
+ return
162
+ if not isinstance(layout_ty, LayoutType):
163
+ return
164
+ assert buffer_ty.dtype == layout_ty.dtype
165
+ return OpaqueTensorType(layout_ty)
166
+
167
+ return typer
168
+
169
+
170
+ @lower_builtin(OpaqueTensor, types.Array, LayoutType)
171
+ def impl_opaque_tensor(context: BaseContext, builder: llir.IRBuilder, sig, args):
172
+ tensor_ty: OpaqueTensorType = sig.return_type
173
+ buffer, buffer_ty = args[0], sig.args[0]
174
+ layout, layout_ty = args[1], sig.args[1]
175
+
176
+ ptr = cgutils.create_struct_proxy(buffer_ty)(context, builder, buffer).data
177
+ # Future release of numba-cuda may have support for address spaces.
178
+ # It is not supported to pass a non generic pointer to device function
179
+ # call.
180
+ if ptr.type.addrspace != 0:
181
+ ptr = builder.addrspacecast(ptr, llir.PointerType(ptr.type.pointee), "generic")
182
+
183
+ voidPtrTy = llir.PointerType(llir.IntType(8))
184
+ ptr = builder.bitcast(ptr, voidPtrTy)
185
+
186
+ opaque_tensor = cgutils.create_struct_proxy(tensor_ty)(context, builder)
187
+ opaque_tensor.ptr = ptr
188
+
189
+ layout_struct = cgutils.create_struct_proxy(layout_ty)(context, builder, layout)
190
+
191
+ if tensor_ty.layout.layout.dynamic_shape_size > 0:
192
+ opaque_tensor.shape = layout_struct.shape
193
+
194
+ if tensor_ty.layout.layout.dynamic_strides_size > 0:
195
+ opaque_tensor.strides = layout_struct.strides
196
+
197
+ return opaque_tensor._getvalue()
198
+
199
+
200
+ make_attribute_wrapper(OpaqueTensorType, "ptr", "ptr")
201
+ overload_type_attribute(OpaqueTensorType, "layout", "layout")