Refactor index manipulation kernels around position-indexed subblocks - #526
Refactor index manipulation kernels around position-indexed subblocks#526lkdvos wants to merge 1 commit into
Conversation
Index manipulations now run through a single kernel that operates on subblocks addressed by position: `StridedSubblocks` (sector-independent views into the flat data of a `TensorMap`) or `TreeSubblocks` (any `AbstractTensorMap`, through `subblock`), both carrying an optional lazy conjugation. `TreeTransformer`s store only the mapping between subblock positions and recoupling coefficients, alongside the subblock structures, and are cached for every tensor type. Adjoint sources and destinations, as well as `conj` in `tensoradd!`, are folded into a conjugation flag, relabeled permutation and levels, and conjugated scalars, so that `AdjointTensorMap` wrappers no longer force the uncached generic path (fixes #516, supersedes #519 and #520). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Benchmark:
|
| main | ld-adjoint | main / ld-adjoint | |
|---|---|---|---|
| indexmanipulations/permute/permute/("ComplexF64", "SU2Irrep", "[48, 48, 48, 48]", "[1.0, 1.0, 1.0, 1.0]", "Any[[1, 3], [2, 4]]") | 1.37 ± 0.22 ms | 1.19 ± 0.23 ms | 1.15 ± 0.29 |
| indexmanipulations/permute/permute/("ComplexF64", "SU2Irrep", "[48, 48, 48, 48]", "[1.0, 1.0, 1.0, 1.0]", "Any[[1, 3], [2, 4]]", "adjoint") | 1.59 ± 0.085 ms | 1.55 ± 0.32 ms | 1.03 ± 0.22 |
| indexmanipulations/permute/permute/("ComplexF64", "SU2Irrep", "[48, 48, 48, 48]", "[1.0, 1.0, 1.0, 1.0]", "Any[[4, 2, 3], [1]]") | 1.34 ± 0.25 ms | 1.22 ± 0.22 ms | 1.1 ± 0.28 |
| indexmanipulations/permute/permute/("ComplexF64", "SU2Irrep", "[48, 48, 48, 48]", "[1.0, 1.0, 1.0, 1.0]", "Any[[4, 2, 3], [1]]", "adjoint") | 1.32 ± 0.12 ms | 0.939 ± 0.15 ms | 1.41 ± 0.26 |
| indexmanipulations/permute/permute/("Float64", "SU2Irrep", "[48, 48, 48, 48]", "[1.0, 1.0, 1.0, 1.0]", "Any[[1, 3], [2, 4]]") | 0.59 ± 0.04 ms | 0.822 ± 0.21 ms | 0.717 ± 0.19 |
| indexmanipulations/permute/permute/("Float64", "SU2Irrep", "[48, 48, 48, 48]", "[1.0, 1.0, 1.0, 1.0]", "Any[[1, 3], [2, 4]]", "adjoint") | 1.27 ± 0.19 ms | 0.775 ± 0.12 ms | 1.64 ± 0.35 |
| indexmanipulations/permute/permute/("Float64", "SU2Irrep", "[48, 48, 48, 48]", "[1.0, 1.0, 1.0, 1.0]", "Any[[4, 2, 3], [1]]") | 0.811 ± 0.089 ms | 0.665 ± 0.059 ms | 1.22 ± 0.17 |
| indexmanipulations/permute/permute/("Float64", "SU2Irrep", "[48, 48, 48, 48]", "[1.0, 1.0, 1.0, 1.0]", "Any[[4, 2, 3], [1]]", "adjoint") | 0.684 ± 0.038 ms | 0.5 ± 0.023 ms | 1.37 ± 0.098 |
| indexmanipulations/permute/permute/("Float64", "SU2Irrep", "[512, 512]", "[1.0, 1.0]", "Any[[2, 1], Any[]]") | 0.144 ± 0.014 ms | 0.0849 ± 0.0089 ms | 1.69 ± 0.24 |
| indexmanipulations/permute/permute/("Float64", "Trivial", "[43408, 1216]", "nothing", "Any[[2, 1], Any[]]") | 0.0453 ± 0.0017 s | 0.0545 ± 0.0086 s | 0.832 ± 0.14 |
| indexmanipulations/permute/permute/("Float64", "Trivial", "[7264, 7264]", "nothing", "Any[[2, 1], Any[]]") | 0.0557 ± 0.0012 s | 0.0548 ± 0.0011 s | 1.02 ± 0.029 |
| indexmanipulations/permute/permute/("Float64", "Z2Irrep", "[43408, 1216]", "[0.5, 0.5]", "Any[[2, 1], Any[]]") | 27.6 ± 1.6 ms | 27.3 ± 0.8 ms | 1.01 ± 0.066 |
| indexmanipulations/permute/permute/("Float64", "Z2Irrep", "[43408, 1216]", "[0.5, 0.5]", "Any[[2, 1], Any[]]", "adjoint") | 27.7 ± 2.1 ms | 27.8 ± 2.7 ms | 0.995 ± 0.12 |
| indexmanipulations/permute/permute/("Float64", "Z2Irrep", "[7264, 7264]", "[0.5, 0.5]", "Any[[2, 1], Any[]]") | 22.9 ± 0.47 ms | 28 ± 0.76 ms | 0.819 ± 0.028 |
| indexmanipulations/permute/permute/("Float64", "Z2Irrep", "[7264, 7264]", "[0.5, 0.5]", "Any[[2, 1], Any[]]", "adjoint") | 23 ± 0.41 ms | 23.1 ± 0.45 ms | 0.994 ± 0.026 |
🤖 Generated with Claude Code
Structural fix for #516 (adjoint
permute!up to 80× slower than the plain path), superseding #519 and #520 and building on #518/#521.What changes
StridedSubblocks: sector-independent, integer-indexed views into the flat data of aTensorMapin canonical fusion-tree order, with an optional lazy conjugation (op = identity/conjas type parameter).TreeSubblocksis the generic counterpart for anyAbstractTensorMap, going throughsubblock.TreeTransformers store only the mapping between subblock positions and recoupling coefficients (plus the two subblock structures), and are cached for every tensor type; the closure-based fallback andTrivialTreeTransformerare gone.TensorMaps and other tensor types.permute!/braid!/transpose!andTO.tensoradd!foldAdjointTensorMapsources and destinations (andconjA) into aconjsrc::Bool, relabeledp/levelsand conjugatedα/β; the flag is resolved into the view type only at the kernel entry, so everything stays type-stable.BraidingTensorsources are converted before the transformer is chosen (the old overload could reach an untyped kernel).Numbers (issue reproducer, adjoint vs plain
permute!): fℤ₂ 1.30× → 1.09×, fℤ₂⊠U₁ 3.57× → 1.05×, SU₂ 4.44× → 1.01×, U₁ 2.23× → 1.11×; plain path unchanged within noise (+32 B from the extraBoolin the cache key).Tests: adjoint source/destination/both for
permute!/transpose!/braid!with accumulation,@tensor conj,BraidingTensorsource, and adot-based isometry check that would catch a wrongly conjugated recoupling matrix for complex sector scalar types. Benchmark suite gainedadjoint = truepermute variants.Follow-ups (not here): keying transformers on sector structure only; passing
conjintotrace_permute!.🤖 Generated with Claude Code