Skip to content

Commit d2287df

Browse files
authored
feat(kernels): ⚡ contract along opt_einsum's dp paths by default (#73)
Assisted-by: Claude Code:claude-opus-5-5
1 parent bd78e23 commit d2287df

13 files changed

Lines changed: 730 additions & 28 deletions

File tree

‎README.md‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -125,7 +125,9 @@ out = src(
125125

126126
On a node with several GPUs, `Resources(devices=4)` splits the sweep among four of them, in the same process.
127127

128-
See [Large problems](https://algorithmiq.github.io/src-method/features/large-problems) for the budgets, the scratch directory, several GPUs and how to read the logged plan.
128+
The contractions follow opt_einsum's cheapest (`dp`) paths. `optimize` selects any other opt_einsum path optimizer, by name or as a `PathOptimizer` instance.
129+
130+
See [Large problems](https://algorithmiq.github.io/src-method/features/large-problems) for the budgets, the scratch directory, several GPUs, contraction paths and how to read the logged plan.
129131

130132
## Installation
131133

‎benches/large/README.md‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,3 +116,52 @@ Speed-up and efficiency are against the single-GPU run:
116116
memory unused was not investigated.
117117
- The stall times the sweep logs include waiting for the GPU's own queue, so they
118118
cannot tell slow transfers from a GPU-bound pass.
119+
120+
### Einsum paths: `auto` against `dp`
121+
122+
The `optimize` argument of `src`, `apply` and `compress` selects the opt_einsum
123+
path optimizer of the kernels, and the planner sizes batches for the same path.
124+
It defaults to `dp` since these measurements; the results above ran with `auto`.
125+
`bench_einsum_dp.py` compares the two: `kernels` times `SiteKernels.env`,
126+
`.sketch` and `.project` on random arrays of the reference shapes, and `sweep`
127+
with `distance` runs the `D_M = 1000` stack, `chi_out = 500`, under the 60 GB
128+
budget, three times per optimizer and compares the outputs. Run on A100-SXM-64GB
129+
(CuPy 14.2.0, NumPy 2.4.4, opt_einsum 3.4.0), complex128.
130+
131+
Full sweep, median of 3 (range), every environment on the device:
132+
133+
| GPUs | `auto` (s) | `dp` (s) | Speed-up | Right-to-left `auto` / `dp` (s) | Planned peak `auto` / `dp` (GB) | QR distance |
134+
|---|---|---|---|---|---|---|
135+
| 1 | 109.7 (109.6-130.6) | 56.2 (56.2-74.1) | 1.95x | 95.4 / 42.3 | 58.8 / 50.3 | 2.7e-13 |
136+
| 4 | 32.4 (32.3-51.7) | 17.4 (17.4-40.1) | 1.86x | 28.6 / 13.6 | 32.3 / 12.2 per GPU | 4.0e-13 |
137+
138+
- The gain is in the right-to-left pass; the left-to-right pass takes the same
139+
time (14.3 against 13.9 s on one GPU). At these shapes `auto`'s sketch and
140+
projection paths cost 3.9x and 2.0x `dp`'s flops. Both optimizers give the
141+
same planned batches.
142+
- The first sweep of each job is 18-23 s slower with either optimizer (the
143+
first-run effect noted above), and is left out of the medians. `dp`'s path
144+
search takes about 1 ms per shape: 0.03-0.10 s per sweep.
145+
146+
Kernels at the reference shapes, `D = (4, 4, 4000, 4)`, physical legs 4,
147+
`eta = 2000`, median of 5 after 2 warm-ups (ms per row or column):
148+
149+
| Kernel | Batch | `auto` | `dp` | Flops `auto` / `dp` |
150+
|---|---|---|---|---|
151+
| project | 8 | 25.2 | 11.8 | 7.9e11 / 4.0e11 |
152+
| project | 16 | 22.1 | 11.0 | 1.6e12 / 7.9e11 |
153+
| project | 32, 64, 128 | over 50 GB | 10.6, 10.5, 10.5 | 9.7e12 / 1.6e12 at 32 |
154+
| sketch | 96 | 10.27 | 10.65 | 4.78e12 / 4.76e12 |
155+
| sketch | 192 | 10.14 | 10.53 | 9.56e12 / 9.51e12 |
156+
| env | 224 | 7.17 | 7.18 | 7.4e12 / 7.4e12 |
157+
158+
- The projection runs at 16-19 TFLOP/s with either optimizer, so the time follows
159+
the flops. From 32 rows `auto` plans an 8.19e9-element intermediate (265-283
160+
GB), which is why it cannot leave 16 rows; `dp` needs 10-25 GB up to 128 rows.
161+
- 384 sketch columns (75 GB) and 256 projection rows with `dp` (50.3 GB) exceed
162+
the 50 GB limit and were not run.
163+
- `dp` is 3.7-3.8% slower on the sketch, where its path has 0.5% fewer flops.
164+
- The outputs of the two optimizers agree to 1.6e-14 or better.
165+
- The reference-size sweep (over an hour) was not run. Taking these per-row and
166+
per-column times for each site predicts its right-to-left pass at about 2060 s
167+
with `dp`, 1.5x faster than the 3112 s measured with `auto`.

0 commit comments

Comments
 (0)