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..7e798c30b 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -579,6 +579,29 @@ Base.@deprecate( return tdst end +""" + add_conj_transform!(tdst, tsrc, p, transformer, α, β, backend, allocator) -> tdst + +Compute `tdst = β * tdst + α * permute(adjoint(tsrc), p)` while reading `tsrc`'s own data +buffer, without materialising an `AdjointTensorMap`. + +`transformer` must come from [`conj_treebraider`](@ref), whose source strides address the +parent's storage; the conjugation itself is folded into the strided reads. +""" +@propagate_inbounds function add_conj_transform!( + tdst::TensorMap, tsrc::TensorMap, p::Index2Tuple, transformer, + α::Number, β::Number, backend, allocator + ) + @boundscheck spacecheck_transform(permute, tdst, tsrc', p) + ntasks = use_threaded_transform(tdst, transformer) ? get_num_transformer_threads() : 1 + scheduler = ntasks == 1 ? SerialScheduler() : DynamicScheduler(; ntasks, split = :roundrobin) + add_transform_kernel!( + tdst.data, tsrc.data, p, transformer, α, β, backend, allocator, scheduler; + conjsrc = true + ) + return tdst +end + function use_threaded_transform(t::TensorMap, transformer) return get_num_transformer_threads() > 1 && length(t.data) > Strided.MINTHREADLENGTH end @@ -587,14 +610,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 +653,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 +681,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 +710,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 +725,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/tensoroperations.jl b/src/tensors/tensoroperations.jl index 147d27b16..01b19171e 100644 --- a/src/tensors/tensoroperations.jl +++ b/src/tensors/tensoroperations.jl @@ -52,9 +52,23 @@ function TO.tensoradd!( return C end if conjA - A′ = adjoint(A) pA′ = adjointtensorindices(A, _canonicalize(pA, C)) - permute!(C, A′, pA′, α, β, backend, allocator) + if C isa TensorMap && A isa TensorMap + # Both operands have a flat data buffer, so the adjoint never has to be + # materialised: `conj_treebraider` yields a memoized transformer whose source + # strides address `A`'s own storage. Materialising `adjoint(A)` here would push + # a perfectly ordinary TensorMap onto the uncached tree-transformation path. + n₁ = numin(A) # == numout(adjoint(A)) + levels = ntuple(identity, numind(A)) + levels′ = ( + TupleTools.getindices(levels, ntuple(identity, n₁)), + TupleTools.getindices(levels, n₁ .+ ntuple(identity, numout(A))), + ) + transformer = conj_treebraider(space(C), space(A), pA′, levels′) + @inbounds add_conj_transform!(C, A, pA′, transformer, α, β, backend, allocator) + else + permute!(C, adjoint(A), pA′, α, β, backend, allocator) + end else permute!(C, A, _canonicalize(pA, C), α, β, backend, allocator) end diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index 86b9c926e..04cfc8c13 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 @@ -166,6 +167,20 @@ 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