Jungfraujoch/common/ParallelFor.h

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#pragma once

#include <algorithm>
#include <atomic>
#include <condition_variable>
#include <deque>
#include <exception>
#include <functional>
#include <mutex>
#include <thread>
#include <vector>

// Two shapes of "run this over a range on several threads", used by the analysis code. Both take the
// worker count from the caller rather than asking the hardware, so a run that was told how many
// threads to use keeps to it.

// How many workers a pass over `n` cheap items should use: enough that each gets at least
// `min_per_thread` of them, and never more than the caller was given. A pass whose data is small can
// otherwise spend more on splitting the work than on doing it. That is not a large machine's problem:
// it is what makes the same code right on an 8-core laptop, a 16-core desktop and a two-socket node,
// none of which should be handed 48 chunks of a few thousand items.
inline size_t ThreadsForWork(size_t n, size_t nthreads, size_t min_per_thread = 32768) {
    if (nthreads <= 1 || n == 0) return 1;
    return std::clamp<size_t>(n / min_per_thread, 1, nthreads);
}

namespace parallel_detail {
    // The threads both helpers below run on. They are made once and kept, because the analysis code
    // repeats some of its passes thousands of times in a run and a thread costs tens of microseconds
    // to create and join - more, on a short pass, than the pass itself.
    class WorkerPool {
    public:
        static WorkerPool &Instance() {
            static WorkerPool pool;
            return pool;
        }

        // True on a thread the pool owns. A parallel pass reached from inside one runs inline instead
        // of queueing: the workers are already occupied by the outer pass, so waiting for one of them
        // to pick up the inner work could wait forever.
        static bool InWorker() { return in_worker; }

        size_t WorkerCount() const { return workers.size(); }

        void Submit(std::function<void()> job) {
            {
                std::lock_guard lock(m);
                queue.push_back(std::move(job));
            }
            cv.notify_one();
        }

    private:
        WorkerPool() {
            const unsigned hw = std::max(1u, std::thread::hardware_concurrency());
            workers.reserve(hw - 1);
            for (unsigned i = 0; i + 1 < hw; i++)      // the submitting thread takes a share too
                workers.emplace_back([this] { in_worker = true; Loop(); });
        }

        ~WorkerPool() {
            {
                std::lock_guard lock(m);
                stop = true;
            }
            cv.notify_all();
            for (auto &t: workers) t.join();
        }

        void Loop() {
            for (;;) {
                std::function<void()> job;
                {
                    std::unique_lock lock(m);
                    cv.wait(lock, [this] { return stop || !queue.empty(); });
                    if (stop) return;
                    job = std::move(queue.front());
                    queue.pop_front();
                }
                job();
            }
        }

        std::mutex m;
        std::condition_variable cv;
        std::deque<std::function<void()> > queue;
        std::vector<std::thread> workers;
        bool stop = false;
        static thread_local bool in_worker;
    };

    inline thread_local bool WorkerPool::in_worker = false;

    // What the tasks of one pass share: the body to call, how many of them are still outstanding, and
    // the first exception any of them threw.
    struct RunState {
        const std::function<void(int)> *body = nullptr;
        std::atomic<int> remaining{0};
        std::mutex done_m;
        std::condition_variable done_cv;
        bool done = false;              // guarded by done_m; see RunOneTask
        std::mutex err_m;
        std::exception_ptr error;
    };

    inline void RunOneTask(RunState &s, int t) {
        try {
            (*s.body)(t);
        } catch (...) {
            std::lock_guard lock(s.err_m);
            if (!s.error) s.error = std::current_exception();
        }
        if (s.remaining.fetch_sub(1, std::memory_order_acq_rel) == 1) {
            // The flag the waiter tests is set UNDER done_m, and the counter is not that flag. If the
            // waiter watched the counter it could see zero the instant the decrement above lands -
            // before this thread has taken the lock - find its predicate already true, never block,
            // and return from RunTasks. RunState is a local of that frame, so the lock and the notify
            // below would then run on a destroyed mutex and condition variable, writing pthread state
            // into a stack frame the submitting thread has already reused. Watching a flag set under
            // the lock means completion cannot be observed until this thread has released it.
            std::lock_guard lock(s.done_m);
            s.done = true;
            s.done_cv.notify_all();
        }
    }

    // Call body(t) for every t in [0, ntasks) on the pool and on this thread, and return once they have
    // all finished. An exception from any of them is held until then and rethrown here, so the others
    // still run to completion - which is what waiting on futures used to give.
    inline void RunTasks(int ntasks, const std::function<void(int)> &body) {
        if (ntasks <= 0) return;
        WorkerPool &pool = WorkerPool::Instance();
        if (ntasks == 1 || WorkerPool::InWorker() || pool.WorkerCount() == 0) {
            for (int t = 0; t < ntasks; t++) body(t);
            return;
        }
        RunState s;
        s.body = &body;
        s.remaining.store(ntasks, std::memory_order_relaxed);
        RunState *sp = &s;
        for (int t = 1; t < ntasks; t++)
            pool.Submit([sp, t] { RunOneTask(*sp, t); });
        RunOneTask(s, 0);
        {
            std::unique_lock lock(s.done_m);
            s.done_cv.wait(lock, [sp] { return sp->done; });
        }
        if (s.error) std::rethrow_exception(s.error);
    }
}

// Chunked: each worker gets one contiguous [lo, hi) range and there is no per-item synchronisation.
// Right for millions of cheap uniform items - the CPU stand-in for a flat CUDA grid-stride kernel.
// The split is fixed and deterministic, so a pass whose per-element work is independent gives the
// same answer as the serial loop, bit for bit.
template <typename Fn>
void ParallelChunks(int n, size_t nthreads, Fn fn) {
    if (n <= 0) return;
    const int nt = static_cast<int>(std::max<size_t>(1, std::min(nthreads, static_cast<size_t>(n))));
    if (nt == 1) { fn(0, n); return; }
    const int chunk = (n + nt - 1) / nt;
    parallel_detail::RunTasks(nt, [&](int t) {
        const int lo = t * chunk, hi = std::min(n, lo + chunk);
        if (lo < hi) fn(lo, hi);
    });
}

// Work-stealing per item, off a shared atomic counter: one atomic per item, so use it only where the
// per-item work is heavy and uneven (per-frame fits, per-ring selections) and the atomic amortises.
// For millions of tiny uniform items a per-item atomic is pure contention - use ParallelChunks.
template <typename Fn>
void ParallelFor(int n, size_t nthreads, Fn fn) {
    if (n <= 0) return;
    if (nthreads <= 1 || n == 1) {
        for (int i = 0; i < n; i++) fn(i);
        return;
    }
    const size_t local = std::min(nthreads, static_cast<size_t>(n));
    std::atomic<int> next = 0;
    parallel_detail::RunTasks(static_cast<int>(local), [&](int) {
        for (int i = next.fetch_add(1); i < n; i = next.fetch_add(1)) fn(i);
    });
}