Parallel Loops

TanayLabUtilities.ParallelLoops.parallel_loop_wo_rng Function
parallel_loop_wo_rng(
    body::Function,
    indices::AbstractVector{<:Integer};
    name::AbstractString = ".loop",
    policy::Symbol = :greedy_sticky,
    order::Maybe{AbstractVector{<:Integer}} = nothing,
    weights::Maybe{AbstractVector{<:Integer}} = nothing,
    nested::Bool = false,
    progress::Maybe{AbstractProgress} = nothing,
)::Nothing

Run the body in parallel, passing it the iteration index . The policy is one of:

  • :greedy_sticky - nthreads() sticky worker tasks spawned via @threads :static for _ in 1:nthreads() , each pulling the next index position from an Atomic{Int} counter. Combines the dynamic load balancing of :greedy with the threadid stability of :static : tasks never migrate, so threadid() -indexed scratch is safe inside the body. Pair with a heaviest-first order (or weights ) to attack the tail-latency problem - heavy items dispatch first, light items fill the tail.
  • :greedy , :dynamic , :static - passed straight to @threads . Note that :greedy and :dynamic create non-sticky tasks that may migrate across threads at yield points, so threadid() -indexed scratch is unsafe under those.
  • :serial - run the loop on the calling thread (useful for debugging).

If order is specified, it must be a permutation of the positions of indices - that is, a length- length(indices) vector whose values are a permutation of 1:length(indices) . The body still receives the corresponding indices[order[k]] value at each step k; only the visit order is altered. Typical usage is heaviest-work-first (so the longest items dispatch as early as possible): pass sortperm(weight_per_index; rev = true) . Under :greedy , :dynamic , :greedy_sticky , or :serial , the iteration visits the positions in order as given. Under :static , order is reshuffled in place in a round-robin fashion each contiguous chunk gets one item from each weight quantile, balancing the static partition. Mutates the caller's order array.

If weights is specified, it must be a length- length(indices) vector of non-negative integers giving the estimated work for each position in indices . When order is not given, weights computes it to be sortperm(weights; rev = true) so the heaviest items dispatch first. When progress is also given, each iteration advances the bar by the visited position's weights entry instead of by 1; the caller is responsible for sizing the progress total to sum(weights) so the percentage reflects work done rather than items done. weights and order are independent and may be combined: pass both to keep an explicit visit order while reporting work-weighted progress. Mutually exclusive with progress_chunk (the unweighted-throttling alternative).

If this is invoked from inside another parallel loop and nested is false (the default), policy is ignored and the loop is executed serially. This makes functions safe to compose: parallel at the top level, serial when called from inside another parallel loop. Pass nested = true to opt out of the safety demotion - the loop then uses policy as if at the top level. The caller is responsible for the scratch contract: under :static / :greedy_sticky each spawned sub-task has a stable threadid() , but multiple outer iterations may be inside their nested loops at the same time, so per-thread scratch indexed by inner threadid() must be local to the outer iteration (not a global per-thread array shared across all outer iterations).

The name is used for flame_timed for the whole loop. If the loop is parallel, the whole loop is a line in the serial flame file, and each iteration is a line in the parallel flame file. The serial file therefore gives a view of the elapsed time of the top-level loops, to identify what actually matters, and the parallel file shows the internal breakdown of the computations in each loop. It is sadly impossible to include threaded tasks in the same flamegraph in a meaningful way (at least not without extending the flamegraph format and visualization).

If progress is specified, it is updated for each iteration. If progress_chunk is specified, then indices must be a range 1:n , order must be nothing , and we throttle calls to the progress bar update to one every progress_chunk .

Note

The code inside the loop body should not use any random number generation. Use parallel_loop_with_rng if random number generation is needed.

using Random

size = 10

for policy in (:serial, :greedy, :dynamic, :static, :greedy_sticky)
    results = zeros(Int, size)
    parallel_loop_wo_rng(1:size; policy) do index
        results[index] = index
        return nothing
    end
    @assert results == collect(1:size)
end

# output


TanayLabUtilities.ParallelLoops.parallel_loop_with_rng Function
parallel_loop_with_rng(
    body::Function,
    indices::AbstractVector{<:Integer};
    name::AbstractString = ".loop",
    policy::Symbol = :greedy,
    order::Maybe{AbstractVector{<:Integer}} = nothing,
    weights::Maybe{AbstractVector{<:Integer}} = nothing,
    progress::Maybe{AbstractProgress} = nothing,
    seed::Maybe{Integer} = nothing,
    rng::Maybe{AbstractRNG} = nothing
)::Nothing

Run the body in parallel, passing it the iteration index and a separate rng that is seeded to a reproducible state regardless of the allocation of tasks to threads. A copy of this rng is given to each iteration, after being reset to seed + index for reproducibility. If no seed is specified, it is just sampled rng before the loop starts. If the rng isn't given, then this uses (and sets for each iteration) the default_rng() . In this case passing it to the body is redundant but is still done for consistency.

The policy is passed to @threads if it is one of (the default :greedy , :dynamic , or :static ). If it is :serial , then the loop is not run in parallel (useful for debugging).

If order and/or weights are specified, they are forwarded to parallel_loop_wo_rng - see there for the position-permutation semantics, the in-place round-robin shuffle under :static , and the work-weighted progress accounting. The body still receives the same indices[order[k]] value at each step k, and the seed is seed + indices[order[k]] so reproducibility is independent of the visit order.

If this is invoked from inside another parallel loop and nested is false (the default), policy is ignored and the loop is executed serially. This makes functions safe to compose: parallel at the top level, serial when called from inside another parallel loop. Pass nested = true to opt out of the safety demotion - the loop then uses policy as if at the top level. The caller is responsible for the scratch contract: under :static / :greedy_sticky each spawned sub-task has a stable threadid() , but multiple outer iterations may be inside their nested loops at the same time, so per-thread scratch indexed by inner threadid() must be local to the outer iteration (not a global per-thread array shared across all outer iterations).

The name is used for flame_timed for the whole loop. If the loop is parallel, the whole loop is a line in the serial flame file, and each iteration is a line in the parallel flame file. The serial file therefore gives a view of the elapsed time of the top-level loops, to identify what actually matters, and the parallel file shows the internal breakdown of the computations in each loop. It is sadly impossible to include threaded tasks in the same flamegraph in a meaningful way (at least not without extending the flamegraph format and visualization).

If progress is specified, it is updated for each iteration.

Note

Yes, the TaskLocalRNG is supposed to do this, but, it actually depends on the way tasks are allocated to threads. The implementation here will give the same results regardless of the thread scheduling policy. Sigh.

using Random

size = 10

function collect_rng(rng::AbstractRNG)::Vector{Float64}
    results = zeros(Float64, size)
    parallel_loop_with_rng(1:size; rng) do index, rng
        results[index] = rand(rng)
    end
    @assert results[1] != results[2]
    return results
end

@assert collect_rng(MersenneTwister(1)) == collect_rng(MersenneTwister(1))

function collect_default_rng()::Vector{Float64}
    results = zeros(Float64, size)
    parallel_loop_with_rng(1:size; seed = 123456, policy = :dynamic) do index, _
        results[index] = rand()
    end
    @assert results[1] != results[2]
    return results
end

@assert collect_default_rng() == collect_default_rng()

# output


TanayLabUtilities.ParallelLoops.DebugProgressUnknown Function
DebugProgressUnknown(; kwargs...)::Maybe{ProgressUnknown}

Same as ProgressUnknown in ProgressMeter , but returns nothing if debug is not enabled for the modules calling DebugProgressUnknown . Use this for progress bars where the total amount of work is not known in advance (for example, a planning phase that walks a data set to discover the set of properties to act on).

Index