File size: 2,372 Bytes
7873319 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
/*
* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
*
* NVIDIA CORPORATION and its licensors retain all intellectual property
* and proprietary rights in and to this software, related documentation
* and any modifications thereto. Any use, reproduction, disclosure or
* distribution of this software and related documentation without an express
* license agreement from NVIDIA CORPORATION is strictly prohibited.
*/
// This file was taken from the tev image viewer and is re-released here
// under the NVIDIA Source Code License with permission from the author.
#include <neural-graphics-primitives/common.h>
#include <neural-graphics-primitives/thread_pool.h>
#include <chrono>
NGP_NAMESPACE_BEGIN
using namespace std;
ThreadPool::ThreadPool()
: ThreadPool{thread::hardware_concurrency()} {}
ThreadPool::ThreadPool(size_t maxNumThreads, bool force) {
if (!force) {
maxNumThreads = min((size_t)thread::hardware_concurrency(), maxNumThreads);
}
startThreads(maxNumThreads);
}
ThreadPool::~ThreadPool() {
shutdownThreads(mThreads.size());
}
void ThreadPool::startThreads(size_t num) {
mNumThreads += num;
for (size_t i = mThreads.size(); i < mNumThreads; ++i) {
mThreads.emplace_back([this, i] {
while (true) {
unique_lock<mutex> lock{mTaskQueueMutex};
// look for a work item
while (i < mNumThreads && mTaskQueue.empty()) {
// if there are none wait for notification
mWorkerCondition.wait(lock);
}
if (i >= mNumThreads) {
break;
}
function<void()> task{move(mTaskQueue.front())};
mTaskQueue.pop_front();
// Unlock the lock, so we can process the task without blocking other threads
lock.unlock();
task();
}
});
}
}
void ThreadPool::shutdownThreads(size_t num) {
auto numToClose = min(num, mNumThreads);
{
lock_guard<mutex> lock{mTaskQueueMutex};
mNumThreads -= numToClose;
}
// Wake up all the threads to have them quit
mWorkerCondition.notify_all();
for (auto i = 0u; i < numToClose; ++i) {
mThreads.back().join();
mThreads.pop_back();
}
}
void ThreadPool::setNThreads(size_t num) {
if (mNumThreads > num) {
shutdownThreads(mNumThreads - num);
} else if (mNumThreads < num) {
startThreads(num - mNumThreads);
}
}
void ThreadPool::flushQueue() {
lock_guard<mutex> lock{mTaskQueueMutex};
mTaskQueue.clear();
}
NGP_NAMESPACE_END
|