pyqcd 0.1.1__tar.gz
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.
- pyqcd-0.1.1/LICENSE +21 -0
- pyqcd-0.1.1/MANIFEST.in +28 -0
- pyqcd-0.1.1/PKG-INFO +133 -0
- pyqcd-0.1.1/README.md +100 -0
- pyqcd-0.1.1/build-kupl.sh +19 -0
- pyqcd-0.1.1/build-nv.sh +19 -0
- pyqcd-0.1.1/cpp/cuda/einsum/AGENTS.md +201 -0
- pyqcd-0.1.1/cpp/cuda/einsum/include/einsum_cuda.h +117 -0
- pyqcd-0.1.1/cpp/cuda/einsum/include/einsum_define.h +99 -0
- pyqcd-0.1.1/cpp/cuda/einsum/include/nucleon2pt.h +180 -0
- pyqcd-0.1.1/cpp/cuda/einsum/include/pattern_3pt.h +119 -0
- pyqcd-0.1.1/cpp/cuda/einsum/make-nv.sh +230 -0
- pyqcd-0.1.1/cpp/cuda/einsum/python/pyqcdeinsum.h +20 -0
- pyqcd-0.1.1/cpp/cuda/einsum/src/apply_contract.cu +1896 -0
- pyqcd-0.1.1/cpp/cuda/einsum/src/apply_plan.cu +197 -0
- pyqcd-0.1.1/cpp/cuda/einsum/src/batched_gemm.cu +435 -0
- pyqcd-0.1.1/cpp/cuda/einsum/src/e4_fused.cu +998 -0
- pyqcd-0.1.1/cpp/cuda/einsum/src/nucleon2pt.cu +1003 -0
- pyqcd-0.1.1/cpp/cuda/einsum/src/pack.cu +433 -0
- pyqcd-0.1.1/cpp/cuda/einsum/src/pattern_3pt.cu +977 -0
- pyqcd-0.1.1/cpp/cuda/einsum/src/sandwich.cu +368 -0
- pyqcd-0.1.1/cpp/kupl/einsum/AGENTS.md +222 -0
- pyqcd-0.1.1/cpp/kupl/einsum/include/einsum_define.h +99 -0
- pyqcd-0.1.1/cpp/kupl/einsum/include/einsum_mpi.h +17 -0
- pyqcd-0.1.1/cpp/kupl/einsum/include/einsum_plan.h +40 -0
- pyqcd-0.1.1/cpp/kupl/einsum/include/einsum_runtime.h +90 -0
- pyqcd-0.1.1/cpp/kupl/einsum/include/nucleon2pt.h +131 -0
- pyqcd-0.1.1/cpp/kupl/einsum/include/pattern_3pt.h +114 -0
- pyqcd-0.1.1/cpp/kupl/einsum/make.sh +238 -0
- pyqcd-0.1.1/cpp/kupl/einsum/src/apply_contract.cpp +1984 -0
- pyqcd-0.1.1/cpp/kupl/einsum/src/apply_plan.cpp +206 -0
- pyqcd-0.1.1/cpp/kupl/einsum/src/mpi_runtime.cpp +555 -0
- pyqcd-0.1.1/cpp/kupl/einsum/src/nucleon2pt.cpp +1137 -0
- pyqcd-0.1.1/cpp/kupl/einsum/src/pattern_3pt.cpp +1163 -0
- pyqcd-0.1.1/cpp/kupl/einsum/src/plan_core.cpp +649 -0
- pyqcd-0.1.1/cpp/kupl/einsum/src/runtime.cpp +456 -0
- pyqcd-0.1.1/cpp/kupl/einsum/tests/test_runtime.cpp +52 -0
- pyqcd-0.1.1/env.sh +69 -0
- pyqcd-0.1.1/install.sh +70 -0
- pyqcd-0.1.1/pyproject.toml +49 -0
- pyqcd-0.1.1/pyqcd/__init__.py +3 -0
- pyqcd-0.1.1/pyqcd/einsum/__init__.py +97 -0
- pyqcd-0.1.1/pyqcd/einsum/_api.py +1155 -0
- pyqcd-0.1.1/pyqcd/einsum/_execute.py +189 -0
- pyqcd-0.1.1/pyqcd/einsum/_expr.py +336 -0
- pyqcd-0.1.1/pyqcd/einsum/_kernel_cuda.py +240 -0
- pyqcd-0.1.1/pyqcd/einsum/_kernel_kupl.py +243 -0
- pyqcd-0.1.1/pyqcd/einsum/_native.py +209 -0
- pyqcd-0.1.1/pyqcd/einsum/_native_plan.py +559 -0
- pyqcd-0.1.1/pyqcd/einsum/_pathinfo.py +77 -0
- pyqcd-0.1.1/pyqcd/einsum/_plan_py.py +917 -0
- pyqcd-0.1.1/pyqcd/einsum/_sharing.py +206 -0
- pyqcd-0.1.1/pyqcd/einsum/_specialize.py +116 -0
- pyqcd-0.1.1/pyqcd/einsum/_version.py +4 -0
- pyqcd-0.1.1/pyqcd/einsum/backends/__init__.py +39 -0
- pyqcd-0.1.1/pyqcd/einsum/backends/dispatch.py +357 -0
- pyqcd-0.1.1/pyqcd/einsum/blas.py +69 -0
- pyqcd-0.1.1/pyqcd/einsum/cuda/__init__.py +68 -0
- pyqcd-0.1.1/pyqcd/einsum/cuda/_program.py +129 -0
- pyqcd-0.1.1/pyqcd/einsum/cuda/einsum.pyi +59 -0
- pyqcd-0.1.1/pyqcd/einsum/cuda/einsum.pyx +223 -0
- pyqcd-0.1.1/pyqcd/einsum/cuda/einsum_api.pxd +31 -0
- pyqcd-0.1.1/pyqcd/einsum/cuda/mpi.py +380 -0
- pyqcd-0.1.1/pyqcd/einsum/helpers.py +44 -0
- pyqcd-0.1.1/pyqcd/einsum/kupl/__init__.py +12 -0
- pyqcd-0.1.1/pyqcd/einsum/kupl/einsum.py +650 -0
- pyqcd-0.1.1/pyqcd/einsum/kupl/mpi.py +374 -0
- pyqcd-0.1.1/pyqcd/einsum/kupl/multiprocess.py +753 -0
- pyqcd-0.1.1/pyqcd/einsum/parser.py +238 -0
- pyqcd-0.1.1/pyqcd/einsum/paths.py +146 -0
- pyqcd-0.1.1/pyqcd.egg-info/PKG-INFO +133 -0
- pyqcd-0.1.1/pyqcd.egg-info/SOURCES.txt +76 -0
- pyqcd-0.1.1/pyqcd.egg-info/dependency_links.txt +1 -0
- pyqcd-0.1.1/pyqcd.egg-info/not-zip-safe +1 -0
- pyqcd-0.1.1/pyqcd.egg-info/requires.txt +7 -0
- pyqcd-0.1.1/pyqcd.egg-info/top_level.txt +1 -0
- pyqcd-0.1.1/setup.cfg +4 -0
- pyqcd-0.1.1/setup.py +239 -0
pyqcd-0.1.1/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 ZhangXin8069
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
pyqcd-0.1.1/MANIFEST.in
ADDED
|
@@ -0,0 +1,28 @@
|
|
|
1
|
+
# Files shipped inside the sdist. Metadata lives in pyproject.toml; the build
|
|
2
|
+
# mechanics (native staging, optional Cython extension) live in setup.py.
|
|
3
|
+
|
|
4
|
+
include LICENSE
|
|
5
|
+
include README.md
|
|
6
|
+
include pyproject.toml
|
|
7
|
+
include setup.py
|
|
8
|
+
include MANIFEST.in
|
|
9
|
+
|
|
10
|
+
# Scripts setup.py invokes to compile the native C/CUDA backends.
|
|
11
|
+
include build-kupl.sh
|
|
12
|
+
include build-nv.sh
|
|
13
|
+
include env.sh
|
|
14
|
+
include install.sh
|
|
15
|
+
|
|
16
|
+
# Native backend sources, headers and their Makefile drivers. The build dirs
|
|
17
|
+
# under cpp/ hold compiled artefacts and are pruned below.
|
|
18
|
+
recursive-include cpp *.h *.hpp *.c *.cc *.cpp *.cu *.cuh *.sh *.md
|
|
19
|
+
prune cpp/kupl/einsum/build
|
|
20
|
+
prune cpp/cuda/einsum/build
|
|
21
|
+
|
|
22
|
+
# Pure-Python package plus the Cython sources of the optional CUDA extension.
|
|
23
|
+
recursive-include pyqcd/einsum *.py *.pyx *.pxd *.pyi
|
|
24
|
+
|
|
25
|
+
# Build artefacts must never reach the sdist.
|
|
26
|
+
prune build
|
|
27
|
+
recursive-exclude pyqcd/einsum __pycache__
|
|
28
|
+
global-exclude *.so *.so.* *.o *.obj *.a *.lib *.log *.pyc
|
pyqcd-0.1.1/PKG-INFO
ADDED
|
@@ -0,0 +1,133 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: pyqcd
|
|
3
|
+
Version: 0.1.1
|
|
4
|
+
Summary: Lattice QCD toolkit with a native contraction engine and an opt_einsum-compatible API.
|
|
5
|
+
Author: ZhangXin8069
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Homepage, https://gitlab.itp.ac.cn/zhangxin8069/PyQCD
|
|
8
|
+
Project-URL: Repository, https://gitlab.itp.ac.cn/zhangxin8069/PyQCD
|
|
9
|
+
Project-URL: Issues, https://gitlab.itp.ac.cn/zhangxin8069/PyQCD/-/issues
|
|
10
|
+
Project-URL: Documentation, https://py-qcd.readthedocs.io/
|
|
11
|
+
Keywords: lattice-qcd,einsum,tensor-contraction,opt-einsum,cuda,mpi
|
|
12
|
+
Classifier: Development Status :: 3 - Alpha
|
|
13
|
+
Classifier: Intended Audience :: Science/Research
|
|
14
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
15
|
+
Classifier: Operating System :: POSIX :: Linux
|
|
16
|
+
Classifier: Programming Language :: C++
|
|
17
|
+
Classifier: Programming Language :: Python :: 3
|
|
18
|
+
Classifier: Programming Language :: Python :: 3 :: Only
|
|
19
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
20
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
21
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
22
|
+
Classifier: Topic :: Scientific/Engineering :: Physics
|
|
23
|
+
Classifier: Typing :: Typed
|
|
24
|
+
Requires-Python: >=3.10
|
|
25
|
+
Description-Content-Type: text/markdown
|
|
26
|
+
License-File: LICENSE
|
|
27
|
+
Requires-Dist: numpy>=1.23
|
|
28
|
+
Provides-Extra: cuda
|
|
29
|
+
Requires-Dist: cupy; extra == "cuda"
|
|
30
|
+
Provides-Extra: mpi
|
|
31
|
+
Requires-Dist: mpi4py; extra == "mpi"
|
|
32
|
+
Dynamic: license-file
|
|
33
|
+
|
|
34
|
+
# PyQCD
|
|
35
|
+
|
|
36
|
+
PyQCD is a lattice-QCD toolkit in Python with optional native acceleration.
|
|
37
|
+
It provides building blocks for propagator and operator contractions, gradient
|
|
38
|
+
flow, renormalisation and gluon TMD-PDF extraction, plus a self-contained
|
|
39
|
+
contraction engine with an **`opt_einsum`-compatible API**.
|
|
40
|
+
|
|
41
|
+
## Install
|
|
42
|
+
|
|
43
|
+
```bash
|
|
44
|
+
pip install pyqcd
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
The native kernels are **optional**: without a C++ compiler (or without CUDA)
|
|
48
|
+
the package installs anyway and the contraction engine runs its pure-Python
|
|
49
|
+
planner — same results, only slower.
|
|
50
|
+
|
|
51
|
+
Set `PYQCD_EINSUM_REQUIRE_NATIVE=1` to make a missing toolchain a hard error.
|
|
52
|
+
|
|
53
|
+
## Documentation
|
|
54
|
+
|
|
55
|
+
* English — https://py-qcd.readthedocs.io/en/latest/usage.html
|
|
56
|
+
* 中文 — https://py-qcd.readthedocs.io/zh-cn/latest/usage.html
|
|
57
|
+
|
|
58
|
+
## Quick start
|
|
59
|
+
|
|
60
|
+
```python
|
|
61
|
+
import numpy as np
|
|
62
|
+
import pyqcd.einsum as pe
|
|
63
|
+
|
|
64
|
+
a = np.arange(6.0).reshape(2, 3)
|
|
65
|
+
b = np.arange(12.0).reshape(3, 4)
|
|
66
|
+
|
|
67
|
+
pe.contract("ij,jk->ik", a, b) # like a @ b
|
|
68
|
+
path, info = pe.contract_path("ij,jk->ik", a, b) # (path, PathInfo)
|
|
69
|
+
int(info.opt_cost)
|
|
70
|
+
|
|
71
|
+
expr = pe.contract_expression("ij,jk->ik", (2, 3), (3, 4))
|
|
72
|
+
expr(a, b) # reusable plan
|
|
73
|
+
```
|
|
74
|
+
|
|
75
|
+
The public surface mirrors `opt_einsum`: `contract`, `contract_path`,
|
|
76
|
+
`contract_expression`, `shared_intermediates`, `PathInfo`, `paths`, `backends`,
|
|
77
|
+
`helpers`, `parser`, `blas`.
|
|
78
|
+
|
|
79
|
+
## What is inside
|
|
80
|
+
|
|
81
|
+
| Subpackage | Content |
|
|
82
|
+
|---|---|
|
|
83
|
+
| `pyqcd.lattice` | lattice constants, DeGrand–Rossi `gamma`, SU(2) Clebsch–Gordan |
|
|
84
|
+
| `pyqcd.gauge` | Wilson/Polyakov loops, Clover field strength, topological charge |
|
|
85
|
+
| `pyqcd.smear` | HYP and Stout smearing |
|
|
86
|
+
| `pyqcd.operator` | Clover / dual field strength, flowed gluon OPE operators |
|
|
87
|
+
| `pyqcd.vertex` | VdV/VVV vertices, momentum smearing, eigenvector compression |
|
|
88
|
+
| `pyqcd.contraction` | Wick contraction, baryon operators, distillation contractions |
|
|
89
|
+
| `pyqcd.renorm` | gradient flow, Z_R, hybrid scheme, NLO matching, TMD extraction |
|
|
90
|
+
| `pyqcd.analysis` | resampling, fits, ratio / effective-mass / energy helpers |
|
|
91
|
+
| `pyqcd.pipeline` | centralised configuration and multi-step pipeline |
|
|
92
|
+
| `pyqcd.parallel` | resource detection and MPI meta-task planning |
|
|
93
|
+
| `pyqcd.einsum` | contraction engine (this package's headline feature) |
|
|
94
|
+
|
|
95
|
+
## Contraction engine
|
|
96
|
+
|
|
97
|
+
* Pure-Python planning (`greedy` / `optimal` / `branch` / `dp`) is always
|
|
98
|
+
available.
|
|
99
|
+
* Two optional native backends: **KUPL** (CPU) and **CUDA** (GPU). Both export
|
|
100
|
+
the same C ABI and share one planning core.
|
|
101
|
+
* Fallback is a hard invariant: a missing library, a failed native call or a
|
|
102
|
+
rejected shape silently degrades to the Python path — results never change.
|
|
103
|
+
* A small set of lattice-QCD equation shapes gets a specialised execution path;
|
|
104
|
+
it is enabled by default and can be turned off per call
|
|
105
|
+
(`contract(..., specialize=False)`) or per process
|
|
106
|
+
(`PYQCD_EINSUM_SPECIALIZE=0`).
|
|
107
|
+
|
|
108
|
+
Introspection:
|
|
109
|
+
|
|
110
|
+
```python
|
|
111
|
+
import pyqcd.einsum as pe
|
|
112
|
+
|
|
113
|
+
pe.build_info() # backend, mpi, mpi_label, native_library
|
|
114
|
+
pe.mpi_enabled() # whether the KUPL library has real MPI support
|
|
115
|
+
pe.specializations() # registered equation shapes
|
|
116
|
+
```
|
|
117
|
+
|
|
118
|
+
MPI sharding exists only in the KUPL backend and only shards a single output
|
|
119
|
+
axis; the CUDA backend has no MPI implementation.
|
|
120
|
+
|
|
121
|
+
## Tests
|
|
122
|
+
|
|
123
|
+
Run from a source checkout:
|
|
124
|
+
|
|
125
|
+
```bash
|
|
126
|
+
python -m pyqcd.testing # integration suite
|
|
127
|
+
bash pyqcd/testing/simple/run-all.sh # per-function API parity vs opt_einsum
|
|
128
|
+
bash pyqcd/testing/einsum/simple/run-all.sh # einsum sub-system suite
|
|
129
|
+
```
|
|
130
|
+
|
|
131
|
+
## License
|
|
132
|
+
|
|
133
|
+
MIT
|
pyqcd-0.1.1/README.md
ADDED
|
@@ -0,0 +1,100 @@
|
|
|
1
|
+
# PyQCD
|
|
2
|
+
|
|
3
|
+
PyQCD is a lattice-QCD toolkit in Python with optional native acceleration.
|
|
4
|
+
It provides building blocks for propagator and operator contractions, gradient
|
|
5
|
+
flow, renormalisation and gluon TMD-PDF extraction, plus a self-contained
|
|
6
|
+
contraction engine with an **`opt_einsum`-compatible API**.
|
|
7
|
+
|
|
8
|
+
## Install
|
|
9
|
+
|
|
10
|
+
```bash
|
|
11
|
+
pip install pyqcd
|
|
12
|
+
```
|
|
13
|
+
|
|
14
|
+
The native kernels are **optional**: without a C++ compiler (or without CUDA)
|
|
15
|
+
the package installs anyway and the contraction engine runs its pure-Python
|
|
16
|
+
planner — same results, only slower.
|
|
17
|
+
|
|
18
|
+
Set `PYQCD_EINSUM_REQUIRE_NATIVE=1` to make a missing toolchain a hard error.
|
|
19
|
+
|
|
20
|
+
## Documentation
|
|
21
|
+
|
|
22
|
+
* English — https://py-qcd.readthedocs.io/en/latest/usage.html
|
|
23
|
+
* 中文 — https://py-qcd.readthedocs.io/zh-cn/latest/usage.html
|
|
24
|
+
|
|
25
|
+
## Quick start
|
|
26
|
+
|
|
27
|
+
```python
|
|
28
|
+
import numpy as np
|
|
29
|
+
import pyqcd.einsum as pe
|
|
30
|
+
|
|
31
|
+
a = np.arange(6.0).reshape(2, 3)
|
|
32
|
+
b = np.arange(12.0).reshape(3, 4)
|
|
33
|
+
|
|
34
|
+
pe.contract("ij,jk->ik", a, b) # like a @ b
|
|
35
|
+
path, info = pe.contract_path("ij,jk->ik", a, b) # (path, PathInfo)
|
|
36
|
+
int(info.opt_cost)
|
|
37
|
+
|
|
38
|
+
expr = pe.contract_expression("ij,jk->ik", (2, 3), (3, 4))
|
|
39
|
+
expr(a, b) # reusable plan
|
|
40
|
+
```
|
|
41
|
+
|
|
42
|
+
The public surface mirrors `opt_einsum`: `contract`, `contract_path`,
|
|
43
|
+
`contract_expression`, `shared_intermediates`, `PathInfo`, `paths`, `backends`,
|
|
44
|
+
`helpers`, `parser`, `blas`.
|
|
45
|
+
|
|
46
|
+
## What is inside
|
|
47
|
+
|
|
48
|
+
| Subpackage | Content |
|
|
49
|
+
|---|---|
|
|
50
|
+
| `pyqcd.lattice` | lattice constants, DeGrand–Rossi `gamma`, SU(2) Clebsch–Gordan |
|
|
51
|
+
| `pyqcd.gauge` | Wilson/Polyakov loops, Clover field strength, topological charge |
|
|
52
|
+
| `pyqcd.smear` | HYP and Stout smearing |
|
|
53
|
+
| `pyqcd.operator` | Clover / dual field strength, flowed gluon OPE operators |
|
|
54
|
+
| `pyqcd.vertex` | VdV/VVV vertices, momentum smearing, eigenvector compression |
|
|
55
|
+
| `pyqcd.contraction` | Wick contraction, baryon operators, distillation contractions |
|
|
56
|
+
| `pyqcd.renorm` | gradient flow, Z_R, hybrid scheme, NLO matching, TMD extraction |
|
|
57
|
+
| `pyqcd.analysis` | resampling, fits, ratio / effective-mass / energy helpers |
|
|
58
|
+
| `pyqcd.pipeline` | centralised configuration and multi-step pipeline |
|
|
59
|
+
| `pyqcd.parallel` | resource detection and MPI meta-task planning |
|
|
60
|
+
| `pyqcd.einsum` | contraction engine (this package's headline feature) |
|
|
61
|
+
|
|
62
|
+
## Contraction engine
|
|
63
|
+
|
|
64
|
+
* Pure-Python planning (`greedy` / `optimal` / `branch` / `dp`) is always
|
|
65
|
+
available.
|
|
66
|
+
* Two optional native backends: **KUPL** (CPU) and **CUDA** (GPU). Both export
|
|
67
|
+
the same C ABI and share one planning core.
|
|
68
|
+
* Fallback is a hard invariant: a missing library, a failed native call or a
|
|
69
|
+
rejected shape silently degrades to the Python path — results never change.
|
|
70
|
+
* A small set of lattice-QCD equation shapes gets a specialised execution path;
|
|
71
|
+
it is enabled by default and can be turned off per call
|
|
72
|
+
(`contract(..., specialize=False)`) or per process
|
|
73
|
+
(`PYQCD_EINSUM_SPECIALIZE=0`).
|
|
74
|
+
|
|
75
|
+
Introspection:
|
|
76
|
+
|
|
77
|
+
```python
|
|
78
|
+
import pyqcd.einsum as pe
|
|
79
|
+
|
|
80
|
+
pe.build_info() # backend, mpi, mpi_label, native_library
|
|
81
|
+
pe.mpi_enabled() # whether the KUPL library has real MPI support
|
|
82
|
+
pe.specializations() # registered equation shapes
|
|
83
|
+
```
|
|
84
|
+
|
|
85
|
+
MPI sharding exists only in the KUPL backend and only shards a single output
|
|
86
|
+
axis; the CUDA backend has no MPI implementation.
|
|
87
|
+
|
|
88
|
+
## Tests
|
|
89
|
+
|
|
90
|
+
Run from a source checkout:
|
|
91
|
+
|
|
92
|
+
```bash
|
|
93
|
+
python -m pyqcd.testing # integration suite
|
|
94
|
+
bash pyqcd/testing/simple/run-all.sh # per-function API parity vs opt_einsum
|
|
95
|
+
bash pyqcd/testing/einsum/simple/run-all.sh # einsum sub-system suite
|
|
96
|
+
```
|
|
97
|
+
|
|
98
|
+
## License
|
|
99
|
+
|
|
100
|
+
MIT
|
|
@@ -0,0 +1,19 @@
|
|
|
1
|
+
#!/usr/bin/env bash
|
|
2
|
+
# Build the KUPL shared library without involving CUDA or Cython.
|
|
3
|
+
set -euo pipefail
|
|
4
|
+
|
|
5
|
+
ROOT=$(cd "$(dirname "${BASH_SOURCE[0]:-$0}")" && pwd)
|
|
6
|
+
PYQCD_REQUESTED_BUILD_DIR=${PYQCD_EINSUM_BUILD_DIR:-}
|
|
7
|
+
export PYQCD_EINSUM_BACKEND=kupl
|
|
8
|
+
# shellcheck source=/dev/null
|
|
9
|
+
source "$ROOT/env.sh"
|
|
10
|
+
case "$PYQCD_REQUESTED_BUILD_DIR" in
|
|
11
|
+
""|"$PYQCD_EINSUM_KUPL_BUILD_DIR"|"$PYQCD_EINSUM_CUDA_BUILD_DIR")
|
|
12
|
+
export PYQCD_EINSUM_BUILD_DIR="$PYQCD_EINSUM_KUPL_BUILD_DIR"
|
|
13
|
+
;;
|
|
14
|
+
*) export PYQCD_EINSUM_BUILD_DIR="$PYQCD_REQUESTED_BUILD_DIR" ;;
|
|
15
|
+
esac
|
|
16
|
+
|
|
17
|
+
printf '=== PyQCD einsum build (kupl) ===\n'
|
|
18
|
+
bash "$ROOT/cpp/kupl/einsum/make.sh" "$@"
|
|
19
|
+
printf '=== PyQCD einsum build (kupl): SUCCESS ===\n'
|
pyqcd-0.1.1/build-nv.sh
ADDED
|
@@ -0,0 +1,19 @@
|
|
|
1
|
+
#!/usr/bin/env bash
|
|
2
|
+
# Build the CUDA shared library. A missing source set is an explicit failure.
|
|
3
|
+
set -euo pipefail
|
|
4
|
+
|
|
5
|
+
ROOT=$(cd "$(dirname "${BASH_SOURCE[0]:-$0}")" && pwd)
|
|
6
|
+
PYQCD_REQUESTED_BUILD_DIR=${PYQCD_EINSUM_BUILD_DIR:-}
|
|
7
|
+
export PYQCD_EINSUM_BACKEND=cuda
|
|
8
|
+
# shellcheck source=/dev/null
|
|
9
|
+
source "$ROOT/env.sh"
|
|
10
|
+
case "$PYQCD_REQUESTED_BUILD_DIR" in
|
|
11
|
+
""|"$PYQCD_EINSUM_KUPL_BUILD_DIR"|"$PYQCD_EINSUM_CUDA_BUILD_DIR")
|
|
12
|
+
export PYQCD_EINSUM_BUILD_DIR="$PYQCD_EINSUM_CUDA_BUILD_DIR"
|
|
13
|
+
;;
|
|
14
|
+
*) export PYQCD_EINSUM_BUILD_DIR="$PYQCD_REQUESTED_BUILD_DIR" ;;
|
|
15
|
+
esac
|
|
16
|
+
|
|
17
|
+
printf '=== PyQCD einsum build (cuda) ===\n'
|
|
18
|
+
bash "$ROOT/cpp/cuda/einsum/make-nv.sh" "$@"
|
|
19
|
+
printf '=== PyQCD einsum build (cuda): SUCCESS ===\n'
|
|
@@ -0,0 +1,201 @@
|
|
|
1
|
+
# cpp/cuda/einsum — CUDA 后端(与 KUPL 同一套 C ABI)
|
|
2
|
+
|
|
3
|
+
本目录产出 `libpyqcdeinsum_cuda.so`(soname 同名),导出面与 `cpp/kupl/einsum`
|
|
4
|
+
(产物 `libpyqcdeinsum_kupl.so`)相同的冻结四入口与
|
|
5
|
+
compiled-context 三入口。冻结契约见 `pyqcd/einsum/AGENTS.md` §5 与
|
|
6
|
+
`../../kupl/einsum/include/einsum_define.h`;本文件只登记本目录的落地细节,
|
|
7
|
+
冲突时以那两处为准。
|
|
8
|
+
|
|
9
|
+
## 1. 文件与归属
|
|
10
|
+
|
|
11
|
+
| 文件 | 职责 | 负责人 |
|
|
12
|
+
|---|---|---|
|
|
13
|
+
| `include/einsum_cuda.h` | pack/reduce/GEMM/stream 内部接口 | 原生组 |
|
|
14
|
+
| `src/apply_plan.cu` | `applyPlanQcd` 薄适配与会话钩子 | 原生组 |
|
|
15
|
+
| `src/apply_contract.cu` | version-2 计划、compiled context、stream pool、workspace | 原生组 |
|
|
16
|
+
| `src/pack.cu` | pack/transpose/reduce/hadamard CUDA kernel | 原生组 |
|
|
17
|
+
| `src/batched_gemm.cu` | cuBLAS K-last batched GEMM | 原生组 |
|
|
18
|
+
| `make-nv.sh` | nvcc 直编 + 原子软链发布 | 装配组(本组) |
|
|
19
|
+
| `include/einsum_define.h` | 冻结 C ABI 的**逐字节镜像**(不属本组,勿改) | 原生组(规划) |
|
|
20
|
+
| `include/einsum_permute.h`、`src/permute.cu` | `qcd_cuda_permute` / `qcd_cuda_hadamard` | A 组 |
|
|
21
|
+
| `include/einsum_batched_gemm.h`、`src/batched_gemm.cu` | `qcd_cuda_batched_gemm` | B 组 |
|
|
22
|
+
| `python/pyqcdeinsum.h` | Cython 侧的平坦 C 声明 | API 组 |
|
|
23
|
+
| `build/`、`libpyqcdeinsum_cuda.so` | 产物与稳定软链(不提交) | — |
|
|
24
|
+
|
|
25
|
+
本组只写第一行三件 + 本文件;`src/permute.cu`、`src/batched_gemm.cu`、
|
|
26
|
+
两个 kernel 头文件、`include/einsum_define.h`、`cpp/kupl/**` 一律不碰。
|
|
27
|
+
|
|
28
|
+
## 2. 不变量
|
|
29
|
+
|
|
30
|
+
1. **导出面**:`-Xcompiler -fvisibility=hidden` + `QCD_EINSUM_API` +
|
|
31
|
+
`-Xlinker --version-script`(`build/einsum_cuda_exports.map`)三重收敛,
|
|
32
|
+
`nm -D --defined-only` 只保留冻结四入口与 `applyCompileQcd` /
|
|
33
|
+
`applyCompiledQcd` / `applyDestroyQcd`;`qcd_cuda_*` 是内部符号。
|
|
34
|
+
2. **规划核心唯一真源**:`../../kupl/einsum/src/plan_core.cpp` 与
|
|
35
|
+
`../../kupl/einsum/include/einsum_plan.h` 被**直接编译**进本库,不复制、
|
|
36
|
+
不 fork。`applyPlanQcd` 只是 thin adapter(校验地址 → `PlanSpec` →
|
|
37
|
+
`pyqcd_einsum::plan_contract`),因此 CUDA 与 KUPL 的路径/代价由构造决定
|
|
38
|
+
逐位相同。
|
|
39
|
+
3. **ABI 镜像守卫**:`include/einsum_define.h` 必须与 KUPL 版本逐字节相同,
|
|
40
|
+
`make-nv.sh:44-49` 用 `cmp -s` 在编译前检查,漂移即报错退出。两库 soname
|
|
41
|
+
已经分开(`libpyqcdeinsum_cuda.so` / `libpyqcdeinsum_kupl.so`),动态链接器
|
|
42
|
+
不会替调用方发现 ABI 漂移,只能靠这道关卡。
|
|
43
|
+
4. **弱符号**:`qcd_cuda_permute` / `qcd_cuda_hadamard` /
|
|
44
|
+
`qcd_cuda_batched_gemm` 在 `einsum_cuda.h` 中以 `__attribute__((weak))`
|
|
45
|
+
声明。缺实现时链接成功、运行期符号为 NULL,派发侧返回
|
|
46
|
+
`QCD_STATUS_INTERNAL`(3)。**不用桩函数伪造成功**。
|
|
47
|
+
5. **不分配**:`apply_plan.cu` 不申请设备内存、不建流、不拷贝张量数据,
|
|
48
|
+
所有缓冲区由调用方持有;规划路径连 GPU 都不需要。
|
|
49
|
+
6. **状态码**:`0/1/2/3` 与 `einsum_define.h` 一致;kernel 返回值是合法枚举
|
|
50
|
+
时原样透传,否则归一为 `3`。异常绝不跨过 C 边界。
|
|
51
|
+
7. **fail-closed**:`applyContractQcd` 先校验**全部**步骤,任何一条不合法都
|
|
52
|
+
返回 `1` 且一个 kernel 都不启动。
|
|
53
|
+
|
|
54
|
+
## 3. 构建
|
|
55
|
+
|
|
56
|
+
```bash
|
|
57
|
+
bash cpp/cuda/einsum/make-nv.sh # 产物 cpp/cuda/einsum/libpyqcdeinsum_cuda.so
|
|
58
|
+
bash cpp/cuda/einsum/make-nv.sh --clean # 删除 build/ 与公开软链
|
|
59
|
+
bash cpp/cuda/einsum/make-nv.sh --require-kernels # 缺 kernel 实现时硬失败(CI 用)
|
|
60
|
+
```
|
|
61
|
+
|
|
62
|
+
实际命令(`nvcc` 12.4 / sm_89,架构按
|
|
63
|
+
`PYQCD_CUDA_ARCHITECTURES > CMAKE_CUDA_ARCHITECTURES > PYQCD_SM_ARCH >` 脚本内
|
|
64
|
+
`DEFAULT_ARCHITECTURES="6.0 7.0"` 解析,`make-nv.sh:126-134`,**不探测 `nvidia-smi`**):
|
|
65
|
+
|
|
66
|
+
```bash
|
|
67
|
+
nvcc -O3 -std=c++17 --shared -Xcompiler -fPIC -Xcompiler -fvisibility=hidden \
|
|
68
|
+
-gencode=arch=compute_89,code=sm_89 \
|
|
69
|
+
-I cpp/kupl/einsum/include -I cpp/cuda/einsum/include \
|
|
70
|
+
cpp/cuda/einsum/src/*.cu cpp/kupl/einsum/src/plan_core.cpp \
|
|
71
|
+
-lcublas -lcudart \
|
|
72
|
+
-Xlinker --version-script=build/einsum_cuda_exports.map \
|
|
73
|
+
-Xlinker -soname -Xlinker libpyqcdeinsum_cuda.so \
|
|
74
|
+
-o build/libpyqcdeinsum_cuda.so
|
|
75
|
+
```
|
|
76
|
+
|
|
77
|
+
- `-I` 顺序把 KUPL 的 include 放在前面:`einsum_plan.h` 与 `einsum_define.h`
|
|
78
|
+
都取自共享目录,本目录的 `include/einsum_define.h` 只是镜像。
|
|
79
|
+
- 先链接到 `build/libpyqcdeinsum_cuda.so.tmp.$$`,`mv -Tf` 落位;公开路径用临时
|
|
80
|
+
软链 + `mv -Tf` 原子替换(与 `cpp/kupl/einsum/make.sh`、PyQCU 一致)。
|
|
81
|
+
- 失败时打印 nvcc 原始输出(同时留档 `build/build.log`)并非零退出。
|
|
82
|
+
- `build/` 可用 `PYQCD_EINSUM_BUILD_DIR` 重定向。
|
|
83
|
+
- **缺 kernel 实现时**:默认仍产出"只有规划 ABI"的库,并在 stderr 打印缺哪
|
|
84
|
+
三个符号、`applyContractQcd` 将返回 `3`,退出码 **0**(上层
|
|
85
|
+
`build-nv.sh` / `_native.py` 需要这个库来规划);`--require-kernels`(或
|
|
86
|
+
`PYQCD_EINSUM_REQUIRE_KERNELS=1`)把这同一种状态变成**明确错误 + 退出码 1**。
|
|
87
|
+
- 本机实测的架构行为、失败模式与复现命令见 §7(2026-11-08 实测)。
|
|
88
|
+
|
|
89
|
+
## 4. 执行与并行
|
|
90
|
+
|
|
91
|
+
- `applyContractQcd` 解析 version-2 执行计划;left `[B,M,K]`、right `[B,N,K]`,
|
|
92
|
+
K 为两侧最低连续轴,cuBLAS 使用 transpose-B。
|
|
93
|
+
- compiled context 首次执行分配并复用中间、packed 与 GEMM workspace;重复调用
|
|
94
|
+
不重新解析计划或 cudaMalloc。
|
|
95
|
+
- stream pool 由 `PYQCD_EINSUM_CUDA_STREAMS` 控制,最多 32 条且受 batch 数限制。
|
|
96
|
+
pack 后记录 event,工作流分片并行,所有流完成后在默认流汇合。
|
|
97
|
+
- 支持 complex128、complex64、float64、float32,混合精度 fail closed。
|
|
98
|
+
|
|
99
|
+
## 5. 已实现 / 未实现边界(如实登记)
|
|
100
|
+
|
|
101
|
+
已实现:
|
|
102
|
+
|
|
103
|
+
- `applyPlanQcd`:与 KUPL **完全相同**的校验与语义,含 0 算子、1 算子、
|
|
104
|
+
`mem_hi < 0` 无限内存、容量不足、下标越界、尺寸为负等全部 fail-closed 分支。
|
|
105
|
+
- `applyContractQcd`、compiled-context 三入口、pack/reduce/hadamard、K-last
|
|
106
|
+
batched GEMM 已实现并通过 GPU 数值对照。
|
|
107
|
+
- `applyInitQcd` / `applyEndQcd`:强定义兼容钩子。
|
|
108
|
+
|
|
109
|
+
**未实现(返回 `1`,不假装支持)**:
|
|
110
|
+
|
|
111
|
+
- 非连续 view 由 Python 绑定显式转连续或回退;不允许静默错误读取 stride。
|
|
112
|
+
- 尚未支持:混合精度、单 operand 重复轴、MPI rank 分片和跨设备 collective。
|
|
113
|
+
- 当前 stream pool 只在单一当前 CUDA device 内工作;跨设备生产拓扑尚需 MPI。
|
|
114
|
+
- `applyContractMpiQcd` 已导出但与 KUPL 不同,当前为返回 `1` 的显式 stub;CUDA MPI
|
|
115
|
+
rank 分片尚未实现,禁止声称支持。
|
|
116
|
+
|
|
117
|
+
## 6. 验收(实测 2026-11-06)
|
|
118
|
+
|
|
119
|
+
```bash
|
|
120
|
+
bash cpp/cuda/einsum/make-nv.sh # 退出码 0
|
|
121
|
+
nm -D --defined-only cpp/cuda/einsum/libpyqcdeinsum_cuda.so | grep Qcd
|
|
122
|
+
```
|
|
123
|
+
|
|
124
|
+
```
|
|
125
|
+
applyPlanQcd / applyContractQcd / applyCompileQcd / applyCompiledQcd /
|
|
126
|
+
applyDestroyQcd / applyInitQcd / applyEndQcd
|
|
127
|
+
```
|
|
128
|
+
|
|
129
|
+
- `readelf -d`:`NEEDED libcublas.so.12`、`libcudart.so.12`;`SONAME
|
|
130
|
+
libpyqcdeinsum_cuda.so`。`qcd_cuda_*` 在 `.symtab` 中是 `t`(内部)。
|
|
131
|
+
- ctypes 对照 KUPL:**8 算子** 310 条表达式(10 条定长 + 300 条随机),
|
|
132
|
+
`path` / `opt_cost` / `largest_intermediate` 逐位相同,差异 **0**;
|
|
133
|
+
同时与 `refer/git-rep/opt_einsum` 的 `contract_path(..., optimize='greedy')`
|
|
134
|
+
对照 **310/310** 一致。
|
|
135
|
+
- V100/P100 四种 dtype 与 NumPy 对照通过;compiled context 重复执行
|
|
136
|
+
`max|Δ| ≈ 1.9e-15`;streams=4 时 batch=4/6 使用四流,batch=2 自动限制两流。
|
|
137
|
+
- T0 distillation GPU 开发窗口:E1 `3.14x`、E2 `4.12x`、E4 `3.10x`;E5 由
|
|
138
|
+
专用 sandwich fast path 继续优化。
|
|
139
|
+
|
|
140
|
+
## 7. 本机实测构建矩阵(实测日期 2026-11-08)
|
|
141
|
+
|
|
142
|
+
只登记**本机实跑**结果。主机 Intel i7-13620H / WSL2 x86_64、nvcc CUDA 12.4.131、
|
|
143
|
+
RTX 4060 Laptop(compute capability **8.9**)。跨后端总表与逐条复现命令见
|
|
144
|
+
`docs/pyqcd_einsum_build_matrix_20261108.md`。
|
|
145
|
+
|
|
146
|
+
### 7.1 四条实测事实
|
|
147
|
+
|
|
148
|
+
1. **默认架构产物在本机不可运行。** `PYQCD_CUDA_ARCHITECTURES=8.9 bash build-nv.sh`
|
|
149
|
+
成功;默认走 `DEFAULT_ARCHITECTURES="6.0 7.0"`(`make-nv.sh:20`),且
|
|
150
|
+
`make-nv.sh:146` 只写 `-gencode=arch=compute_XX,code=sm_XX`、**不生成 PTX**,
|
|
151
|
+
于是默认 fatbin 里没有任何 8.x 目标:构建退出码 0,但在 sm_89 上无法执行。
|
|
152
|
+
本机复核命令(`cuobjdump` 不在本机 PATH,故本次未用它复核):
|
|
153
|
+
|
|
154
|
+
```bash
|
|
155
|
+
cuobjdump --list-elf cpp/cuda/einsum/libpyqcdeinsum_cuda.so
|
|
156
|
+
cuobjdump --list-ptx cpp/cuda/einsum/libpyqcdeinsum_cuda.so # 预期为空:无 PTX
|
|
157
|
+
```
|
|
158
|
+
|
|
159
|
+
2. **MPI 与 `CXX` 冲突(KUPL 侧事实,本目录不受影响)。** KUPL 的 `make.sh` 默认
|
|
160
|
+
`PYQCD_EINSUM_ENABLE_MPI=auto`,本机 `mpicxx` 存在 → 自动开 MPI;但 conda 导出的
|
|
161
|
+
`CXX=x86_64-conda-linux-gnu-c++` 覆盖了脚本默认的 `mpicxx`,构建在
|
|
162
|
+
`mpi_runtime.cpp:7:10: fatal error: mpi.h: No such file or directory` 处失败。
|
|
163
|
+
`make-nv.sh` 没有 MPI 分支,因此
|
|
164
|
+
`PYQCD_CUDA_ARCHITECTURES=8.9 bash build-nv.sh` 不受该冲突影响;`install.sh` /
|
|
165
|
+
`setup.py` 的 `auto` 分支在有 nvcc 且 `cpp/cuda/einsum/src/*.cu` 存在时选 CUDA
|
|
166
|
+
(`install.sh:29-35`、`setup.py:29-31`),本机默认安装路径恰好绕开该失败。
|
|
167
|
+
|
|
168
|
+
3. **OpenBLAS 探测打印 `not found`,实际由运行期 `dlopen` 解析。** 构建日志是
|
|
169
|
+
`OpenBLAS: not found`(无 `pkg-config`、无系统 `cblas.h`),但本库不链接 BLAS:
|
|
170
|
+
先试 `RTLD_DEFAULT`(进程内已有符号),再按 `libcblas.so.3` → `libcblas.so` →
|
|
171
|
+
`libopenblas.so.0` → `libopenblas.so` 的顺序 `dlopen`
|
|
172
|
+
(`cpp/kupl/einsum/src/apply_contract.cpp:287-290`)。本机
|
|
173
|
+
`ctypes.CDLL("libcblas.so.3")` 与 `ctypes.CDLL("libopenblas.so.0")` 均成功。
|
|
174
|
+
|
|
175
|
+
4. **`-march=native` 默认与 PyQCU 相反(对照登记,本目录不涉及)。** `make-nv.sh` 用
|
|
176
|
+
`-O3 -std=c++17`(`make-nv.sh:179`)编译,不含任何 `-march`;`-march=native`
|
|
177
|
+
只出现在 KUPL 的 `make.sh:37-41`,`PYQCD_EINSUM_NATIVE_ARCH` 默认 `1`,而 PyQCU
|
|
178
|
+
对应开关默认 OFF。CUDA 产物的跨机风险来自第 1 条的架构目标,不是主机 ISA。
|
|
179
|
+
|
|
180
|
+
### 7.2 本目录的可用组合
|
|
181
|
+
|
|
182
|
+
| 组合 | 命令 | 状态(2026-11-08) |
|
|
183
|
+
|---|---|---|
|
|
184
|
+
| 脚本默认(`6.0 7.0`,无 PTX) | `bash build-nv.sh` | 构建退出码 0;sm_89 上不可运行(7.1-1) |
|
|
185
|
+
| 显式 `8.9` | `PYQCD_CUDA_ARCHITECTURES=8.9 bash build-nv.sh` | 实测通过 |
|
|
186
|
+
| 显式 `8.0` | `PYQCD_CUDA_ARCHITECTURES=8.0 bash build-nv.sh` | 未实测 |
|
|
187
|
+
| 多架构 | `PYQCD_CUDA_ARCHITECTURES="8.0 8.9" bash build-nv.sh` | 未实测(逗号/分号/空格均可分隔,`make-nv.sh:135-136`) |
|
|
188
|
+
| CUDA + MPI | — | 不支持(`applyContractMpiQcd` 是返回 1 的 stub) |
|
|
189
|
+
|
|
190
|
+
产物核对:
|
|
191
|
+
|
|
192
|
+
```bash
|
|
193
|
+
readelf -d cpp/cuda/einsum/libpyqcdeinsum_cuda.so | grep -E 'SONAME|NEEDED'
|
|
194
|
+
# SONAME libpyqcdeinsum_cuda.so;NEEDED libcublas.so.12 / libcudart.so.12 / libstdc++ / libm / libgcc_s / libc
|
|
195
|
+
nm -D --defined-only cpp/cuda/einsum/libpyqcdeinsum_cuda.so | awk '{print $NF}' | grep Qcd | sort
|
|
196
|
+
# 8 个入口(含 fail-closed 的 applyContractMpiQcd stub)
|
|
197
|
+
```
|
|
198
|
+
|
|
199
|
+
`setup.py` 会把这个文件复制到 `pyqcd/einsum/cuda/libpyqcdeinsum_cuda.so`
|
|
200
|
+
(`setup.py:42-53`),Cython 扩展以 `-Wl,-rpath,$ORIGIN` 解析它(`setup.py:119-122`);
|
|
201
|
+
磁盘上的 `cpp/cuda/einsum/libpyqcdeinsum_cuda.so` 是指向 `build/` 的软链。
|
|
@@ -0,0 +1,117 @@
|
|
|
1
|
+
#ifndef PYQCD_EINSUM_CUDA_H
|
|
2
|
+
#define PYQCD_EINSUM_CUDA_H
|
|
3
|
+
|
|
4
|
+
#include "einsum_define.h"
|
|
5
|
+
|
|
6
|
+
#include <cublas_v2.h>
|
|
7
|
+
#include <cuda_runtime.h>
|
|
8
|
+
|
|
9
|
+
#include <cstdint>
|
|
10
|
+
|
|
11
|
+
constexpr int32_t QCD_CUDA_MAX_RANK = 64;
|
|
12
|
+
constexpr int32_t QCD_CUDA_MAX_PACKED_RANK = 128;
|
|
13
|
+
|
|
14
|
+
struct qcd_cuda_copy_spec {
|
|
15
|
+
int32_t rank;
|
|
16
|
+
int64_t elements;
|
|
17
|
+
int64_t shape[QCD_CUDA_MAX_PACKED_RANK];
|
|
18
|
+
int64_t src_stride[QCD_CUDA_MAX_PACKED_RANK];
|
|
19
|
+
};
|
|
20
|
+
|
|
21
|
+
struct qcd_cuda_hadamard_spec {
|
|
22
|
+
int32_t rank;
|
|
23
|
+
int64_t elements;
|
|
24
|
+
int64_t shape[QCD_CUDA_MAX_RANK];
|
|
25
|
+
int64_t left_stride[QCD_CUDA_MAX_RANK];
|
|
26
|
+
int64_t right_stride[QCD_CUDA_MAX_RANK];
|
|
27
|
+
};
|
|
28
|
+
|
|
29
|
+
struct qcd_cuda_reduce_spec {
|
|
30
|
+
int32_t output_rank;
|
|
31
|
+
int32_t reduction_rank;
|
|
32
|
+
int64_t output_elements;
|
|
33
|
+
int64_t reduction_elements;
|
|
34
|
+
int64_t output_shape[QCD_CUDA_MAX_RANK];
|
|
35
|
+
int64_t output_src_stride[QCD_CUDA_MAX_RANK];
|
|
36
|
+
int64_t reduction_extent[QCD_CUDA_MAX_RANK];
|
|
37
|
+
int64_t reduction_src_stride[QCD_CUDA_MAX_RANK];
|
|
38
|
+
};
|
|
39
|
+
|
|
40
|
+
struct qcd_cuda_sandwich_spec {
|
|
41
|
+
int64_t prefix;
|
|
42
|
+
int64_t left;
|
|
43
|
+
int64_t contracted_left;
|
|
44
|
+
int64_t contracted_right;
|
|
45
|
+
int64_t right;
|
|
46
|
+
int64_t lane;
|
|
47
|
+
};
|
|
48
|
+
|
|
49
|
+
struct qcd_cuda_e4_spec {
|
|
50
|
+
int64_t z;
|
|
51
|
+
int64_t y;
|
|
52
|
+
int64_t x;
|
|
53
|
+
int64_t c;
|
|
54
|
+
int64_t t;
|
|
55
|
+
int64_t a;
|
|
56
|
+
int64_t b;
|
|
57
|
+
int64_t n;
|
|
58
|
+
int64_t input_g;
|
|
59
|
+
int64_t output_g;
|
|
60
|
+
int64_t d;
|
|
61
|
+
int64_t e;
|
|
62
|
+
int64_t f;
|
|
63
|
+
int64_t output_elements;
|
|
64
|
+
};
|
|
65
|
+
|
|
66
|
+
#if defined(_WIN32)
|
|
67
|
+
#define QCD_CUDA_LOCAL
|
|
68
|
+
#define QCD_CUDA_WEAK
|
|
69
|
+
#else
|
|
70
|
+
#define QCD_CUDA_LOCAL __attribute__((visibility("hidden")))
|
|
71
|
+
#define QCD_CUDA_WEAK __attribute__((weak, visibility("hidden")))
|
|
72
|
+
#endif
|
|
73
|
+
|
|
74
|
+
extern "C" QCD_CUDA_WEAK int qcd_cuda_permute(
|
|
75
|
+
int dtype_id, const void* src, void* dst, qcd_cuda_copy_spec spec,
|
|
76
|
+
cudaStream_t stream);
|
|
77
|
+
|
|
78
|
+
extern "C" QCD_CUDA_WEAK int qcd_cuda_hadamard(
|
|
79
|
+
int dtype_id, const void* left, const void* right, void* dst,
|
|
80
|
+
qcd_cuda_hadamard_spec spec, cudaStream_t stream);
|
|
81
|
+
|
|
82
|
+
extern "C" QCD_CUDA_WEAK int qcd_cuda_reduce(
|
|
83
|
+
int dtype_id, const void* src, void* dst, qcd_cuda_reduce_spec spec,
|
|
84
|
+
cudaStream_t stream);
|
|
85
|
+
|
|
86
|
+
extern "C" QCD_CUDA_WEAK int qcd_cuda_sandwich(
|
|
87
|
+
int dtype_id,
|
|
88
|
+
const void* left_matrix,
|
|
89
|
+
const void* tensor,
|
|
90
|
+
const void* right_matrix,
|
|
91
|
+
void* dst,
|
|
92
|
+
qcd_cuda_sandwich_spec spec,
|
|
93
|
+
cudaStream_t stream);
|
|
94
|
+
|
|
95
|
+
extern "C" QCD_CUDA_WEAK int qcd_cuda_e4_fused(
|
|
96
|
+
int dtype_id,
|
|
97
|
+
const void* a,
|
|
98
|
+
const void* b,
|
|
99
|
+
const void* gamma,
|
|
100
|
+
const void* c,
|
|
101
|
+
void* dst,
|
|
102
|
+
qcd_cuda_e4_spec spec,
|
|
103
|
+
cudaStream_t stream);
|
|
104
|
+
|
|
105
|
+
extern "C" QCD_CUDA_WEAK int qcd_cuda_batched_gemm(
|
|
106
|
+
int dtype_id,
|
|
107
|
+
cublasHandle_t handle,
|
|
108
|
+
const void* left,
|
|
109
|
+
const void* right,
|
|
110
|
+
void* dst,
|
|
111
|
+
int64_t m,
|
|
112
|
+
int64_t n,
|
|
113
|
+
int64_t k,
|
|
114
|
+
int64_t batch,
|
|
115
|
+
cudaStream_t stream);
|
|
116
|
+
|
|
117
|
+
#endif
|