Cached treetwister for twist!

docs/src/Changelog.md

@@ -23,6 +23,8 @@ When releasing a new version, move the "Unreleased" changes to a new version sec
 ### Added
 
 - Enzyme rules for `planartrace!`, with a pullback that also supports adjoint tensors
+- `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
 
@@ -36,6 +38,8 @@ When releasing a new version, move the "Unreleased" changes to a new version sec
 
 ### Performance
 
+- `twist!` and `twist` on a `TensorMap`, a `DiagonalTensorMap` or their adjoints 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, diagonal or not, into a single kernel launch ([#595](https://github.com/QuantumKitHub/TensorKit.jl/issues/595))
+
 ## [0.17.3](https://github.com/QuantumKitHub/TensorKit.jl/compare/v0.17.2...v0.17.3) - 2026-10-30
 
 ### Added

docs/src/lib/fusiontrees.md

@@ -49,4 +49,5 @@ Base.transpose(src::Union{FusionTreePair, FusionTreeBlock}, p::Index2Tuple)
 braid(src::Union{FusionTreePair, FusionTreeBlock}, p::Index2Tuple, levels::Index2Tuple)
 permute(src::Union{FusionTreePair, FusionTreeBlock}, p::Index2Tuple)
 planar_trace(src::Union{FusionTreePair, FusionTreeBlock}, p::Index2Tuple, q::Index2Tuple)
+twist(src::Union{FusionTreePair, FusionTreeBlock}, inds)
 ```

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)},

ext/TensorKitEnzymeExt/utility.jl

@@ -151,6 +151,7 @@ end
 @inline EnzymeRules.inactive(::typeof(TensorKit.treetransposer), ::Any...) = nothing
 @inline EnzymeRules.inactive(::typeof(TensorKit.treebraider), ::Any...) = nothing
 @inline EnzymeRules.inactive(::typeof(TensorKit.treetracer), ::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

ext/TensorKitGPUArraysExt.jl

@@ -394,6 +394,14 @@ end
     return device_transformer(A, host, p, q)
 end
 
+@cached function TensorKit.treetwister(
+        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)}, Vdst, Vsrc, p, conjsrc, twistmask, inv)
+    return device_transformer(A, host, linearize(p))
+end
+
 const DEFAULT_DEVICECACHE_SIZE = 1000
 
 function __init__()
@@ -404,7 +412,9 @@ function __init__()
             Symbol("TensorKitGPUArraysExt.GLOBAL_TREETRANSPOSER_CACHE") => GLOBAL_TREETRANSPOSER_CACHE,
             Symbol("TensorKitGPUArraysExt.GLOBAL_DIAGONAL_TREEBRAIDER_CACHE") => GLOBAL_DIAGONAL_TREEBRAIDER_CACHE,
             Symbol("TensorKitGPUArraysExt.GLOBAL_DIAGONAL_TREETRANSPOSER_CACHE") => GLOBAL_DIAGONAL_TREETRANSPOSER_CACHE,
+            Symbol("TensorKitGPUArraysExt.GLOBAL_DIAGONAL_TREETWISTER_CACHE") => GLOBAL_DIAGONAL_TREETWISTER_CACHE,
             Symbol("TensorKitGPUArraysExt.GLOBAL_TREETRACER_CACHE") => GLOBAL_TREETRACER_CACHE,
+            Symbol("TensorKitGPUArraysExt.GLOBAL_TREETWISTER_CACHE") => GLOBAL_TREETWISTER_CACHE,
         )
     )
 
@@ -413,7 +423,9 @@ function __init__()
     resize!(GLOBAL_TREETRANSPOSER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE)
     resize!(GLOBAL_DIAGONAL_TREEBRAIDER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE)
     resize!(GLOBAL_DIAGONAL_TREETRANSPOSER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE)
+    resize!(GLOBAL_DIAGONAL_TREETWISTER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE)
     resize!(GLOBAL_TREETRACER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE)
+    resize!(GLOBAL_TREETWISTER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE)
 
     return nothing
 end
@@ -694,6 +706,14 @@ end
     return diagonal_device_transformer(A, host, dstdiagonal)
 end
 
+@cached function diagonal_treetwister(
+        A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple,
+        conjsrc::Bool, twistmask::NTuple{N, Bool}, inv::Bool, dstdiagonal::Bool
+    ) where {TA <: AnyGPUVector, N}
+    host = TensorKit.treetwister(Vector{scalartype(A)}, Vdst, Vsrc, p, conjsrc, twistmask, inv)
+    return diagonal_device_transformer(A, host, dstdiagonal)
+end
+
 function _add_diagonal!(tdst::GPUDiagonalTensorMap, tsrc, conjsrc, transformer, α, β)
     _launch_unique!(tdst.data, tsrc.data, conjsrc ? conj : identity, transformer, α, β, Val(1))
     return tdst
@@ -729,6 +749,18 @@ Base.@propagate_inbounds function TensorKit._transpose!(
     return _add_diagonal!(tdst, tsrc, conjsrc, transformer, α, β)
 end
 
+Base.@propagate_inbounds function TensorKit._twist!(
+        tdst::Union{GPUTensorMap, GPUDiagonalTensorMap}, tsrc::GPUDiagonalTensorMap,
+        p::Index2Tuple, conjsrc::Bool, twistmask::Tuple{Vararg{Bool}}, inv::Bool, α, β, backend, allocator
+    )
+    Vdst, Vsrc = space(tdst), space(tsrc)
+    @boundscheck TensorKit.spacecheck_transform(permute, Vdst, Vsrc, p, conjsrc)
+    transformer = diagonal_treetwister(
+        storagetype(tdst), Vdst, Vsrc, p, conjsrc, twistmask, inv, tdst isa DiagonalTensorMap
+    )
+    return _add_diagonal!(tdst, tsrc, conjsrc, transformer, α, β)
+end
+
 function TensorKit.add_transform_kernel!(
         dst::AnyGPUVector, src::AnyGPUVector, p, conjsrc::Bool,
         device::DeviceUniqueTreeTransformer, α, β, backend, allocator, ntasks::Int

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}

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
 

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}

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))