SMAUG
Simulating Machine Learning Applications on gem5-Aladdin
thread_pool.cpp
1 #include "smaug/utility/thread_pool.h"
2 #include "smaug/utility/utils.h"
3 #include "smaug/core/globals.h"
4 
5 
6 namespace smaug {
7 
8 ThreadPool::ThreadPool(int nthreads) : workers(nthreads) {}
9 
10 ThreadPool::~ThreadPool() {
11  // Shutdown the thread pool and free all resources.
12  for (int i = 0; i < workers.size(); i++) {
13  WorkerThread* worker = &workers[i];
14  gem5::wakeCpu(worker->cpuid);
15  pthread_mutex_lock(&worker->statusMutex);
16  worker->exit = true;
17  pthread_cond_signal(&worker->wakeupCond);
18  pthread_mutex_unlock(&worker->statusMutex);
19  }
20  for (int i = 0; i < workers.size(); i++) {
21  WorkerThread* worker = &workers[i];
22  pthread_join(worker->thread, NULL);
23  pthread_mutex_destroy(&worker->statusMutex);
24  pthread_cond_destroy(&worker->wakeupCond);
25  pthread_cond_destroy(&worker->statusCond);
26  }
27 }
28 
29 void* ThreadPool::workerLoop(void* args) {
30  ThreadInitArgs* initArgs = reinterpret_cast<ThreadInitArgs*>(args);
31  WorkerThread* worker = initArgs->worker;
32  // Notify the main thread about this thread's cpuid. This can only be done
33  // after the thread context is created.
34  pthread_mutex_lock(&initArgs->cpuidMutex);
35  initArgs->cpuid = gem5::getCpuId();
36  worker->status = Idle;
37  pthread_cond_signal(&initArgs->cpuidCond);
38  pthread_mutex_unlock(&initArgs->cpuidMutex);
39 
40  do {
41  gem5::quiesce();
42  pthread_mutex_lock(&worker->statusMutex);
43  while (!worker->valid && !worker->exit)
44  pthread_cond_wait(&worker->wakeupCond, &worker->statusMutex);
45  if (worker->valid) {
46  worker->status = Running;
47  pthread_mutex_unlock(&worker->statusMutex);
48 
49  // Run the function.
50  worker->func(worker->args);
51 
52  pthread_mutex_lock(&worker->statusMutex);
53  worker->status = Idle;
54  worker->valid = false;
55  pthread_cond_signal(&worker->statusCond);
56  }
57  bool exitThread = worker->exit;
58  pthread_mutex_unlock(&worker->statusMutex);
59  if (exitThread)
60  break;
61  } while (true);
62 
63  pthread_exit(NULL);
64 }
65 
67  // Initialize the CPU ID for each worker thread.
68  for (int i = 0; i < workers.size(); i++) {
69  WorkerThread* worker = &workers[i];
70  ThreadInitArgs initArgs(worker);
71  pthread_create(
72  &worker->thread, NULL, &ThreadPool::workerLoop, &initArgs);
73 
74  // Fill in the CPU ID of the worker thread.
75  pthread_mutex_lock(&initArgs.cpuidMutex);
76  while (initArgs.cpuid == -1 || worker->status == Uninitialized)
77  pthread_cond_wait(&initArgs.cpuidCond, &initArgs.cpuidMutex);
78  worker->cpuid = initArgs.cpuid;
79  pthread_mutex_unlock(&initArgs.cpuidMutex);
80  assert(worker->status != Uninitialized &&
81  "Worker thread did not successfully initialize!");
82  }
83 }
84 
85 int ThreadPool::dispatchThread(WorkerThreadFunc func, void* args) {
86  for (int i = 0; i < workers.size(); i++) {
87  WorkerThread* worker = &workers[i];
88  pthread_mutex_lock(&worker->statusMutex);
89  if (worker->status == Idle && !worker->valid) {
90  worker->func = func;
91  worker->args = args;
92  worker->valid = true;
93  gem5::wakeCpu(worker->cpuid);
94  pthread_cond_signal(&worker->wakeupCond);
95  pthread_mutex_unlock(&worker->statusMutex);
96  return i;
97  }
98  pthread_mutex_unlock(&worker->statusMutex);
99  }
100  return -1;
101 }
102 
104  // There is no need to call wakeCpu here. If the CPU is quiesced, then
105  // it cannot possibly be running anything, so its status will be Idle, and
106  // this will move on to the next CPU.
107  for (int i = 0; i < workers.size(); i++) {
108  WorkerThread* worker = &workers[i];
109  pthread_mutex_lock(&worker->statusMutex);
110  while (worker->status == Running || worker->valid == true)
111  pthread_cond_wait(&worker->statusCond, &worker->statusMutex);
112  pthread_mutex_unlock(&worker->statusMutex);
113  }
114 }
115 
116 } // namespace smaug
smaug::ThreadPool::WorkerThread::wakeupCond
pthread_cond_t wakeupCond
The main thread signals this condition variable to wake up the thread and have it check for work (ind...
Definition: thread_pool.h:82
smaug::ThreadPool::WorkerThread::args
void * args
User-provided arguments.
Definition: thread_pool.h:64
smaug::gem5::getCpuId
int getCpuId()
Returns the logical CPU number.
Definition: utils.cpp:71
smaug::ThreadPool::joinThreadPool
void joinThreadPool()
Wait for all threads in the pool to finish work.
Definition: thread_pool.cpp:103
smaug::ThreadPool::workerLoop
static void * workerLoop(void *args)
The main event loop executed by all worker threads.
Definition: thread_pool.cpp:29
smaug::ThreadPool::WorkerThread::statusMutex
pthread_mutex_t statusMutex
This mutex protects all of the subsequent fields of this struct.
Definition: thread_pool.h:72
smaug::gem5::wakeCpu
void wakeCpu(int id)
Wakes up a quiesced CPU.
Definition: utils.cpp:65
smaug::ThreadPool::WorkerThread::exit
bool exit
Set to true to inform the worker thread to terminate.
Definition: thread_pool.h:74
smaug::ThreadPool::ThreadPool
ThreadPool(int nthreads)
Create a ThreadPool with N threads.
Definition: thread_pool.cpp:8
smaug::ThreadPool::workers
std::vector< WorkerThread > workers
Worker threads.
Definition: thread_pool.h:120
globals.h
SMAUG Global variables.
smaug::ThreadPool::initThreadPool
void initThreadPool()
Initialize the thread pool.
Definition: thread_pool.cpp:66
smaug::ThreadPool::WorkerThread::statusCond
pthread_cond_t statusCond
The worker thread signals this condition variable to inform the main thread of a change in status (us...
Definition: thread_pool.h:87
smaug::ThreadPool::WorkerThread::valid
bool valid
Set to true if the func and args are valid and need to be executed.
Definition: thread_pool.h:76
smaug::ThreadPool::WorkerThread::cpuid
int cpuid
The gem5 simulation CPU ID assigned to this worker thread.
Definition: thread_pool.h:89
smaug
The smaug namespace is the parent namespace of all C++ code in SMAUG.
Definition: backend.cpp:38
smaug::ThreadPool::dispatchThread
int dispatchThread(WorkerThreadFunc func, void *args)
Dispatch the function to a worker in the thread pool.
Definition: thread_pool.cpp:85
smaug::gem5::quiesce
void quiesce()
Puts the CPU to sleep.
Definition: utils.cpp:59
smaug::ThreadPool::WorkerThread::thread
pthread_t thread
pthread handle.
Definition: thread_pool.h:67
smaug::ThreadPool::ThreadInitArgs
Definition: thread_pool.h:103
smaug::ThreadPool::WorkerThread
All state and metadata for a worker thread.
Definition: thread_pool.h:61