[leburgel]

Draft, mainly as a placeholder to demonstrate the benefits and one potential way of implementing this.

On the CPU, permute, braid and transpose call tensoradd! once per fusion-tree subblock, which costs about 0.5–0.7 µs per call independent of the subblock size. For tensors with many small subblocks, such as fermionic U(1)-symmetric PEPS tensors, this dominates: permute! becomes 25–125× slower than copying the data. This is the same per-subblock pattern as #595, and the CPU counterpart of #508.

This PR collects the subblocks with at most 1024 entries into groups of equal size and strides when the tree transformer is built, and caches them with it: per group the relative positions of the entries in the destination and the permuted source, per subblock the offsets and the coefficient. add_transform_kernel! then transforms every group in a single loop, while larger subblocks still go through tensoradd! one by one. For sectors with multiple fusion channels, the recoupling of small subblocks is applied term by term in the same loops instead of through pack, mul! and unpack. With multiple transformer threads, the tasks handle disjoint ranges of destination offsets.

permute! into a preallocated destination, one thread, best of 30, in ms (p1 = ((1, 4, 2), (5, 3, 6)) on Vχ ⊗ VD ⊗ VD' ← Vχ ⊗ VD ⊗ VD'):

sector, legs	subblocks	largest	copy!	main	PR
fℤ₂ × U(1), Vχ charges ±2, deg. 4; VD ±1, deg. 1	9632	16	0.187	4.711	0.261
fℤ₂ × U(1), Vχ charges ±8, deg. 2; VD ±1, deg. 1	40736	4	0.204	25.30	0.329
U(1), Vχ charges ±2, deg. 16; VD ±1, deg. 1	301	256	0.079	0.277	0.112
U(1), Vχ charges ±2, deg. 32; VD ±1, deg. 2	301	16384	14.58	17.25	17.40
SU(2), Vχ j ≤ 1, deg. 2; VD j ≤ 1/2, deg. 1	120	4	0.000	0.104	0.005
SU(2), Vχ j ≤ 1, deg. 16; VD j ≤ 1/2, deg. 2	120	4096	0.62	2.18–2.46	1.69–2.20
fℤ₂ × SU(2), deg. 2 / 1	977	4	0.003	0.953	0.024
fℤ₂ × SU(2), deg. 4 / 2	977	256	0.313	2.514	1.342

Tensors whose subblocks are all larger than 1024 entries take the same path as before. The existing permutation tests only have subblocks of up to 12 entries, so I added one testset with subblocks on both sides of the threshold.

The groups add about 25 bytes per small subblock to each cached transformer; the GPU extension builds them as well but does not use them.

Reproducer

using TensorKit
V(qmax, d) = Vect[FermionParity ⊠ U1Irrep]([(p, q) => d for p in 0:1 for q in -qmax:qmax]...)
W = V(8, 2) ⊗ V(1, 1) ⊗ V(1, 1)' ← V(8, 2) ⊗ V(1, 1) ⊗ V(1, 1)'
t = randn(ComplexF64, W)
p = ((1, 4, 2), (5, 3, 6))
tdst, tcopy = permute(t, p; copy = true), similar(t)
besttime(f) = (f(); minimum(_ -> @elapsed(f()), 1:20))
println(length(fusiontrees(t)), " subblocks of ", dim(t) ÷ length(fusiontrees(t)), " entries on average")
println("permute!: ", round(1e3 * besttime(() -> permute!(tdst, t, p)); digits = 3), " ms")
println("copy!:    ", round(1e3 * besttime(() -> copy!(tcopy, t)); digits = 3), " ms")

main:

40736 subblocks of 4 entries on average
permute!: 24.778 ms
copy!:    0.199 ms

this PR:

40736 subblocks of 4 entries on average
permute!: 0.342 ms
copy!:    0.2 ms

🤖 Generated with Claude Code

[codecov]

Codecov Report

❌ Patch coverage is 96.29630% with 3 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/tensors/indexmanipulations.jl	92.30%	3 Missing ⚠️

Files with missing lines	Coverage Δ
src/tensors/treetransformers.jl	97.45% <100.00%> (-1.27%)	⬇️
src/tensors/indexmanipulations.jl	89.78% <92.30%> (-0.50%)	⬇️

... and 23 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[lkdvos]

Do you happen to have an idea about whether the speedup is mostly because of the pre-computed linear indices, or just from the chunking, i.e. reducing the overhead? I recently attempted to make the Strided calls a lot less overhead, especially for small tensors, so I'm hoping that just the chunking could suffice.

See also the work in #584, where I also introduced some chunking (there the main motivation was multithreading, since for many small blocks the contention on the single atomic Int is quite large).

[leburgel]

I tried to disentangle the two by running the different steps separately on the same transformer data (Strided 2.6.6, so including your recent changes; ComplexF64, single thread, best of 30 × 8 interleaved rounds, using the tensors from the PR description). Times are in ns per subblock:

tensor	largest subblock	tensoradd! (main)	bare Strided call	nested loop per subblock	groups, nested loop	groups, index lists (this PR)
fℤ₂ × U(1), 9632 subblocks	16	552	459	66	66	27
fℤ₂ × U(1), 40736 subblocks	4	543	476	104	99	8
U(1), 301 subblocks	64	686	567	98	92	67
U(1), 301 subblocks	256	1080	807	390	383	386
U(1), 301 subblocks	1024	2869	2060	1662	1631	1726

It turns out that both matter, but most of the time saved comes from avoiding the per-call overhead. For the five tensors with subblocks of at most 16 entries (two of which are shown above), a single Strided call still costs 440–570 ns, and simply replacing it by a plain loop over each subblock already makes things 3–7× faster. Grouping the subblocks by shape doesn't do much by itself, but it's precisely what makes the precomputed index lists cheap, since we only need one list per shape. Those index lists then give another 2.5–12×. Once the subblocks reach 256 entries, all of the hand-written loops perform about the same. A second run gave the same picture. I haven't tried QuantumKitHub/Strided.jl#84 here.

I also ran #584 next to main and this PR on the four fℤ₂ × U(1) tensors from the description (permute! into a preallocated destination, ms, best of 30 within a run and of three runs, one run at a time on a shared node):

	main	#584	this PR
1 thread	9.7 / 14.5 / 33.6 / 29.4	5.8 / 13.9 / 18.6 / 23.9	0.28 / 0.71 / 0.69 / 0.37
4 transformer threads	2.4 / 5.0 / 6.9 / 8.9	2.1 / 4.4 / 6.4 / 7.8	0.11 / 0.27 / 0.26 / 0.18

The node was still busy, so these are noisier than the numbers above, but #584 stays within 2× of main here, since it still calls tensoradd! once per subblock. For the multithreaded case, I split the destination offsets into 4 × ntasks chunks, which is similar to what you did in #584. I also added a test for this, since codecov flagged it as uncovered.