From 683b1b5d38918e2d3871a40ff955b1b137bee63a Mon Sep 17 00:00:00 2001 From: Sam Haves Date: Mon, 2 Feb 2026 06:30:25 -0800 Subject: [PATCH] Create threadpool for src/user and primarily compiler usage. PiperOrigin-RevId: 864328236 Change-Id: Iffaa9b8e6eacdb038020af9d6ad82d7cb7a6c941 --- src/user/CMakeLists.txt | 2 + src/user/user_threadpool.cc | 85 +++++++++++++ src/user/user_threadpool.h | 85 +++++++++++++ test/user/user_threadpool_test.cc | 197 ++++++++++++++++++++++++++++++ 4 files changed, 369 insertions(+) create mode 100644 src/user/user_threadpool.cc create mode 100644 src/user/user_threadpool.h create mode 100644 test/user/user_threadpool_test.cc diff --git a/src/user/CMakeLists.txt b/src/user/CMakeLists.txt index 6bd7399a..bd93c028 100644 --- a/src/user/CMakeLists.txt +++ b/src/user/CMakeLists.txt @@ -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 diff --git a/src/user/user_threadpool.cc b/src/user/user_threadpool.cc new file mode 100644 index 00000000..1fcfa1aa --- /dev/null +++ b/src/user/user_threadpool.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 +#include +#include +#include +#include + +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 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 task) { + std::unique_lock 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 lock(m_); + cv_in_.wait(lock, [&]() { return !queue_.empty(); }); + std::function task = std::move(queue_.front()); + queue_.pop(); + cv_in_.notify_one(); + return task; + }(); + if (task == nullptr) { + { + std::unique_lock lock(m_); + ++ctr_; + cv_ext_.notify_one(); + } + break; + } + task(); + + { + std::unique_lock lock(m_); + ++ctr_; + cv_ext_.notify_one(); + } + } +} + +} // namespace mujoco::user diff --git a/src/user/user_threadpool.h b/src/user/user_threadpool.h new file mode 100644 index 00000000..99354136 --- /dev/null +++ b/src/user/user_threadpool.h @@ -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 +#include +#include +#include +#include +#include +#include + + +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 task); + + // return number of tasks completed + std::uint64_t GetCount() { + std::lock_guard lock(m_); + return ctr_; + } + + // reset count to zero + void ResetCount() { + std::lock_guard lock(m_); + ctr_ = 0; + } + + // wait for count, then return + void WaitCount(int value) { + std::unique_lock 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 threads_; + std::mutex m_; + std::condition_variable cv_in_; + std::condition_variable cv_ext_; + std::queue> queue_; + std::uint64_t ctr_; +}; + +} // namespace mujoco::user + +#endif // MUJOCO_SRC_USER_USER_THREADPOOL_H_ diff --git a/test/user/user_threadpool_test.cc b/test/user/user_threadpool_test.cc new file mode 100644 index 00000000..80b85c0c --- /dev/null +++ b/test/user/user_threadpool_test.cc @@ -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 +#include +#include +#include + +#include +#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 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 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 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 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 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 worker_ids; + + for (int i = 0; i < kNumTasks; ++i) { + pool.Schedule([&ids_mutex, &worker_ids]() { + int id = ThreadPool::WorkerId(); + std::lock_guard 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 fast_tasks_completed{0}; + constexpr int kNumFastTasks = 10; + + pool.Schedule([&block_mutex, &block_cv, &unblock]() { + std::unique_lock 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 lock(block_mutex); + unblock = true; + } + block_cv.notify_one(); + + pool.WaitCount(kNumFastTasks + 1); +} +} // namespace +} // namespace mujoco