From 93baf98ad1ede769476587652e711b5477b56456 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Wed, 28 Oct 2026 09:41:38 -0400 Subject: [PATCH 1/5] Implement `twist!` through a cached `treetwister` and the transform kernel `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 --- docs/src/Changelog.md | 2 + docs/src/lib/fusiontrees.md | 1 + ext/TensorKitEnzymeExt/utility.jl | 1 + ext/TensorKitGPUArraysExt.jl | 9 +++++ src/auxiliary/caches.jl | 2 +- src/fusiontrees/braiding_manipulations.jl | 17 ++++++++ src/tensors/indexmanipulations.jl | 48 ++++++++++++++--------- src/tensors/treetransformers.jl | 8 ++++ test/cuda/tensors.jl | 16 ++++++++ 9 files changed, 84 insertions(+), 20 deletions(-) diff --git a/docs/src/Changelog.md b/docs/src/Changelog.md index a6ab14c3b..a30130890 100644 --- a/docs/src/Changelog.md +++ b/docs/src/Changelog.md @@ -23,6 +23,7 @@ When releasing a new version, move the "Unreleased" changes to a new version sec ### Added - Re-export `NamedSector` and `@NamedSector` from TensorKitSectors +- `twist(src, inds; inv = false)` for fusion tree pairs and blocks, returning `src` together with the scalar twist factor of the legs `inds`, which can be given as indices or as a `Bool` mask ### Changed @@ -44,6 +45,7 @@ When releasing a new version, move the "Unreleased" changes to a new version sec - Recoupling matrices are stored such that they can leverage BLAS, also for mixed complex tensor with real recoupling coefficients. - Constructing a `FusionTreeBlock` only takes the fusion product of the shorter side, and `fusionblocks` groups the cached fusion trees of the space instead of constructing a block for every combination of sectors (internal) - Fusion tree block manipulations of (named) product sectors are computed as Kronecker products of the manipulations of their parts, splitting off the `UniqueFusion` sectors, such that they share the caches of the individual sectors ([#592](https://github.com/QuantumKitHub/TensorKit.jl/pull/592)) +- `twist!` and `twist` on a `TensorMap` use a cached tree transformer and the index manipulation kernel instead of a loop over all fusion trees, which removes the per-subblock bookkeeping and turns a twist of a GPU tensor into a single kernel launch ([#595](https://github.com/QuantumKitHub/TensorKit.jl/issues/595)) ## [0.17.2](https://github.com/QuantumKitHub/TensorKit.jl/compare/v0.17.1...v0.17.2) - 2026-09-20 diff --git a/docs/src/lib/fusiontrees.md b/docs/src/lib/fusiontrees.md index 57033eca6..e1d72a1c5 100644 --- a/docs/src/lib/fusiontrees.md +++ b/docs/src/lib/fusiontrees.md @@ -48,4 +48,5 @@ repartition(src::Union{FusionTreePair, FusionTreeBlock}, N::Int) Base.transpose(src::Union{FusionTreePair, FusionTreeBlock}, p::Index2Tuple) braid(src::Union{FusionTreePair, FusionTreeBlock}, p::Index2Tuple, levels::Index2Tuple) permute(src::Union{FusionTreePair, FusionTreeBlock}, p::Index2Tuple) +twist(src::Union{FusionTreePair, FusionTreeBlock}, inds) ``` diff --git a/ext/TensorKitEnzymeExt/utility.jl b/ext/TensorKitEnzymeExt/utility.jl index bb9a21db0..ddc734fdb 100644 --- a/ext/TensorKitEnzymeExt/utility.jl +++ b/ext/TensorKitEnzymeExt/utility.jl @@ -150,6 +150,7 @@ end @inline EnzymeRules.inactive(::typeof(TensorKit.product_transform), ::Any...) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.treetransposer), ::Any...) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.treebraider), ::Any...) = nothing +@inline EnzymeRules.inactive(::typeof(TensorKit.treetwister), ::Any...) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.insertleftunit), ::HomSpace, ::Any) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.insertrightunit), ::HomSpace, ::Any) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.removeunit), ::HomSpace, ::Any) = nothing diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl index 5e013865d..2388e43ba 100644 --- a/ext/TensorKitGPUArraysExt.jl +++ b/ext/TensorKitGPUArraysExt.jl @@ -287,6 +287,13 @@ end return device_transformer(A, host, linearize(p)) end +@cached function TensorKit.treetwister( + A::Type{TA}, V::TensorMapSpace, twistmask::NTuple{N, Bool}, inv::Bool + ) where {TA <: AnyGPUVector, N} + host = TensorKit.treetwister(Vector{scalartype(A)}, V, twistmask, inv) + return device_transformer(A, host, linearize((codomainind(V), domainind(V)))) +end + const DEFAULT_DEVICECACHE_SIZE = 1000 function __init__() @@ -295,12 +302,14 @@ function __init__() TensorKit.GLOBAL_CACHES, ( Symbol("TensorKitGPUArraysExt.GLOBAL_TREEBRAIDER_CACHE") => GLOBAL_TREEBRAIDER_CACHE, Symbol("TensorKitGPUArraysExt.GLOBAL_TREETRANSPOSER_CACHE") => GLOBAL_TREETRANSPOSER_CACHE, + Symbol("TensorKitGPUArraysExt.GLOBAL_TREETWISTER_CACHE") => GLOBAL_TREETWISTER_CACHE, ) ) # device-caches should probably have smaller sizes: resize!(GLOBAL_TREEBRAIDER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) resize!(GLOBAL_TREETRANSPOSER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) + resize!(GLOBAL_TREETWISTER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) return nothing end diff --git a/src/auxiliary/caches.jl b/src/auxiliary/caches.jl index d6753ead9..5525f1714 100644 --- a/src/auxiliary/caches.jl +++ b/src/auxiliary/caches.jl @@ -50,7 +50,7 @@ end # category of the miss-path (construction) timer section of an `@cached` function function _cached_category(fname::Symbol) - return fname in (:fsbraid, :fstranspose, :treebraider, :treetransposer) ? + return fname in (:fsbraid, :fstranspose, :treebraider, :treetransposer, :treetwister) ? "symmetry" : "bookkeeping" end diff --git a/src/fusiontrees/braiding_manipulations.jl b/src/fusiontrees/braiding_manipulations.jl index 0ad6353ed..b737f354f 100644 --- a/src/fusiontrees/braiding_manipulations.jl +++ b/src/fusiontrees/braiding_manipulations.jl @@ -367,6 +367,23 @@ function permute(src::Union{FusionTreePair, FusionTreeBlock}, p::Index2Tuple) return braid(src, p, (levels1, levels2)) end +""" + twist(src::Union{FusionTreePair, FusionTreeBlock}, inds; inv::Bool = false) -> (src, θ) + +Apply a twist to the uncoupled legs `inds` of the fusion-splitting tree pair or block `src`, +where `inds` is either a collection of leg indices or an `NTuple{N, Bool}` mask over all legs. +The twist is diagonal in the fusion tree basis, so this returns `src` itself together with the +scalar factor `θ`, which is shared by all trees in a block. If `inv = true`, use the inverse twist. +""" +function twist(src::Union{FusionTreePair{I}, FusionTreeBlock{I}}, inds; inv::Bool = false) where {I} + uncoupled₁, uncoupled₂ = src isa FusionTreeBlock ? src.uncoupled : (src[1].uncoupled, src[2].uncoupled) + uncoupled = (uncoupled₁..., uncoupled₂...) + N = length(uncoupled) + mask = inds isa NTuple{N, Bool} ? inds : ntuple(i -> i in inds, N) + θ = prod(map((m, a) -> m ? twist(a) : one(twist(a)), mask, uncoupled); init = one(sectorscalartype(I))) + return src, inv ? θ' : θ +end + """ flip((f₁, f₂)::FusionTreePair, i::Int; inv::Bool = false) -> SingletonDict{FusionTreePair, <:Number} diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index f6623c197..38fb5dc38 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -67,14 +67,7 @@ function twist!(t::AbstractTensorMap, inds; inv::Bool = false) (scalartype(t) <: Real && !(sectorscalartype(sectortype(t)) <: Real)) && throw(ArgumentError("Can't in-place twist a real tensor with complex sector type")) has_shared_twist(t, inds) && return t - - N₁ = numout(t) - for (f₁, f₂) in fusiontrees(t) - θ = prod(i -> i <= N₁ ? twist(f₁.uncoupled[i]) : twist(f₂.uncoupled[i - N₁]), inds) - inv && (θ = θ') - scale!(t[f₁, f₂], θ) - end - return t + return _twist!(t, t, inds, inv) end """ @@ -92,8 +85,25 @@ function twist(t::AbstractTensorMap, inds; inv::Bool = false, copy::Bool = false return copy ? Base.copy(t) : t end tdst = similar(t, promote_twist(t)) - copy!(tdst, t) - return twist!(tdst, inds; inv) + return _twist!(tdst, t, inds, inv) +end + +# `tdst = twist(tsrc, inds; inv)`, where `tdst` and `tsrc` are allowed to be the same tensor +function _twist!(tdst::TensorMap, tsrc::TensorMap, inds, inv::Bool) + V = space(tsrc) + twistmask = ntuple(i -> i in inds, numind(V)) + transformer = treetwister(storagetype(tdst), V, twistmask, inv) + return add_transform!( + tdst, tsrc, (codomainind(V), domainind(V)), false, transformer, + One(), Zero(), TO.DefaultBackend(), TO.DefaultAllocator() + ) +end +function _twist!(tdst::AbstractTensorMap, tsrc::AbstractTensorMap, inds, inv::Bool) + for (f₁, f₂) in fusiontrees(tsrc) + _, θ = twist((f₁, f₂), inds; inv) + scale!(tdst[f₁, f₂], tsrc[f₁, f₂], θ) + end + return tdst end # ========================= @@ -605,6 +615,7 @@ end Vdst, Vsrc = space(tdst), space(tsrc) @boundscheck spacecheck_transform(permute, Vdst, Vsrc, p, conjsrc) all(has_array_view, (tdst, tsrc)) && return _dense_transform!(tdst, tsrc, p, conjsrc, α, β, backend, allocator) + _is_trivial_transform(tsrc, p, conjsrc) && return add!(tdst, tsrc, α, β) transformer = treebraider(storagetype(tdst), Vdst, Vsrc, p, conjsrc, levels) return add_transform!(tdst, tsrc, p, conjsrc, transformer, α, β, backend, allocator) end @@ -616,6 +627,7 @@ end Vdst, Vsrc = space(tdst), space(tsrc) @boundscheck spacecheck_transform(permute, Vdst, Vsrc, p, conjsrc) all(has_array_view, (tdst, tsrc)) && return _dense_transform!(tdst, tsrc, p, conjsrc, α, β, backend, allocator) + _is_trivial_transform(tsrc, p, conjsrc) && return add!(tdst, tsrc, α, β) transformer = treetransposer(storagetype(tdst), Vdst, Vsrc, p, conjsrc) return add_transform!(tdst, tsrc, p, conjsrc, transformer, α, β, backend, allocator) end @@ -626,18 +638,16 @@ function add_transform!( tdst::AbstractTensorMap, tsrc::AbstractTensorMap, p::Index2Tuple, conjsrc::Bool, transformer, α::Number, β::Number, backend, allocator ) - if !conjsrc && p[1] === codomainind(tsrc) && p[2] === domainind(tsrc) - add!(tdst, tsrc, α, β) - else - p2 = (linearize(p), ()) # only the linear permutation matters for the array kernels - ntasks = get_num_transformer_threads(tdst) - dst, src = _transform_subblocks(tdst, tsrc, transformer) - add_transform_kernel!(dst, src, p2, conjsrc, transformer, α, β, backend, allocator, ntasks) - end - + p2 = (linearize(p), ()) # only the linear permutation matters for the array kernels + ntasks = get_num_transformer_threads(tdst) + dst, src = _transform_subblocks(tdst, tsrc, transformer) + add_transform_kernel!(dst, src, p2, conjsrc, transformer, α, β, backend, allocator, ntasks) return tdst end +_is_trivial_transform(tsrc, p::Index2Tuple, conjsrc::Bool) = + !conjsrc && p[1] === codomainind(tsrc) && p[2] === domainind(tsrc) + # TensorMaps address their flat data directly, other tensor types go through `subblock` _transform_subblocks(tdst::TensorMap, tsrc::TensorMap, transformer) = StridedSubblocks(tdst, transformer.structure_dst), StridedSubblocks(tsrc, transformer.structure_src) diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index b274c6f8c..d7dd57b31 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -202,4 +202,12 @@ end return TreeTransformer(A, fusiontreetransform, p, Vdst, Vsrc, conjsrc) end +# a twist maps every tree pair onto itself, so this is a `UniqueTreeTransformer` for all sectors +@cached function treetwister( + A::Type{TA}, V::TensorMapSpace, twistmask::NTuple{N, Bool}, inv::Bool + )::UniqueTreeTransformer{recoupling_scalartype(A, sectorscalartype(sectortype(V))), N} where {TA, N} + fusiontreetwister(f) = twist(f, twistmask; inv) + return UniqueTreeTransformer(A, fusiontreetwister, (codomainind(V), domainind(V)), V, V, false) +end + # default cachestyle is GlobalLRUCache diff --git a/test/cuda/tensors.jl b/test/cuda/tensors.jl index 59196d52c..980567e29 100644 --- a/test/cuda/tensors.jl +++ b/test/cuda/tensors.jl @@ -6,6 +6,7 @@ ad = adapt(Array) const CUDAExt = Base.get_extension(TensorKit, :TensorKitCUDAExt) @assert !isnothing(CUDAExt) const CuTensorMap = getglobal(CUDAExt, :CuTensorMap) +const GPUArraysExt = Base.get_extension(TensorKit, :TensorKitGPUArraysExt) Random.seed!(123456) @@ -308,6 +309,21 @@ for V in spacelist @test ht3 ≈ adapt(Array, dt3) end end + hasbraiding && @timedtestset "Twists: test via CPU" begin + t = cuRAND.rand(ComplexF64, V1 ⊗ V2 ⊗ V3 ← (V4 ⊗ V5)') + ht = adapt(Array, t) + for inds in (1, (2, 4), [1, 3, 5]), inv in (false, true) + @test adapt(Array, twist(t, inds; inv)) ≈ twist(ht, inds; inv) + @test adapt(Array, twist!(copy(t), inds; inv)) ≈ twist!(copy(ht), inds; inv) + end + @test twist(twist(t, 2), 2; inv = true) ≈ t + if !TensorKit.has_shared_twist(t, (1, 3, 5)) + twistmask = (true, false, true, false, true) + twister = TensorKit.treetwister(storagetype(t), space(t), twistmask, false) + @test twister isa getglobal(GPUArraysExt, :DeviceUniqueTreeTransformer) + @test twister === TensorKit.treetwister(storagetype(t), space(t), twistmask, false) + end + end symmetricbraiding && @timedtestset "Full trace: test self-consistency" begin t = cuRAND.rand(ComplexF64, V1 ⊗ V2' ⊗ V2 ⊗ V1') t2 = permute(t, ((1, 2), (4, 3))) From 866ee15035acd84c7ad1c5b348d879fb237ff983 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Wed, 28 Oct 2026 11:26:43 -0400 Subject: [PATCH 2/5] Handle adjoint operands of `twist!` and `twist` through the twister 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 --- docs/src/Changelog.md | 2 +- ext/TensorKitGPUArraysExt.jl | 7 ++++--- src/tensors/indexmanipulations.jl | 23 +++++++++++++++-------- src/tensors/treetransformers.jl | 14 +++++++++++--- test/cuda/tensors.jl | 14 ++++++++++---- test/tensors/indexmanipulations.jl | 20 ++++++++++++++++++++ 6 files changed, 61 insertions(+), 19 deletions(-) diff --git a/docs/src/Changelog.md b/docs/src/Changelog.md index a30130890..b87fb5e64 100644 --- a/docs/src/Changelog.md +++ b/docs/src/Changelog.md @@ -45,7 +45,7 @@ When releasing a new version, move the "Unreleased" changes to a new version sec - Recoupling matrices are stored such that they can leverage BLAS, also for mixed complex tensor with real recoupling coefficients. - Constructing a `FusionTreeBlock` only takes the fusion product of the shorter side, and `fusionblocks` groups the cached fusion trees of the space instead of constructing a block for every combination of sectors (internal) - Fusion tree block manipulations of (named) product sectors are computed as Kronecker products of the manipulations of their parts, splitting off the `UniqueFusion` sectors, such that they share the caches of the individual sectors ([#592](https://github.com/QuantumKitHub/TensorKit.jl/pull/592)) -- `twist!` and `twist` on a `TensorMap` use a cached tree transformer and the index manipulation kernel instead of a loop over all fusion trees, which removes the per-subblock bookkeeping and turns a twist of a GPU tensor into a single kernel launch ([#595](https://github.com/QuantumKitHub/TensorKit.jl/issues/595)) +- `twist!` and `twist` on a `TensorMap` or its adjoint use a cached tree transformer and the index manipulation kernel instead of a loop over all fusion trees, which removes the per-subblock bookkeeping and turns a twist of a GPU tensor into a single kernel launch ([#595](https://github.com/QuantumKitHub/TensorKit.jl/issues/595)) ## [0.17.2](https://github.com/QuantumKitHub/TensorKit.jl/compare/v0.17.1...v0.17.2) - 2026-09-20 diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl index 2388e43ba..eecb63827 100644 --- a/ext/TensorKitGPUArraysExt.jl +++ b/ext/TensorKitGPUArraysExt.jl @@ -288,10 +288,11 @@ end end @cached function TensorKit.treetwister( - A::Type{TA}, V::TensorMapSpace, twistmask::NTuple{N, Bool}, inv::Bool + A::Type{TA}, V::TensorMapSpace, conjsrc::Bool, twistmask::NTuple{N, Bool}, inv::Bool ) where {TA <: AnyGPUVector, N} - host = TensorKit.treetwister(Vector{scalartype(A)}, V, twistmask, inv) - return device_transformer(A, host, linearize((codomainind(V), domainind(V)))) + host = TensorKit.treetwister(Vector{scalartype(A)}, V, conjsrc, twistmask, inv) + _, p = TensorKit.twistpermutation(V, conjsrc) + return device_transformer(A, host, linearize(p)) end const DEFAULT_DEVICECACHE_SIZE = 1000 diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index 38fb5dc38..9771806dc 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -88,15 +88,22 @@ function twist(t::AbstractTensorMap, inds; inv::Bool = false, copy::Bool = false return _twist!(tdst, t, inds, inv) end +const TensorMapOrAdjoint = Union{TensorMap, AdjointTensorMap{<:Any, <:Any, <:Any, <:Any, <:TensorMap}} + # `tdst = twist(tsrc, inds; inv)`, where `tdst` and `tsrc` are allowed to be the same tensor -function _twist!(tdst::TensorMap, tsrc::TensorMap, inds, inv::Bool) - V = space(tsrc) - twistmask = ntuple(i -> i in inds, numind(V)) - transformer = treetwister(storagetype(tdst), V, twistmask, inv) - return add_transform!( - tdst, tsrc, (codomainind(V), domainind(V)), false, transformer, - One(), Zero(), TO.DefaultBackend(), TO.DefaultAllocator() - ) +function _twist!(tdst::TensorMapOrAdjoint, tsrc::TensorMapOrAdjoint, inds, inv::Bool) + twistmask = ntuple(i -> i in inds, numind(tdst)) + # the twist acts on the legs of `tdst`, which is the inverse twist of the legs of its parent + if tdst isa AdjointTensorMap + tp = parent(tdst) + twistmask = TupleTools.getindices(twistmask, adjointtensorindices(tp, allind(tp))) + inv = !inv + end + p = (codomainind(tsrc), domainind(tsrc)) + tdst′, tsrc′, p′, _, conjsrc, α, β = unwrap_adjoints(tdst, tsrc, p, nothing, false, One(), Zero()) + transformer = treetwister(storagetype(tdst′), space(tsrc′), conjsrc, twistmask, inv) + add_transform!(tdst′, tsrc′, p′, conjsrc, transformer, α, β, TO.DefaultBackend(), TO.DefaultAllocator()) + return tdst end function _twist!(tdst::AbstractTensorMap, tsrc::AbstractTensorMap, inds, inv::Bool) for (f₁, f₂) in fusiontrees(tsrc) diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index d7dd57b31..e2c46304e 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -202,12 +202,20 @@ end return TreeTransformer(A, fusiontreetransform, p, Vdst, Vsrc, conjsrc) end -# a twist maps every tree pair onto itself, so this is a `UniqueTreeTransformer` for all sectors +# a twist maps every tree pair onto itself, so this is a `UniqueTreeTransformer` for all sectors; +# it twists the legs `twistmask` of `conjsrc ? V' : V`, reading the subblocks of a source with space `V` @cached function treetwister( - A::Type{TA}, V::TensorMapSpace, twistmask::NTuple{N, Bool}, inv::Bool + A::Type{TA}, V::TensorMapSpace, conjsrc::Bool, twistmask::NTuple{N, Bool}, inv::Bool )::UniqueTreeTransformer{recoupling_scalartype(A, sectorscalartype(sectortype(V))), N} where {TA, N} + Vdst, p = twistpermutation(V, conjsrc) fusiontreetwister(f) = twist(f, twistmask; inv) - return UniqueTreeTransformer(A, fusiontreetwister, (codomainind(V), domainind(V)), V, V, false) + return UniqueTreeTransformer(A, fusiontreetwister, p, Vdst, V, conjsrc) +end + +# destination space of a twist of `conjsrc ? V' : V`, and the permutation of the legs of `V` it applies +function twistpermutation(V::TensorMapSpace, conjsrc::Bool) + conjsrc || return V, (codomainind(V), domainind(V)) + return V', adjointtensorindices(V', (codomainind(V'), domainind(V'))) end # default cachestyle is GlobalLRUCache diff --git a/test/cuda/tensors.jl b/test/cuda/tensors.jl index 980567e29..ec69dc496 100644 --- a/test/cuda/tensors.jl +++ b/test/cuda/tensors.jl @@ -313,15 +313,21 @@ for V in spacelist t = cuRAND.rand(ComplexF64, V1 ⊗ V2 ⊗ V3 ← (V4 ⊗ V5)') ht = adapt(Array, t) for inds in (1, (2, 4), [1, 3, 5]), inv in (false, true) - @test adapt(Array, twist(t, inds; inv)) ≈ twist(ht, inds; inv) - @test adapt(Array, twist!(copy(t), inds; inv)) ≈ twist!(copy(ht), inds; inv) + ref = twist(ht, inds; inv) + @test adapt(Array, twist(t, inds; inv)) ≈ ref + @test adapt(Array, twist!(copy(t), inds; inv)) ≈ ref + # adjoint operands + @test adapt(Array, twist(t', inds; inv)) ≈ twist(copy(ht'), inds; inv) + t2 = copy(t) + twist!(t2', inds; inv) + @test adapt(Array, t2)' ≈ twist(copy(ht'), inds; inv) end @test twist(twist(t, 2), 2; inv = true) ≈ t if !TensorKit.has_shared_twist(t, (1, 3, 5)) twistmask = (true, false, true, false, true) - twister = TensorKit.treetwister(storagetype(t), space(t), twistmask, false) + twister = TensorKit.treetwister(storagetype(t), space(t), false, twistmask, false) @test twister isa getglobal(GPUArraysExt, :DeviceUniqueTreeTransformer) - @test twister === TensorKit.treetwister(storagetype(t), space(t), twistmask, false) + @test twister === TensorKit.treetwister(storagetype(t), space(t), false, twistmask, false) end end symmetricbraiding && @timedtestset "Full trace: test self-consistency" begin diff --git a/test/tensors/indexmanipulations.jl b/test/tensors/indexmanipulations.jl index 0cc276c6d..cfc576457 100644 --- a/test/tensors/indexmanipulations.jl +++ b/test/tensors/indexmanipulations.jl @@ -173,6 +173,26 @@ for V in spacelist end end end + hasbraiding && @timedtestset "Twists: adjoint operands" begin + W = V1 ⊗ V2 ← V3 ⊗ V4 + for T in (Float64, ComplexF64) + T <: Real && !(sectorscalartype(I) <: Real) && continue + t = rand(T, W) + tref = copy(t') # genuine TensorMap on W' + for inds in (1, (2, 3), [1, 4], (1, 2, 3, 4)), inv in (false, true) + ref = twist(tref, inds; inv) + # adjoint source + @test twist(t', inds; inv) ≈ ref + # adjoint source and destination, in place + t2 = copy(t) + @test twist!(t2', inds; inv) ≈ ref + @test t2 ≈ ref' + # adjoint destination + @test TensorKit._twist!(similar(t)', tref, inds, inv) ≈ ref + @test twist(twist(t', inds), inds; inv = true) ≈ tref + end + end + end symmetricbraiding && @timedtestset "Permutations: BraidingTensor source" begin τ = BraidingTensor(V1, V2) p = ((2, 1), (3, 4)) From cef98d88f37354bf07bb4d243068ae26496d964d Mon Sep 17 00:00:00 2001 From: lkdvos Date: Wed, 28 Oct 2026 13:26:55 -0400 Subject: [PATCH 3/5] =?UTF-8?q?Make=20`twist!`=20mirror=20`braid!`=20with?= =?UTF-8?q?=20a=20public=20`twist!(tdst,=20tsrc,=20inds,=20=CE=B1,=20?= =?UTF-8?q?=CE=B2,=20...)`?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `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 --- docs/src/Changelog.md | 1 + ext/TensorKitEnzymeExt/indexmanipulations.jl | 138 ++++++++++++++++++ ext/TensorKitGPUArraysExt.jl | 6 +- .../indexmanipulations.jl | 87 +++++++++++ src/pullbacks/indexmanipulations.jl | 11 ++ src/tensors/indexmanipulations.jl | 106 +++++++++----- src/tensors/treetransformers.jl | 23 ++- test/chainrules/linalg.jl | 2 + test/cuda/tensors.jl | 5 +- test/enzyme/flip.jl | 9 ++ test/mooncake/indexmanipulations.jl | 12 ++ test/tensors/indexmanipulations.jl | 13 +- 12 files changed, 354 insertions(+), 59 deletions(-) diff --git a/docs/src/Changelog.md b/docs/src/Changelog.md index b87fb5e64..7cc5da908 100644 --- a/docs/src/Changelog.md +++ b/docs/src/Changelog.md @@ -24,6 +24,7 @@ When releasing a new version, move the "Unreleased" changes to a new version sec - Re-export `NamedSector` and `@NamedSector` from TensorKitSectors - `twist(src, inds; inv = false)` for fusion tree pairs and blocks, returning `src` together with the scalar twist factor of the legs `inds`, which can be given as indices or as a `Bool` mask +- `twist!(tdst, tsrc, inds, α = 1, β = 0, [backend], [allocator]; inv = false)` computing `tdst = β * tdst + α * twist(tsrc, inds; inv)`, in line with `permute!` and `braid!`, together with Mooncake and Enzyme rules; `twist` now also accepts `backend` and `allocator` keywords ### Changed diff --git a/ext/TensorKitEnzymeExt/indexmanipulations.jl b/ext/TensorKitEnzymeExt/indexmanipulations.jl index 50a1da480..9908ede18 100644 --- a/ext/TensorKitEnzymeExt/indexmanipulations.jl +++ b/ext/TensorKitEnzymeExt/indexmanipulations.jl @@ -186,6 +186,144 @@ function EnzymeRules.forward( end end +function EnzymeRules.augmented_primal( + config::EnzymeRules.RevConfigWidth{1}, + func::Const{typeof(TK.twist!)}, + ::Type{RT}, + C::Annotation{<:AbstractTensorMap}, + A::Annotation{<:AbstractTensorMap}, + inds::Const, + α::Annotation{<:Number}, + β::Annotation{<:Number}, + ba::Const...; + inv::Bool = false + ) where {RT} + C_cache = !isa(β, Const) ? copy(C.val) : nothing + A_cache = EnzymeRules.overwritten(config)[3] ? copy(A.val) : nothing + # if we need to compute Δa, it is faster to allocate an intermediate twisted A + # 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) + add!(C.val, Ap, α.val, β.val) + Ap + else + bavs = map(a -> a.val, ba) + TK.twist!(C.val, A.val, inds.val, α.val, β.val, bavs...; inv) + nothing + end + cache = (C_cache, A_cache, Ap) + primal = EnzymeRules.needs_primal(config) ? C.val : nothing + shadow = EnzymeRules.needs_shadow(config) ? C.dval : nothing + return EnzymeRules.AugmentedReturn(primal, shadow, cache) +end +function EnzymeRules.reverse( + config::EnzymeRules.RevConfigWidth{1}, + func::Const{typeof(TK.twist!)}, + ::Type{RT}, + cache, + C::Annotation{<:AbstractTensorMap}, + A::Annotation{<:AbstractTensorMap}, + inds::Const, + α::Annotation{<:Number}, + β::Annotation{<:Number}, + ba::Const...; + inv::Bool = false + ) where {RT} + C_cache, A_cache, Ap = cache + Cval = something(C_cache, C.val) + Aval = something(A_cache, A.val) + bavs = map(a -> a.val, ba) + # ΔA + if !isa(A, Const) && !isa(C, Const) + TK.twist_pb(A.dval, Aval, C.dval, C.val, inds.val, α.val, bavs...; inv) + end + Δαr = pullback_dα(α, C, Ap) + Δβr = pullback_dβ(β, C, Cval) + !isa(C, Const) && pullback_dC!(C.dval, β.val) + return nothing, nothing, nothing, Δαr, Δβr, map(Returns(nothing), ba)... +end +function EnzymeRules.forward( + config::EnzymeRules.FwdConfigWidth{1}, + func::Const{typeof(TK.twist!)}, + ::Type{RT}, + C::Annotation{<:AbstractTensorMap}, + A::Annotation{<:AbstractTensorMap}, + inds::Annotation, + α::Annotation{<:Number}, + β::Annotation{<:Number}, + ba::Const...; + inv::Bool = false + ) where {RT} + bavs = map(a -> a.val, ba) + if !isa(C, Const) + if isa(β, Const) + scale!(C.dval, β.val) + else + add!(C.dval, C.val, β.dval, β.val) + end + !isa(α, Const) && TK.twist!(C.dval, A.val, inds.val, α.dval, One(), bavs...; inv) + !isa(A, Const) && TK.twist!(C.dval, A.dval, inds.val, α.val, One(), bavs...; inv) + end + TK.twist!(C.val, A.val, inds.val, α.val, β.val, bavs...; inv) + if EnzymeRules.needs_primal(config) && EnzymeRules.needs_shadow(config) + return C + elseif EnzymeRules.needs_primal(config) + return C.val + elseif EnzymeRules.needs_shadow(config) + return C.dval + else + return nothing + end +end + +# in-place twist: the source and destination are the same tensor +function EnzymeRules.augmented_primal( + config::EnzymeRules.RevConfigWidth{1}, + func::Const{typeof(TK.twist!)}, + ::Type{RT}, + t::Annotation{<:AbstractTensorMap}, + inds::Const; + inv::Bool = false + ) where {RT} + TK.twist!(t.val, inds.val; inv) + primal = EnzymeRules.needs_primal(config) ? t.val : nothing + shadow = EnzymeRules.needs_shadow(config) ? t.dval : nothing + return EnzymeRules.AugmentedReturn(primal, shadow, nothing) +end +function EnzymeRules.reverse( + config::EnzymeRules.RevConfigWidth{1}, + func::Const{typeof(TK.twist!)}, + ::Type{RT}, + cache, + t::Annotation{<:AbstractTensorMap}, + inds::Const; + inv::Bool = false + ) where {RT} + !isa(t, Const) && TK.twist!(t.dval, inds.val; inv = !inv) + return (nothing, nothing) +end +function EnzymeRules.forward( + config::EnzymeRules.FwdConfigWidth{1}, + func::Const{typeof(TK.twist!)}, + ::Type{RT}, + t::Annotation{<:AbstractTensorMap}, + inds::Annotation; + inv::Bool = false + ) where {RT} + TK.twist!(t.val, inds.val; inv) + !isa(t, Const) && TK.twist!(t.dval, inds.val; inv) + if EnzymeRules.needs_primal(config) && EnzymeRules.needs_shadow(config) + return t + elseif EnzymeRules.needs_primal(config) + return t.val + elseif EnzymeRules.needs_shadow(config) + return t.dval + else + return nothing + end +end + function EnzymeRules.augmented_primal( config::EnzymeRules.RevConfigWidth{1}, func::Const{typeof(flip)}, diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl index eecb63827..e8018dfc5 100644 --- a/ext/TensorKitGPUArraysExt.jl +++ b/ext/TensorKitGPUArraysExt.jl @@ -288,10 +288,10 @@ end end @cached function TensorKit.treetwister( - A::Type{TA}, V::TensorMapSpace, conjsrc::Bool, twistmask::NTuple{N, Bool}, inv::Bool + A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, + conjsrc::Bool, twistmask::NTuple{N, Bool}, inv::Bool ) where {TA <: AnyGPUVector, N} - host = TensorKit.treetwister(Vector{scalartype(A)}, V, conjsrc, twistmask, inv) - _, p = TensorKit.twistpermutation(V, conjsrc) + host = TensorKit.treetwister(Vector{scalartype(A)}, Vdst, Vsrc, p, conjsrc, twistmask, inv) return device_transformer(A, host, linearize(p)) end diff --git a/ext/TensorKitMooncakeExt/indexmanipulations.jl b/ext/TensorKitMooncakeExt/indexmanipulations.jl index bedc0af72..cd7dd63be 100644 --- a/ext/TensorKitMooncakeExt/indexmanipulations.jl +++ b/ext/TensorKitMooncakeExt/indexmanipulations.jl @@ -160,6 +160,93 @@ function Mooncake.rrule!!( return t_Δt, twist_pullback end +# both are needed for correctly capturing every dispatch +@is_primitive( + DefaultCtx, + ReverseMode, + Tuple{ + typeof(TK.twist!), + AbstractTensorMap, + AbstractTensorMap, Any, + Number, Number, Vararg{Any}, + } +) +@is_primitive( + DefaultCtx, + ReverseMode, + Tuple{ + typeof(Core.kwcall), @NamedTuple{inv::Bool}, typeof(TK.twist!), + AbstractTensorMap, + AbstractTensorMap, Any, + Vararg{Any}, + } +) + +function Mooncake.rrule!!( + ::CoDual{typeof(TK.twist!)}, + C_ΔC::CoDual{<:AbstractTensorMap}, + A_ΔA::CoDual{<:AbstractTensorMap}, inds_Δinds::CoDual, + α_Δα::CoDual{<:Number}, β_Δβ::CoDual{<:Number}, + ba_Δba::CoDual... + ) + C_ΔC, pullback = twist!_rrule(C_ΔC, A_ΔA, inds_Δinds, false, (α_Δα, β_Δβ, ba_Δba...)) + twist!_pullback(Δ::NoRData) = (NoRData(), pullback(Δ)...) + return C_ΔC, twist!_pullback +end +function Mooncake.rrule!!( + ::CoDual{typeof(Core.kwcall)}, kwargs_Δkwargs::CoDual{@NamedTuple{inv::Bool}}, ::CoDual{typeof(TK.twist!)}, + C_ΔC::CoDual{<:AbstractTensorMap}, + A_ΔA::CoDual{<:AbstractTensorMap}, inds_Δinds::CoDual, + args_Δargs::CoDual... + ) + inv = primal(kwargs_Δkwargs).inv + C_ΔC, pullback = twist!_rrule(C_ΔC, A_ΔA, inds_Δinds, inv, args_Δargs) + twist!_pullback(Δ::NoRData) = (NoRData(), NoRData(), NoRData(), pullback(Δ)...) + return C_ΔC, twist!_pullback +end + +# `args_Δargs` holds the optional `α`, `β`, backend and allocator, where a missing `α` or `β` is constant +function twist!_rrule(C_ΔC, A_ΔA, inds_Δinds, inv, args_Δargs) + # prepare arguments + C, ΔC = arrayify(C_ΔC) + A, ΔA = arrayify(A_ΔA) + inds = primal(inds_Δinds) + args = primal.(args_Δargs) + α = length(args) >= 1 ? args[1] : One() + β = length(args) >= 2 ? args[2] : Zero() + ba = length(args) >= 3 ? args[3:end] : () + + C_cache = copy(C) + + # if we need to compute Δa, it is faster to allocate an intermediate twisted A + # and store that instead of repeating the twist in the pullback each time. + # effectively, we replace `twist!` by `add ∘ twist`. + Ap = if _needs_tangent(α) + Ap = twist(A, inds; inv, copy = true) + add!(C, Ap, α, β) + Ap + else + TK.twist!(C, A, inds, α, β, ba...; inv) + nothing + end + + function pullback(::NoRData) + copy!(C, C_cache) + + TK.twist_pb(ΔA, A, ΔC, C, inds, α, ba...; inv) + ΔAr = NoRData() + + Δαr = length(args) >= 1 ? pullback_dα(α, ΔC, Ap) : nothing + Δβr = length(args) >= 2 ? pullback_dβ(β, ΔC, C) : nothing + ΔCr = pullback_dC!(ΔC, β) # this typically returns NoRData() + + Δargsr = (Δαr, Δβr, map(Returns(NoRData()), ba)...)[1:length(args)] + return ΔCr, ΔAr, NoRData(), Δargsr... + end + + return C_ΔC, pullback +end + # both are needed for correctly capturing every dispatch @is_primitive DefaultCtx ReverseMode Tuple{typeof(flip), AbstractTensorMap, Any} @is_primitive DefaultCtx ReverseMode Tuple{typeof(Core.kwcall), @NamedTuple{inv::Bool}, typeof(flip), AbstractTensorMap, Any} diff --git a/src/pullbacks/indexmanipulations.jl b/src/pullbacks/indexmanipulations.jl index 809cad7a6..0cfe57c37 100644 --- a/src/pullbacks/indexmanipulations.jl +++ b/src/pullbacks/indexmanipulations.jl @@ -15,6 +15,17 @@ for transform in (:permute, :transpose) return end end +function twist_pb(ΔA, A, ΔC, C, inds, α, ba...; inv::Bool = false) + TC = promote_twist(VectorInterface.promote_scale(ΔC, α), sectortype(ΔC)) + if scalartype(ΔA) <: Real && !(TC <: Real) + ΔAc = similar(ΔC, TC) + twist!(ΔAc, ΔC, inds, conj(α), Zero(), ba...; inv = !inv) + add!(ΔA, real(ΔAc)) + else + twist!(ΔA, ΔC, inds, conj(α), One(), ba...; inv = !inv) + end + return +end function braid_pb(ΔA, A, ΔC, C, p, levels, α, ba...) ip = invperm(linearize(p)) pΔA = TO.repartition(ip, numout(A)) diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index 9771806dc..710680904 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -57,60 +57,60 @@ end Apply a twist to the `i`th index of `t`, or all indices in `inds`, storing the result in `t`. If `inv=true`, use the inverse twist. + twist!(tdst, tsrc, inds, α = 1, β = 0, [backend], [allocator]; inv::Bool = false) -> tdst + +Compute `tdst = β * tdst + α * twist(tsrc, inds; inv)`, writing the result into `tdst`. +Optionally specify a `backend` and `allocator` for the underlying array operation. + See [`twist`](@ref) for creating a new tensor. """ function twist!(t::AbstractTensorMap, inds; inv::Bool = false) - if !all(in(allind(t)), inds) - msg = "Can't twist indices $inds of a tensor with only $(numind(t)) indices." - throw(ArgumentError(msg)) - end + spacecheck_transform(twist, space(t), space(t), inds) (scalartype(t) <: Real && !(sectorscalartype(sectortype(t)) <: Real)) && throw(ArgumentError("Can't in-place twist a real tensor with complex sector type")) has_shared_twist(t, inds) && return t - return _twist!(t, t, inds, inv) + return @inbounds twist!(t, t, inds, One(), Zero(); inv) +end +@propagate_inbounds function twist!( + tdst::AbstractTensorMap, tsrc::AbstractTensorMap, inds, + α::Number = One(), β::Number = Zero(), + backend::AbstractBackend = TO.DefaultBackend(), allocator = TO.DefaultAllocator(); + inv::Bool = false + ) + @boundscheck spacecheck_transform(twist, space(tdst), space(tsrc), inds) + @timeit_debug GLOBAL_TIMER "twist!" begin + twistmask = ntuple(i -> i in inds, numind(tsrc)) + p = (codomainind(tsrc), domainind(tsrc)) + tdst′, tsrc′, p′, twistmask′, conjsrc′, α′, β′ = unwrap_adjoints(tdst, tsrc, p, twistmask, false, α, β) + # like `α`, the twist is conjugated by an adjoint destination + inv′ = inv ⊻ (tdst isa AdjointTensorMap) + @inbounds _twist!(tdst′, tsrc′, p′, conjsrc′, twistmask′, inv′, α′, β′, backend, allocator) + end + return tdst end """ - twist(tsrc::AbstractTensorMap, i::Int; inv::Bool = false, copy::Bool = false) -> tdst - twist(tsrc::AbstractTensorMap, inds; inv::Bool = false, copy::Bool = false) -> tdst + twist( + tsrc::AbstractTensorMap, inds; inv::Bool = false, copy::Bool = false, + backend = DefaultBackend(), allocator = DefaultAllocator() + ) -> tdst -Apply a twist to the `i`th index of `tsrc` and return the result as a new tensor. +Apply a twist to the `i`th index of `tsrc`, or all indices in `inds`, and return the result as a new tensor. If `inv = true`, use the inverse twist. If `copy = false`, `tdst` might share data with `tsrc` whenever possible. Otherwise, a copy is always made. +Optionally specify a `backend` and `allocator` for the underlying array operation. See [`twist!`](@ref) for storing the result in place. """ -function twist(t::AbstractTensorMap, inds; inv::Bool = false, copy::Bool = false) - if has_shared_twist(t, inds) - return copy ? Base.copy(t) : t - end - tdst = similar(t, promote_twist(t)) - return _twist!(tdst, t, inds, inv) -end - -const TensorMapOrAdjoint = Union{TensorMap, AdjointTensorMap{<:Any, <:Any, <:Any, <:Any, <:TensorMap}} +function twist( + t::AbstractTensorMap, inds; inv::Bool = false, + copy::Bool = false, backend::AbstractBackend = TO.DefaultBackend(), allocator = TO.DefaultAllocator() + ) + has_shared_twist(t, inds) && return copy ? Base.copy(t) : t -# `tdst = twist(tsrc, inds; inv)`, where `tdst` and `tsrc` are allowed to be the same tensor -function _twist!(tdst::TensorMapOrAdjoint, tsrc::TensorMapOrAdjoint, inds, inv::Bool) - twistmask = ntuple(i -> i in inds, numind(tdst)) - # the twist acts on the legs of `tdst`, which is the inverse twist of the legs of its parent - if tdst isa AdjointTensorMap - tp = parent(tdst) - twistmask = TupleTools.getindices(twistmask, adjointtensorindices(tp, allind(tp))) - inv = !inv - end - p = (codomainind(tsrc), domainind(tsrc)) - tdst′, tsrc′, p′, _, conjsrc, α, β = unwrap_adjoints(tdst, tsrc, p, nothing, false, One(), Zero()) - transformer = treetwister(storagetype(tdst′), space(tsrc′), conjsrc, twistmask, inv) - add_transform!(tdst′, tsrc′, p′, conjsrc, transformer, α, β, TO.DefaultBackend(), TO.DefaultAllocator()) - return tdst -end -function _twist!(tdst::AbstractTensorMap, tsrc::AbstractTensorMap, inds, inv::Bool) - for (f₁, f₂) in fusiontrees(tsrc) - _, θ = twist((f₁, f₂), inds; inv) - scale!(tdst[f₁, f₂], tsrc[f₁, f₂], θ) - end - return tdst + # general case + tdst = similar(t, promote_twist(t)) + return @inbounds twist!(tdst, t, inds, One(), Zero(), backend, allocator; inv) end # ========================= @@ -543,6 +543,22 @@ end Vsrc′, p′ = conjsrc ? (Vsrc', adjointtensorindices(Vsrc, p)) : (Vsrc, p) return spacecheck_transform(f, Vdst, Vsrc′, p′) end +@noinline function spacecheck_transform(::typeof(twist), Vdst::TensorMapSpace, Vsrc::TensorMapSpace, inds) + check_spacetype(Vdst, Vsrc) + all(in(allind(Vsrc)), inds) || + throw(ArgumentError(lazy"Can't twist indices $inds of a tensor with only $(numind(Vsrc)) indices.")) + Vdst == Vsrc || + throw( + SpaceMismatch( + lazy""" + incompatible spaces for `twist(Vsrc, $inds) -> Vdst` + Vsrc = $Vsrc + Vdst = $Vdst + """ + ) + ) + return nothing +end @noinline function spacecheck_transform(::typeof(braid), Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, levels::IndexTuple) check_spacetype(Vdst, Vsrc) braid(Vsrc, p, levels) == Vdst || @@ -579,9 +595,9 @@ Base.@deprecate( # with `p` indexing the legs of `tsrc`. Adjoint wrappers are absorbed into `conjsrc`, `p`, `levels` # and the scalars, such that the actual work always happens on the subblocks of the parent tensors. -# levels are attached to the legs: relabel them from the legs of `t'` to those of `t` +# levels (or a twist mask) are attached to the legs: relabel them from the legs of `t'` to those of `t` _adjoint_levels(t, ::Nothing) = nothing -_adjoint_levels(t, levels::IndexTuple) = TupleTools.getindices(levels, adjointtensorindices(t, allind(t))) +_adjoint_levels(t, levels::Tuple) = TupleTools.getindices(levels, adjointtensorindices(t, allind(t))) _unwrap_source(tsrc::AbstractTensorMap, p, levels, conjsrc::Bool) = (tsrc, p, levels, conjsrc) function _unwrap_source(tsrc::AdjointTensorMap, p, levels, conjsrc::Bool) @@ -639,6 +655,16 @@ end return add_transform!(tdst, tsrc, p, conjsrc, transformer, α, β, backend, allocator) end +# counterpart of `_braid!` for `twist!`, where `p` only relabels the legs of a conjugated source +@propagate_inbounds function _twist!( + tdst, tsrc, p::Index2Tuple, conjsrc::Bool, twistmask::Tuple{Vararg{Bool}}, inv::Bool, α, β, backend, allocator + ) + Vdst, Vsrc = space(tdst), space(tsrc) + @boundscheck spacecheck_transform(permute, Vdst, Vsrc, p, conjsrc) + transformer = treetwister(storagetype(tdst), Vdst, Vsrc, p, conjsrc, twistmask, inv) + return add_transform!(tdst, tsrc, p, conjsrc, transformer, α, β, backend, allocator) +end + # kernel for computing `tdst = β * tdst + α * permutedims(conjsrc ? conj(tsrc) : tsrc, p)` # with fusion tree transformation encoded in `transformer`. function add_transform!( diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index e2c46304e..f0fcafe61 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -202,20 +202,17 @@ end return TreeTransformer(A, fusiontreetransform, p, Vdst, Vsrc, conjsrc) end -# a twist maps every tree pair onto itself, so this is a `UniqueTreeTransformer` for all sectors; -# it twists the legs `twistmask` of `conjsrc ? V' : V`, reading the subblocks of a source with space `V` +# a twist maps every tree pair onto itself, so this is a `UniqueTreeTransformer` for all sectors, +# where `p` only relabels the legs of a conjugated source @cached function treetwister( - A::Type{TA}, V::TensorMapSpace, conjsrc::Bool, twistmask::NTuple{N, Bool}, inv::Bool - )::UniqueTreeTransformer{recoupling_scalartype(A, sectorscalartype(sectortype(V))), N} where {TA, N} - Vdst, p = twistpermutation(V, conjsrc) - fusiontreetwister(f) = twist(f, twistmask; inv) - return UniqueTreeTransformer(A, fusiontreetwister, p, Vdst, V, conjsrc) -end - -# destination space of a twist of `conjsrc ? V' : V`, and the permutation of the legs of `V` it applies -function twistpermutation(V::TensorMapSpace, conjsrc::Bool) - conjsrc || return V, (codomainind(V), domainind(V)) - return V', adjointtensorindices(V', (codomainind(V'), domainind(V'))) + A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool, + twistmask::NTuple{N, Bool}, inv::Bool + )::UniqueTreeTransformer{recoupling_scalartype(A, sectorscalartype(sectortype(Vdst))), N} where {TA, N} + # the twist mask is attached to the legs, so it follows the same relabeling as the levels in `treebraider` + Vsrc′ = conjsrc ? Vsrc' : Vsrc + twistmask′ = conjsrc ? TupleTools.getindices(twistmask, adjointtensorindices(Vsrc′, allind(Vsrc′))) : twistmask + fusiontreetwister(f) = twist(f, twistmask′; inv) + return UniqueTreeTransformer(A, fusiontreetwister, p, Vdst, Vsrc, conjsrc) end # default cachestyle is GlobalLRUCache diff --git a/test/chainrules/linalg.jl b/test/chainrules/linalg.jl index 880412a3c..3e81dc2ad 100644 --- a/test/chainrules/linalg.jl +++ b/test/chainrules/linalg.jl @@ -126,6 +126,7 @@ for V in spacelist test_rrule(transpose, A, ((2, 5, 4), (1, 3)); fkwargs = bakwargs) symmetricbraiding && test_rrule(permute, A, ((1, 3, 2), (5, 4)); fkwargs = bakwargs) + hasbraiding && test_rrule(twist, A, (2, 4); fkwargs = bakwargs) # `repartition` has no rrule of its own: it is differentiated by AD-ing through # its body down to the `transpose` rrule, so it has to be tested end-to-end @@ -139,6 +140,7 @@ for V in spacelist hasbraiding && test_rrule(twist, A, 1) hasbraiding && test_rrule(twist, A, [1, 3]) + hasbraiding && test_rrule(twist, A, [1, 3]; fkwargs = (; inv = true)) hasbraiding && test_rrule(flip, A, 1) hasbraiding && test_rrule(flip, A, [1, 3, 4]) diff --git a/test/cuda/tensors.jl b/test/cuda/tensors.jl index ec69dc496..556dbd2e9 100644 --- a/test/cuda/tensors.jl +++ b/test/cuda/tensors.jl @@ -325,9 +325,10 @@ for V in spacelist @test twist(twist(t, 2), 2; inv = true) ≈ t if !TensorKit.has_shared_twist(t, (1, 3, 5)) twistmask = (true, false, true, false, true) - twister = TensorKit.treetwister(storagetype(t), space(t), false, twistmask, false) + p = (codomainind(t), domainind(t)) + twister = TensorKit.treetwister(storagetype(t), space(t), space(t), p, false, twistmask, false) @test twister isa getglobal(GPUArraysExt, :DeviceUniqueTreeTransformer) - @test twister === TensorKit.treetwister(storagetype(t), space(t), false, twistmask, false) + @test twister === TensorKit.treetwister(storagetype(t), space(t), space(t), p, false, twistmask, false) end end symmetricbraiding && @timedtestset "Full trace: test self-consistency" begin diff --git a/test/enzyme/flip.jl b/test/enzyme/flip.jl index 66efa44f4..39f3b3756 100644 --- a/test/enzyme/flip.jl +++ b/test/enzyme/flip.jl @@ -29,6 +29,15 @@ is_ci = get(ENV, "CI", "false") == "true" EnzymeTestUtils.test_forward(twist!, TA, (A, TA), ([1, 3], Const); atol, rtol) end end + α = randn(T) + β = randn(T) + C = randn!(similar(A, TensorKit.promote_twist(A))) + EnzymeTestUtils.test_reverse(twist!, TA, (copy(C), TA), (A, TA), ((1, 3), Const), (α, Active), (β, Active); atol, rtol, fkwargs = (inv = true,)) + EnzymeTestUtils.test_forward(twist!, TA, (copy(C), TA), (A, TA), ((1, 3), Const), (α, Duplicated), (β, Duplicated); atol, rtol, fkwargs = (inv = false,)) + if !is_ci + EnzymeTestUtils.test_reverse(twist!, TA, (copy(C), TA), (A, TA), ([2, 4], Const), (α, Const), (β, Const); atol, rtol) + EnzymeTestUtils.test_forward(twist!, TA, (copy(C), TA), (A, TA), ([2, 4], Const), (α, Const), (β, Const); atol, rtol) + end EnzymeTestUtils.test_reverse(flip, TA, (A, TA), (1, Const); atol, rtol, fkwargs = (inv = false,)) EnzymeTestUtils.test_reverse(flip, TA, (A, TA), ([1, 3], Const); atol, rtol, fkwargs = (inv = true,)) if !is_ci diff --git a/test/mooncake/indexmanipulations.jl b/test/mooncake/indexmanipulations.jl index 390721d71..a7a637dac 100644 --- a/test/mooncake/indexmanipulations.jl +++ b/test/mooncake/indexmanipulations.jl @@ -82,6 +82,18 @@ eltypes = (Float64, ComplexF64) Mooncake.TestUtils.test_rule(rng, twist!, A, [1, 3]; atol, rtol, mode) end + α = randn(T) + β = randn(T) + C = randn!(similar(A, TensorKit.promote_twist(A))) + Mooncake.TestUtils.test_rule(rng, TensorKit.twist!, C, A, (1, 3), α, β; atol, rtol, mode) + Mooncake.TestUtils.test_rule(rng, Core.kwcall, (; inv = true), TensorKit.twist!, C, A, [2, 4], α, β; atol, rtol, mode) + Mooncake.TestUtils.test_rule(rng, Core.kwcall, (; inv = false), TensorKit.twist!, C, A, 1; atol, rtol, mode) + if !(T <: Real) + Mooncake.TestUtils.test_rule(rng, TensorKit.twist!, C, real(A), (1, 3), α, β; atol, rtol, mode) + Mooncake.TestUtils.test_rule(rng, TensorKit.twist!, C, A, (1, 3), real(α), real(β); atol, rtol, mode) + end + Mooncake.TestUtils.test_rule(rng, Core.kwcall, (; inv = true), twist, A, (1, 3); atol, rtol, mode, is_primitive = false) + Mooncake.TestUtils.test_rule(rng, Core.kwcall, (; inv = false), flip, A, 1; atol, rtol, mode) Mooncake.TestUtils.test_rule(rng, Core.kwcall, (; inv = true), flip, A, [1, 3]; atol, rtol, mode) Mooncake.TestUtils.test_rule(rng, flip, A, 1; atol, rtol, mode) diff --git a/test/tensors/indexmanipulations.jl b/test/tensors/indexmanipulations.jl index cfc576457..c4059e00d 100644 --- a/test/tensors/indexmanipulations.jl +++ b/test/tensors/indexmanipulations.jl @@ -188,9 +188,20 @@ for V in spacelist @test twist!(t2', inds; inv) ≈ ref @test t2 ≈ ref' # adjoint destination - @test TensorKit._twist!(similar(t)', tref, inds, inv) ≈ ref + @test twist!(similar(t)', tref, inds; inv) ≈ ref @test twist(twist(t', inds), inds; inv = true) ≈ tref end + # accumulation into the destination + α = T <: Complex ? T(2.1 + 0.3im) : T(2.1) + β = T <: Complex ? T(-0.7im) : T(-0.7) + tdst = rand(T, space(tref)) + ref = β * tdst + α * twist(tref, (2, 3)) + @test @constinferred(twist!(copy(tdst), tref, (2, 3), α, β)) ≈ ref + @test twist!(copy(tdst), t', (2, 3), α, β) ≈ ref + D = copy(tdst') + twist!(D', t', (2, 3), α, β) + @test D' ≈ ref + @test twist!(copy(tdst), tref, (2, 3), α, β; inv = true) ≈ β * tdst + α * twist(tref, (2, 3); inv = true) end end symmetricbraiding && @timedtestset "Permutations: BraidingTensor source" begin From 378900a3bbaa57ca646d1dc09104c4541fe222a4 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Wed, 28 Oct 2026 18:45:24 -0400 Subject: [PATCH 4/5] Fix `spacecheck_transform` ambiguity and Enzyme crash of in-place `twist!` on Julia 1.10 Co-Authored-By: Claude Opus 5.5 --- ext/TensorKitEnzymeExt/indexmanipulations.jl | 9 +++++---- src/tensors/indexmanipulations.jl | 3 +++ 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/ext/TensorKitEnzymeExt/indexmanipulations.jl b/ext/TensorKitEnzymeExt/indexmanipulations.jl index 9908ede18..a6384333a 100644 --- a/ext/TensorKitEnzymeExt/indexmanipulations.jl +++ b/ext/TensorKitEnzymeExt/indexmanipulations.jl @@ -278,6 +278,7 @@ function EnzymeRules.forward( end # in-place twist: the source and destination are the same tensor +# (calling the wrapper `twist!(t, inds)` here crashes Enzyme on Julia 1.10) function EnzymeRules.augmented_primal( config::EnzymeRules.RevConfigWidth{1}, func::Const{typeof(TK.twist!)}, @@ -286,7 +287,7 @@ function EnzymeRules.augmented_primal( inds::Const; inv::Bool = false ) where {RT} - TK.twist!(t.val, inds.val; inv) + TK.twist!(t.val, t.val, inds.val; inv) primal = EnzymeRules.needs_primal(config) ? t.val : nothing shadow = EnzymeRules.needs_shadow(config) ? t.dval : nothing return EnzymeRules.AugmentedReturn(primal, shadow, nothing) @@ -300,7 +301,7 @@ function EnzymeRules.reverse( inds::Const; inv::Bool = false ) where {RT} - !isa(t, Const) && TK.twist!(t.dval, inds.val; inv = !inv) + !isa(t, Const) && TK.twist!(t.dval, t.dval, inds.val; inv = !inv) return (nothing, nothing) end function EnzymeRules.forward( @@ -311,8 +312,8 @@ function EnzymeRules.forward( inds::Annotation; inv::Bool = false ) where {RT} - TK.twist!(t.val, inds.val; inv) - !isa(t, Const) && TK.twist!(t.dval, inds.val; inv) + TK.twist!(t.val, t.val, inds.val; inv) + !isa(t, Const) && TK.twist!(t.dval, t.dval, inds.val; inv) if EnzymeRules.needs_primal(config) && EnzymeRules.needs_shadow(config) return t elseif EnzymeRules.needs_primal(config) diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index 710680904..227655efc 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -559,6 +559,9 @@ end ) return nothing end +# disambiguation with the generic method +spacecheck_transform(f::typeof(twist), Vdst::TensorMapSpace, Vsrc::TensorMapSpace, inds::Index2Tuple) = + invoke(spacecheck_transform, Tuple{typeof(twist), TensorMapSpace, TensorMapSpace, Any}, f, Vdst, Vsrc, inds) @noinline function spacecheck_transform(::typeof(braid), Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, levels::IndexTuple) check_spacetype(Vdst, Vsrc) braid(Vsrc, p, levels) == Vdst || From 7584f98bc7fd8664fb1eb1df1613b8fbb30e13b5 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 29 Oct 2026 15:46:59 -0400 Subject: [PATCH 5/5] Return `t` from the in-place `twist!` wrapper to avoid an Enzyme crash on Julia 1.10 Enzyme aborts in `fixup_callconv!` when compiling the wrapper while it returned the result of the three-argument `twist!`, which also broke the fermionic contraction rules. The in-place Enzyme rules call the wrapper again. Co-Authored-By: Claude Opus 5.5 --- ext/TensorKitEnzymeExt/indexmanipulations.jl | 9 ++++----- src/tensors/indexmanipulations.jl | 4 +++- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/ext/TensorKitEnzymeExt/indexmanipulations.jl b/ext/TensorKitEnzymeExt/indexmanipulations.jl index a6384333a..9908ede18 100644 --- a/ext/TensorKitEnzymeExt/indexmanipulations.jl +++ b/ext/TensorKitEnzymeExt/indexmanipulations.jl @@ -278,7 +278,6 @@ function EnzymeRules.forward( end # in-place twist: the source and destination are the same tensor -# (calling the wrapper `twist!(t, inds)` here crashes Enzyme on Julia 1.10) function EnzymeRules.augmented_primal( config::EnzymeRules.RevConfigWidth{1}, func::Const{typeof(TK.twist!)}, @@ -287,7 +286,7 @@ function EnzymeRules.augmented_primal( inds::Const; inv::Bool = false ) where {RT} - TK.twist!(t.val, t.val, inds.val; inv) + TK.twist!(t.val, inds.val; inv) primal = EnzymeRules.needs_primal(config) ? t.val : nothing shadow = EnzymeRules.needs_shadow(config) ? t.dval : nothing return EnzymeRules.AugmentedReturn(primal, shadow, nothing) @@ -301,7 +300,7 @@ function EnzymeRules.reverse( inds::Const; inv::Bool = false ) where {RT} - !isa(t, Const) && TK.twist!(t.dval, t.dval, inds.val; inv = !inv) + !isa(t, Const) && TK.twist!(t.dval, inds.val; inv = !inv) return (nothing, nothing) end function EnzymeRules.forward( @@ -312,8 +311,8 @@ function EnzymeRules.forward( inds::Annotation; inv::Bool = false ) where {RT} - TK.twist!(t.val, t.val, inds.val; inv) - !isa(t, Const) && TK.twist!(t.dval, t.dval, inds.val; inv) + TK.twist!(t.val, inds.val; inv) + !isa(t, Const) && TK.twist!(t.dval, inds.val; inv) if EnzymeRules.needs_primal(config) && EnzymeRules.needs_shadow(config) return t elseif EnzymeRules.needs_primal(config) diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index 227655efc..78a7d82fe 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -69,7 +69,9 @@ function twist!(t::AbstractTensorMap, inds; inv::Bool = false) (scalartype(t) <: Real && !(sectorscalartype(sectortype(t)) <: Real)) && throw(ArgumentError("Can't in-place twist a real tensor with complex sector type")) has_shared_twist(t, inds) && return t - return @inbounds twist!(t, t, inds, One(), Zero(); inv) + # returning `t` rather than the result of `twist!` avoids an Enzyme crash on Julia 1.10 + @inbounds twist!(t, t, inds, One(), Zero(); inv) + return t end @propagate_inbounds function twist!( tdst::AbstractTensorMap, tsrc::AbstractTensorMap, inds,