diff --git a/docs/src/Changelog.md b/docs/src/Changelog.md index a6ab14c3b..7cc5da908 100644 --- a/docs/src/Changelog.md +++ b/docs/src/Changelog.md @@ -23,6 +23,8 @@ 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 +- `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 @@ -44,6 +46,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` 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/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/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/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..e8018dfc5 100644 --- a/ext/TensorKitGPUArraysExt.jl +++ b/ext/TensorKitGPUArraysExt.jl @@ -287,6 +287,14 @@ end return device_transformer(A, host, linearize(p)) 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__() @@ -295,12 +303,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/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/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/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 f6623c197..78a7d82fe 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -57,43 +57,62 @@ 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 - - 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 + # 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, + α::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 +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 + + # general case tdst = similar(t, promote_twist(t)) - copy!(tdst, t) - return twist!(tdst, inds; inv) + return @inbounds twist!(tdst, t, inds, One(), Zero(), backend, allocator; inv) end # ========================= @@ -526,6 +545,25 @@ 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 +# 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 || @@ -562,9 +600,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) @@ -605,6 +643,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,28 +655,37 @@ 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 +# 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!( 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..f0fcafe61 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -202,4 +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, +# where `p` only relabels the legs of a conjugated source +@cached function treetwister( + 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 59196d52c..556dbd2e9 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,28 @@ 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) + 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) + 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), space(t), p, false, 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))) 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 0cc276c6d..c4059e00d 100644 --- a/test/tensors/indexmanipulations.jl +++ b/test/tensors/indexmanipulations.jl @@ -173,6 +173,37 @@ 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 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 τ = BraidingTensor(V1, V2) p = ((2, 1), (3, 4))