Refactor index manipulation kernels around position-indexed subblocks - #526
Conversation
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
* Intercept `TO.tensoradd!` instead of `TensorKit.add_transform!` The blockwise `tensoradd` implementations were only reachable through `TensorKit`'s internals: `TO.tensoradd!` delegated to `permute!`, and the `add_transform!` methods defined here caught the kernel. Neither is a contract `TensorKit` owes us, and both stopped holding on TensorKit's index-manipulation refactor (QuantumKitHub/TensorKit.jl#526), where `TO.tensoradd!` calls the braid kernel directly and `add_transform!` gained a `conjsrc` argument. Nothing errors in that case: dense block tensors survive via TensorKit's generic subblock fallback, while sparse ones silently produce zeros, because the fallback writes through blocks a sparse container never materialized. Implement `TO.tensoradd!` for the block tensor types instead. That is the public entry point `@tensor` lowers to, so it is reachable no matter how TensorKit arranges its kernels, and `conjA` is handled explicitly rather than relying on it having been unwrapped into an adjoint beforehand. The mixed block/plain methods take the concrete `TensorMap` so that `(BlockTensorMap, SparseBlockTensorMap)` and the reverse resolve to the general method rather than being ambiguous. The added testset calls `TO.tensoradd!` directly rather than through `@tensor`, so this stays covered independently of TensorKit's routing. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * Pin the planar entry points in the tests `planaradd!`, `planartrace!` and `planarcontract!` are reached through `transpose!`, `trace_permute!` and `contract!`, so they currently land on methods defined here and are not affected by TensorKit#526. They rest on the same undocumented delegation that broke `tensoradd!` though, and there was no planar coverage here at all, so pin them directly. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
All comments should now be addressed again, I think I also managed to simplify slightly more because of @Jutho's comments, in particular the abelian case now no longer needs to refer to the "adjoint spaces" and simply loops over the source fusiontrees and exchanges them on the fly, which at least reads a lot easier. For the non-abelian case I still am doing it that way because swapping splitting and fusiontree in the |
| f = gettokenvalue(fusiontrees(iter.t), i) | ||
| return f => iter.structure[i], i + 1 | ||
| end | ||
| end |
There was a problem hiding this comment.
Reposting as a new comment, so that it does not get lost:
I guess I am wondering to what extent this specialization is necessary. As far as I can tell, subblocks is not really used in any performance critical code, they all go via StridedSubblocks or TreeSubblocks directly. The subblocks function and the associated SubblockIterator seem mostly user convenience functions, for getting and setting data in the tensor. So I don't know if we really need to make this more complicated for negligible performance gain.
There was a problem hiding this comment.
To address this more precisely, this is almost true, in the sense that for example operations that mix diagonal tensors and regular tensors would still end up with subblocks being called, and I do see future work benefiting from the knowledge that subblocks is a performant primitive to build around, e.g. for twist or flip implementations.
|
Ok, I finally managed to make my way through. Looks really great. I have four final questions or suggestions, but otherwise fully approve. |
"Abelian" is ambiguous for sectors: it can refer either to the fusion of two sectors having a unique result, or to the commutativity of the fusion rules. The transformer is selected on `FusionStyle(I) == UniqueFusion()`, so name it after that. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`StridedSubblocks` and `TreeSubblocks` no longer apply `identity`/`conj` to every view. Instead `conjsrc` is threaded through `add_transform_kernel!` into `_add_transform_block!`, where it is handed to `TO.tensoradd!` as its `conjA` argument, at the single-tree call and when packing a multi-tree block. This drops a type parameter from both collections, so the kernel compiles to one instance per (storage, numind) rather than one per conjugation. The runtime flag is free: `flag2op` is union-split, and `conj` of a real-eltype `StridedView` is a type-level no-op, which also makes the previous `scalartype(t) <: Real` guard redundant. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
Thanks for the careful review, I think all of these comments were helpful to further simplify and improve the implementation, which I've hopefully managed to achieve here. In particular, your comment about the As for whether the TensorMap specialization is still worth it, I would argue yes, mostly because there are still other parts of the code that iterate fusiontrees and look at the data, such as I'm hoping all tests pass, and if so might merge this Monday morning (ET) if no further comments, or unless someone merges this first. |
| @noinline _throw_subblock_bounds(iter, i) = throw(BoundsError(iter, i)) | ||
| @noinline _throw_subblock_missing(f) = throw(SectorMismatch(lazy"fusion tree pair $f is not present")) | ||
|
|
||
| @propagate_inbounds function Base.getindex(iter::SubblockIterator{<:TensorMap}, i::Int) |
There was a problem hiding this comment.
This is now relying on iter being constructed as SubblockIterator(t). What happens if someone constructs SubblockIterator(t::TensorMap, fusiontrees(t)). This method is still being called and will fail, no?
There was a problem hiding this comment.
Yes, although that is also true for someone calling SubblockIterator(t::AbstractTensorMap, nothing), in the sense that you really should not be doing that? I can try and restrict the method even further, but since this should all be internals I don't know if that makes too much sense 😄
|
By the way, did you rerun the benchmarks from the original issue with this latest iteration of the PR? |
Issue #516 reproducer, rerun against the final state of this PR
In every case the adjoint path now allocates exactly what the plain path allocates and runs at the same speed. The plain path itself did not move ( Two notes on the baseline: Full suite ( |
* Draft changelog for v0.17.2 Consolidates the Unreleased section (which already included the real entries added by #526/#532 on merge) with entries for the remaining PRs merged since v0.17.1 (#487-#535), and retitles it as 0.17.2. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Bump version to v0.17.2 Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Fix confirmed small bugs from the pre-release audit - isunitspace: require dim(V) == 1 for GenericUnit sectors (#537) - GradedSpace ⊕/supremum: check unit homogeneity of the result (#538) - isconj(::ComplexSpace): return isdual(V) instead of always true (#539) - multi_associator: return a vector, not a scalar, on early-exit for GenericFusion (#540) - split(f, 0): use leftunit(f.coupled) instead of indexing an empty uncoupled tuple (#541) - repartition: return a Pair in the identity branch, matching every other branch (#542) - Mooncake scalar_pullback: accumulate into the tangent instead of overwriting it (#543) - rand/randn/randexp/randisometry(rng, T, space): fix one(domain) typo (#544) - pinv(::DiagonalTensorMap): fix inverted atol/rtol defaulting and empty-tensor throw (#545) - t1 / t2: promote to a float scalartype, matching t1 \ t2 (#546) Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Add changelog entry for the audit bugfixes Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Address fable review findings on the audit bugfixes - split(f, 0): also guard innerlines_extended construction, which still indexed the empty uncoupled tuple for a 0-leg tree - pinv(::DiagonalTensorMap): use eps (not sqrt(eps)) for the default rtol, matching _default_rtol's convention and dense LinearAlgebra.pinv Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Address tuicr review comments on the audit bugfixes - pinv(::DiagonalTensorMap): reuse _default_rtol instead of duplicating its formula - Add regression tests for split(f, 0) on a genuine 0-leg tree and for multi_associator's early-exit branch on a GenericFusion sector Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * fix planar issues after MPSKit test rerun * harden Mooncake scalar pullback --------- Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
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