mirror of
https://github.com/pytorch/pytorch.git
synced 2025-11-11 22:34:53 +08:00
1544c37520
This PR continues to fix clang-tidy warnings for headers in c10/core and c10/util. Pull Request resolved: https://github.com/pytorch/pytorch/pull/115495 Approved by: https://github.com/malfet
121 lines
2.9 KiB
C++
121 lines
2.9 KiB
C++
#pragma once
|
|
|
|
#include <atomic>
|
|
#include <condition_variable>
|
|
#include <cstddef>
|
|
#include <functional>
|
|
#include <mutex>
|
|
#include <queue>
|
|
#include <thread>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
#include <c10/macros/Export.h>
|
|
#include <c10/util/Registry.h>
|
|
#include <c10/util/numa.h>
|
|
#include <c10/util/thread_name.h>
|
|
|
|
namespace c10 {
|
|
|
|
class C10_API TaskThreadPoolBase {
|
|
public:
|
|
virtual void run(std::function<void()> func) = 0;
|
|
|
|
virtual size_t size() const = 0;
|
|
|
|
/**
|
|
* The number of available (i.e. idle) threads in this thread pool.
|
|
*/
|
|
virtual size_t numAvailable() const = 0;
|
|
|
|
/**
|
|
* Check if the current thread is from the thread pool.
|
|
*/
|
|
virtual bool inThreadPool() const = 0;
|
|
|
|
virtual ~TaskThreadPoolBase() noexcept = default;
|
|
|
|
static size_t defaultNumThreads();
|
|
};
|
|
|
|
class C10_API ThreadPool : public c10::TaskThreadPoolBase {
|
|
protected:
|
|
struct task_element_t {
|
|
bool run_with_id;
|
|
// NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
|
|
const std::function<void()> no_id;
|
|
// NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
|
|
const std::function<void(std::size_t)> with_id;
|
|
|
|
explicit task_element_t(std::function<void()> f)
|
|
: run_with_id(false), no_id(std::move(f)), with_id(nullptr) {}
|
|
explicit task_element_t(std::function<void(std::size_t)> f)
|
|
: run_with_id(true), no_id(nullptr), with_id(std::move(f)) {}
|
|
};
|
|
|
|
std::queue<task_element_t> tasks_;
|
|
std::vector<std::thread> threads_;
|
|
mutable std::mutex mutex_;
|
|
std::condition_variable condition_;
|
|
std::condition_variable completed_;
|
|
std::atomic_bool running_;
|
|
bool complete_;
|
|
std::size_t available_;
|
|
std::size_t total_;
|
|
int numa_node_id_;
|
|
|
|
public:
|
|
ThreadPool() = delete;
|
|
|
|
explicit ThreadPool(
|
|
int pool_size,
|
|
int numa_node_id = -1,
|
|
const std::function<void()>& init_thread = nullptr);
|
|
|
|
~ThreadPool() override;
|
|
|
|
size_t size() const override;
|
|
|
|
size_t numAvailable() const override;
|
|
|
|
bool inThreadPool() const override;
|
|
|
|
void run(std::function<void()> func) override;
|
|
|
|
template <typename Task>
|
|
void runTaskWithID(Task task) {
|
|
std::unique_lock<std::mutex> lock(mutex_);
|
|
|
|
// Set task and signal condition variable so that a worker thread will
|
|
// wake up and use the task.
|
|
tasks_.emplace(static_cast<std::function<void(std::size_t)>>(task));
|
|
complete_ = false;
|
|
condition_.notify_one();
|
|
}
|
|
|
|
/// @brief Wait for queue to be empty
|
|
void waitWorkComplete();
|
|
|
|
private:
|
|
// @brief Entry point for pool threads.
|
|
void main_loop(std::size_t index);
|
|
};
|
|
|
|
class C10_API TaskThreadPool : public c10::ThreadPool {
|
|
public:
|
|
explicit TaskThreadPool(int pool_size, int numa_node_id = -1)
|
|
: ThreadPool(pool_size, numa_node_id, [numa_node_id]() {
|
|
setThreadName("CaffeTaskThread");
|
|
NUMABind(numa_node_id);
|
|
}) {}
|
|
};
|
|
|
|
C10_DECLARE_SHARED_REGISTRY(
|
|
ThreadPoolRegistry,
|
|
TaskThreadPoolBase,
|
|
int,
|
|
int,
|
|
bool);
|
|
|
|
} // namespace c10
|