From 67c42cac6bd56eda9b718508fc61e74c66848ec6 Mon Sep 17 00:00:00 2001 From: leburgel Date: Thu, 27 Aug 2026 16:13:10 +0200 Subject: [PATCH 1/2] Add dedicated cached `treebraider` and `treetransposer` for `AdjointTensorMap` --- src/spaces/homspace.jl | 24 ++++++++++++++ src/tensors/indexmanipulations.jl | 28 +++++++++++----- src/tensors/treetransformers.jl | 55 +++++++++++++++++++++++++++---- 3 files changed, 92 insertions(+), 15 deletions(-) diff --git a/src/spaces/homspace.jl b/src/spaces/homspace.jl index 8460f7578..480d2cfa1 100644 --- a/src/spaces/homspace.jl +++ b/src/spaces/homspace.jl @@ -169,6 +169,30 @@ See also [`degeneracystructure`](@ref), [`blockstructure`](@ref). """ subblockstructure(W::HomSpace) = Dictionary(fusiontrees(W), degeneracystructure(W).subblockstructure) +""" + adjoint_subblockstructure(W::HomSpace) -> Dictionary + +Subblock structure of `W'`, expressed as strides into a buffer laid out for `W`. + +This mirrors the relation used by `subblock(::AdjointTensorMap, ...)`: the adjoint shares its +parent's data, with the tree pair swapped and sizes and strides permuted by +`(domainind..., codomainind...)`. Permuting `adjoint(t)` can therefore read `t`'s own buffer +directly, instead of materialising an `AdjointTensorMap` and falling back to the uncached +tree-transformation path. + +See also [`subblockstructure`](@ref). +""" +function adjoint_subblockstructure(W::HomSpace) + N₁, N₂ = numout(W), numin(W) + swap = (ntuple(i -> N₁ + i, N₂)..., ntuple(identity, N₁)...) + structs = subblockstructure(W) + newkeys = map(((f₁, f₂),) -> (f₂, f₁), collect(keys(structs))) + newvals = map(collect(values(structs))) do (sz, str, off) + return (TupleTools.getindices(sz, swap), TupleTools.getindices(str, swap), off) + end + return Dictionary(newkeys, newvals) +end + """ fusionblocks(W::HomSpace) diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index ce4add069..ead110183 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -570,6 +570,17 @@ Base.@deprecate( scheduler = ntasks == 1 ? SerialScheduler() : DynamicScheduler(; ntasks, split = :roundrobin) if tdst isa TensorMap && tsrc isa TensorMap # unpack data fields to avoid specializing add_transform_kernel!(tdst.data, tsrc.data, p, transformer, α, β, backend, allocator, scheduler) + elseif tdst isa TensorMap && tsrc isa AdjointTensorMap && + parent(tsrc) isa TensorMap && + transformer isa Union{AbelianTreeTransformer, GenericTreeTransformer} + # An adjoint shares its parent's buffer. `conj_treebraider` builds a + # transformer whose source strides already address that buffer, so the data + # can be read directly and conjugated on the fly, instead of falling back to + # the uncached per-fusion-block path. + add_transform_kernel!( + tdst.data, parent(tsrc).data, p, transformer, α, β, backend, allocator, + scheduler; conjsrc = true + ) else add_transform_kernel!(tdst, tsrc, p, transformer, α, β, backend, allocator, scheduler) end @@ -587,14 +598,15 @@ function use_threaded_transform(t::AbstractTensorMap, transformer) end function add_transform_kernel!( - tdst, tsrc, p, transformer, α, β, backend, allocator, scheduler + tdst, tsrc, p, transformer, α, β, backend, allocator, scheduler; + conjsrc::Bool = false ) I = sectortype(tdst) if FusionStyle(I) === UniqueFusion() tforeach(fusiontrees(tsrc); scheduler) do (f₁, f₂) (f₁′, f₂′), coeff = transformer((f₁, f₂)) @inbounds TO.tensoradd!( - tdst[f₁′, f₂′], tsrc[f₁, f₂], p, false, α * coeff, β, backend, allocator + tdst[f₁′, f₂′], tsrc[f₁, f₂], p, conjsrc, α * coeff, β, backend, allocator ) end return nothing @@ -629,7 +641,7 @@ function add_transform_kernel!( @inbounds for (i, (f₁, f₂)) in enumerate(fusiontrees(src)) TO.tensoradd!( sreshape(view(buffer_src, :, i), sz_src), tsrc[f₁, f₂], - ptriv, false, One(), Zero(), backend, allocator + ptriv, conjsrc, One(), Zero(), backend, allocator ) end @@ -657,19 +669,19 @@ end # repeated specialization -- this only depends on `numind` and `eltype`. function add_transform_kernel!( data_dst::DenseVector, data_src::DenseVector, p, transformer::AbelianTreeTransformer, - α, β, backend, allocator, scheduler + α, β, backend, allocator, scheduler; conjsrc::Bool = false ) tforeach(transformer.data; scheduler) do (coeff, struct_dst, struct_src) TO.tensoradd!( StridedView(data_dst, struct_dst...), StridedView(data_src, struct_src...), - p, false, α * coeff, β, backend, allocator + p, conjsrc, α * coeff, β, backend, allocator ) end return nothing end function add_transform_kernel!( data_dst::DenseVector, data_src::DenseVector, p, transformer::GenericTreeTransformer, - α, β, backend, allocator, scheduler + α, β, backend, allocator, scheduler; conjsrc::Bool = false ) cp = TO.allocator_checkpoint!(allocator) @@ -686,7 +698,7 @@ function add_transform_kernel!( TO.tensoradd!( StridedView(data_dst, sz_dst, only(structs_dst)...), StridedView(data_src, sz_src, only(structs_src)...), - p, false, α * coeff, β, backend, allocator + p, conjsrc, α * coeff, β, backend, allocator ) else # Multi-tree block: pack → recoupling matmul → unpack. rows, cols = size(U) @@ -701,7 +713,7 @@ function add_transform_kernel!( @inbounds for (i, struct_src_i) in enumerate(structs_src) TO.tensoradd!( sreshape(view(buffer_src, :, i), sz_src), StridedView(data_src, sz_src, struct_src_i...), - ptriv, false, One(), Zero(), backend, allocator + ptriv, conjsrc, One(), Zero(), backend, allocator ) end diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index 86b9c926e..83e8e39e3 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -13,10 +13,10 @@ struct AbelianTreeTransformer{T, N} <: TreeTransformer data::Vector{AbelianTransformerData{T, N}} end -function AbelianTreeTransformer(transform, p, Vdst, Vsrc) +function AbelianTreeTransformer(transform, p, Vdst, Vsrc; srcstructure = subblockstructure(Vsrc)) t₀ = Base.time() permute(Vsrc, p) == Vdst || throw(SpaceMismatch("Incompatible spaces for permuting.")) - fts_src = subblockstructure(Vsrc) + fts_src = srcstructure fts_dst = subblockstructure(Vdst) L = length(fts_src) T = sectorscalartype(sectortype(Vdst)) @@ -51,11 +51,11 @@ struct GenericTreeTransformer{T, N} <: TreeTransformer data::Vector{GenericTransformerData{T, N}} end -function GenericTreeTransformer(transform, p, Vdst, Vsrc) +function GenericTreeTransformer(transform, p, Vdst, Vsrc; srcstructure = subblockstructure(Vsrc)) t₀ = Base.time() permute(Vsrc, p) == Vdst || throw(SpaceMismatch("Incompatible spaces for permuting.")) fusionstructure_dst = subblockstructure(Vdst) - fusionstructure_src = subblockstructure(Vsrc) + fusionstructure_src = srcstructure I = sectortype(Vsrc) T = sectorscalartype(I) N = numind(Vdst) @@ -139,7 +139,8 @@ function treetransformertype(Vdst, Vsrc) end function TreeTransformer( - transform::Function, p, Vdst::HomSpace{S}, Vsrc::HomSpace{S} + transform::Function, p, Vdst::HomSpace{S}, Vsrc::HomSpace{S}; + srcstructure = subblockstructure(Vsrc) ) where {S} permute(Vsrc, p) == Vdst || throw(SpaceMismatch("Incompatible spaces for permuting")) @@ -148,8 +149,8 @@ function TreeTransformer( I === Trivial && return TrivialTreeTransformer() return FusionStyle(I) == UniqueFusion() ? - AbelianTreeTransformer(transform, p, Vdst, Vsrc) : - GenericTreeTransformer(transform, p, Vdst, Vsrc) + AbelianTreeTransformer(transform, p, Vdst, Vsrc; srcstructure) : + GenericTreeTransformer(transform, p, Vdst, Vsrc; srcstructure) end # braid is special because it has levels @@ -159,6 +160,14 @@ end function treebraider(tdst::TensorMap, tsrc::TensorMap, p::Index2Tuple, levels) return treebraider(space(tdst), space(tsrc), p, levels) end +# The adjoint of a plain TensorMap can use a memoized transformer too: it shares its parent's +# data, so `conj_treebraider` supplies one whose source strides address the parent's buffer. +function treebraider( + tdst::TensorMap, tsrc::AdjointTensorMap{<:Any, <:Any, <:Any, <:Any, <:TensorMap}, + p::Index2Tuple, levels + ) + return conj_treebraider(space(tdst), space(parent(tsrc)), p, levels) +end @cached function treebraider( Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, levels )::treetransformertype(Vdst, Vsrc) @@ -166,12 +175,34 @@ end return TreeTransformer(fusiontreebraider, p, Vdst, Vsrc) end +# Transformer for braiding `adjoint(t)` while reading `t`'s own data buffer: the tree +# transformation is that of `Vparent'`, but the source strides address the parent's storage +# (see `adjoint_subblockstructure`). Kept as a separate cached function so that its entries +# cannot collide with the plain `treebraider` ones for the same pair of spaces. +@cached function conj_treebraider( + Vdst::TensorMapSpace, Vparent::TensorMapSpace, p::Index2Tuple, levels + )::treetransformertype(Vdst, Vparent') + fusiontreebraider(f) = braid(f, p, levels) + return TreeTransformer( + fusiontreebraider, p, Vdst, Vparent'; + srcstructure = adjoint_subblockstructure(Vparent) + ) +end + function treetransposer(::AbstractTensorMap, ::AbstractTensorMap, p::Index2Tuple) return fusiontreetransform(f) = transpose(f, p) end function treetransposer(tdst::TensorMap, tsrc::TensorMap, p::Index2Tuple) return treetransposer(space(tdst), space(tsrc), p) end +# As for `treebraider`: the adjoint of a plain TensorMap shares its parent's buffer, so it can +# use a memoized transformer whose source strides address that buffer. +function treetransposer( + tdst::TensorMap, tsrc::AdjointTensorMap{<:Any, <:Any, <:Any, <:Any, <:TensorMap}, + p::Index2Tuple + ) + return conj_treetransposer(space(tdst), space(parent(tsrc)), p) +end @cached function treetransposer( Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple )::treetransformertype(Vdst, Vsrc) @@ -179,6 +210,16 @@ end return TreeTransformer(fusiontreetransform, p, Vdst, Vsrc) end +@cached function conj_treetransposer( + Vdst::TensorMapSpace, Vparent::TensorMapSpace, p::Index2Tuple + )::treetransformertype(Vdst, Vparent') + fusiontreetransform(f) = transpose(f, p) + return TreeTransformer( + fusiontreetransform, p, Vdst, Vparent'; + srcstructure = adjoint_subblockstructure(Vparent) + ) +end + # default cachestyle is GlobalLRUCache # Sorting based on cost model From 48f926b613293a6d3e2412da37f95de0d75e0af6 Mon Sep 17 00:00:00 2001 From: leburgel Date: Thu, 27 Aug 2026 16:33:11 +0200 Subject: [PATCH 2/2] Transposer was actually never slow, don't touch it --- src/tensors/treetransformers.jl | 18 ------------------ 1 file changed, 18 deletions(-) diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index 83e8e39e3..7a4c75511 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -195,14 +195,6 @@ end function treetransposer(tdst::TensorMap, tsrc::TensorMap, p::Index2Tuple) return treetransposer(space(tdst), space(tsrc), p) end -# As for `treebraider`: the adjoint of a plain TensorMap shares its parent's buffer, so it can -# use a memoized transformer whose source strides address that buffer. -function treetransposer( - tdst::TensorMap, tsrc::AdjointTensorMap{<:Any, <:Any, <:Any, <:Any, <:TensorMap}, - p::Index2Tuple - ) - return conj_treetransposer(space(tdst), space(parent(tsrc)), p) -end @cached function treetransposer( Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple )::treetransformertype(Vdst, Vsrc) @@ -210,16 +202,6 @@ end return TreeTransformer(fusiontreetransform, p, Vdst, Vsrc) end -@cached function conj_treetransposer( - Vdst::TensorMapSpace, Vparent::TensorMapSpace, p::Index2Tuple - )::treetransformertype(Vdst, Vparent') - fusiontreetransform(f) = transpose(f, p) - return TreeTransformer( - fusiontreetransform, p, Vdst, Vparent'; - srcstructure = adjoint_subblockstructure(Vparent) - ) -end - # default cachestyle is GlobalLRUCache # Sorting based on cost model