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/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl index 5b8e43a97..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 @@ -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 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 e11a1edf0..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,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 diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index 2bc2ec4ac..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 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))