From 0dbfa3b0d005a40b6953943bd51a970451cf5a42 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Wed, 28 Oct 2026 08:42:36 -0400 Subject: [PATCH 1/3] Fold the twist of fermionic contraction operands into their permutation The operand that `blas_contract!` twists for fermionic sectors is always copied first, so the twist is now folded into the recoupling coefficients of that copy through a cached `treepermuter`, instead of a separate `twist!` pass over the copied operand (#595). Co-Authored-By: Claude Opus 5.5 --- docs/src/Changelog.md | 1 + ext/TensorKitEnzymeExt/utility.jl | 1 + ext/TensorKitGPUArraysExt.jl | 10 ++++++++++ src/auxiliary/caches.jl | 2 +- src/tensors/tensoroperations.jl | 27 +++++++++++++++++++++++---- src/tensors/treetransformers.jl | 14 ++++++++++++++ 6 files changed, 50 insertions(+), 5 deletions(-) diff --git a/docs/src/Changelog.md b/docs/src/Changelog.md index aa3d1b297..602d2fff4 100644 --- a/docs/src/Changelog.md +++ b/docs/src/Changelog.md @@ -39,6 +39,7 @@ 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)) +- Contractions of fermionic tensors fold the twist of one of the operands into the permutation that copies it, instead of twisting the copy in a separate pass ([#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 diff --git a/ext/TensorKitEnzymeExt/utility.jl b/ext/TensorKitEnzymeExt/utility.jl index 48a2bff02..7d506fd42 100644 --- a/ext/TensorKitEnzymeExt/utility.jl +++ b/ext/TensorKitEnzymeExt/utility.jl @@ -152,6 +152,7 @@ end @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.treepermuter), ::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 5b8e43a97..53557aacd 100644 --- a/ext/TensorKitGPUArraysExt.jl +++ b/ext/TensorKitGPUArraysExt.jl @@ -387,6 +387,14 @@ end return device_transformer(A, host, linearize(p)) end +@cached function TensorKit.treepermuter( + A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, + conjsrc::Bool, twistmask::NTuple{N, Bool} + ) where {TA <: AnyGPUVector, N} + host = TensorKit.treepermuter(Vector{scalartype(A)}, Vdst, Vsrc, p, conjsrc, twistmask) + return device_transformer(A, host, linearize(p)) +end + @cached function TensorKit.treetracer( A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, q::Index2Tuple ) where {TA <: AnyGPUVector} @@ -410,6 +418,7 @@ function __init__() TensorKit.GLOBAL_CACHES, ( Symbol("TensorKitGPUArraysExt.GLOBAL_TREEBRAIDER_CACHE") => GLOBAL_TREEBRAIDER_CACHE, Symbol("TensorKitGPUArraysExt.GLOBAL_TREETRANSPOSER_CACHE") => GLOBAL_TREETRANSPOSER_CACHE, + Symbol("TensorKitGPUArraysExt.GLOBAL_TREEPERMUTER_CACHE") => GLOBAL_TREEPERMUTER_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, @@ -421,6 +430,7 @@ function __init__() # 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_TREEPERMUTER_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) diff --git a/src/auxiliary/caches.jl b/src/auxiliary/caches.jl index 5525f1714..353b85edc 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, :treetwister) ? + return fname in (:fsbraid, :fstranspose, :treebraider, :treetransposer, :treepermuter, :treetwister) ? "symmetry" : "bookkeeping" end diff --git a/src/tensors/tensoroperations.jl b/src/tensors/tensoroperations.jl index e11a1edf0..f5dc297c1 100644 --- a/src/tensors/tensoroperations.jl +++ b/src/tensors/tensoroperations.jl @@ -390,6 +390,19 @@ function TO.isblasdestination(A::AbstractTensorMap, ipAB::Index2Tuple) return scalartype(A) <: LinearAlgebra.BlasFloat && has_shared_permute(A, ipAB) end +# `tdst = twist(permute(tsrc, p), inds)`, with the twist folded into the coefficients of the permutation +function permute_twist!(tdst::TensorMap, tsrc::AbstractTensorMap, p::Index2Tuple, inds, backend, allocator) + tsrc′, p′, _, conjsrc = _unwrap_source(tsrc, p, nothing, false) + # other tensor types, such as diagonal GPU tensors, rely on the dispatch of `tensoradd!` + if !(tsrc′ isa TensorMap) + TO.tensoradd!(tdst, tsrc, p, false, One(), Zero(), backend, allocator) + return twist!(tdst, inds) + end + twistmask = ntuple(in(inds), numind(tdst)) + transformer = treepermuter(storagetype(tdst), space(tdst), space(tsrc′), p′, conjsrc, twistmask) + return add_transform!(tdst, tsrc′, p′, conjsrc, transformer, One(), Zero(), backend, allocator) +end + function blas_contract!( C::AbstractTensorMap, A::AbstractTensorMap, pA::Index2Tuple, @@ -425,8 +438,11 @@ function blas_contract!( # Bring A in the correct form for BLAS contraction if copyA Anew = @timeit_debug GLOBAL_TIMER "alloc: buffers" TO.tensoralloc_add(TC, A, pA, false, Val(true), allocator) - Anew = TO.tensoradd!(Anew, A, pA, false, One(), Zero(), backend, allocator) - twistA && twist!(Anew, filter(!isdual ∘ Base.Fix1(space, Anew), domainind(Anew))) + if twistA + permute_twist!(Anew, A, pA, filter(!isdual ∘ Base.Fix1(space, Anew), domainind(Anew)), backend, allocator) + else + Anew = TO.tensoradd!(Anew, A, pA, false, One(), Zero(), backend, allocator) + end else Anew = permute(A, pA) end @@ -435,8 +451,11 @@ function blas_contract!( # Bring B in the correct form for BLAS contraction if copyB Bnew = @timeit_debug GLOBAL_TIMER "alloc: buffers" TO.tensoralloc_add(TC, B, pB, false, Val(true), allocator) - Bnew = TO.tensoradd!(Bnew, B, pB, false, One(), Zero(), backend, allocator) - twistB && twist!(Bnew, filter(isdual ∘ Base.Fix1(space, Bnew), codomainind(Bnew))) + if twistB + permute_twist!(Bnew, B, pB, filter(isdual ∘ Base.Fix1(space, Bnew), codomainind(Bnew)), backend, allocator) + else + Bnew = TO.tensoradd!(Bnew, B, pB, false, One(), Zero(), backend, allocator) + end else Bnew = permute(B, pB) end diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index 2bc2ec4ac..d28c1d3cb 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -202,6 +202,20 @@ end return TreeTransformer(A, fusiontreetransform, p, Vdst, Vsrc, conjsrc) end +# permutation for symmetric braidings followed by a twist of the destination legs in `twistmask`, +# which is diagonal in the fusion tree basis and therefore folded into the recoupling coefficients +@cached function treepermuter( + A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool, + twistmask::NTuple{N, Bool} + )::treetransformertype(A, Vdst, Vsrc) where {TA, N} + p′ = conjsrc ? adjointtensorindices(Vsrc, p) : p + function fusiontreepermuter(f) + f′, U = permute(f, p′) + return f′, U * last(twist(f′, twistmask)) + end + return TreeTransformer(A, fusiontreepermuter, p, Vdst, Vsrc, conjsrc) +end + """ TraceTransformer{T, N₁, N₂} From e16b5eb9deb4694bda92a75a4880ff095a70568b Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 5 Nov 2026 10:33:27 -0500 Subject: [PATCH 2/3] Generalize `treebraider` with a twist mask instead of a separate `treepermuter` `blas_contract!` now computes the twist masks of both operands in `contraction_twistmasks` and always copies through `permute_twist!`, with an all-false mask when no twist is needed. Co-Authored-By: Claude Opus 5.5 --- ext/TensorKitEnzymeExt/utility.jl | 1 - ext/TensorKitGPUArraysExt.jl | 26 +++++--------- src/auxiliary/caches.jl | 2 +- src/tensors/indexmanipulations.jl | 13 ++++--- src/tensors/tensoroperations.jl | 56 +++++++++++-------------------- src/tensors/treetransformers.jl | 27 +++++---------- 6 files changed, 46 insertions(+), 79 deletions(-) diff --git a/ext/TensorKitEnzymeExt/utility.jl b/ext/TensorKitEnzymeExt/utility.jl index 7d506fd42..48a2bff02 100644 --- a/ext/TensorKitEnzymeExt/utility.jl +++ b/ext/TensorKitEnzymeExt/utility.jl @@ -152,7 +152,6 @@ end @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.treepermuter), ::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 53557aacd..78dafc2fd 100644 --- a/ext/TensorKitGPUArraysExt.jl +++ b/ext/TensorKitGPUArraysExt.jl @@ -374,9 +374,9 @@ end # Cache the device transformer directly, constructing the cached host version only on a miss. @cached function TensorKit.treebraider( A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, - conjsrc::Bool, levels::IndexTuple - ) where {TA <: AnyGPUVector} - host = TensorKit.treebraider(Vector{scalartype(A)}, Vdst, Vsrc, p, conjsrc, levels) + conjsrc::Bool, levels::IndexTuple, twistmask::NTuple{N, Bool} + ) where {TA <: AnyGPUVector, N} + host = TensorKit.treebraider(Vector{scalartype(A)}, Vdst, Vsrc, p, conjsrc, levels, twistmask) return device_transformer(A, host, linearize(p)) end @@ -387,14 +387,6 @@ end return device_transformer(A, host, linearize(p)) end -@cached function TensorKit.treepermuter( - A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, - conjsrc::Bool, twistmask::NTuple{N, Bool} - ) where {TA <: AnyGPUVector, N} - host = TensorKit.treepermuter(Vector{scalartype(A)}, Vdst, Vsrc, p, conjsrc, twistmask) - return device_transformer(A, host, linearize(p)) -end - @cached function TensorKit.treetracer( A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, q::Index2Tuple ) where {TA <: AnyGPUVector} @@ -418,7 +410,6 @@ function __init__() TensorKit.GLOBAL_CACHES, ( Symbol("TensorKitGPUArraysExt.GLOBAL_TREEBRAIDER_CACHE") => GLOBAL_TREEBRAIDER_CACHE, Symbol("TensorKitGPUArraysExt.GLOBAL_TREETRANSPOSER_CACHE") => GLOBAL_TREETRANSPOSER_CACHE, - Symbol("TensorKitGPUArraysExt.GLOBAL_TREEPERMUTER_CACHE") => GLOBAL_TREEPERMUTER_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, @@ -430,7 +421,6 @@ function __init__() # 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_TREEPERMUTER_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) @@ -702,9 +692,9 @@ end @cached function diagonal_treebraider( A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, - conjsrc::Bool, levels::IndexTuple, dstdiagonal::Bool - ) where {TA <: AnyGPUVector} - host = TensorKit.treebraider(Vector{scalartype(A)}, Vdst, Vsrc, p, conjsrc, levels) + conjsrc::Bool, levels::IndexTuple, twistmask::NTuple{N, Bool}, dstdiagonal::Bool + ) where {TA <: AnyGPUVector, N} + host = TensorKit.treebraider(Vector{scalartype(A)}, Vdst, Vsrc, p, conjsrc, levels, twistmask) return diagonal_device_transformer(A, host, dstdiagonal) end @@ -737,12 +727,12 @@ end Base.@propagate_inbounds function TensorKit._braid!( tdst::Union{GPUTensorMap, GPUDiagonalTensorMap}, tsrc::GPUDiagonalTensorMap, - p::Index2Tuple, conjsrc::Bool, levels::IndexTuple, α, β, backend, allocator + p::Index2Tuple, conjsrc::Bool, levels::IndexTuple, twistmask::Tuple{Vararg{Bool}}, α, β, backend, allocator ) Vdst, Vsrc = space(tdst), space(tsrc) @boundscheck TensorKit.spacecheck_transform(permute, Vdst, Vsrc, p, conjsrc) transformer = diagonal_treebraider( - storagetype(tdst), Vdst, Vsrc, p, conjsrc, levels, tdst isa DiagonalTensorMap + storagetype(tdst), Vdst, Vsrc, p, conjsrc, levels, twistmask, tdst isa DiagonalTensorMap ) return _add_diagonal!(tdst, tsrc, conjsrc, transformer, α, β) end diff --git a/src/auxiliary/caches.jl b/src/auxiliary/caches.jl index 353b85edc..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, :treepermuter, :treetwister) ? + return fname in (:fsbraid, :fstranspose, :treebraider, :treetransposer, :treetwister) ? "symmetry" : "bookkeeping" end diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index 78a7d82fe..1b768de52 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -241,7 +241,7 @@ See also [`permute`](@ref) for creating a new tensor. @boundscheck spacecheck_transform(permute, space(tdst), space(tsrc), p) @timeit_debug GLOBAL_TIMER "permute!/braid!" begin tdst′, tsrc′, p′, _, conjsrc′, α′, β′ = unwrap_adjoints(tdst, tsrc, p, nothing, false, α, β) - @inbounds _braid!(tdst′, tsrc′, p′, conjsrc′, allind(tsrc′), α′, β′, backend, allocator) + @inbounds _braid!(tdst′, tsrc′, p′, conjsrc′, allind(tsrc′), _notwist(tdst′), α′, β′, backend, allocator) end return tdst end @@ -331,7 +331,7 @@ See also [`braid`](@ref) for creating a new tensor. @boundscheck spacecheck_transform(braid, space(tdst), space(tsrc), p, levels) @timeit_debug GLOBAL_TIMER "permute!/braid!" begin tdst′, tsrc′, p′, levels′, conjsrc′, α′, β′ = unwrap_adjoints(tdst, tsrc, p, levels, false, α, β) - @inbounds _braid!(tdst′, tsrc′, p′, conjsrc′, levels′, α′, β′, backend, allocator) + @inbounds _braid!(tdst′, tsrc′, p′, conjsrc′, levels′, _notwist(tdst′), α′, β′, backend, allocator) end return tdst end @@ -637,14 +637,17 @@ function _dense_transform!(tdst, tsrc, p::Index2Tuple, conjsrc::Bool, α, β, ba return tdst end +_notwist(t) = ntuple(Returns(false), numind(t)) + +# `twistmask` selects the destination legs that are twisted after braiding @propagate_inbounds function _braid!( - tdst, tsrc, p::Index2Tuple, conjsrc::Bool, levels::IndexTuple, α, β, backend, allocator + tdst, tsrc, p::Index2Tuple, conjsrc::Bool, levels::IndexTuple, twistmask::Tuple{Vararg{Bool}}, α, β, backend, allocator ) 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) + !any(twistmask) && _is_trivial_transform(tsrc, p, conjsrc) && return add!(tdst, tsrc, α, β) + transformer = treebraider(storagetype(tdst), Vdst, Vsrc, p, conjsrc, levels, twistmask) return add_transform!(tdst, tsrc, p, conjsrc, transformer, α, β, backend, allocator) end diff --git a/src/tensors/tensoroperations.jl b/src/tensors/tensoroperations.jl index f5dc297c1..fd6d5c23c 100644 --- a/src/tensors/tensoroperations.jl +++ b/src/tensors/tensoroperations.jl @@ -48,7 +48,7 @@ function TO.tensoradd!( backend, allocator ) tdst, tsrc, p, _, conjA′, α′, β′ = unwrap_adjoints(C, A, _canonicalize(pA, C), nothing, conjA, α, β) - _braid!(tdst, tsrc, p, conjA′, allind(tsrc), α′, β′, backend, allocator) + _braid!(tdst, tsrc, p, conjA′, allind(tsrc), _notwist(tdst), α′, β′, backend, allocator) return C end @@ -390,17 +390,21 @@ function TO.isblasdestination(A::AbstractTensorMap, ipAB::Index2Tuple) return scalartype(A) <: LinearAlgebra.BlasFloat && has_shared_permute(A, ipAB) end -# `tdst = twist(permute(tsrc, p), inds)`, with the twist folded into the coefficients of the permutation -function permute_twist!(tdst::TensorMap, tsrc::AbstractTensorMap, p::Index2Tuple, inds, backend, allocator) +# `tdst = twist(permute(tsrc, p), twistmask)`, with the twist folded into the coefficients of the permutation +function permute_twist!(tdst::AbstractTensorMap, tsrc::AbstractTensorMap, p::Index2Tuple, twistmask, backend, allocator) tsrc′, p′, _, conjsrc = _unwrap_source(tsrc, p, nothing, false) - # other tensor types, such as diagonal GPU tensors, rely on the dispatch of `tensoradd!` - if !(tsrc′ isa TensorMap) - TO.tensoradd!(tdst, tsrc, p, false, One(), Zero(), backend, allocator) - return twist!(tdst, inds) - end - twistmask = ntuple(in(inds), numind(tdst)) - transformer = treepermuter(storagetype(tdst), space(tdst), space(tsrc′), p′, conjsrc, twistmask) - return add_transform!(tdst, tsrc′, p′, conjsrc, transformer, One(), Zero(), backend, allocator) + return _braid!(tdst, tsrc′, p′, conjsrc, allind(tsrc′), twistmask, One(), Zero(), backend, allocator) +end + +# fermionic contractions need a twist of the contracted legs of either `A` or `B`: +# twist the one that is copied anyway, or the smallest if neither or both are copied +function contraction_twistmasks(A, pA, B, pB, copyA, copyB) + twists = map(isdual ∘ Base.Fix1(space, B), pB[1]) + notwistA, notwistB = map(Returns(false), linearize(pA)), map(Returns(false), linearize(pB)) + (BraidingStyle(sectortype(B)) isa Fermionic && any(twists)) || return notwistA, notwistB + twistA = copyA ⊻ copyB ? copyA : dim(A) < dim(B) + return twistA ? ((map(Returns(false), pA[1])..., twists...), notwistB) : + (notwistA, (twists..., map(Returns(false), pB[2])...)) end function blas_contract!( @@ -419,30 +423,14 @@ function blas_contract!( copyA = !(TO.isblascontractable(A, pA) && scalartype(A) === TC) copyB = !(TO.isblascontractable(B, pB) && scalartype(B) === TC) - if bstyle isa Fermionic && any(isdual ∘ Base.Fix1(space, B), pB[1]) - # twist smallest object if neither or both already have to be permuted - # otherwise twist the one that already is copied - if !(copyA ⊻ copyB) - twistA = dim(A) < dim(B) - else - twistA = copyA - end - twistB = !twistA - copyA |= twistA - copyB |= twistB - else - twistA = false - twistB = false - end + twistmaskA, twistmaskB = contraction_twistmasks(A, pA, B, pB, copyA, copyB) + copyA |= any(twistmaskA) + copyB |= any(twistmaskB) # Bring A in the correct form for BLAS contraction if copyA Anew = @timeit_debug GLOBAL_TIMER "alloc: buffers" TO.tensoralloc_add(TC, A, pA, false, Val(true), allocator) - if twistA - permute_twist!(Anew, A, pA, filter(!isdual ∘ Base.Fix1(space, Anew), domainind(Anew)), backend, allocator) - else - Anew = TO.tensoradd!(Anew, A, pA, false, One(), Zero(), backend, allocator) - end + permute_twist!(Anew, A, pA, twistmaskA, backend, allocator) else Anew = permute(A, pA) end @@ -451,11 +439,7 @@ function blas_contract!( # Bring B in the correct form for BLAS contraction if copyB Bnew = @timeit_debug GLOBAL_TIMER "alloc: buffers" TO.tensoralloc_add(TC, B, pB, false, Val(true), allocator) - if twistB - permute_twist!(Bnew, B, pB, filter(isdual ∘ Base.Fix1(space, Bnew), codomainind(Bnew)), backend, allocator) - else - Bnew = TO.tensoradd!(Bnew, B, pB, false, One(), Zero(), backend, allocator) - end + permute_twist!(Bnew, B, pB, twistmaskB, backend, allocator) else Bnew = permute(B, pB) end diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index d28c1d3cb..d50f0c157 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -182,15 +182,20 @@ function TreeTransformer( GenericTreeTransformer(A, transform, p, Vdst, Vsrc, conjsrc) end -# braid is special because it has levels +# braid is special because it has levels; the destination legs in `twistmask` are twisted afterwards, +# which is diagonal in the fusion tree basis and therefore folded into the recoupling coefficients @cached function treebraider( - A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool, levels::IndexTuple - )::treetransformertype(A, Vdst, Vsrc) where {TA} + A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool, levels::IndexTuple, + twistmask::NTuple{N, Bool} + )::treetransformertype(A, Vdst, Vsrc) where {TA, N} Vsrc′, p′ = conjsrc ? (Vsrc', adjointtensorindices(Vsrc, p)) : (Vsrc, p) # levels are attached to the legs, so they follow the same relabeling as the permutation levels′ = conjsrc ? TupleTools.getindices(levels, adjointtensorindices(Vsrc′, allind(Vsrc′))) : levels levels″ = (TupleTools.getindices(levels′, codomainind(Vsrc′)), TupleTools.getindices(levels′, domainind(Vsrc′))) - fusiontreebraider(f) = braid(f, p′, levels″) + function fusiontreebraider(f) + f′, U = braid(f, p′, levels″) + return any(twistmask) ? (f′, U * last(twist(f′, twistmask))) : (f′, U) + end return TreeTransformer(A, fusiontreebraider, p, Vdst, Vsrc, conjsrc) end @@ -202,20 +207,6 @@ end return TreeTransformer(A, fusiontreetransform, p, Vdst, Vsrc, conjsrc) end -# permutation for symmetric braidings followed by a twist of the destination legs in `twistmask`, -# which is diagonal in the fusion tree basis and therefore folded into the recoupling coefficients -@cached function treepermuter( - A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool, - twistmask::NTuple{N, Bool} - )::treetransformertype(A, Vdst, Vsrc) where {TA, N} - p′ = conjsrc ? adjointtensorindices(Vsrc, p) : p - function fusiontreepermuter(f) - f′, U = permute(f, p′) - return f′, U * last(twist(f′, twistmask)) - end - return TreeTransformer(A, fusiontreepermuter, p, Vdst, Vsrc, conjsrc) -end - """ TraceTransformer{T, N₁, N₂} From 6b51f3ee1ba61270299992513ddbbafce16e4931 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 5 Nov 2026 11:09:58 -0500 Subject: [PATCH 3/3] Use a valid space for the GPU `SectorVector` adapt test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `V1 ⊗ V2 ← V3` has incompatible units for `IsingBimodule`; reuse the space of the `Adapt` testset. Co-Authored-By: Claude Opus 5.5 --- test/amd/tensors.jl | 2 +- test/cuda/tensors.jl | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/test/amd/tensors.jl b/test/amd/tensors.jl index d0c9ca014..453933d72 100644 --- a/test/amd/tensors.jl +++ b/test/amd/tensors.jl @@ -107,7 +107,7 @@ for V in spacelist v_cpu = @constinferred adapt(Array, v_gpu) @test v_cpu isa TensorKit.SectorVector{T, I, Vector{T}} @test parent(v_cpu) == parent(v) - t = AMDGPU.randn(T, V1 ⊗ V2 ← V3) + t = AMDGPU.randn(T, V1 ⊗ V2 ⊗ V3 ← (V4 ⊗ V5)') s = svd_vals(t) s_cpu = adapt(Array, s) @test parent(s_cpu) == collect(parent(s)) diff --git a/test/cuda/tensors.jl b/test/cuda/tensors.jl index 750d7384a..f4c178549 100644 --- a/test/cuda/tensors.jl +++ b/test/cuda/tensors.jl @@ -126,7 +126,7 @@ for V in spacelist v_cpu = @constinferred adapt(Array, v_gpu) @test v_cpu isa TensorKit.SectorVector{T, I, Vector{T}} @test parent(v_cpu) == parent(v) - t = CUDA.randn(T, V1 ⊗ V2 ← V3) + t = CUDA.randn(T, V1 ⊗ V2 ⊗ V3 ← (V4 ⊗ V5)') s = svd_vals(t) s_cpu = adapt(Array, s) @test parent(s_cpu) == collect(parent(s))