1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
|
#include "caffe2/utils/threadpool/pthreadpool.h"
#include "caffe2/utils/threadpool/pthreadpool-cpp.h"
#include "caffe2/utils/threadpool/ThreadPool.h"
#if defined(USE_PTHREADPOOL)
namespace caffe2 {
namespace {
static thread_local bool using_new_threadpool{false};
}
WithCastToNewThreadPool::WithCastToNewThreadPool(bool use_new_threadpool) {
use_new_threadpool_ = using_new_threadpool;
using_new_threadpool = use_new_threadpool;
}
WithCastToNewThreadPool::~WithCastToNewThreadPool() {
using_new_threadpool = use_new_threadpool_;
}
}
#endif
//
// External API
//
void legacy_pthreadpool_compute_1d(
legacy_pthreadpool_t threadpool,
legacy_pthreadpool_function_1d_t function,
void* argument,
size_t range) {
if (threadpool == nullptr) {
/* No thread pool provided: execute function sequentially on the calling
* thread */
for (size_t i = 0; i < range; i++) {
function(argument, i);
}
return;
}
#if defined(USE_PTHREADPOOL)
if (caffe2::using_new_threadpool) {
pthreadpool_parallelize_1d(threadpool, function, argument, range, 0u);
} else {
reinterpret_cast<caffe2::ThreadPool*>(threadpool)
->run(
[function, argument](int threadId, size_t workId) {
function(argument, workId);
},
range);
}
#else
reinterpret_cast<caffe2::ThreadPool*>(threadpool)
->run(
[function, argument](int threadId, size_t workId) {
function(argument, workId);
},
range);
#endif
}
void legacy_pthreadpool_parallelize_1d(
const legacy_pthreadpool_t threadpool,
const legacy_pthreadpool_function_1d_t function,
void* const argument,
const size_t range,
uint32_t) {
legacy_pthreadpool_compute_1d(threadpool, function, argument, range);
}
size_t legacy_pthreadpool_get_threads_count(legacy_pthreadpool_t threadpool) {
// The current fix only useful when XNNPACK calls legacy_pthreadpool_get_threads_count with nullptr.
if (threadpool == nullptr) {
return 1;
}
return reinterpret_cast<caffe2::ThreadPool*>(threadpool)->getNumThreads();
}
legacy_pthreadpool_t legacy_pthreadpool_create(size_t threads_count) {
std::mutex thread_pool_creation_mutex_;
std::lock_guard<std::mutex> guard(thread_pool_creation_mutex_);
return reinterpret_cast<legacy_pthreadpool_t>(caffe2::ThreadPool::createThreadPool(threads_count));
}
void legacy_pthreadpool_destroy(legacy_pthreadpool_t pthreadpool) {
if (pthreadpool) {
caffe2::ThreadPool* threadpool =
reinterpret_cast<caffe2::ThreadPool*>(pthreadpool);
delete threadpool;
}
}
|
