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.
- nvmath/__init__.pxd +0 -0
- nvmath/__init__.py +45 -0
- nvmath/_internal/__init__.py +0 -0
- nvmath/_internal/attribute_ifc_factory.py +330 -0
- nvmath/_internal/layout.py +70 -0
- nvmath/_internal/templates.py +130 -0
- nvmath/_internal/threadsafe.py +106 -0
- nvmath/_internal/utils.py +43 -0
- nvmath/_internal/workspace.py +490 -0
- nvmath/_utils.py +147 -0
- nvmath/bindings/__init__.py +60 -0
- nvmath/bindings/_internal/__init__.pxd +0 -0
- nvmath/bindings/_internal/__init__.py +0 -0
- nvmath/bindings/_internal/common_types.pxd +31 -0
- nvmath/bindings/_internal/cublas.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/_internal/cublas.pxd +530 -0
- nvmath/bindings/_internal/cublasLt.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/_internal/cublasLt.pxd +59 -0
- nvmath/bindings/_internal/cublasMp.pxd +52 -0
- nvmath/bindings/_internal/cudss.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/_internal/cudss.pxd +54 -0
- nvmath/bindings/_internal/cufft.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/_internal/cufft.pxd +70 -0
- nvmath/bindings/_internal/cufftMp.pxd +77 -0
- nvmath/bindings/_internal/curand.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/_internal/curand.pxd +42 -0
- nvmath/bindings/_internal/cusolver.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/_internal/cusolver.pxd +15 -0
- nvmath/bindings/_internal/cusolverDn.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/_internal/cusolverDn.pxd +406 -0
- nvmath/bindings/_internal/cusolverMp.pxd +71 -0
- nvmath/bindings/_internal/cusolverSp.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/_internal/cusolverSp.pxd +75 -0
- nvmath/bindings/_internal/cusparse.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/_internal/cusparse.pxd +471 -0
- nvmath/bindings/_internal/cusparseLt.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/_internal/cusparseLt.pxd +48 -0
- nvmath/bindings/_internal/cutensor.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/_internal/cutensor.pxd +58 -0
- nvmath/bindings/_internal/mathdx.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/_internal/mathdx.pxd +116 -0
- nvmath/bindings/_internal/nvshmem.pxd +29 -0
- nvmath/bindings/_internal/utils.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/_internal/utils.pxd +174 -0
- nvmath/bindings/_internal/utils.pyi +10 -0
- nvmath/bindings/cublas.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cublas.pxd +558 -0
- nvmath/bindings/cublas.pyi +812 -0
- nvmath/bindings/cublasLt.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cublasLt.pxd +109 -0
- nvmath/bindings/cublasLt.pyi +1461 -0
- nvmath/bindings/cublasMp.pxd +85 -0
- nvmath/bindings/cublasMp.pyi +267 -0
- nvmath/bindings/cudss.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cudss.pxd +98 -0
- nvmath/bindings/cudss.pyi +443 -0
- nvmath/bindings/cufft.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cufft.pxd +118 -0
- nvmath/bindings/cufft.pyi +301 -0
- nvmath/bindings/cufftMp.pxd +124 -0
- nvmath/bindings/cufftMp.pyi +326 -0
- nvmath/bindings/curand.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/curand.pxd +71 -0
- nvmath/bindings/curand.pyi +189 -0
- nvmath/bindings/cusolver.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cusolver.pxd +62 -0
- nvmath/bindings/cusolver.pyi +320 -0
- nvmath/bindings/cusolverDn.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cusolverDn.pxd +430 -0
- nvmath/bindings/cusolverDn.pyi +422 -0
- nvmath/bindings/cusolverMp.pxd +98 -0
- nvmath/bindings/cusolverMp.pyi +114 -0
- nvmath/bindings/cusolverSp.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cusolverSp.pxd +95 -0
- nvmath/bindings/cusolverSp.pyi +70 -0
- nvmath/bindings/cusparse.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cusparse.pxd +546 -0
- nvmath/bindings/cusparse.pyi +1017 -0
- nvmath/bindings/cusparseLt.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cusparseLt.pxd +99 -0
- nvmath/bindings/cusparseLt.pyi +252 -0
- nvmath/bindings/cutensor.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cutensor.pxd +98 -0
- nvmath/bindings/cutensor.pyi +324 -0
- nvmath/bindings/cycublas.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cycublas.pxd +664 -0
- nvmath/bindings/cycublasLt.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cycublasLt.pxd +1045 -0
- nvmath/bindings/cycublasMp.pxd +171 -0
- nvmath/bindings/cycudss.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cycudss.pxd +277 -0
- nvmath/bindings/cycufft.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cycufft.pxd +333 -0
- nvmath/bindings/cycufftMp.pxd +342 -0
- nvmath/bindings/cycurand.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cycurand.pxd +141 -0
- nvmath/bindings/cycusolver.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cycusolver.pxd +137 -0
- nvmath/bindings/cycusolverDn.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cycusolverDn.pxd +443 -0
- nvmath/bindings/cycusolverMp.pxd +107 -0
- nvmath/bindings/cycusolverSp.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cycusolverSp.pxd +93 -0
- nvmath/bindings/cycusparse.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cycusparse.pxd +679 -0
- nvmath/bindings/cycusparseLt.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cycusparseLt.pxd +135 -0
- nvmath/bindings/cycutensor.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cycutensor.pxd +189 -0
- nvmath/bindings/cymathdx.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/cymathdx.pxd +552 -0
- nvmath/bindings/cynvshmem.pxd +118 -0
- nvmath/bindings/mathdx.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/mathdx.pxd +182 -0
- nvmath/bindings/mathdx.pyi +1562 -0
- nvmath/bindings/nvpl/__init__.pxd +0 -0
- nvmath/bindings/nvpl/__init__.py +13 -0
- nvmath/bindings/nvpl/_internal/__init__.pxd +0 -0
- nvmath/bindings/nvpl/_internal/__init__.py +0 -0
- nvmath/bindings/nvpl/_internal/blas.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/nvpl/_internal/blas.pxd +237 -0
- nvmath/bindings/nvpl/_internal/fft.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/nvpl/_internal/fft.pxd +36 -0
- nvmath/bindings/nvpl/blas.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/nvpl/blas.pxd +131 -0
- nvmath/bindings/nvpl/blas.pyi +168 -0
- nvmath/bindings/nvpl/cyblas.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/nvpl/cyblas.pxd +280 -0
- nvmath/bindings/nvpl/cyfft.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/nvpl/cyfft.pxd +93 -0
- nvmath/bindings/nvpl/fft.cp314t-win_amd64.pyd +0 -0
- nvmath/bindings/nvpl/fft.pxd +100 -0
- nvmath/bindings/nvpl/fft.pyi +168 -0
- nvmath/bindings/nvshmem.pxd +54 -0
- nvmath/bindings/nvshmem.pyi +191 -0
- nvmath/device/__init__.py +38 -0
- nvmath/device/_deprecated.py +33 -0
- nvmath/device/common.py +315 -0
- nvmath/device/common_backend.py +131 -0
- nvmath/device/common_cuda.py +201 -0
- nvmath/device/common_numba.py +300 -0
- nvmath/device/common_numba_cuda_mlir.py +202 -0
- nvmath/device/common_opaque_tensor.py +201 -0
- nvmath/device/cublasdx.py +1606 -0
- nvmath/device/cublasdx_backend.py +860 -0
- nvmath/device/cublasdx_numba.py +1534 -0
- nvmath/device/cublasdx_numba_cuda_mlir.py +208 -0
- nvmath/device/cufftdx.py +373 -0
- nvmath/device/cufftdx_backend.py +220 -0
- nvmath/device/cufftdx_numba.py +140 -0
- nvmath/device/cufftdx_numba_cuda_mlir.py +79 -0
- nvmath/device/curand_kernel.py +9147 -0
- nvmath/device/cusolverdx.py +2708 -0
- nvmath/device/cusolverdx_backend.py +440 -0
- nvmath/device/cusolverdx_numba.py +567 -0
- nvmath/device/cusolverdx_numba_cuda_mlir.py +604 -0
- nvmath/device/cusolverdx_overload_backend.py +1029 -0
- nvmath/device/llvm_array.py +29 -0
- nvmath/device/random.py +441 -0
- nvmath/device/random_helpers.py +23 -0
- nvmath/device/random_states.py +187 -0
- nvmath/device/types.py +138 -0
- nvmath/device/vector_types_numba.py +259 -0
- nvmath/distributed/__init__.py +200 -0
- nvmath/distributed/_internal/__init__.py +0 -0
- nvmath/distributed/_internal/nccl.py +86 -0
- nvmath/distributed/_internal/nvshmem.py +307 -0
- nvmath/distributed/_internal/symmetric_memory.py +35 -0
- nvmath/distributed/_internal/tensor_ifc.py +70 -0
- nvmath/distributed/_internal/tensor_ifc_cupy.py +68 -0
- nvmath/distributed/_internal/tensor_ifc_host_device.py +172 -0
- nvmath/distributed/_internal/tensor_ifc_numpy.py +46 -0
- nvmath/distributed/_internal/tensor_ifc_torch.py +162 -0
- nvmath/distributed/_internal/tensor_wrapper.py +81 -0
- nvmath/distributed/_utils.py +167 -0
- nvmath/distributed/distribution/__init__.py +30 -0
- nvmath/distributed/distribution/_configuration.py +39 -0
- nvmath/distributed/distribution/distributions.py +1024 -0
- nvmath/distributed/distribution/redistribute.py +1284 -0
- nvmath/distributed/fft/__init__.py +7 -0
- nvmath/distributed/fft/_configuration.py +82 -0
- nvmath/distributed/fft/fft.py +2742 -0
- nvmath/distributed/linalg/__init__.py +22 -0
- nvmath/distributed/linalg/_internal/__init__.py +3 -0
- nvmath/distributed/linalg/_internal/epilog_protocol.py +586 -0
- nvmath/distributed/linalg/_internal/matmul_desc_ifc.py +28 -0
- nvmath/distributed/linalg/advanced/__init__.py +8 -0
- nvmath/distributed/linalg/advanced/_configuration.py +171 -0
- nvmath/distributed/linalg/advanced/matmulmod.py +3573 -0
- nvmath/distributed/linalg/generic/__init__.py +8 -0
- nvmath/distributed/linalg/generic/_caching.py +66 -0
- nvmath/distributed/linalg/generic/_configuration.py +61 -0
- nvmath/distributed/linalg/generic/_factorization.py +172 -0
- nvmath/distributed/linalg/generic/_initialization.py +966 -0
- nvmath/distributed/linalg/generic/_problem_spec.py +511 -0
- nvmath/distributed/linalg/generic/solvermod.py +1368 -0
- nvmath/distributed/process_group.py +408 -0
- nvmath/fft/__init__.py +7 -0
- nvmath/fft/_configuration.py +189 -0
- nvmath/fft/_exec_utils.py +82 -0
- nvmath/fft/_helpers.py +237 -0
- nvmath/fft/fft.py +3122 -0
- nvmath/internal/__init__.pxd +3 -0
- nvmath/internal/__init__.py +10 -0
- nvmath/internal/_bindings.cp314t-win_amd64.pyd +0 -0
- nvmath/internal/_bindings.pxd +18 -0
- nvmath/internal/_device_utils.py +45 -0
- nvmath/internal/_layout/__init__.pxd +3 -0
- nvmath/internal/_layout/__init__.py +7 -0
- nvmath/internal/_layout/_layout.cp314t-win_amd64.pyd +0 -0
- nvmath/internal/_layout/_layout.pxd +1303 -0
- nvmath/internal/_layout/_layout.pyi +1145 -0
- nvmath/internal/enum_utils.py +142 -0
- nvmath/internal/formatters.py +87 -0
- nvmath/internal/mem_limit.py +51 -0
- nvmath/internal/memory.cp314t-win_amd64.pyd +0 -0
- nvmath/internal/memory.pxd +13 -0
- nvmath/internal/memory.pyi +50 -0
- nvmath/internal/ndbuffer/__init__.pxd +3 -0
- nvmath/internal/ndbuffer/__init__.py +9 -0
- nvmath/internal/ndbuffer/_copy_kernel.cp314t-win_amd64.pyd +0 -0
- nvmath/internal/ndbuffer/_copy_kernel.pxd +10 -0
- nvmath/internal/ndbuffer/_jit.cp314t-win_amd64.pyd +0 -0
- nvmath/internal/ndbuffer/_jit.pxd +7 -0
- nvmath/internal/ndbuffer/_ndbuffer.cp314t-win_amd64.pyd +0 -0
- nvmath/internal/ndbuffer/_ndbuffer.pxd +40 -0
- nvmath/internal/ndbuffer/_ndbuffer.pyi +463 -0
- nvmath/internal/ndbuffer/copy_kernel/args.h +34 -0
- nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/array_view.h +52 -0
- nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/elementwise.h +68 -0
- nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/grid_indexer.h +69 -0
- nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/transposed.h +242 -0
- nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/type_utils.h +39 -0
- nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/utils.h +132 -0
- nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/vec.h +159 -0
- nvmath/internal/ndbuffer/copy_kernel/elementwise.h +53 -0
- nvmath/internal/ndbuffer/copy_kernel/transposed.h +58 -0
- nvmath/internal/package_ifc.py +168 -0
- nvmath/internal/package_ifc_cuda.py +57 -0
- nvmath/internal/package_ifc_cupy.py +67 -0
- nvmath/internal/package_ifc_torch.py +69 -0
- nvmath/internal/package_wrapper.py +14 -0
- nvmath/internal/tensor_ifc.py +179 -0
- nvmath/internal/tensor_ifc_cupy.py +234 -0
- nvmath/internal/tensor_ifc_ndbuffer.py +147 -0
- nvmath/internal/tensor_ifc_numpy.py +184 -0
- nvmath/internal/tensor_ifc_torch.py +178 -0
- nvmath/internal/tensor_wrapper.py +160 -0
- nvmath/internal/typemaps.py +113 -0
- nvmath/internal/utils.py +805 -0
- nvmath/linalg/__init__.py +56 -0
- nvmath/linalg/_internal/__init__.py +3 -0
- nvmath/linalg/_internal/algo_cap_ifc.py +82 -0
- nvmath/linalg/_internal/algo_config_ifc.py +43 -0
- nvmath/linalg/_internal/batch.py +234 -0
- nvmath/linalg/_internal/enum_to_tuples.py +64 -0
- nvmath/linalg/_internal/epilog_protocol.py +766 -0
- nvmath/linalg/_internal/layout.py +624 -0
- nvmath/linalg/_internal/matmul_desc_ifc.py +28 -0
- nvmath/linalg/_internal/matmul_pref_ifc.py +27 -0
- nvmath/linalg/_internal/matrix_layout_ifc.py +26 -0
- nvmath/linalg/_internal/solver_utils.py +432 -0
- nvmath/linalg/_internal/typemaps.py +144 -0
- nvmath/linalg/_internal/utils.py +157 -0
- nvmath/linalg/advanced/__init__.py +8 -0
- nvmath/linalg/advanced/_algorithmmod.py +170 -0
- nvmath/linalg/advanced/_configuration.py +351 -0
- nvmath/linalg/advanced/helpers/__init__.py +5 -0
- nvmath/linalg/advanced/helpers/matmul.py +1316 -0
- nvmath/linalg/advanced/matmulmod.py +3734 -0
- nvmath/linalg/generic/__init__.py +53 -0
- nvmath/linalg/generic/_configuration/__init__.py +39 -0
- nvmath/linalg/generic/_configuration/layout.py +263 -0
- nvmath/linalg/generic/_configuration/match.py +734 -0
- nvmath/linalg/generic/_configuration/qualifiers.py +493 -0
- nvmath/linalg/generic/_configuration/solver_configuration.py +59 -0
- nvmath/linalg/generic/_configuration/wrap.py +217 -0
- nvmath/linalg/generic/_dtype.py +15 -0
- nvmath/linalg/generic/matmulmod.py +2094 -0
- nvmath/linalg/generic/solvermod.py +1301 -0
- nvmath/memory.py +279 -0
- nvmath/sparse/__init__.py +38 -0
- nvmath/sparse/_internal/__init__.py +21 -0
- nvmath/sparse/_internal/common_utils.py +147 -0
- nvmath/sparse/_internal/cudss_config_ifc.py +702 -0
- nvmath/sparse/_internal/cudss_data_ifc.py +399 -0
- nvmath/sparse/_internal/cudss_utils.py +506 -0
- nvmath/sparse/_internal/cusparse_utils.py +382 -0
- nvmath/sparse/_internal/sparse_bsc_ifc.py +303 -0
- nvmath/sparse/_internal/sparse_bsr_ifc.py +305 -0
- nvmath/sparse/_internal/sparse_coo_ifc.py +256 -0
- nvmath/sparse/_internal/sparse_csc_ifc.py +268 -0
- nvmath/sparse/_internal/sparse_csr_ifc.py +288 -0
- nvmath/sparse/_internal/sparse_dia_ifc.py +242 -0
- nvmath/sparse/_internal/sparse_format_helpers.py +601 -0
- nvmath/sparse/_internal/sparse_tensor_ifc.py +133 -0
- nvmath/sparse/_internal/sparse_ust_ifc.py +141 -0
- nvmath/sparse/_internal/utils.py +56 -0
- nvmath/sparse/advanced/__init__.py +7 -0
- nvmath/sparse/advanced/_configuration.py +227 -0
- nvmath/sparse/advanced/direct_solver.py +2069 -0
- nvmath/sparse/generic/__init__.py +7 -0
- nvmath/sparse/generic/_configuration.py +129 -0
- nvmath/sparse/generic/_helpers.py +137 -0
- nvmath/sparse/generic/_thunks.py +21 -0
- nvmath/sparse/generic/matmulmod.py +2353 -0
- nvmath/sparse/ust/__init__.py +7 -0
- nvmath/sparse/ust/_converters.py +422 -0
- nvmath/sparse/ust/_cpp.py +28 -0
- nvmath/sparse/ust/_drawer.py +565 -0
- nvmath/sparse/ust/_emitter.py +1033 -0
- nvmath/sparse/ust/_jit.py +188 -0
- nvmath/sparse/ust/_kernel.py +282 -0
- nvmath/sparse/ust/_utils.py +149 -0
- nvmath/sparse/ust/interfaces/__init__.py +0 -0
- nvmath/sparse/ust/interfaces/torch_interface.py +476 -0
- nvmath/sparse/ust/tensor.py +1016 -0
- nvmath/sparse/ust/tensor_format.py +957 -0
- nvmath/tensor/__init__.py +6 -0
- nvmath/tensor/_configuration.py +120 -0
- nvmath/tensor/_internal/__init__.py +3 -0
- nvmath/tensor/_internal/cutensor_config_ifc.py +279 -0
- nvmath/tensor/_internal/cutensor_utils.py +230 -0
- nvmath/tensor/_internal/data.py +43 -0
- nvmath/tensor/_internal/einsum_parser.py +444 -0
- nvmath/tensor/_internal/typemaps.py +96 -0
- nvmath/tensor/contract.py +1900 -0
- nvmath_python-1.0.0.dist-info/METADATA +134 -0
- nvmath_python-1.0.0.dist-info/RECORD +332 -0
- nvmath_python-1.0.0.dist-info/WHEEL +5 -0
- nvmath_python-1.0.0.dist-info/licenses/LICENSE +177 -0
- 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")
|