Copybara import of the project:

-- 3a95b62f59e81bfef0f076afb173ecc14b27943d by Levi Burner <leviburner@gmail.com>: rollout prototype native threadpool for comparing to python threads -- efd8be1124ac839b902de45973a3ca8b9f2215e6 by Levi Burner <leviburner@gmail.com>: copy mjpcs threadpool into python bindings -- 75603eea3e8362e354a9675e8a6cd14e56ec3d28 by Levi Burner <leviburner@gmail.com>: rollout use threadpool as translation unit -- 06b90febd021663f6cc81fd7895e4d6e2008ed97 by Levi Burner <leviburner@gmail.com>: rollout add chunk_divisor parameter -- 298ab2f3c0d6e12530832c3cdbf784dd92d54806 by Levi Burner <leviburner@gmail.com>: rollout add native threading test -- 169cf9978e7abad6edd1392b8e6aab995e4f8f10 by Levi Burner <leviburner@gmail.com>: rollout exchange chunk_divisor arg for chunk_size -- 265af851d74432d261277d3dbda11cdef1841bc8 by Levi Burner <leviburner@gmail.com>: rollout fix cosmetics -- 1e8bffa88bf36190501b334bef31147e23db39f7 by Levi Burner <leviburner@gmail.com>: make native rollout a class instead of a function -- ba788214b047577f58c41ce0ab6c62c277cd8b0d by Levi Burner <leviburner@gmail.com>: rollout update docs and changelog -- e4cb7732319e04cba2ab2c2ad848c659f6309808 by Levi Burner <leviburner@gmail.com>: rollout don't register atexit handler for Rollout objects -- 5a08d2efdbbbb01d4b1231ff9a36a1dc44f4d9ee by Levi Burner <leviburner@gmail.com>: rollout nthread kwarg, rename shutdown_pool to close, fixups -- f622378543596a208339af0208fa3a70bf2a8007 by Levi Burner <leviburner@gmail.com>: rollout add missing .close() calls -- 50f3ebca43c53eac03f03943c34bb1e46967bd4f by Levi Burner <leviburner@gmail.com>: rollout return immediately COPYBARA_INTEGRATE_REVIEW=https://github.com/google-deepmind/mujoco/pull/2282 from aftersomemath:rollout-threaded 50f3ebca43c53eac03f03943c34bb1e46967bd4f PiperOrigin-RevId: 706744277 Change-Id: I1ab2263b7d6ce30cf1908aec8fd5f2eb976a19e6
2024-12-16 09:56:12 -08:00
parent b26d6f0466
commit a7eb6efd4e
8 changed files with 761 additions and 203 deletions
@@ -0,0 +1,87 @@
+// Copyright 2024 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 "threadpool.h"
+
+#include <condition_variable>
+#include <functional>
+#include <mutex>
+#include <thread>
+#include <utility>
+
+#include <absl/base/attributes.h>
+
+namespace mujoco::python {
+
+ABSL_CONST_INIT 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::python