Repository navigation
Conversation
…ernel `twist!` looped over all fusion trees and scaled every subblock separately, with a cost that grows with the number of subblocks rather than the data size, and one kernel launch per subblock for GPU tensors (#595). - Add `twist(src, inds; inv)` for fusion tree pairs and blocks, returning the scalar twist factor, with `inds` given as indices or as a `Bool` mask. - Add the cached `treetwister`, a `UniqueTreeTransformer` for all sectors since a twist maps every tree pair onto itself, keyed on a fixed-length mask. - `twist!` and `twist` on a `TensorMap` go through `add_transform!` (in place for `twist!`, a single pass for `twist`); other tensor types keep the loop. - Move the shortcut for trivial permutations from `add_transform!` into `_braid!` and `_transpose!`, so `add_transform!` always applies the transformer. - Cache a device `treetwister` in the GPUArrays extension. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Adjoint destinations and sources are unwrapped as in the other index manipulations: the twist mask of an adjoint destination is relabeled onto the legs of its parent with the inverse twist, and `treetwister` takes a `conjsrc` flag, with the destination space and the permutation of the source legs following from the source space. An in-place twist of an adjoint therefore uses the same transformer as the equivalent twist of its parent. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… α, β, ...)` `twist!(t, inds)` and `twist(t, inds)` are now thin wrappers around the three-argument form, which goes through `unwrap_adjoints` and a `_twist!` kernel entry like `_braid!`. Adds `twist_pb` with Mooncake and Enzyme rules, and backend/allocator keywords for `twist`. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
kshyatt
reviewed
Oct 28, 2026
| # and store that instead of repeating the twist in the pullback each time. | ||
| # effectively, we replace `twist!` by `add ∘ twist`. | ||
| Ap = if !isa(α, Const) | ||
| Ap = twist(A.val, inds.val; inv, copy = true) |
Member
There was a problem hiding this comment.
similar to permute and transform, is there any chance this could secretly be aliasing A.val?
Member
Author
There was a problem hiding this comment.
That's what the copy=true is for right?
kshyatt
reviewed
Oct 28, 2026
| shadow = EnzymeRules.needs_shadow(config) ? C.dval : nothing | ||
| return EnzymeRules.AugmentedReturn(primal, shadow, cache) | ||
| end | ||
| function EnzymeRules.reverse( |
Member
There was a problem hiding this comment.
Was Enzyme not able to make its own rules for these? Kind of sad if so...
Member
Author
There was a problem hiding this comment.
These changes make twist go through the same path as permute and transpose etc, so even if it was able to do this I think this path is way more friendly on both compile time and possibly runtime?
…ist!` on Julia 1.10 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Addresses the standalone
twist!part of #595.twist!looped overfusiontrees(t)and calledscale!on every subblock, so its cost grew with the number of subblocks rather than with the data size, and a twist of a GPU tensor launched one kernel per subblock. This PR routestwist!/twistthrough the existing index manipulation kernel with a cached tree transformer.Changes
twist(src::Union{FusionTreePair, FusionTreeBlock}, inds; inv = false) -> (src, θ): fusion-tree level twist, next topermute/braid, withindsgiven as indices or as anNTuple{N, Bool}mask (indices are converted to a fixed-length mask, so all forms are inferable).twist!(tdst, tsrc, inds, α = 1, β = 0, [backend], [allocator]; inv = false)computingtdst = β * tdst + α * twist(tsrc, inds; inv), laid out likebraid!: it unwraps adjoints throughunwrap_adjointsand calls a_twist!kernel entry mirroring_braid!.twist!(t, inds)andtwist(t, inds)are thin wrappers around it, andtwistnow also acceptsbackend/allocatorkeywords.treetwister(A, Vdst, Vsrc, p, conjsrc, twistmask, inv), with the same arguments astreebraider(the mask taking the place of the levels): always aUniqueTreeTransformer, also for non-abelian sectors, since a twist maps every tree pair onto itself with a scalar factor (no recoupling matrices / pack-matmul-unpack needed).twist!(t, inds)runs throughadd_transform!(aliasing is allowed for trivial permutations in TensorOperations), andtwistis a single out-of-place pass instead ofcopy!followed bytwist!.TensorMap/AdjointTensorMapdestinations and sources go through the twister. The twist mask is relabeled like the levels ofbraid!, and an adjoint destination flipsinv(just as it conjugatesα).twist_pbpullback with Mooncake and Enzyme rules for the three-argumenttwist!, plus Enzyme rules for the in-placetwist!(t, inds). The existing ChainRulestwistrrule already forwards keywords and is now also tested withinvandbackend/allocator.add_transform!into_braid!/_transpose!(which also skips the transformer lookup there), soadd_transform!always applies its transformer.treetwister, registered and sized like the other device caches, so a GPUtwist!is a single batched kernel launch.Benchmarks (CPU, single thread)
fℤ₂ × U(1) tensors
Vχ ⊗ VD ⊗ VD' ← Vχ ⊗ VD ⊗ VD', twisted on indices(2, 3, 5, 6), minimum of 20 runs:twist!beforetwist!aftercopyBuilding the twister on a cache miss costs ~3× a cached call in the smallest case. There is still a per-subblock cost from the individual
tensoradd!calls on the CPU, which is left for a follow-up. GPU timings were not measured.Testing
tensors/indexmanipulations.jl,tensors/contractions.jl,chainrules/linalg.jl,mooncake/indexmanipulations.jl,enzyme/flip.jlandcuda/tensors.jl(RTX A6000) pass.test/cuda/tensors.jlcomparingtwist/twist!onCuTensorMaps (and their adjoints) against the host results, and checking that the cached device transformer is used.test/tensors/indexmanipulations.jl, covering adjoint sources, adjoint destinations and in-place twists of adjoints, plus α/β accumulation.Follow-up
Folding the twist into the permutations that precede a contraction (the second part of #595) will be done on top of this, through a cached contraction plan holding the transformers.
🤖 Generated with Claude Code