thread.hpp

Go to the documentation of this file.

#pragma once
#include "barretenberg/common/compiler_hints.hpp"
#include <atomic>
#include <barretenberg/env/hardware_concurrency.hpp>
#include <barretenberg/numeric/bitop/get_msb.hpp>
#include <functional>
#include <iostream>
#include <ranges>
#include <vector>
 
namespace bb {
#ifdef __wasm__
// Fixed number of workers in WASM environment
constexpr size_t PARALLEL_FOR_MAX_NESTING = 1;
#else
constexpr size_t PARALLEL_FOR_MAX_NESTING = 2;
#endif
 
// Useful for programatically benching different thread counts
// Note this is threadsafe and affects parallel_for's just in that thread if so.
void set_parallel_for_concurrency(size_t num_cores);
size_t get_num_cpus();
 
// For algorithms that need to be divided amongst power of 2 threads.
inline size_t get_num_cpus_pow2()
{
    return static_cast<size_t>(1ULL << numeric::get_msb(get_num_cpus()));
}
 
void parallel_for(size_t num_iterations, const std::function<void(size_t)>& func);
void parallel_for_range(size_t num_points,
                        const std::function<void(size_t, size_t)>& func,
                        size_t no_multhreading_if_less_or_equal = 0);
 
void parallel_for_heuristic(size_t num_points,
                            const std::function<void(size_t, size_t, size_t)>& func,
                            size_t heuristic_cost);
 
template <typename Func>
    requires std::invocable<Func, std::size_t>
void parallel_for_heuristic(size_t num_points, const Func& func, size_t heuristic_cost)
{
    parallel_for_heuristic(
        num_points,
        [&](size_t start_idx, size_t end_idx, BB_UNUSED size_t chunk_index) {
            for (size_t i = start_idx; i < end_idx; i++) {
                func(i);
            }
        },
        heuristic_cost);
}
 
template <typename Func, typename Accum>
    requires std::invocable<Func, std::size_t, Accum&>
std::vector<Accum> parallel_for_heuristic(size_t num_points,
                                          const Accum& initial_accum,
                                          const Func& func,
                                          size_t heuristic_cost)
{
    // thread-safe accumulators
    std::vector<Accum> accumulators(get_num_cpus(), initial_accum);
    parallel_for_heuristic(
        num_points,
        [&](size_t start_idx, size_t end_idx, size_t chunk_index) {
            for (size_t i = start_idx; i < end_idx; i++) {
                func(i, accumulators[chunk_index]);
            }
        },
        heuristic_cost);
    return accumulators;
}
 
const size_t DEFAULT_MIN_ITERS_PER_THREAD = 1 << 4;
 
struct MultithreadData {
    size_t num_threads;
    // index bounds for each thread
    std::vector<size_t> start;
    std::vector<size_t> end;
};
 
MultithreadData calculate_thread_data(size_t num_iterations,
                                      size_t min_iterations_per_thread = DEFAULT_MIN_ITERS_PER_THREAD);
 
size_t calculate_num_threads(size_t num_iterations, size_t min_iterations_per_thread = DEFAULT_MIN_ITERS_PER_THREAD);
 
namespace thread_heuristics {
// Maximum observed parallel_for overhead in nanoseconds (rounded up from 388us measurement)
constexpr size_t PARALLEL_FOR_COST = 400000;
// Rough cost of operations (the operation costs are derives in basics_bench and the units are nanoseconds)
// Field element (16 byte) addition cost
constexpr size_t FF_ADDITION_COST = 4;
// Field element (16 byte) multiplication cost
constexpr size_t FF_MULTIPLICATION_COST = 21;
// Field element (16 byte) inversion cost
constexpr size_t FF_INVERSION_COST = 7000;
// Group element projective addition number
constexpr size_t GE_ADDITION_COST = 350;
// Group element projective doubling number
constexpr size_t GE_DOUBLING_COST = 194;
// Group element scalar multiplication cost
constexpr size_t SM_COST = 50000;
// Field element (16 byte) sequential copy number
constexpr size_t FF_COPY_COST = 3;
// Fine default if something looks 'chunky enough that I don't want to calculate'
constexpr size_t ALWAYS_MULTITHREAD = 100000;
} // namespace thread_heuristics
 
struct ThreadChunk {
    size_t thread_index;
    size_t total_threads;
    auto range(size_t size, size_t offset = 0) const
    {
        if (total_threads == 0 || thread_index >= total_threads) {
            return std::views::iota(size_t{ 0 }, size_t{ 0 });
        }
        // Calculate base chunk size and remainder
        size_t chunk_size = size / total_threads;
        size_t remainder = size % total_threads;
 
        if (thread_index < remainder) {
            // Threads with index < remainder get chunk_size + 1 elements
            size_t start = thread_index * (chunk_size + 1);
            size_t end = start + chunk_size + 1;
            return std::views::iota(start + offset, end + offset);
        }
        // Threads with index >= remainder get chunk_size elements
        size_t start = remainder * (chunk_size + 1) + (thread_index - remainder) * chunk_size;
        size_t end = start + chunk_size;
        return std::views::iota(start + offset, end + offset);
    }
};
 
template <typename Func>
    requires std::invocable<Func, ThreadChunk>
void parallel_for(const Func& func)
{
    size_t total_threads = get_num_cpus();
    parallel_for(total_threads, [&](size_t thread_index) {
        func(ThreadChunk{ .thread_index = thread_index, .total_threads = total_threads });
    });
}
 
// Overload that allows specifying the number of threads explicitly while still using ThreadChunk
template <typename Func>
    requires std::invocable<Func, ThreadChunk>
void parallel_for(size_t num_threads, const Func& func)
{
    parallel_for(num_threads, [&](size_t thread_index) {
        func(ThreadChunk{ .thread_index = thread_index, .total_threads = num_threads });
    });
}
 
// parallel_for_heuristic variant that uses ThreadChunk for work distribution.
// Parallelizes only when the estimated total work exceeds the parallel_for overhead.
template <typename Func>
    requires std::invocable<Func, ThreadChunk>
void parallel_for_heuristic(size_t num_points, const Func& func, size_t heuristic_cost)
{
    const size_t num_cpus = get_num_cpus();
    const size_t chunk_size = (num_points / num_cpus) + (num_points % num_cpus == 0 ? 0 : 1);
    const size_t offset_cost = (num_points - chunk_size) * heuristic_cost;
 
    if (offset_cost < thread_heuristics::PARALLEL_FOR_COST) {
        func(ThreadChunk{ .thread_index = 0, .total_threads = 1 });
        return;
    }
    parallel_for(num_cpus, [&](size_t thread_index) {
        func(ThreadChunk{ .thread_index = thread_index, .total_threads = num_cpus });
    });
}
 
} // namespace bb