Skip to content

Refactor index manipulation kernels around position-indexed subblocks - #526

Merged
lkdvos merged 16 commits into
mainfrom
ld-adjoint
Sep 20, 2026
Merged

lkdvos merged 16 commits into
mainfrom
ld-adjoint

Conversation

@lkdvos

@lkdvos lkdvos commented Sep 3, 2026

Copy link
Copy Markdown
Member

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 a TensorMap in canonical fusion-tree order, with an optional lazy conjugation (op = identity/conj as type parameter). TreeSubblocks is the generic counterpart for any AbstractTensorMap, going through subblock.
  • 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 and TrivialTreeTransformer are gone.
  • One kernel serves abelian and generic transformers, TensorMaps and other tensor types.
  • permute!/braid!/transpose! and TO.tensoradd! fold AdjointTensorMap sources and destinations (and conjA) into a conjsrc::Bool, relabeled p/levels and conjugated α/β; the flag is resolved into the view type only at the kernel entry, so everything stays type-stable.
  • BraidingTensor sources 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 extra Bool in the cache key).

Tests: adjoint source/destination/both for permute!/transpose!/braid! with accumulation, @tensor conj, BraidingTensor source, and a dot-based isometry check that would catch a wrongly conjugated recoupling matrix for complex sector scalar types. Benchmark suite gained adjoint = true permute variants.

Follow-ups (not here): keying transformers on sector structure only; passing conj into trace_permute!.

🤖 Generated with Claude Code

@codecov

codecov Bot commented Sep 3, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 85.79545% with 25 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
src/tensors/blockiterators.jl 45.83% 13 Missing ⚠️
src/tensors/tensor.jl 56.25% 7 Missing ⚠️
src/tensors/treetransformers.jl 96.66% 2 Missing ⚠️
ext/TensorKitEnzymeExt/utility.jl 0.00% 1 Missing ⚠️
src/tensors/braidingtensor.jl 50.00% 1 Missing ⚠️
src/tensors/indexmanipulations.jl 98.52% 1 Missing ⚠️
Files with missing lines Coverage Δ
src/TensorKit.jl 17.24% <ø> (ø)
src/tensors/abstracttensor.jl 54.85% <100.00%> (ø)
src/tensors/tensoroperations.jl 96.44% <100.00%> (-0.10%) ⬇️
ext/TensorKitEnzymeExt/utility.jl 20.68% <0.00%> (-0.25%) ⬇️
src/tensors/braidingtensor.jl 87.58% <50.00%> (+0.65%) ⬆️
src/tensors/indexmanipulations.jl 91.13% <98.52%> (+1.56%) ⬆️
src/tensors/treetransformers.jl 96.34% <96.66%> (+1.49%) ⬆️
src/tensors/tensor.jl 81.28% <56.25%> (-2.32%) ⬇️
src/tensors/blockiterators.jl 43.69% <45.83%> (ø)

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@lkdvos

lkdvos commented Sep 6, 2026

Copy link
Copy Markdown
Member Author

Benchmark: main vs ld-adjoint (rusty, dedicated rome node, --threads=4)

benchpkg TensorKit --rev=main,ld-adjoint --bench-on=ld-adjoint --filter=indexmanipulations, i.e. both revisions ran this branch's suite including the new adjoint = true variants (permute!(C, A', p)). Ratio is main / ld-adjoint, so > 1 means this branch is faster.

Adjoint sources (the #516 case): 1.03–1.69× faster than main, and at parity with the plain TensorMap path on this branch (e.g. SU₂ Float64 [[1,3],[2,4]]: 1.27 ms → 0.78 ms next to 0.82 ms plain; on main the adjoint variant was 2.2× slower than plain).

Plain path: unchanged within noise. Two entries show a < 1 ratio (SU₂ Float64 [[1,3],[2,4]] 0.72 ± 0.19 and ℤ₂ [7264, 7264] 0.82 ± 0.03); both were re-measured locally against main at 4 threads and are at parity (0.63–0.67 vs 0.67 ms, and 49–50 vs 50 ms). The ℤ₂ adjoint variant executes the identical kernel (real eltype ⇒ op = identity) and is at parity in the table as well, and the untouched Trivial path shows the same kind of scatter (0.83 ± 0.14), so these are between-process memory-placement effects on the 400 MB transposes rather than code differences.

Full table
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

Comment thread src/tensors/blockiterators.jl
Comment thread src/tensors/indexmanipulations.jl Outdated
@lkdvos
lkdvos marked this pull request as ready for review September 10, 2026 13:52
@lkdvos
lkdvos requested a review from Jutho September 10, 2026 13:52
Comment thread src/tensors/blockiterators.jl Outdated
Comment thread src/tensors/indexmanipulations.jl
Comment thread src/tensors/indexmanipulations.jl Outdated
Comment thread src/tensors/indexmanipulations.jl Outdated
Comment thread src/tensors/indexmanipulations.jl Outdated
lkdvos added a commit to QuantumKitHub/BlockTensorKit.jl that referenced this pull request Sep 15, 2026
* 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>
Comment thread src/tensors/indexmanipulations.jl
Comment thread src/tensors/indexmanipulations.jl Outdated
Comment thread src/tensors/blockiterators.jl Outdated
Comment thread src/tensors/tensor.jl Outdated
Comment thread src/tensors/tensoroperations.jl Outdated
Comment thread src/tensors/indexmanipulations.jl Outdated
Comment thread ext/TensorKitEnzymeExt/utility.jl Outdated
Comment thread src/tensors/blockiterators.jl Outdated
Comment thread src/tensors/treetransformers.jl Outdated
Comment thread src/tensors/tensor.jl Outdated
Comment thread src/tensors/treetransformers.jl Outdated
Comment thread src/tensors/treetransformers.jl Outdated
@lkdvos

lkdvos commented Sep 17, 2026

Copy link
Copy Markdown
Member Author

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 FusionBlock requires rebuilding the block array and then sorting it (in order to have a unique cached value for the unitary generated there), so there's a lot more subtle annoyances to get rid of this. I think in this case that is just complexity that is required... :(

@lkdvos
lkdvos enabled auto-merge (squash) September 17, 2026 18:13
@lkdvos lkdvos linked an issue Sep 17, 2026 that may be closed by this pull request
Comment thread src/tensors/treetransformers.jl
Comment thread src/tensors/treetransformers.jl Outdated
Comment thread src/tensors/indexmanipulations.jl Outdated
Comment thread src/tensors/tensor.jl Outdated
f = gettokenvalue(fusiontrees(iter.t), i)
return f => iter.structure[i], i + 1
end
end

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@Jutho

Jutho commented Sep 18, 2026

Copy link
Copy Markdown
Member

Ok, I finally managed to make my way through. Looks really great. I have four final questions or suggestions, but otherwise fully approve.

@lkdvos
lkdvos disabled auto-merge September 18, 2026 22:38
lkdvos and others added 3 commits September 19, 2026 15:20
"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>
@lkdvos

lkdvos commented Sep 19, 2026

Copy link
Copy Markdown
Member Author

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 SubblockIterator specialization made me realize that the TreeSubblocks were functionally the same structure as that, and so I managed to merge these implementations, just adding the "token-based" access to the original. This has further reduced some code at otherwise zero additional cost, which is always nice 😄.

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 flip and twist, and while I don't want to delay this PR further, it would definitely be useful to rewrite these in terms of the subblocks as well, further bypassing some symmetry overhead.

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.

Comment thread src/tensors/tensor.jl
@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)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 😄

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ok, fair enough.

@Jutho

Jutho commented Sep 19, 2026

Copy link
Copy Markdown
Member

By the way, did you rerun the benchmarks from the original issue with this latest iteration of the PR?

@lkdvos

lkdvos commented Sep 20, 2026

Copy link
Copy Markdown
Member Author

Issue #516 reproducer, rerun against the final state of this PR

main @ 7a187c4 vs ld-adjoint @ 6b2dd79, rusty rome node, julia 1.12.7, single-threaded, @benchmark samples=500 seconds=10. The ratio is in-place adjoint permute! over plain TensorMap permute! within the same revision, so 1.00 means an adjoint source costs nothing extra.

sector type ratio on main ratio here adjoint allocs, main → here
fℤ₂ 1.31× 1.09× 134 → 68
fℤ₂ ⊠ U₁ 3.82× 1.01× 675 → 320
U₁ 2.01× 1.11× 299 → 145
SU₂ 3.92× 1.01× 7096 → 1926
fℤ₂ ⊠ U₁ ⊠ SU₃ 19.06× 0.99× 7873 → 486

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 (SU₂ 0.183 → 0.187 ms, fℤ₂ ⊠ U₁ ⊠ SU₃ 0.0296 → 0.0301 ms), so this is the adjoint side coming down to the baseline rather than the baseline drifting up. In absolute terms the adjoint column goes SU₂ 0.718 → 0.189 ms and fℤ₂ ⊠ U₁ ⊠ SU₃ 0.564 → 0.030 ms.

Two notes on the baseline: main already carries #518 and #521, so the comparison is against the post-hash-fix numbers from the issue thread, not the original 82×. And the U₁ row is not in the issue script — I added it with a guessed space so the U₁ figure quoted in the PR description has a counterpart.

Full suite (linalg + indexmanipulations + tensornetworks) on both revisions is still running; I'll post the table when it lands.

@lkdvos
lkdvos merged commit bfca561 into main Sep 20, 2026
65 of 67 checks passed
@lkdvos
lkdvos deleted the ld-adjoint branch September 20, 2026 13:03
lkdvos added a commit that referenced this pull request Sep 21, 2026
* 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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Permuting an AdjointTensorMap seems much slower than it should be

3 participants