Add threading primitives.

PiperOrigin-RevId: 562649644
Change-Id: I9c35b270e4b3b50cc7eb5a152f24d71def8e2fcd

src/thread/CMakeLists.txt

@@ -0,0 +1,23 @@
+# Copyright 2023 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
+#
+#     https://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.
+
+set(MUJOCO_THREAD_SRCS
+    lockless_queue.h
+    task.cc
+    task.h
+    thread_pool.cc
+    thread_pool.h
+)
+
+target_sources(mujoco PRIVATE ${MUJOCO_THREAD_SRCS})

src/thread/lockless_queue.h

@@ -0,0 +1,154 @@
+// Copyright 2023 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.
+// IWYU pragma: private, include "third_party/mujoco/include/mujoco.h"
+// IWYU pragma: friend "third_party/(py/)?mujoco/.*"
+
+#ifndef MUJOCO_SRC_THREAD_LOCKLESS_QUEUE_H_
+#define MUJOCO_SRC_THREAD_LOCKLESS_QUEUE_H_
+
+#include <atomic>
+#include <climits>
+#include <cstddef>
+#include <thread>
+
+namespace mujoco {
+
+// A Lockless Queue allows for sending information quickly between different
+// threads. This is a Multi-Producer Multi-Consumer Lockless Queue allowing for
+// multiple threads to be adding items to the queue while multiple threads are
+// consuming items from the queue. Internally it uses a Ring Buffer for storage
+// so it will not grow as items are added. Push will block if the Queue is full
+// and Pop will block if it is empty.
+//
+// For a basic overview of this category of structures:
+// https://www.linuxjournal.com/content/lock-free-multi-producer-multi-consumer-queue-ring-buffer
+template  <typename T, size_t buffer_capacity>
+class LocklessQueue {
+ public:
+  bool full() const {
+    return full_internal(
+        convert_to_index(read_cursor_), convert_to_index(write_cursor_));
+  }
+
+  bool empty() const {
+    return maximum_read_cursor_ == read_cursor_;
+  }
+
+  // Push an element into the queue.
+  void push(const T& input) {
+    // Reserve a slot in the queue
+    size_t current_write_cursor;
+    size_t dummy_current_write_cursor;
+    size_t next_write_cursor;
+    size_t current_write_index;
+    size_t current_read_index;
+    do {
+      // Check if the queue is full.
+      do {
+        current_write_cursor = write_cursor_.load();
+        current_write_index = convert_to_index(current_write_cursor);
+        next_write_cursor = get_next_cursor(current_write_cursor);
+
+        current_read_index = convert_to_index(read_cursor_.load());
+      } while (full_internal(current_read_index, current_write_index));
+
+      // Once it's not full, attempt to grab a slot to write.
+      dummy_current_write_cursor = current_write_cursor;
+    } while (!write_cursor_.compare_exchange_weak(
+                 dummy_current_write_cursor, next_write_cursor));
+
+    // Write the entry.
+    buffer_[current_write_index].store(input);
+
+    // Increment maximum read cursor. Note here it has to wait if the compare
+    // and exchange fails as another thread might not have completed its write.
+    do {
+      dummy_current_write_cursor = current_write_cursor;
+    } while (!maximum_read_cursor_.compare_exchange_weak(
+                 dummy_current_write_cursor, next_write_cursor));
+  }
+
+  // Pop an element from the queue.
+  T pop() {
+    size_t current_read_cursor;
+    size_t dummy_current_read_cursor;
+    size_t current_read_index;
+    size_t next_read_cursor;
+    size_t current_maximum_read_cursor;
+    size_t current_maximum_read_index;
+    bool empty = false;
+    T result;
+    do {
+      // Wait until the queue has an element
+      do {
+        if (empty) {
+          std::this_thread::yield();
+        }
+        current_read_cursor = read_cursor_.load();
+        current_maximum_read_cursor = maximum_read_cursor_.load();
+
+        current_read_index = convert_to_index(current_read_cursor);
+        current_maximum_read_index = convert_to_index(
+            current_maximum_read_cursor);
+
+        empty = empty_internal(
+            current_read_index, current_maximum_read_index);
+      } while (empty);
+
+      next_read_cursor = get_next_cursor(current_read_cursor);
+
+      // Attempt to grab the element, if unsuccessful then wait for the next
+      // element to arrive.
+      result = buffer_[current_read_index].load();
+      dummy_current_read_cursor = current_read_cursor;
+    } while (!read_cursor_.compare_exchange_weak(
+                 dummy_current_read_cursor, next_read_cursor));
+
+    return result;
+  }
+
+ private:
+  size_t convert_to_index(size_t input) const {
+    return input % internal_buffer_capacity_;
+  }
+
+  size_t get_next_cursor(size_t input) const {
+    return (input + 1) % cursor_max_;
+  }
+
+  size_t get_next_index(size_t input) const {
+    return convert_to_index(get_next_cursor(input));
+  }
+
+  bool full_internal(size_t read_index, size_t write_index) const {
+    return get_next_index(write_index) == read_index;
+  }
+
+  bool empty_internal(size_t read_index, size_t write_index) const {
+    return read_index == write_index;
+  }
+
+  const size_t internal_buffer_capacity_ = buffer_capacity + 1;
+  const size_t cursor_max_ = UINT_MAX - (UINT_MAX % internal_buffer_capacity_);
+
+  std::atomic<size_t> read_cursor_ = 0;
+  std::atomic<size_t> write_cursor_ = 0;
+  std::atomic<size_t> maximum_read_cursor_ = 0;
+
+  std::atomic<T> buffer_[(buffer_capacity + 1)];
+};
+
+}  // namespace mujoco
+
+#endif  // MUJOCO_SRC_THREAD_LOCKLESS_QUEUE_H_

src/thread/task.cc

@@ -0,0 +1,24 @@
+// Copyright 2023 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 "thread/task.h"
+
+#include <mujoco/mjthread.h>
+#include <mujoco/mujoco.h>
+
+// waits for a task to complete
+void mju_taskJoin(mjTask* task) {
+  mujoco::Task* task_ptr = static_cast<mujoco::Task*>(static_cast<void*>(task));
+  task_ptr->Join();
+}

src/thread/task.h

@@ -0,0 +1,69 @@
+// Copyright 2023 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.
+// IWYU pragma: private, include "third_party/mujoco/include/mujoco.h"
+// IWYU pragma: friend "third_party/(py/)?mujoco/.*"
+
+#ifndef MUJOCO_SRC_THREAD_TASK_H_
+#define MUJOCO_SRC_THREAD_TASK_H_
+
+#ifdef __cplusplus
+
+#include <atomic>
+#include <functional>
+#include <thread>
+
+namespace mujoco {
+
+class Task {
+ public:
+  enum Status {
+    QUEUED,
+    COMPLETE,
+  };
+
+  static void Initialize(
+      Task* task,
+      std::function<void*(void*)> start_routine,
+      void* args) {
+    // instantiate a task at the pointer passed in
+    new(task) Task();
+    task->start_routine_ = start_routine;
+    task->args_ = args;
+    task->status_ = Status::QUEUED;
+  }
+
+  void Execute() {
+    args_ = start_routine_(args_);
+    status_ = Status::COMPLETE;
+  }
+
+  void Join() {
+    while (status_ != Status::COMPLETE) {
+      std::this_thread::yield();
+    }
+  }
+
+ private:
+  std::function<void*(void*)> start_routine_;
+
+  void* args_;
+
+  std::atomic<Status> status_ = Status::QUEUED;
+};
+
+}  // namespace mujoco
+
+#endif  // __cplusplus
+
+#endif  // MUJOCO_SRC_THREAD_TASK_H_

src/thread/thread_pool.cc

@@ -0,0 +1,52 @@
+// Copyright 2023 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 "thread/thread_pool.h"
+
+#include <cstddef>
+
+#include <mujoco/mjthread.h>
+#include <mujoco/mujoco.h>
+#include "thread/task.h"
+
+static constexpr size_t kMaxThreads = 128;
+
+// create a thread pool
+mjThreadPool* mju_threadPoolCreate(size_t number_of_threads) {
+  mujoco::ThreadPool<kMaxThreads>* thread_pool =
+      new mujoco::ThreadPool<kMaxThreads>(number_of_threads);
+  return static_cast<mjThreadPool*>(static_cast<void*>(thread_pool));
+}
+
+// start a task in the threadpool
+void mju_threadPoolEnqueue(
+    mjThreadPool* thread_pool, mjTask* task, mjStartRoutine start_routine,
+    void* args) {
+  mujoco::ThreadPool<kMaxThreads>* thread_pool_ptr =
+      static_cast<mujoco::ThreadPool<kMaxThreads>*>(
+          static_cast<void*>(thread_pool));
+  thread_pool_ptr->Enqueue(
+      static_cast<mujoco::Task*>(static_cast<void*>(task)), start_routine,
+      args);
+}
+
+// shutdown the threadpool and free the memory
+void mju_threadPoolDestroy(mjThreadPool* thread_pool) {
+  mujoco::ThreadPool<kMaxThreads>* thread_pool_ptr =
+      static_cast<mujoco::ThreadPool<kMaxThreads>*>(
+          static_cast<void*>(thread_pool));
+  thread_pool_ptr->Shutdown();
+  delete thread_pool_ptr;
+}
+

src/thread/thread_pool.h

@@ -0,0 +1,103 @@
+// Copyright 2023 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.
+// IWYU pragma: private, include "third_party/mujoco/include/mujoco.h"
+// IWYU pragma: friend "third_party/(py/)?mujoco/.*"
+
+#ifndef MUJOCO_SRC_THREAD_THREAD_POOL_H_
+#define MUJOCO_SRC_THREAD_THREAD_POOL_H_
+
+#ifdef __cplusplus
+
+#include <atomic>
+#include <cstddef>
+#include <functional>
+#include <thread>
+
+#include "thread/lockless_queue.h"
+#include "thread/task.h"
+
+namespace mujoco {
+
+static constexpr size_t kThreadPoolQueueSize = 640;
+
+template <size_t max_number_of_threads>
+class ThreadPool {
+ public:
+  ThreadPool(size_t number_of_threads)
+      : number_of_threads_(number_of_threads) {
+    for (int i = 0; i < number_of_threads_; ++i) {
+      threads_[i] = std::thread(ThreadPoolWorker, static_cast<void*>(this));
+    }
+  }
+
+  // start a task in the threadpool
+  void Enqueue(
+      Task* task, std::function<void*(void*)> start_routine, void* args) {
+    Task::Initialize(task, start_routine, args);
+    lockless_queue_.push(static_cast<void*>(task));
+  }
+
+  // shutdown the threadpool
+  void Shutdown() {
+    if (shutdown_) {
+      return;
+    }
+
+    shutdown_ = true;
+    Task shutdown_tasks[number_of_threads_];
+    for (int i = 0; i < number_of_threads_; ++i) {
+      Enqueue(&shutdown_tasks[i], ShutdownFunction, nullptr);
+    }
+
+    for (int i = 0; i < number_of_threads_; ++i) {
+      threads_[i].join();
+    }
+  }
+
+  ~ThreadPool() { Shutdown(); }
+
+ private:
+  // method executed by running threads
+  static void ThreadPoolWorker(void* arg) {
+    ThreadPool<max_number_of_threads>* thread_pool =
+        static_cast<ThreadPool<max_number_of_threads>*>(arg);
+    while (!thread_pool->shutdown_) {
+      Task* task = static_cast<Task*>(thread_pool->lockless_queue_.pop());
+      task->Execute();
+    }
+  }
+
+  // shutdown function passed to running threads to ensure cleans shutdown
+  static void* ShutdownFunction(void* args) {
+    return NULL;
+  }
+
+  // is the thread pool is being shut down
+  std::atomic<bool> shutdown_ = false;
+
+  // actual number of running threads in the threadpool
+  const size_t number_of_threads_;
+
+  // OS threads that are running in this pool
+  std::thread threads_[max_number_of_threads];
+
+  // queue of tasks to execute
+  LocklessQueue<void*, kThreadPoolQueueSize> lockless_queue_;
+};
+
+}  // namespace mujoco
+
+#endif  // __cplusplus
+
+#endif  // MUJOCO_SRC_THREAD_THREAD_POOL_H_