Add dedicated cached treebraider for AdjointTensorMap

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)
 

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
 

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,13 +160,35 @@ 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)
     fusiontreebraider(f) = braid(f, p, levels)
     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