diff --git a/benchmark/TensorKitBenchmarks/indexmanipulations/benchparams.toml b/benchmark/TensorKitBenchmarks/indexmanipulations/benchparams.toml index c19aa0f66..549072d23 100644 --- a/benchmark/TensorKitBenchmarks/indexmanipulations/benchparams.toml +++ b/benchmark/TensorKitBenchmarks/indexmanipulations/benchparams.toml @@ -10,3 +10,17 @@ I = "Z2Irrep" p = [[[2, 1], []]] dims = [[7264, 7264], [43408, 1216]] sigmas = [[0.5, 0.5]] + +[[permute]] +T = ["Float64"] +I = "SU2Irrep" +p = [[[2, 1], []]] +dims = [[512, 512]] +sigmas = [[1.0, 1.0]] + +[[permute]] +T = ["Float64", "ComplexF64"] +I = "SU2Irrep" +p = [[[1, 3], [2, 4]], [[4, 2, 3], [1]]] +dims = [[48, 48, 48, 48]] +sigmas = [[1.0, 1.0, 1.0, 1.0]] diff --git a/src/auxiliary/auxiliary.jl b/src/auxiliary/auxiliary.jl index e5ed7cc90..3d3ff0fe4 100644 --- a/src/auxiliary/auxiliary.jl +++ b/src/auxiliary/auxiliary.jl @@ -92,3 +92,49 @@ end @noinline function _throw_ambiguous_levels(l) throw(ArgumentError(lazy"ambiguous braid: two indices with equal level $l have to cross")) end + +""" + taskforeach(f, items, ntasks::Int) -> Nothing + taskforeach(f, items, resources::Vector) -> Nothing + +Apply `f(item)` (respectively `f(item, resource)`) to all elements of `items`, +distributing the work over at most `ntasks` workers (respectively one worker per entry +of `resources`), with dynamic load balancing through an atomic counter. The calling +thread acts as one of the workers, so at most `ntasks - 1` tasks are spawned, and none +at all for a single worker. +""" +function taskforeach(f, items, ntasks::Int) + items′ = items isa AbstractArray ? items : collect(items) + n = length(items′) + counter = Threads.Atomic{Int}(1) + Threads.@sync begin + for _ in 2:min(ntasks, n) + Threads.@spawn _taskforeach_worker(f, items′, counter, n) + end + _taskforeach_worker(f, items′, counter, n) + end + return nothing +end +function taskforeach(f, items, resources::Vector) + isempty(resources) && return nothing + items′ = items isa AbstractArray ? items : collect(items) + n = length(items′) + counter = Threads.Atomic{Int}(1) + Threads.@sync begin + for j in 2:min(length(resources), n) + local r = resources[j] + Threads.@spawn _taskforeach_worker(Base.Fix2(f, r), items′, counter, n) + end + _taskforeach_worker(Base.Fix2(f, first(resources)), items′, counter, n) + end + return nothing +end + +function _taskforeach_worker(f, items, counter, n::Int) + while true + i = Threads.atomic_add!(counter, 1) + i > n && break + f(@inbounds(items[i])) + end + return nothing +end diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index ce4add069..2f1d6bf2a 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -567,11 +567,10 @@ Base.@deprecate( TO.tensoradd!(tdst[], tsrc[], p, false, α, β, backend, allocator) else ntasks = use_threaded_transform(tdst, transformer) ? get_num_transformer_threads() : 1 - 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) + add_transform_kernel!(tdst.data, tsrc.data, p, transformer, α, β, backend, allocator, ntasks) else - add_transform_kernel!(tdst, tsrc, p, transformer, α, β, backend, allocator, scheduler) + add_transform_kernel!(tdst, tsrc, p, transformer, α, β, backend, allocator, ntasks) end end end @@ -587,11 +586,11 @@ 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, ntasks::Int ) I = sectortype(tdst) if FusionStyle(I) === UniqueFusion() - tforeach(fusiontrees(tsrc); scheduler) do (f₁, f₂) + taskforeach(fusiontrees(tsrc), ntasks) do (f₁, f₂) (f₁′, f₂′), coeff = transformer((f₁, f₂)) @inbounds TO.tensoradd!( tdst[f₁′, f₂′], tsrc[f₁, f₂], p, false, α * coeff, β, backend, allocator @@ -599,56 +598,24 @@ function add_transform_kernel!( end return nothing end - cp = TO.allocator_checkpoint!(allocator) - - # buffers have to be created without race condition: err on the side of caution with a lock - buffer_lock = Threads.ReentrantLock() - - OhMyThreads.@tasks for src in fusionblocks(tsrc) - # setup - OhMyThreads.@set scheduler = scheduler - dst, U = transformer(src) - if length(src) == 1 # Degenerate block with a single tree: no matmul needed. - (f₁, f₂) = only(fusiontrees(src)) - (f₁′, f₂′) = only(fusiontrees(dst)) - @inbounds TO.tensoradd!( - tdst[f₁′, f₂′], tsrc[f₁, f₂], p, false, α * only(U), β, backend, allocator - ) - else # Multi-tree block: pack → recoupling matmul → unpack. - rows, cols = size(U) - sz_src = size(tsrc[first(fusiontrees(src))...]) - blocksize = prod(sz_src) - buffer = @lock buffer_lock TO.tensoralloc(storagetype(tdst), blocksize * (rows + cols), Val(true), allocator) - ptriv = (ntuple(identity, length(sz_src)), ()) - buffer_dst = StridedView(buffer, (blocksize, rows), (1, blocksize), 0) - buffer_src = StridedView(buffer, (blocksize, cols), (1, blocksize), blocksize * rows) - - # 1. Extract: copy each source block into column i of buffer_src as a flat vector, - # using a trivial permutation so the layout is canonical before the matmul. - @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 - ) - end + fblocks = fusionblocks(tsrc) + bufsize = buffersize(tsrc, fblocks) - # 2. Recoupling: buffer_dst = buffer_src * U^T (each output tree is a linear - # combination of input trees weighted by the recoupling coefficients). - U′ = Adapt.adapt(storagetype(tdst), StridedView(U)) - mul!(buffer_dst, buffer_src, transpose(U′)) - - # 3. Insert: scatter column i of buffer_dst into the destination, applying the - # actual index permutation p in the same tensoradd! call. - @inbounds for (i, (f₃, f₄)) in enumerate(fusiontrees(dst)) - TO.tensoradd!( - tdst[f₃, f₄], sreshape(view(buffer_dst, :, i), sz_src), - p, false, α, β, backend, allocator - ) - end - @lock buffer_lock TO.tensorfree!(buffer, allocator) - end + # One max-sized workspace per task (a single one that is reused by all blocks when + # serial), allocated on the calling thread before any task spawns, so that also + # allocators that are not thread-safe can be used. + cp = TO.allocator_checkpoint!(allocator) + buffers = [ + TO.tensoralloc(storagetype(tdst), bufsize, Val(true), allocator) + for _ in 1:clamp(length(fblocks), 1, ntasks) + ] + taskforeach(fblocks, buffers) do src, buffer + _add_transform_block!( + tdst, tsrc, p, src, transformer, buffer, α, β, backend, allocator + ) end + foreach(Base.Fix2(TO.tensorfree!, allocator), buffers) TO.allocator_reset!(allocator, cp) return nothing end @@ -657,9 +624,9 @@ 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, ntasks::Int ) - tforeach(transformer.data; scheduler) do (coeff, struct_dst, struct_src) + taskforeach(transformer.data, ntasks) do (coeff, struct_dst, struct_src) TO.tensoradd!( StridedView(data_dst, struct_dst...), StridedView(data_src, struct_src...), p, false, α * coeff, β, backend, allocator @@ -669,58 +636,117 @@ function add_transform_kernel!( end function add_transform_kernel!( data_dst::DenseVector, data_src::DenseVector, p, transformer::GenericTreeTransformer, - α, β, backend, allocator, scheduler + α, β, backend, allocator, ntasks::Int ) + bufsize = buffersize(transformer) + + # One max-sized workspace per task (a single one that is reused by all blocks when + # serial), allocated on the calling thread before any task spawns, so that also + # allocators that are not thread-safe can be used. cp = TO.allocator_checkpoint!(allocator) + buffers = [ + TO.tensoralloc(typeof(data_dst), bufsize, Val(true), allocator) + for _ in 1:clamp(length(transformer.data), 1, ntasks) + ] + taskforeach(transformer.data, buffers) do subtransformer, buffer + _add_transform_block!( + data_dst, data_src, p, subtransformer, buffer, α, β, backend, allocator + ) + end + foreach(Base.Fix2(TO.tensorfree!, allocator), buffers) + TO.allocator_reset!(allocator, cp) + return nothing +end + +function _add_transform_block!( + tdst, tsrc, p, src::FusionTreeBlock, transformer, buffer, + α, β, backend, allocator + ) + dst, U = transformer(src) + + if length(src) == 1 # Degenerate block with a single tree: no matmul needed. + (f₁, f₂) = only(fusiontrees(src)) + (f₁′, f₂′) = only(fusiontrees(dst)) + @inbounds TO.tensoradd!( + tdst[f₁′, f₂′], tsrc[f₁, f₂], p, false, α * only(U), β, backend, allocator + ) + else # Multi-tree block: pack → recoupling matmul → unpack. + rows, cols = size(U) + sz_src = size(tsrc[first(fusiontrees(src))...]) + blocksize = prod(sz_src) + # the buffer was sized assuming a square recoupling matrix + rows == cols || throw(DimensionMismatch(lazy"recoupling matrix is not square: $(size(U))")) + ptriv = (ntuple(identity, length(sz_src)), ()) + buffer_dst = StridedView(buffer, (blocksize, rows), (1, blocksize), 0) + buffer_src = StridedView(buffer, (blocksize, cols), (1, blocksize), blocksize * rows) + + # 1. Extract: copy each source block into column i of buffer_src as a flat vector, + # using a trivial permutation so the layout is canonical before the matmul. + @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 + ) + end - # buffers have to be created without race condition: err on the side of caution with a lock - buffer_lock = Threads.ReentrantLock() + # 2. Recoupling: buffer_dst = α * buffer_src * U^T (each output tree is a linear + # combination of input trees weighted by the recoupling coefficients). + U′ = _adapt_recoupling(storagetype(tdst), U) + mul!(buffer_dst, buffer_src, transpose(U′), α, Zero()) - OhMyThreads.@tasks for subtransformer in transformer.data - # setup - OhMyThreads.@set scheduler = scheduler - U, (sz_dst, structs_dst), (sz_src, structs_src) = subtransformer + # 3. Insert: scatter column i of buffer_dst into the destination, applying the + # actual index permutation p in the same tensoradd! call. + @inbounds for (i, (f₃, f₄)) in enumerate(fusiontrees(dst)) + TO.tensoradd!( + tdst[f₃, f₄], sreshape(view(buffer_dst, :, i), sz_src), + p, false, One(), β, backend, allocator + ) + end + end + return nothing +end - if length(U) == 1 # Degenerate block with a single tree: no matmul needed. - coeff = only(U) +function _add_transform_block!( + data_dst::DenseVector, data_src::DenseVector, p, + ((U, (sz_dst, structs_dst), (sz_src, structs_src)))::GenericTransformerData, + buffer, α, β, backend, allocator + ) + if length(U) == 1 # Degenerate block with a single tree: no matmul needed. + coeff = only(U) + TO.tensoradd!( + StridedView(data_dst, sz_dst, only(structs_dst)...), + StridedView(data_src, sz_src, only(structs_src)...), + p, false, α * coeff, β, backend, allocator + ) + else # Multi-tree block: pack → recoupling matmul → unpack. + rows, cols = size(U) + blocksize = prod(sz_src) + ptriv = (ntuple(identity, length(sz_src)), ()) + buffer_dst = StridedView(buffer, (blocksize, rows), (1, blocksize), 0) + buffer_src = StridedView(buffer, (blocksize, cols), (1, blocksize), blocksize * rows) + + # 1. Extract: copy each source block into column i of buffer_src as a flat vector, + # using a trivial permutation so the layout is canonical before the matmul. + @inbounds for (i, struct_src_i) in enumerate(structs_src) TO.tensoradd!( - StridedView(data_dst, sz_dst, only(structs_dst)...), - StridedView(data_src, sz_src, only(structs_src)...), - p, false, α * coeff, β, backend, allocator + sreshape(view(buffer_src, :, i), sz_src), StridedView(data_src, sz_src, struct_src_i...), + ptriv, false, One(), Zero(), backend, allocator ) - else # Multi-tree block: pack → recoupling matmul → unpack. - rows, cols = size(U) - blocksize = prod(sz_src) - buffer = @lock buffer_lock TO.tensoralloc(typeof(data_dst), blocksize * (rows + cols), Val(true), allocator) - ptriv = (ntuple(identity, length(sz_src)), ()) - buffer_dst = StridedView(buffer, (blocksize, rows), (1, blocksize), 0) - buffer_src = StridedView(buffer, (blocksize, cols), (1, blocksize), blocksize * rows) - - # 1. Extract: copy each source block into column i of buffer_src as a flat vector, - # using a trivial permutation so the layout is canonical before the matmul. - @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 - ) - end + end - # 2. Recoupling: buffer_dst = buffer_src * U^T (each output tree is a linear - # combination of input trees weighted by the recoupling coefficients). - U′ = Adapt.adapt(typeof(data_dst), StridedView(U)) - mul!(buffer_dst, buffer_src, transpose(U′)) - - # 3. Insert: scatter column i of buffer_dst into the destination, applying the - # actual index permutation p in the same tensoradd! call. - @inbounds for (i, struct_dst_i) in enumerate(structs_dst) - TO.tensoradd!( - StridedView(data_dst, sz_dst, struct_dst_i...), sreshape(view(buffer_dst, :, i), sz_src), - p, false, α, β, backend, allocator - ) - end - @lock buffer_lock TO.tensorfree!(buffer, allocator) + # 2. Recoupling: buffer_dst = α * buffer_src * U^T (each output tree is a linear + # combination of input trees weighted by the recoupling coefficients). + U′ = _adapt_recoupling(typeof(data_dst), U) + mul!(buffer_dst, buffer_src, transpose(U′), α, Zero()) + + # 3. Insert: scatter column i of buffer_dst into the destination, applying the + # actual index permutation p in the same tensoradd! call. + @inbounds for (i, struct_dst_i) in enumerate(structs_dst) + TO.tensoradd!( + StridedView(data_dst, sz_dst, struct_dst_i...), sreshape(view(buffer_dst, :, i), sz_src), + p, false, One(), β, backend, allocator + ) end end - TO.allocator_reset!(allocator, cp) return nothing end diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index 86b9c926e..4a6401a36 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -51,6 +51,27 @@ struct GenericTreeTransformer{T, N} <: TreeTransformer data::Vector{GenericTransformerData{T, N}} end +""" + buffersize(transformer::GenericTreeTransformer) -> Int + buffersize(t::AbstractTensorMap, fblocks) -> Int + +Compute the workspace size required to pack, recouple and unpack the largest multi-tree +block, i.e. `prod(sz_src) * (rows + cols)` where `(rows, cols) = size(U)` is the size of +the recoupling matrix. The second form computes this from the fusion-tree blocks of a +tensor instead, using the fact that the recoupling matrix is square. +""" +function buffersize(transformer::GenericTreeTransformer) + return maximum(transformer.data; init = 0) do (U, _, (sz_src, _)) + return length(U) == 1 ? 0 : prod(sz_src) * sum(size(U)) + end +end +function buffersize(t::AbstractTensorMap, fblocks) + return maximum(fblocks; init = 0) do src + n = length(src) + return n == 1 ? 0 : 2 * n * length(t[first(fusiontrees(src))...]) + end +end + function GenericTreeTransformer(transform, p, Vdst, Vsrc) t₀ = Base.time() permute(Vsrc, p) == Vdst || throw(SpaceMismatch("Incompatible spaces for permuting.")) @@ -67,41 +88,19 @@ function GenericTreeTransformer(transform, p, Vdst, Vsrc) data = Vector{GenericTransformerData{T, N}}(undef, nblocks) nthreads = get_num_manipulation_threads() - if nthreads > 1 - counter = Threads.Atomic{Int}(1) - Threads.@sync for _ in 1:min(nthreads, nblocks) - Threads.@spawn begin - while true - local_counter = Threads.atomic_add!(counter, 1) - local_counter > nblocks && break - fs_src = fblocks[local_counter] - fs_dst, U = transform(fs_src) - sz_src, newstructs_src = repack_transformer_structure(fusionstructure_src, fusiontrees(fs_src)) - sz_dst, newstructs_dst = repack_transformer_structure(fusionstructure_dst, fusiontrees(fs_dst)) - data[local_counter] = U, (sz_dst, newstructs_dst), (sz_src, newstructs_src) - - @debug( - lazy"Created recoupling block for uncoupled: $(fs_src.uncoupled)", - sz = size(U), sparsity = count(!iszero, U) / length(U) - ) - end - end - end - transformer = GenericTreeTransformer{T, N}(data) - else - for (i, fs_src) in enumerate(fblocks) - fs_dst, U = transform(fs_src) - sz_src, newstructs_src = repack_transformer_structure(fusionstructure_src, fusiontrees(fs_src)) - sz_dst, newstructs_dst = repack_transformer_structure(fusionstructure_dst, fusiontrees(fs_dst)) - data[i] = U, (sz_dst, newstructs_dst), (sz_src, newstructs_src) - - @debug( - lazy"Created recoupling block for uncoupled: $(fs_src.uncoupled)", - sz = size(U), sparsity = count(!iszero, U) / length(U) - ) - end - transformer = GenericTreeTransformer{T, N}(data) + taskforeach(1:nblocks, nthreads) do i + fs_src = fblocks[i] + fs_dst, U = transform(fs_src) + sz_src, newstructs_src = repack_transformer_structure(fusionstructure_src, fusiontrees(fs_src)) + sz_dst, newstructs_dst = repack_transformer_structure(fusionstructure_dst, fusiontrees(fs_dst)) + data[i] = U, (sz_dst, newstructs_dst), (sz_src, newstructs_src) + + @debug( + lazy"Created recoupling block for uncoupled: $(fs_src.uncoupled)", + sz = size(U), sparsity = count(!iszero, U) / length(U) + ) end + transformer = GenericTreeTransformer{T, N}(data) # sort by (approximate) weight to facilitate multi-threading strategies sort!(transformer) @@ -191,6 +190,14 @@ function Base.sort!( return transformer end +# For CPU arrays the recoupling matrix can be used as is, also when the scalar types +# do not match, since Strided handles mixed-eltype mul! without the copy that +# Adapt.adapt would make (which additionally dispatches dynamically). Other storage +# types (e.g. GPU arrays) do require the conversion. +# TODO: transformers with dedicated storagetypes +_adapt_recoupling(::Type{<:Array}, U::Matrix) = StridedView(U) +_adapt_recoupling(::Type{A}, U::Matrix) where {A} = Adapt.adapt(A, StridedView(U)) + function _transformer_weight((coeff, struct_dst, struct_src)::AbelianTransformerData) return prod(struct_dst[1]) end