Jungfraujoch/tests/ParallelForTest.cpp

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

#include <catch2/catch_all.hpp>

#include <algorithm>
#include <cmath>
#include <mutex>
#include <numeric>
#include <stdexcept>
#include <thread>
#include <vector>

#include "../common/CUDAWrapper.h"
#include "../common/ParallelFor.h"
#include "../common/ThreadAffinity.h"

namespace {
    const std::vector<size_t> THREAD_COUNTS = {0, 1, 2, 3, 8, 64};
    const std::vector<int> SIZES = {0, 1, 2, 7, 64, 1000};
}

// The split is fixed and contiguous, which is what lets a pass whose per-element work is independent
// give the serial answer bit for bit. Every element has to land in exactly one slice, the slices have
// to tile [0, n) in order, and there must never be an empty one.
TEST_CASE("ParallelChunks_SlicesTileTheRange", "[ParallelFor]") {
    for (int n: SIZES) {
        for (size_t nthreads: THREAD_COUNTS) {
            CAPTURE(n, nthreads);
            std::mutex m;
            std::vector<std::pair<int, int>> slices;
            std::vector<int> visits(static_cast<size_t>(std::max(n, 1)), 0);

            ParallelChunks(n, nthreads, [&](int lo, int hi) {
                {
                    std::scoped_lock lock(m);
                    slices.emplace_back(lo, hi);
                }
                for (int i = lo; i < hi; i++)
                    visits[static_cast<size_t>(i)]++;   // slices are disjoint, so no lock needed
            });

            std::sort(slices.begin(), slices.end());
            int expected_lo = 0;
            for (const auto &[lo, hi]: slices) {
                CHECK(lo == expected_lo);
                CHECK(hi > lo);                          // never an empty slice
                expected_lo = hi;
            }
            CHECK(expected_lo == n);                     // and they reach the end
            CHECK(slices.size() <= static_cast<size_t>(std::max(n, 0)));
            for (int i = 0; i < n; i++)
                CHECK(visits[static_cast<size_t>(i)] == 1);
        }
    }
}

// Work stealing, so a worker takes whatever is next rather than a fixed slice - but every item still
// has to be done exactly once, whatever order they come out in.
TEST_CASE("ParallelFor_VisitsEveryItemOnce", "[ParallelFor]") {
    for (int n: SIZES) {
        for (size_t nthreads: THREAD_COUNTS) {
            CAPTURE(n, nthreads);
            std::vector<std::atomic<int>> visits(static_cast<size_t>(std::max(n, 1)));
            for (auto &v: visits)
                v.store(0);

            ParallelFor(n, nthreads, [&](int i) { visits[static_cast<size_t>(i)].fetch_add(1); });

            for (int i = 0; i < n; i++)
                CHECK(visits[static_cast<size_t>(i)].load() == 1);
        }
    }
}

// One thread means the caller's loop, in order - the property the serial fallbacks rely on.
TEST_CASE("ParallelFor_IsSerialAndInOrderForOneThread", "[ParallelFor]") {
    for (size_t nthreads: {size_t{0}, size_t{1}}) {
        CAPTURE(nthreads);
        std::vector<int> order;
        ParallelFor(16, nthreads, [&](int i) { order.push_back(i); });
        std::vector<int> expected(16);
        std::iota(expected.begin(), expected.end(), 0);
        CHECK(order == expected);
    }
}

// Splitting the work must not change the answer. Each element is written by exactly one worker, so
// the result has to match the serial loop element for element, at every thread count.
TEST_CASE("ParallelFor_SplitDoesNotChangeTheResult", "[ParallelFor]") {
    constexpr int N = 1000;
    std::vector<double> serial(N);
    for (int i = 0; i < N; i++)
        serial[static_cast<size_t>(i)] = std::sin(i * 0.001) * 1e6 + i;

    for (size_t nthreads: THREAD_COUNTS) {
        CAPTURE(nthreads);
        std::vector<double> chunked(N, 0.0), stolen(N, 0.0);
        ParallelChunks(N, nthreads, [&](int lo, int hi) {
            for (int i = lo; i < hi; i++)
                chunked[static_cast<size_t>(i)] = std::sin(i * 0.001) * 1e6 + i;
        });
        ParallelFor(N, nthreads, [&](int i) {
            stolen[static_cast<size_t>(i)] = std::sin(i * 0.001) * 1e6 + i;
        });
        CHECK(chunked == serial);       // bit for bit, not approximately
        CHECK(stolen == serial);
    }
}

// The blocks depend on n alone: the same tiling of [0, n), block b always the same range, whatever the
// thread count.
TEST_CASE("ParallelBlocks_PartitionIgnoresTheThreadCount", "[ParallelFor]") {
    for (int n: {0, 1, 7, 4095, 4096, 100000, 5000000}) {
        const int nb = ReductionBlocks(n);
        CHECK(nb <= MAX_REDUCTION_BLOCKS);
        std::vector<std::pair<int, int>> reference;
        for (size_t nthreads: THREAD_COUNTS) {
            CAPTURE(n, nthreads);
            std::vector<std::pair<int, int>> range(static_cast<size_t>(nb), {-1, -1});
            ParallelBlocks(n, nthreads, [&](int b, int lo, int hi) { range[static_cast<size_t>(b)] = {lo, hi}; });
            int expected_lo = 0;
            for (const auto &[lo, hi]: range) {
                CHECK(lo == expected_lo);
                CHECK(hi > lo);
                expected_lo = hi;
            }
            CHECK(expected_lo == n);
            if (reference.empty()) reference = range;
            CHECK(range == reference);
        }
    }
}

// What the blocks are for: a floating-point sum folded per block and added up in block order has the
// same bits at any thread count. The terms span many decades, so any change of grouping would show.
TEST_CASE("ParallelBlocks_SumIsTheSameAtAnyThreadCount", "[ParallelFor]") {
    constexpr int N = 1000003;
    constexpr int NCELL = 37;
    std::vector<double> term(N);
    for (int i = 0; i < N; i++)
        term[static_cast<size_t>(i)] = std::exp(std::sin(i * 0.37) * 20.0) * (i % 3 == 0 ? -1.0 : 1.0);

    auto blocked_sums = [&](size_t nthreads) {
        const int nb = ReductionBlocks(N);
        std::vector<double> part(static_cast<size_t>(nb), 0.0);
        std::vector<std::vector<double>> cell_part(static_cast<size_t>(nb), std::vector<double>(NCELL, 0.0));
        ParallelBlocks(N, nthreads, [&](int b, int lo, int hi) {
            for (int i = lo; i < hi; i++) {
                part[static_cast<size_t>(b)] += term[static_cast<size_t>(i)];
                cell_part[static_cast<size_t>(b)][static_cast<size_t>(i % NCELL)] += term[static_cast<size_t>(i)];
            }
        });
        std::vector<double> out(NCELL + 1, 0.0);
        for (int b = 0; b < nb; b++) {
            out[NCELL] += part[static_cast<size_t>(b)];
            for (int c = 0; c < NCELL; c++)
                out[static_cast<size_t>(c)] += cell_part[static_cast<size_t>(b)][static_cast<size_t>(c)];
        }
        return out;
    };

    const std::vector<double> one = blocked_sums(1);
    for (size_t nthreads: {size_t{3}, size_t{7}, size_t{32}}) {
        CAPTURE(nthreads);
        CHECK(blocked_sums(nthreads) == one);   // bit for bit, not approximately
    }
}

// A negative or zero count is a no-op rather than an error - callers pass a computed size.
TEST_CASE("ParallelFor_DoesNothingForAnEmptyRange", "[ParallelFor]") {
    int calls = 0;
    for (int n: {0, -1, -1000}) {
        ParallelChunks(n, 8, [&](int, int) { calls++; });
        ParallelFor(n, 8, [&](int) { calls++; });
        ParallelBlocks(n, 8, [&](int, int, int) { calls++; });
    }
    CHECK(calls == 0);
}

// An exception thrown in a worker reaches the caller rather than terminating: the futures are waited
// on, so the other workers finish first and only then does it propagate.
TEST_CASE("ParallelFor_PropagatesAnException", "[ParallelFor]") {
    CHECK_THROWS_AS(ParallelChunks(64, 4, [](int lo, int) {
        if (lo == 0) throw std::runtime_error("from a chunk");
    }), std::runtime_error);

    CHECK_THROWS_AS(ParallelFor(64, 4, [](int i) {
        if (i == 0) throw std::runtime_error("from an item");
    }), std::runtime_error);
}

#ifdef __linux__
// The merging engine's ScopedGpuPin keeps its thread by the card only while it lives: a thread kept to
// one CPU before is kept to that CPU after, and the pool's workers keep the process's CPUs while a
// narrowed thread hands them work.
TEST_CASE("ScopedGpuPin_GivesTheThreadItsCpusBack", "[ParallelFor]") {
    const std::vector<int> all = GetThreadCpus();
    REQUIRE(all.size() >= 2);
    std::vector<int> after_pin;
    std::vector<std::vector<int>> seen_by_task(4);
    std::thread([&] {
        SetThreadCpus({all.front()});
        {
            const ScopedGpuPin pin(0);
            ParallelChunks(4, 4, [&](int lo, int) { seen_by_task[lo] = GetThreadCpus(); });
        }
        after_pin = GetThreadCpus();
    }).join();
    CHECK(after_pin == std::vector<int>{all.front()});
    for (int t = 1; t < 4; t++)
        CHECK(seen_by_task[t] == all);
}
#endif

namespace {
    std::vector<int> CpuRange(int lo, int hi) {   // [lo, hi]
        std::vector<int> v(hi - lo + 1);
        std::iota(v.begin(), v.end(), lo);
        return v;
    }

    std::vector<int> Concat(std::vector<int> a, const std::vector<int> &b) {
        a.insert(a.end(), b.begin(), b.end());
        return a;
    }
}

// Synthetic machines: which CPUs a GPU run keeps to, given the CPUs it may use, the package of every
// CPU and the package of every card.
TEST_CASE("CpusOfGpuPackages_KeepsToTheCardsPackages", "[ParallelFor]") {
    // Four packages of 12 CPUs, four cards on packages 2 and 3.
    std::vector<int> four(48);
    for (int c = 0; c < 48; c++)
        four[c] = c / 12;
    CHECK(CpusOfGpuPackages(CpuRange(0, 47), four, {2, 2, 3, 3}) == CpuRange(24, 47));

    // Two packages of 12 cores with hyper-threading, numbered as Linux does: CPUs 0-11 and 24-35 on
    // package 0, 12-23 and 36-47 on package 1.
    std::vector<int> two(48);
    for (int c = 0; c < 48; c++)
        two[c] = (c / 12) % 2;
    const std::vector<int> package1 = Concat(CpuRange(12, 23), CpuRange(36, 47));
    CHECK(CpusOfGpuPackages(CpuRange(0, 47), two, {1}) == package1);
    CHECK(CpusOfGpuPackages(CpuRange(0, 47), two, {1, 1}) == package1);

    // Cards on both packages: nothing to keep to.
    CHECK(CpusOfGpuPackages(CpuRange(0, 47), two, {0, 1}).empty());

    // No card (a CPU build, or none visible), or a card whose package is not known.
    CHECK(CpusOfGpuPackages(CpuRange(0, 47), two, {}).empty());
    CHECK(CpusOfGpuPackages(CpuRange(0, 47), two, {1, -1}).empty());

    // One package, as on a desktop.
    CHECK(CpusOfGpuPackages(CpuRange(0, 31), std::vector<int>(32, 0), {0}).empty());
}

// A process started on fewer CPUs (taskset, numactl, a cpuset) keeps within them.
TEST_CASE("CpusOfGpuPackages_RespectsARestrictedMask", "[ParallelFor]") {
    std::vector<int> four(48);
    for (int c = 0; c < 48; c++)
        four[c] = c / 12;
    // Already on one package - with the cards or without them - stays where it is.
    CHECK(CpusOfGpuPackages(CpuRange(0, 11), four, {2, 3}).empty());
    CHECK(CpusOfGpuPackages(CpuRange(24, 35), four, {2, 3}).empty());
    // Already on the cards' packages only.
    CHECK(CpusOfGpuPackages(CpuRange(24, 47), four, {2, 3}).empty());
    // Over packages 1 and 2, cards on 2 and 3: the part on package 2.
    CHECK(CpusOfGpuPackages(CpuRange(18, 29), four, {2, 3}) == CpuRange(24, 29));
}

// Without a card nothing is confined, so the process keeps the CPUs it started with.
TEST_CASE("ConfineToGpuPackages_DoesNothingWithoutACard", "[ParallelFor]") {
    const unsigned before = AvailableCpus();
    const std::vector<int> cpus = GetThreadCpus();
    CHECK(ConfineToGpuPackages({}).empty());
    CHECK(AvailableCpus() == before);
    CHECK(GetThreadCpus() == cpus);
}