Transform small subblocks in groups in the CPU add_transform!

src/tensors/indexmanipulations.jl

@@ -655,9 +655,25 @@ function add_transform_kernel!(
         dst::TransformSubblocks, src::TransformSubblocks, p, conjsrc::Bool,
         transformer::TreeTransformer, α, β, backend, allocator, ntasks::Int
     )
+    return _add_transform_blocks!(dst, src, p, conjsrc, transformer, transformer.data, α, β, backend, allocator, ntasks)
+end
+# on the CPU, the small subblocks of `TensorMap`s are transformed in groups, the others one by one
+function add_transform_kernel!(
+        dst::StridedSubblocks{<:CPUStorage}, src::StridedSubblocks{<:CPUStorage}, p, conjsrc::Bool,
+        transformer::TreeTransformer, α, β, backend, allocator, ntasks::Int
+    )
+    @timeit_debug GLOBAL_TIMER "dense: subblock groups" _add_subblockgroups!(
+        dst.data, src.data, conjsrc, transformer.groups, α, β, ntasks
+    )
+    large = view(transformer.data, transformer.large)
+    return _add_transform_blocks!(dst, src, p, conjsrc, transformer, large, α, β, backend, allocator, ntasks)
+end
+
+function _add_transform_blocks!(dst, src, p, conjsrc::Bool, transformer, data, α, β, backend, allocator, ntasks::Int)
+    isempty(data) && return nothing
     bufsize = buffersize(transformer)
     if bufsize == 0 # no recoupling needed: every block consists of a single tree
-        taskforeach(transformer.data, ntasks) do (U, inds_dst, inds_src)
+        taskforeach(data, ntasks) do (U, inds_dst, inds_src)
             _add_transform_block!(dst, src, p, conjsrc, U, inds_dst, inds_src, nothing, α, β, backend, allocator)
         end
     else
@@ -667,9 +683,9 @@ function add_transform_kernel!(
         cp = TO.allocator_checkpoint!(allocator)
         @timeit_debug GLOBAL_TIMER "alloc: buffers" buffers = [
             TO.tensoralloc(storagetype(dst), bufsize, Val(true), allocator)
-                for _ in 1:min(length(transformer.data), ntasks)
+                for _ in 1:min(length(data), ntasks)
         ]
-        taskforeach(transformer.data, buffers) do (U, inds_dst, inds_src), buffer
+        taskforeach(data, buffers) do (U, inds_dst, inds_src), buffer
             _add_transform_block!(dst, src, p, conjsrc, U, inds_dst, inds_src, buffer, α, β, backend, allocator)
         end
         foreach(Base.Fix2(TO.tensorfree!, allocator), buffers)
@@ -721,6 +737,49 @@ function _add_transform_block!(
     return nothing
 end
 
+function _add_subblockgroups!(A, B, conjsrc::Bool, groups, α, β, ntasks::Int)
+    return conjsrc ? _add_subblockgroups!(A, B, conj, groups, α, β, ntasks) :
+        _add_subblockgroups!(A, B, identity, groups, α, β, ntasks)
+end
+function _add_subblockgroups!(A, B, op, groups, α, β, ntasks::Int)
+    if ntasks == 1
+        foreach(group -> _add_subblockgroup!(A, B, op, group, eachindex(group.coeff), α, β), groups)
+    else
+        # tasks handle disjoint ranges of destination offsets, and thus disjoint subblocks
+        nchunks = 4 * ntasks
+        taskforeach(1:nchunks, ntasks) do c
+            lo, hi = div((c - 1) * length(A), nchunks), div(c * length(A), nchunks)
+            for group in groups
+                terms = searchsortedfirst(group.offset_dst, lo):(searchsortedfirst(group.offset_dst, hi) - 1)
+                _add_subblockgroup!(A, B, op, group, terms, α, β)
+            end
+        end
+    end
+    return nothing
+end
+
+# the metadata of the groups is built together with the transformer, which ensures the indices are inbounds
+function _add_subblockgroup!(A, B, op, group::SubblockGroup, terms, α, β)
+    (; rel_dst, rel_src, offset_dst, offset_src, coeff) = group
+    @inbounds for n in terms
+        o₁, o₂, a = offset_dst[n], offset_src[n], α * coeff[n]
+        if group.accumulate
+            @simd for k in eachindex(rel_dst)
+                A[o₁ + rel_dst[k]] += a * op(B[o₂ + rel_src[k]])
+            end
+        elseif iszero(β)
+            @simd for k in eachindex(rel_dst)
+                A[o₁ + rel_dst[k]] = a * op(B[o₂ + rel_src[k]])
+            end
+        else
+            @simd for k in eachindex(rel_dst)
+                A[o₁ + rel_dst[k]] = a * op(B[o₂ + rel_src[k]]) + β * A[o₁ + rel_dst[k]]
+            end
+        end
+    end
+    return nothing
+end
+
 # computes `buffer_dst = buffer_src * transpose(U)`, where both are column-major views into `buffer`
 _recouple!(buffer, buffer_dst, buffer_src, U) =
     (mul!(buffer_dst, buffer_src, transpose(StridedView(U))); nothing)

src/tensors/treetransformers.jl

@@ -33,6 +33,30 @@ function recoupling_scalartype(::Type{A}, ::Type{Tₛ}) where {A, Tₛ <: Number
     return Tₛ <: Real ? real(T) : complex(T)
 end
 
+"""
+    SubblockGroup{T}
+
+Small subblocks of equal size and strides, which are transformed by a single loop over all of them
+rather than by one `tensoradd!` call per subblock. Entry `k` of a subblock lies at `rel_dst[k]`
+and `rel_src[k]` relative to the offsets of the destination and source subblock, where the source
+strides are permuted such that both refer to the same entry. Term `n` then computes
+`dst[offset_dst[n] .+ rel_dst] = α * coeff[n] * src[offset_src[n] .+ rel_src] + b * dst[offset_dst[n] .+ rel_dst]`,
+with `b = β` for the first term into every destination subblock and `b = 1` for the groups of
+further recoupling terms, which have `accumulate` set and are applied after all others.
+The terms are sorted by `offset_dst`.
+"""
+struct SubblockGroup{T}
+    rel_dst::Vector{Int}
+    rel_src::Vector{Int}
+    offset_dst::Vector{Int}
+    offset_src::Vector{Int}
+    coeff::Vector{T}
+    accumulate::Bool
+end
+
+# subblocks up to this length are transformed in `SubblockGroup`s
+const SUBBLOCKGROUP_MAXLENGTH = 1024
+
 # (coefficient, destination position, source position)
 const UniqueTransformerData{T} = Tuple{T, Int, Int}
 
@@ -42,12 +66,16 @@ const UniqueTransformerData{T} = Tuple{T, Int, Int}
 Tree transformation for `UniqueFusion` sectors, where every source subblock maps onto a single
 destination subblock with a scalar coefficient, stored as `(coeff, idst, isrc)`. The subblock
 structures of the destination and source spaces are kept alongside, such that the
-[`StridedSubblocks`](@ref) of both tensors can be created without further lookups.
+[`StridedSubblocks`](@ref) of both tensors can be created without further lookups. The small
+subblocks are also collected in [`SubblockGroup`](@ref)s, and `large` holds the positions in
+`data` of the others.
 """
 struct UniqueTreeTransformer{T, N} <: TreeTransformer
     data::Vector{UniqueTransformerData{T}}
     structure_dst::Vector{StridedStructure{N}}
     structure_src::Vector{StridedStructure{N}}
+    groups::Vector{SubblockGroup{T}}
+    large::Vector{Int}
 end
 
 # (recoupling matrix, destination positions, source positions)
@@ -60,12 +88,16 @@ Tree transformation for sectors with multiple fusion channels, where the subbloc
 [`FusionTreeBlock`](@ref) map onto the subblocks of the transformed block through a recoupling
 matrix, stored as `(U, inds_dst, inds_src)`, where `U[j, i]` maps source `i` onto destination `j`.
 The subblock structures of the destination and source spaces are kept alongside, such that the
-[`StridedSubblocks`](@ref) of both tensors can be created without further lookups.
+[`StridedSubblocks`](@ref) of both tensors can be created without further lookups. The blocks of
+small subblocks are also collected in [`SubblockGroup`](@ref)s, and `large` holds the positions in
+`data` of the others.
 """
 struct GenericTreeTransformer{T, N} <: TreeTransformer
     data::Vector{GenericTransformerData{T}}
     structure_dst::Vector{StridedStructure{N}}
     structure_src::Vector{StridedStructure{N}}
+    groups::Vector{SubblockGroup{T}}
+    large::Vector{Int}
 end
 
 function UniqueTreeTransformer(::Type{A}, transform, p, Vdst::HomSpace, Vsrc::HomSpace, conjsrc::Bool) where {A}
@@ -85,7 +117,10 @@ function UniqueTreeTransformer(::Type{A}, transform, p, Vdst::HomSpace, Vsrc::Ho
 
     structure_dst = degeneracystructure(Vdst).subblockstructure
     structure_src = degeneracystructure(Vsrc).subblockstructure
-    transformer = UniqueTreeTransformer(data, structure_dst, structure_src)
+    groups, large = @timeit_debug GLOBAL_TIMER "bookkeeping: subblock groups" subblockgroups(
+        T, data, structure_dst, structure_src, linearize(p)
+    )
+    transformer = UniqueTreeTransformer(data, structure_dst, structure_src, groups, large)
 
     Δt = Base.time() - t₀
     @debug(lazy"Treetransformer for $Vsrc to $Vdst via $p", conjsrc, nblocks = length(data), Δt)
@@ -137,7 +172,10 @@ function GenericTreeTransformer(::Type{A}, transform, p, Vdst::HomSpace, Vsrc::H
 
     # sort by (approximate) weight to facilitate multi-threading strategies
     @timeit_debug GLOBAL_TIMER "bookkeeping: sort" Base.permute!(data, sortperm(weights; rev = true))
-    transformer = GenericTreeTransformer{T, numind(Vdst)}(data, structure_dst, structure_src)
+    groups, large = @timeit_debug GLOBAL_TIMER "bookkeeping: subblock groups" subblockgroups(
+        T, data, structure_dst, structure_src, linearize(p)
+    )
+    transformer = GenericTreeTransformer{T, numind(Vdst)}(data, structure_dst, structure_src, groups, large)
 
     Δt = Base.time() - t₀
     @debug(
@@ -151,6 +189,59 @@ function GenericTreeTransformer(::Type{A}, transform, p, Vdst::HomSpace, Vsrc::H
     return transformer
 end
 
+"""
+    subblockgroups(T, data, structure_dst, structure_src, p) -> (groups, large)
+
+Collect the terms of the transformer `data` that act on subblocks with at most
+`SUBBLOCKGROUP_MAXLENGTH` entries into [`SubblockGroup`](@ref)s by size and strides, with the
+accumulating groups last, and return them together with the positions in `data` of the other entries.
+"""
+function subblockgroups(::Type{T}, data, structure_dst::Vector{StridedStructure{N}}, structure_src, p) where {T, N}
+    terms = Tuple{T, Int, Int, Bool}[]
+    large = Int[]
+    for (n, (U, inds_dst, inds_src)) in enumerate(data)
+        if prod(first(structure_dst[first(inds_dst)])) > SUBBLOCKGROUP_MAXLENGTH
+            push!(large, n)
+        else
+            _push_terms!(terms, U, inds_dst, inds_src)
+        end
+    end
+    sort!(terms; by = term -> structure_dst[term[2]][3])
+
+    # strides of dimensions of size 1 are irrelevant, and set to zero to obtain fewer groups
+    groups = Dict{Tuple{Bool, NTuple{3, NTuple{N, Int}}}, SubblockGroup{T}}()
+    for (coeff, idst, isrc, accumulate) in terms
+        sz, strides_dst, offset_dst = structure_dst[idst]
+        _, strides_src, offset_src = structure_src[isrc]
+        strides_dst = map((n, s) -> n == 1 ? 0 : s, sz, strides_dst)
+        strides_src′ = map((n, s) -> n == 1 ? 0 : s, sz, TupleTools.getindices(strides_src, p))
+        group = get!(groups, (accumulate, (sz, strides_dst, strides_src′))) do
+            rel_dst, rel_src = Int[], Int[]
+            for I in CartesianIndices(sz)
+                push!(rel_dst, 1 + sum((Tuple(I) .- 1) .* strides_dst; init = 0))
+                push!(rel_src, 1 + sum((Tuple(I) .- 1) .* strides_src′; init = 0))
+            end
+            return SubblockGroup(rel_dst, rel_src, Int[], Int[], T[], accumulate)
+        end
+        push!(group.offset_dst, offset_dst)
+        push!(group.offset_src, offset_src)
+        push!(group.coeff, coeff)
+    end
+    return sort!(collect(values(groups)); by = group -> group.accumulate), large
+end
+
+_push_terms!(terms, coeff::Number, idst::Int, isrc::Int) = push!(terms, (coeff, idst, isrc, false))
+function _push_terms!(terms, U::AbstractMatrix, inds_dst, inds_src)
+    for (i, idst) in enumerate(inds_dst)
+        js = findall(!iszero, view(U, i, :))
+        isempty(js) && push!(js, 1) # such that β is still applied
+        for j in js
+            push!(terms, (U[i, j], idst, inds_src[j], j != first(js)))
+        end
+    end
+    return terms
+end
+
 """
     buffersize(transformer::TreeTransformer) -> Int
 

test/tensors/indexmanipulations.jl

@@ -235,3 +235,37 @@ for V in spacelist
     end
     TensorKit.empty_globalcaches!()
 end
+
+@timedtestset "Permutations with small and large subblocks" begin
+    # subblocks up to `SUBBLOCKGROUP_MAXLENGTH` entries are transformed in groups, larger ones one by one
+    V = Vect[SU2Irrep](0 => 6, 1 // 2 => 6, 1 => 1)
+    t = rand(ComplexF64, V ⊗ V ⊗ V ← V)
+    @test maximum(length ∘ last, subblocks(t)) > TensorKit.SUBBLOCKGROUP_MAXLENGTH
+    a = convert(Array, t)
+    for p in (((1, 4), (2, 3)), ((3, 1, 4), (2,)))
+        @test convert(Array, permute(t, p)) ≈ permutedims(a, (p[1]..., p[2]...))
+    end
+end
+
+@timedtestset "Permutations with several transformer tasks" begin
+    # the tasks split the subblock groups by destination offset; they also run on a single thread,
+    # so the number is set directly rather than through `set_num_transformer_threads`, which caps it
+    V1 = Vect[FermionParity ⊠ U1Irrep]((p, q) => 3 for p in 0:1 for q in -2:2)
+    V2 = Vect[SU2Irrep](0 => 8, 1 // 2 => 8, 1 => 1)
+    ntasks = TensorKit.TRANSFORMER_THREADS[]
+    for V in (V1, V2)
+        t = rand(ComplexF64, V ⊗ V ⊗ V ← V)
+        @test dim(t) > TensorKit.Strided.MINTHREADLENGTH # below, a single task is used
+        for p in (((1, 4), (2, 3)), ((3, 1, 4), (2,)))
+            tdst = rand(ComplexF64, permute(space(t), p))
+            refs = (permute(t, p), permute!(copy(tdst), t, p, 0.7, 0.3))
+            try
+                TensorKit.TRANSFORMER_THREADS[] = 3
+                @test permute(t, p) ≈ refs[1]
+                @test permute!(copy(tdst), t, p, 0.7, 0.3) ≈ refs[2]
+            finally
+                TensorKit.TRANSFORMER_THREADS[] = ntasks
+            end
+        end
+    end
+end