From 5ea47423088964ffdf88bf99caeec2d06d92701d Mon Sep 17 00:00:00 2001 From: leburgel Date: Mon, 26 Oct 2026 12:31:58 +0100 Subject: [PATCH 1/3] Transform small subblocks in shape groups on the CPU instead of one tensoradd! per subblock Co-Authored-By: Claude Opus 5.5 --- src/tensors/indexmanipulations.jl | 63 +++++++++++++++++++- src/tensors/treetransformers.jl | 95 +++++++++++++++++++++++++++++-- 2 files changed, 151 insertions(+), 7 deletions(-) diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index f6623c197..8a77c226e 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -655,9 +655,23 @@ 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 + ) + _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 +681,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 +735,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) diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index b274c6f8c..831a8a593 100644 --- a/src/tensors/treetransformers.jl +++ b/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,8 @@ 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 = 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 +170,8 @@ 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 = 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 +185,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 From 9464cd0889a7c9bfbe69eed1456469276355daae Mon Sep 17 00:00:00 2001 From: leburgel Date: Mon, 26 Oct 2026 13:02:39 +0100 Subject: [PATCH 2/3] Time the subblock groups and test permutations across the group threshold Co-Authored-By: Claude Opus 5.5 --- src/tensors/indexmanipulations.jl | 4 +++- src/tensors/treetransformers.jl | 8 ++++++-- test/tensors/indexmanipulations.jl | 11 +++++++++++ 3 files changed, 20 insertions(+), 3 deletions(-) diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index 8a77c226e..84465ede1 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -662,7 +662,9 @@ function add_transform_kernel!( dst::StridedSubblocks{<:CPUStorage}, src::StridedSubblocks{<:CPUStorage}, p, conjsrc::Bool, transformer::TreeTransformer, α, β, backend, allocator, ntasks::Int ) - _add_subblockgroups!(dst.data, src.data, conjsrc, transformer.groups, α, β, ntasks) + @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 diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index 831a8a593..6e2d7883a 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -117,7 +117,9 @@ function UniqueTreeTransformer(::Type{A}, transform, p, Vdst::HomSpace, Vsrc::Ho structure_dst = degeneracystructure(Vdst).subblockstructure structure_src = degeneracystructure(Vsrc).subblockstructure - groups, large = subblockgroups(T, data, structure_dst, structure_src, linearize(p)) + 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₀ @@ -170,7 +172,9 @@ 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)) - groups, large = subblockgroups(T, data, structure_dst, structure_src, linearize(p)) + 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₀ diff --git a/test/tensors/indexmanipulations.jl b/test/tensors/indexmanipulations.jl index 0cc276c6d..7f6e3af22 100644 --- a/test/tensors/indexmanipulations.jl +++ b/test/tensors/indexmanipulations.jl @@ -235,3 +235,14 @@ 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 From ed76496201996b73ae2db6053dfa67ad11627dd8 Mon Sep 17 00:00:00 2001 From: leburgel Date: Tue, 27 Oct 2026 18:38:56 +0100 Subject: [PATCH 3/3] Test the subblock groups with several transformer tasks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The test suite runs with one transformer task, so the path that splits the groups over tasks by destination offset was not covered. Compare `permute` and `permute!` with β ≠ 0 using three tasks against the serial result, for a fermionic U(1) tensor (unique fusion) and an SU(2) tensor with subblocks on both sides of the group threshold (generic fusion, accumulating groups). Co-Authored-By: Claude Opus 5.5 --- test/tensors/indexmanipulations.jl | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/test/tensors/indexmanipulations.jl b/test/tensors/indexmanipulations.jl index 7f6e3af22..c8d5f1e05 100644 --- a/test/tensors/indexmanipulations.jl +++ b/test/tensors/indexmanipulations.jl @@ -246,3 +246,26 @@ end @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