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
nvmath/internal/utils.py
ADDED
|
@@ -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"
|