Create threadpool for src/user and primarily compiler usage.
PiperOrigin-RevId: 864328236 Change-Id: Iffaa9b8e6eacdb038020af9d6ad82d7cb7a6c941
This commit is contained in:
Sam Haves
committed by
Copybara-Service
Copybara-Service
parent
b7a46d04b6
commit
683b1b5d38
@@ -29,6 +29,8 @@ set(MUJOCO_USER_SRCS
|
||||
user_objects.h
|
||||
user_resource.cc
|
||||
user_resource.h
|
||||
user_threadpool.cc
|
||||
user_threadpool.h
|
||||
user_util.cc
|
||||
user_util.h
|
||||
user_vfs.cc
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
// Copyright 2026 DeepMind Technologies Limited
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include "user/user_threadpool.h"
|
||||
|
||||
#include <condition_variable>
|
||||
#include <functional>
|
||||
#include <mutex>
|
||||
#include <thread>
|
||||
#include <utility>
|
||||
|
||||
namespace mujoco::user {
|
||||
|
||||
constinit thread_local int ThreadPool::worker_id_ = -1;
|
||||
|
||||
// ThreadPool constructor
|
||||
ThreadPool::ThreadPool(int num_threads) : ctr_(0) {
|
||||
for (int i = 0; i < num_threads; i++) {
|
||||
threads_.push_back(std::thread(&ThreadPool::WorkerThread, this, i));
|
||||
}
|
||||
}
|
||||
|
||||
// ThreadPool destructor
|
||||
ThreadPool::~ThreadPool() {
|
||||
{
|
||||
std::unique_lock<std::mutex> lock(m_);
|
||||
for (int i = 0; i < threads_.size(); i++) {
|
||||
queue_.push(nullptr);
|
||||
}
|
||||
cv_in_.notify_all();
|
||||
}
|
||||
for (auto& thread : threads_) {
|
||||
thread.join();
|
||||
}
|
||||
}
|
||||
|
||||
// ThreadPool scheduler
|
||||
void ThreadPool::Schedule(std::function<void()> task) {
|
||||
std::unique_lock<std::mutex> lock(m_);
|
||||
queue_.push(std::move(task));
|
||||
cv_in_.notify_one();
|
||||
}
|
||||
|
||||
// ThreadPool worker
|
||||
void ThreadPool::WorkerThread(int i) {
|
||||
worker_id_ = i;
|
||||
while (true) {
|
||||
auto task = [&]() {
|
||||
std::unique_lock<std::mutex> lock(m_);
|
||||
cv_in_.wait(lock, [&]() { return !queue_.empty(); });
|
||||
std::function<void()> task = std::move(queue_.front());
|
||||
queue_.pop();
|
||||
cv_in_.notify_one();
|
||||
return task;
|
||||
}();
|
||||
if (task == nullptr) {
|
||||
{
|
||||
std::unique_lock<std::mutex> lock(m_);
|
||||
++ctr_;
|
||||
cv_ext_.notify_one();
|
||||
}
|
||||
break;
|
||||
}
|
||||
task();
|
||||
|
||||
{
|
||||
std::unique_lock<std::mutex> lock(m_);
|
||||
++ctr_;
|
||||
cv_ext_.notify_one();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace mujoco::user
|
||||
@@ -0,0 +1,85 @@
|
||||
// Copyright 2026 DeepMind Technologies Limited
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#ifndef MUJOCO_SRC_USER_USER_THREADPOOL_H_
|
||||
#define MUJOCO_SRC_USER_USER_THREADPOOL_H_
|
||||
|
||||
#include <condition_variable>
|
||||
#include <cstdint>
|
||||
#include <functional>
|
||||
#include <mutex>
|
||||
#include <queue>
|
||||
#include <thread>
|
||||
#include <vector>
|
||||
|
||||
|
||||
namespace mujoco::user {
|
||||
|
||||
// ThreadPool class
|
||||
class ThreadPool {
|
||||
public:
|
||||
// constructor
|
||||
explicit ThreadPool(int num_threads);
|
||||
|
||||
// destructor
|
||||
~ThreadPool();
|
||||
|
||||
int NumThreads() const { return threads_.size(); }
|
||||
|
||||
// returns an ID between 0 and NumThreads() - 1. must be called within
|
||||
// worker thread (returns -1 if not).
|
||||
static int WorkerId() { return worker_id_; }
|
||||
|
||||
// ----- methods ----- //
|
||||
// set task for threadpool
|
||||
void Schedule(std::function<void()> task);
|
||||
|
||||
// return number of tasks completed
|
||||
std::uint64_t GetCount() {
|
||||
std::lock_guard<std::mutex> lock(m_);
|
||||
return ctr_;
|
||||
}
|
||||
|
||||
// reset count to zero
|
||||
void ResetCount() {
|
||||
std::lock_guard<std::mutex> lock(m_);
|
||||
ctr_ = 0;
|
||||
}
|
||||
|
||||
// wait for count, then return
|
||||
void WaitCount(int value) {
|
||||
std::unique_lock<std::mutex> lock(m_);
|
||||
cv_ext_.wait(lock, [&]() { return ctr_ >= value; });
|
||||
}
|
||||
|
||||
private:
|
||||
// ----- methods ----- //
|
||||
|
||||
// execute task with available thread
|
||||
void WorkerThread(int i);
|
||||
|
||||
constinit static thread_local int worker_id_;
|
||||
|
||||
// ----- members ----- //
|
||||
std::vector<std::thread> threads_;
|
||||
std::mutex m_;
|
||||
std::condition_variable cv_in_;
|
||||
std::condition_variable cv_ext_;
|
||||
std::queue<std::function<void()>> queue_;
|
||||
std::uint64_t ctr_;
|
||||
};
|
||||
|
||||
} // namespace mujoco::user
|
||||
|
||||
#endif // MUJOCO_SRC_USER_USER_THREADPOOL_H_
|
||||
@@ -0,0 +1,197 @@
|
||||
// Copyright 2026 DeepMind Technologies Limited
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include "src/user/user_threadpool.h"
|
||||
|
||||
#include <atomic>
|
||||
#include <condition_variable>
|
||||
#include <mutex>
|
||||
#include <set>
|
||||
|
||||
#include <gtest/gtest.h>
|
||||
#include "test/fixture.h"
|
||||
|
||||
namespace mujoco {
|
||||
namespace {
|
||||
|
||||
using user::ThreadPool;
|
||||
using ThreadPoolTest = MujocoTest;
|
||||
|
||||
TEST(ThreadPoolTest, ConstructorCreatesCorrectNumberOfThreads) {
|
||||
ThreadPool pool(4);
|
||||
EXPECT_EQ(pool.NumThreads(), 4);
|
||||
}
|
||||
|
||||
TEST(ThreadPoolTest, SingleThread) {
|
||||
ThreadPool pool(1);
|
||||
EXPECT_EQ(pool.NumThreads(), 1);
|
||||
}
|
||||
|
||||
TEST(ThreadPoolTest, ScheduleAndWait) {
|
||||
ThreadPool pool(4);
|
||||
std::atomic<int> counter{0};
|
||||
|
||||
for (int i = 0; i < 10; ++i) {
|
||||
pool.Schedule([&counter]() { counter++; });
|
||||
}
|
||||
|
||||
pool.WaitCount(10);
|
||||
EXPECT_EQ(counter.load(), 10);
|
||||
}
|
||||
|
||||
TEST(ThreadPoolTest, GetCountReturnsCompletedTasks) {
|
||||
ThreadPool pool(2);
|
||||
|
||||
pool.Schedule([]() {});
|
||||
pool.Schedule([]() {});
|
||||
pool.WaitCount(2);
|
||||
|
||||
EXPECT_GE(pool.GetCount(), 2);
|
||||
}
|
||||
|
||||
TEST(ThreadPoolTest, ResetCountClearsCounter) {
|
||||
ThreadPool pool(2);
|
||||
|
||||
pool.Schedule([]() {});
|
||||
pool.WaitCount(1);
|
||||
EXPECT_GE(pool.GetCount(), 1);
|
||||
|
||||
pool.ResetCount();
|
||||
EXPECT_EQ(pool.GetCount(), 0);
|
||||
}
|
||||
|
||||
TEST(ThreadPoolTest, WorkerIdIsValid) {
|
||||
ThreadPool pool(4);
|
||||
std::atomic<int> valid_ids{0};
|
||||
|
||||
for (int i = 0; i < 10; ++i) {
|
||||
pool.Schedule([&valid_ids, &pool]() {
|
||||
int id = ThreadPool::WorkerId();
|
||||
if (id >= 0 && id < pool.NumThreads()) {
|
||||
valid_ids++;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
pool.WaitCount(10);
|
||||
EXPECT_EQ(valid_ids.load(), 10);
|
||||
}
|
||||
|
||||
TEST(ThreadPoolTest, WorkerIdOutsidePoolReturnsNegativeOne) {
|
||||
EXPECT_EQ(ThreadPool::WorkerId(), -1);
|
||||
}
|
||||
|
||||
TEST(ThreadPoolTest, ParallelExecution) {
|
||||
ThreadPool pool(4);
|
||||
std::atomic<int> sum{0};
|
||||
const int n = 1000;
|
||||
|
||||
for (int i = 0; i < n; ++i) {
|
||||
pool.Schedule([&sum]() { sum++; });
|
||||
}
|
||||
|
||||
pool.WaitCount(n);
|
||||
EXPECT_EQ(sum.load(), n);
|
||||
}
|
||||
|
||||
TEST(ThreadPoolTest, MultipleWaitCounts) {
|
||||
ThreadPool pool(2);
|
||||
std::atomic<int> counter{0};
|
||||
|
||||
for (int i = 0; i < 5; ++i) {
|
||||
pool.Schedule([&counter]() { counter++; });
|
||||
}
|
||||
pool.WaitCount(5);
|
||||
EXPECT_EQ(counter.load(), 5);
|
||||
|
||||
pool.ResetCount();
|
||||
|
||||
for (int i = 0; i < 3; ++i) {
|
||||
pool.Schedule([&counter]() { counter++; });
|
||||
}
|
||||
pool.WaitCount(3);
|
||||
EXPECT_EQ(counter.load(), 8);
|
||||
}
|
||||
|
||||
TEST(ThreadPoolTest, ScopedPoolDestruction) {
|
||||
std::atomic<int> counter{0};
|
||||
{
|
||||
ThreadPool pool(2);
|
||||
for (int i = 0; i < 5; ++i) {
|
||||
pool.Schedule([&counter]() { counter++; });
|
||||
}
|
||||
pool.WaitCount(5);
|
||||
}
|
||||
EXPECT_EQ(counter.load(), 5);
|
||||
}
|
||||
|
||||
TEST(ThreadPoolTest, OnlyConfiguredThreadsAreUsed) {
|
||||
constexpr int kNumThreads = 4;
|
||||
constexpr int kNumTasks = 100;
|
||||
|
||||
ThreadPool pool(kNumThreads);
|
||||
std::mutex ids_mutex;
|
||||
std::set<int> worker_ids;
|
||||
|
||||
for (int i = 0; i < kNumTasks; ++i) {
|
||||
pool.Schedule([&ids_mutex, &worker_ids]() {
|
||||
int id = ThreadPool::WorkerId();
|
||||
std::lock_guard<std::mutex> lock(ids_mutex);
|
||||
worker_ids.insert(id);
|
||||
});
|
||||
}
|
||||
|
||||
pool.WaitCount(kNumTasks);
|
||||
|
||||
EXPECT_LE(worker_ids.size(), kNumThreads);
|
||||
for (int id : worker_ids) {
|
||||
EXPECT_GE(id, 0);
|
||||
EXPECT_LT(id, kNumThreads);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(ThreadPoolTest, BlockingTaskDoesNotPreventOtherTasks) {
|
||||
constexpr int kNumThreads = 4;
|
||||
ThreadPool pool(kNumThreads);
|
||||
|
||||
std::mutex block_mutex;
|
||||
std::condition_variable block_cv;
|
||||
bool unblock = false;
|
||||
|
||||
std::atomic<int> fast_tasks_completed{0};
|
||||
constexpr int kNumFastTasks = 10;
|
||||
|
||||
pool.Schedule([&block_mutex, &block_cv, &unblock]() {
|
||||
std::unique_lock<std::mutex> lock(block_mutex);
|
||||
block_cv.wait(lock, [&unblock]() { return unblock; });
|
||||
});
|
||||
|
||||
for (int i = 0; i < kNumFastTasks; ++i) {
|
||||
pool.Schedule([&fast_tasks_completed]() { fast_tasks_completed++; });
|
||||
}
|
||||
|
||||
pool.WaitCount(kNumFastTasks);
|
||||
|
||||
EXPECT_EQ(fast_tasks_completed.load(), kNumFastTasks);
|
||||
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(block_mutex);
|
||||
unblock = true;
|
||||
}
|
||||
block_cv.notify_one();
|
||||
|
||||
pool.WaitCount(kNumFastTasks + 1);
|
||||
}
|
||||
} // namespace
|
||||
} // namespace mujoco
|
||||
Reference in New Issue
Block a user