mx/Jungfraujoch
2
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 68 Wiki Activity

Run the parallel helpers on a pool instead of starting threads per chunk

Browse Source 
ParallelFor and ParallelChunks started fresh OS threads on every call, one per chunk,
through std::async. There are 57 call sites and several sit inside iterative fits, so
one run of the heaviest crystal created 13 497 threads and a 900-frame dataset 7 320.
Both now run on a persistent pool. The contracts are unchanged: ParallelChunks keeps
the same worker count and the same fixed split, so a reduction sums term for term as
before, and ParallelFor keeps stealing per item.

Two things in the pool are worth knowing. The caller is one of the hands - it claims
its own region's tasks and then waits only on tasks already running - so a region
entered from inside another region cannot deadlock at any depth, which a shared-queue
pool would. And a task wakes one worker rather than the whole pool: on a large machine
notify_all wakes every idle thread to find nothing, once per region, tens of thousands
of times a run.

Two hand-rolled copies of the same pattern now use it, in FrenchWilson and in the two
histogram passes of ComputeAsuGroups.

Be clear about what this buys today: nothing measurable. Thread creations drop 13 497
-> 770 and entering a parallel region goes from 1.2-2.2 ms to 112 us, an 11-20x cut,
but wall clock on 48 threads is level with before, inside the +-5 % this machine's
run-to-run placement is worth. What it removes is a cost that grows with the thread
count - measured, entry is linear in it - and the machine this is heading for has four
times the threads of the one it was measured on, where the same 335 regions a run
would cost about 1.8 s of pure thread creation.

ComputeAsuGroups' histogram also changes. It is an nthreads x n_groups table, 936 MB
at -N 48 on the heaviest crystal and allocated five times a run, and the prefix over
it walked DOWN a column - a 19.5 MB stride, so a cache and TLB miss per step, 234 M of
them, serially. Both passes now walk rows and split over group ranges. The counts are
integers, so the result is bit-identical. This one is reasoning, not measurement: at 48
threads it sits under this machine's noise and could not be shown either way.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
jungfrau
2026-08-17 18:05:35 -04:00
co-authored by Claude Opus 5
parent 65eb86ab66
commit a22a51372e
3 changed files with 187 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/ParallelFor.h
+155 -22
View File
@@ -5,7 +5,11 @@
#include <algorithm>
#include <atomic>
#include <future>
#include <condition_variable>
#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
@@ -13,16 +17,151 @@
// 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. Both helpers below start a
// thread per chunk, which costs tens of microseconds, so a pass that is repeated - or one whose data
// is small - can easily 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.
// `min_per_thread` of them, and never more than the caller was given. Entering a parallel region below
// costs a condition-variable wakeup rather than a thread, but a pass over a few thousand items still
// finishes before the last worker has woken up, so the floor stays. 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 {
// One parallel region: `ntasks` invocations of one body, each claimed and run by exactly one thread.
struct Region {
const std::function<void(int)> *body;
int ntasks;
int next = 0; // next unclaimed task, guarded by the pool's mutex
std::atomic<int> done{0};
std::exception_ptr error; // the first task to throw, guarded by the pool's mutex
};
// The threads both helpers below run on: started once, parked on a condition variable between regions.
// The analysis code enters tens of thousands of parallel regions in a run - several from inside
// iterative fits - and a thread per chunk per call was tens of microseconds of clone and first-touch
// each time, paid again on every iteration.
class Pool {
public:
static Pool &Instance() {
static Pool pool;
return pool;
}
// Run body(0) ... body(ntasks-1), each exactly once, on the pool's threads and on this one, and
// return when they have all finished. An exception from a task is rethrown here, as it was when
// these were futures.
void Run(int ntasks, const std::function<void(int)> &body) {
Region region{&body, ntasks};
{
std::lock_guard lock(mutex_);
regions_.push_back(&region);
}
// One wakeup per task the caller is not going to run itself. notify_all instead would wake every
// parked thread for a region of two chunks - on a node with 192 of them that is 190 threads woken
// to find nothing and go back to sleep, per region, tens of thousands of times in a run.
for (int i = 1; i < ntasks; i++)
work_.notify_one();
// The calling thread is one of the hands and keeps claiming until nothing of this region is
// left unclaimed. That is what makes a region entered from inside a region safe: a region never
// waits for a pool thread to come free, only for tasks that are already running.
for (int i = Claim(region); i >= 0; i = Claim(region))
Execute(region, i);
std::unique_lock lock(mutex_);
done_.wait(lock, [&] { return region.done.load() == ntasks; });
if (region.error)
std::rethrow_exception(region.error);
}
private:
Pool() {
const unsigned n = std::max(1u, std::thread::hardware_concurrency());
threads_.reserve(n);
for (unsigned i = 0; i < n; i++)
threads_.emplace_back([this] { Work(); });
}
~Pool() {
{
std::lock_guard lock(mutex_);
stop_ = true;
}
work_.notify_all();
for (auto &t : threads_)
t.join();
}
// Claim one of `region`'s tasks, or -1 if it has none left. A region leaves the list the moment its
// last task is claimed, so no pool thread can look at it after Run returns and it goes out of scope.
int Claim(Region &region) {
std::lock_guard lock(mutex_);
if (region.next == region.ntasks)
return -1;
const int i = region.next++;
if (region.next == region.ntasks)
regions_.erase(std::find(regions_.begin(), regions_.end(), &region));
return i;
}
// Claim one task of the newest region that still has one, called with the mutex held. Newest first,
// so an inner region is finished before more hands go to the outer one waiting on it.
Region *ClaimAny(int &i) {
for (size_t at = regions_.size(); at-- > 0;) {
Region *region = regions_[at];
if (region->next < region->ntasks) {
i = region->next++;
if (region->next == region->ntasks)
regions_.erase(regions_.begin() + at);
return region;
}
}
return nullptr;
}
// Run one task and count it. The last one to finish wakes whoever is waiting in Run; it reads ntasks
// before the count, because after the count the region may already be gone.
void Execute(Region &region, int i) {
try {
(*region.body)(i);
} catch (...) {
std::lock_guard lock(mutex_);
if (!region.error)
region.error = std::current_exception();
}
const int ntasks = region.ntasks;
if (region.done.fetch_add(1) + 1 == ntasks) {
std::lock_guard lock(mutex_);
done_.notify_all();
}
}
void Work() {
std::unique_lock lock(mutex_);
for (;;) {
int i = 0;
Region *region = ClaimAny(i);
if (!region) {
if (stop_)
return;
work_.wait(lock);
continue;
}
lock.unlock();
Execute(*region, i);
lock.lock();
}
}
std::mutex mutex_;
std::condition_variable work_, done_;
std::vector<Region *> regions_;
std::vector<std::thread> threads_;
bool stop_ = false;
};
} // namespace parallel_detail
// 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
@@ -33,14 +172,11 @@ void ParallelChunks(int n, size_t nthreads, Fn fn) {
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;
std::vector<std::future<void>> futures;
futures.reserve(nt);
for (int t = 0; t < nt; t++) {
const std::function<void(int)> body = [&fn, n, chunk](int t) {
const int lo = t * chunk, hi = std::min(n, lo + chunk);
if (lo >= hi) break;
futures.emplace_back(std::async(std::launch::async, [&fn, lo, hi] { fn(lo, hi); }));
}
for (auto &f : futures) f.get();
if (lo < hi) fn(lo, hi);
};
parallel_detail::Pool::Instance().Run(nt, body);
}
// Work-stealing per item, off a shared atomic counter: one atomic per item, so use it only where the
@@ -53,14 +189,11 @@ void ParallelFor(int n, size_t nthreads, Fn fn) {
for (int i = 0; i < n; i++) fn(i);
return;
}
const size_t local = std::min(nthreads, static_cast<size_t>(n));
const int nt = static_cast<int>(std::min(nthreads, static_cast<size_t>(n)));
std::atomic<int> next = 0;
std::vector<std::future<void>> futures;
futures.reserve(local);
for (size_t t = 0; t < local; t++)
futures.emplace_back(std::async(std::launch::async, [&] {
for (int i = next.fetch_add(1); i < n; i = next.fetch_add(1))
fn(i);
}));
for (auto &f : futures) f.get();
const std::function<void(int)> body = [&fn, &next, n](int) {
for (int i = next.fetch_add(1); i < n; i = next.fetch_add(1))
fn(i);
};
parallel_detail::Pool::Instance().Run(nt, body);
}
image_analysis/scale_merge/FrenchWilson.cpp
+3 -18
View File
@@ -5,10 +5,10 @@
#include <algorithm>
#include <cmath>
#include <future>
#include <limits>
#include <vector>
#include "../../common/ParallelFor.h"
#include "../../common/ResolutionShells.h"
#include "gemmi/symmetry.hpp"
@@ -146,10 +146,7 @@ void ApplyFrenchWilson(std::vector<MergedReflection> &merged, int32_t space_grou
};
// Each reflection's amplitudes depend only on itself and the shell priors above, so the loop is
// data-parallel over contiguous chunks and gives the same result whatever the worker count.
const int n = static_cast<int>(merged.size());
const int nt = std::clamp(opts.num_threads, 1, n);
const int chunk = (n + nt - 1) / nt;
auto do_chunk = [&](int lo, int hi) {
ParallelChunks(static_cast<int>(merged.size()), std::max(1, opts.num_threads), [&](int lo, int hi) {
std::vector<double> logw(opts.integration_points);
for (int i = lo; i < hi; ++i) {
MergedReflection &r = merged[i];
@@ -157,17 +154,5 @@ void ApplyFrenchWilson(std::vector<MergedReflection> &merged, int32_t space_grou
fw_one(r, r.I_plus, r.sigma_plus, r.F_plus, r.sigmaF_plus, logw);
fw_one(r, r.I_minus, r.sigma_minus, r.F_minus, r.sigmaF_minus, logw);
}
};
if (nt == 1) {
do_chunk(0, n);
return;
}
std::vector<std::future<void>> futures;
futures.reserve(nt);
for (int t = 0; t < nt; ++t) {
const int lo = t * chunk, hi = std::min(n, lo + chunk);
if (lo >= hi) break;
futures.emplace_back(std::async(std::launch::async, [&do_chunk, lo, hi] { do_chunk(lo, hi); }));
}
for (auto &f : futures) f.get();
});
}
image_analysis/scale_merge/RotationScaleMerge.cpp
+29 -23
View File
@@ -1013,40 +1013,46 @@ int RotationScaleMerge::ComputeAsuGroups(const HKLKeyGenerator &keygen) {
std::vector<std::vector<int32_t>> hist(nt, std::vector<int32_t>(n_groups, 0));
// Pass 1 (parallel): per-chunk group histogram over the flat group_ids.
std::vector<std::future<void>> f1;
for (int t = 0; t < nt; ++t) {
const int lo = t * chunk, hi = std::min(n, lo + chunk);
if (lo >= hi) break;
f1.emplace_back(std::async(std::launch::async, [&, t, lo, hi] {
ParallelChunks(nt, nthreads, [&](int tlo, int thi) {
for (int t = tlo; t < thi; ++t) {
auto &h = hist[t];
const int lo = t * chunk, hi = std::min(n, lo + chunk);
for (int i = lo; i < hi; ++i) { const int g = group_ids[i]; if (g >= 0) ++h[g]; }
}));
}
for (auto &f : f1) f.get();
}
});
// CSR starts + convert hist[t][g] into chunk t's write base for group g (exclusive prefix over t).
// Both passes over the nt x n_groups table run a range of groups at a time with t OUTERMOST, so
// each step is the next entry of a row. Taking one group all the way down its column instead - as
// this did - strides n_groups apart on every step and misses cache on all of them, which is a cost
// that grows with the thread count: 234 million such steps at 48 threads on the largest crystal
// here, twice that on a node with 96. Integer counts, so the split changes no result.
std::vector<int32_t> gstart(n_groups), gcount(n_groups);
ParallelChunks(n_groups, nthreads, [&](int glo, int ghi) {
for (int t = 0; t < nt; ++t)
for (int g = glo; g < ghi; ++g) gcount[g] += hist[t][g];
});
int acc = 0;
for (int g = 0; g < n_groups; ++g) {
int base = acc;
gstart[g] = acc;
for (int t = 0; t < nt; ++t) { const int c = hist[t][g]; hist[t][g] = base; base += c; }
gcount[g] = base - acc;
acc = base;
}
for (int g = 0; g < n_groups; ++g) { gstart[g] = acc; acc += gcount[g]; }
ParallelChunks(n_groups, nthreads, [&](int glo, int ghi) {
std::vector<int32_t> base(gstart.begin() + glo, gstart.begin() + ghi);
for (int t = 0; t < nt; ++t)
for (int g = glo; g < ghi; ++g) {
const int c = hist[t][g];
hist[t][g] = base[g - glo];
base[g - glo] += c;
}
});
// Pass 2 (parallel): each chunk fills its obs into gperm at its per-group base (stable).
std::vector<int32_t> gperm(acc);
std::vector<std::future<void>> f2;
for (int t = 0; t < nt; ++t) {
const int lo = t * chunk, hi = std::min(n, lo + chunk);
if (lo >= hi) break;
f2.emplace_back(std::async(std::launch::async, [&, t, lo, hi] {
ParallelChunks(nt, nthreads, [&](int tlo, int thi) {
for (int t = tlo; t < thi; ++t) {
std::vector<int32_t> fill = hist[t];
const int lo = t * chunk, hi = std::min(n, lo + chunk);
for (int i = lo; i < hi; ++i) { const int g = group_ids[i]; if (g >= 0) gperm[fill[g]++] = i; }
}));
}
for (auto &f : f2) f.get();
}
});
gpu_->SetGroups(n_groups, group_ids.data(), gperm.data(), acc, gstart.data(), gcount.data());
}
#endif