Fold the twist of fermionic contraction operands into their permutation

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
 

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
 
@@ -692,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
 
@@ -727,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

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
 

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,6 +390,23 @@ function TO.isblasdestination(A::AbstractTensorMap, ipAB::Index2Tuple)
     return scalartype(A) <: LinearAlgebra.BlasFloat && has_shared_permute(A, ipAB)
 end
 
+# `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)
+    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!(
         C::AbstractTensorMap,
         A::AbstractTensorMap, pA::Index2Tuple,
@@ -406,27 +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)
-        Anew = TO.tensoradd!(Anew, A, pA, false, One(), Zero(), backend, allocator)
-        twistA && twist!(Anew, filter(!isdual ∘ Base.Fix1(space, Anew), domainind(Anew)))
+        permute_twist!(Anew, A, pA, twistmaskA, backend, allocator)
     else
         Anew = permute(A, pA)
     end
@@ -435,8 +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)
-        Bnew = TO.tensoradd!(Bnew, B, pB, false, One(), Zero(), backend, allocator)
-        twistB && twist!(Bnew, filter(isdual ∘ Base.Fix1(space, Bnew), codomainind(Bnew)))
+        permute_twist!(Bnew, B, pB, twistmaskB, backend, allocator)
     else
         Bnew = permute(B, pB)
     end

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
 

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

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