agent
40a3535393
ParallelFor runs a counted loop across a fixed set of threads created up front and hands out indices from an atomic counter, so the calling thread shares the work. A thread-start hook runs once per helper thread.
74 lines
1.5 KiB
C++
74 lines
1.5 KiB
C++
#include "ParallelFor.hpp"
|
|
|
|
namespace tp {
|
|
|
|
ParallelFor::ParallelFor(size_t threads, ThreadStart onStart) :
|
|
mOnStart(std::move(onStart)) {
|
|
for (size_t index = 1; index < threads; index++)
|
|
mThreads.emplace_back([this, index]() {
|
|
if (mOnStart) mOnStart(index);
|
|
work(0);
|
|
});
|
|
}
|
|
|
|
ParallelFor::~ParallelFor() {
|
|
{
|
|
std::lock_guard lock(mMutex);
|
|
mStopping = true;
|
|
}
|
|
mStart.notify_all();
|
|
for (std::thread& thread : mThreads)
|
|
thread.join();
|
|
}
|
|
|
|
void ParallelFor::drain() {
|
|
while (true) {
|
|
const size_t index = mNext.fetch_add(1, std::memory_order_relaxed);
|
|
if (index >= mCount) return;
|
|
(*mBody)(index);
|
|
}
|
|
}
|
|
|
|
void ParallelFor::work(size_t seen) {
|
|
while (true) {
|
|
{
|
|
std::unique_lock lock(mMutex);
|
|
mStart.wait(lock, [&]() { return mStopping || (mGeneration != seen && mBody); });
|
|
if (mStopping) return;
|
|
seen = mGeneration;
|
|
mBusy++;
|
|
}
|
|
drain();
|
|
bool last = false;
|
|
{
|
|
std::lock_guard lock(mMutex);
|
|
last = --mBusy == 0;
|
|
}
|
|
if (last) mDone.notify_one();
|
|
}
|
|
}
|
|
|
|
void ParallelFor::run(size_t count, const std::function<void(size_t)>& body) {
|
|
if (count == 0) return;
|
|
if (mThreads.empty() || count == 1) {
|
|
for (size_t index = 0; index < count; index++)
|
|
body(index);
|
|
return;
|
|
}
|
|
|
|
{
|
|
std::lock_guard lock(mMutex);
|
|
mBody = &body;
|
|
mCount = count;
|
|
mNext.store(0, std::memory_order_relaxed);
|
|
mGeneration++;
|
|
}
|
|
mStart.notify_all();
|
|
drain();
|
|
|
|
std::unique_lock lock(mMutex);
|
|
mDone.wait(lock, [this]() { return mBusy == 0; });
|
|
mBody = nullptr;
|
|
mCount = 0;
|
|
}
|
|
}
|