10 #if defined(EIGEN_USE_THREADS) && !defined(EIGEN_CXX11_TENSOR_TENSOR_DEVICE_THREAD_POOL_H) 11 #define EIGEN_CXX11_TENSOR_TENSOR_DEVICE_THREAD_POOL_H 17 #ifndef EIGEN_USE_SIMPLE_THREAD_POOL 18 template <
typename Env>
using ThreadPoolTempl = NonBlockingThreadPoolTempl<Env>;
19 typedef NonBlockingThreadPool ThreadPool;
21 template <
typename Env>
using ThreadPoolTempl = SimpleThreadPoolTempl<Env>;
22 typedef SimpleThreadPool ThreadPool;
30 Barrier(
unsigned int count) : state_(count << 1), notified_(false) {
31 eigen_assert(((count << 1) >> 1) == count);
34 eigen_assert((state_>>1) == 0);
38 unsigned int v = state_.fetch_sub(2, std::memory_order_acq_rel) - 2;
40 eigen_assert(((v + 2) & ~1) != 0);
43 std::unique_lock<std::mutex> l(mu_);
44 eigen_assert(!notified_);
50 unsigned int v = state_.fetch_or(1, std::memory_order_acq_rel);
51 if ((v >> 1) == 0)
return;
52 std::unique_lock<std::mutex> l(mu_);
60 std::condition_variable cv_;
61 std::atomic<unsigned int> state_;
71 struct Notification : Barrier {
72 Notification() : Barrier(1) {};
78 template <
typename Function,
typename... Args>
struct FunctionWrapperWithNotification
80 static void run(Notification* n, Function f, Args... args) {
88 template <
typename Function,
typename... Args>
struct FunctionWrapperWithBarrier
90 static void run(Barrier* b, Function f, Args... args) {
98 template <
typename SyncType>
99 static EIGEN_STRONG_INLINE
void wait_until_ready(SyncType* n) {
107 struct ThreadPoolDevice {
109 ThreadPoolDevice(ThreadPoolInterface* pool,
int num_cores) : pool_(pool), num_threads_(num_cores) { }
111 EIGEN_STRONG_INLINE
void* allocate(
size_t num_bytes)
const {
112 return internal::aligned_malloc(num_bytes);
115 EIGEN_STRONG_INLINE
void deallocate(
void* buffer)
const {
116 internal::aligned_free(buffer);
119 EIGEN_STRONG_INLINE
void memcpy(
void* dst,
const void* src,
size_t n)
const {
120 ::memcpy(dst, src, n);
122 EIGEN_STRONG_INLINE
void memcpyHostToDevice(
void* dst,
const void* src,
size_t n)
const {
125 EIGEN_STRONG_INLINE
void memcpyDeviceToHost(
void* dst,
const void* src,
size_t n)
const {
129 EIGEN_STRONG_INLINE
void memset(
void* buffer,
int c,
size_t n)
const {
130 ::memset(buffer, c, n);
133 EIGEN_STRONG_INLINE
int numThreads()
const {
137 EIGEN_STRONG_INLINE
size_t firstLevelCacheSize()
const {
138 return l1CacheSize();
141 EIGEN_STRONG_INLINE
size_t lastLevelCacheSize()
const {
143 return l3CacheSize() / num_threads_;
146 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
int majorDeviceVersion()
const {
151 template <
class Function,
class... Args>
152 EIGEN_STRONG_INLINE Notification* enqueue(Function&& f, Args&&... args)
const {
153 Notification* n =
new Notification();
154 pool_->Schedule(std::bind(&FunctionWrapperWithNotification<Function, Args...>::run, n, f, args...));
158 template <
class Function,
class... Args>
159 EIGEN_STRONG_INLINE
void enqueue_with_barrier(Barrier* b,
161 Args&&... args)
const {
162 pool_->Schedule(std::bind(
163 &FunctionWrapperWithBarrier<Function, Args...>::run, b, f, args...));
166 template <
class Function,
class... Args>
167 EIGEN_STRONG_INLINE
void enqueueNoNotification(Function&& f, Args&&... args)
const {
168 pool_->Schedule(std::bind(f, args...));
173 EIGEN_STRONG_INLINE
int currentThreadId()
const {
174 return pool_->CurrentThreadId();
182 void parallelFor(Index n,
const TensorOpCost& cost,
183 std::function<Index(Index)> block_align,
184 std::function<
void(Index, Index)> f)
const {
185 typedef TensorCostModel<ThreadPoolDevice> CostModel;
186 if (n <= 1 || numThreads() == 1 ||
187 CostModel::numThreads(n, cost, static_cast<int>(numThreads())) == 1) {
198 double block_size_f = 1.0 / CostModel::taskSize(1, cost);
199 Index block_size = numext::mini(n, numext::maxi<Index>(1, block_size_f));
200 const Index max_block_size =
201 numext::mini(n, numext::maxi<Index>(1, 2 * block_size_f));
203 Index new_block_size = block_align(block_size);
204 eigen_assert(new_block_size >= block_size);
205 block_size = numext::mini(n, new_block_size);
207 Index block_count = divup(n, block_size);
210 double max_efficiency =
211 static_cast<double>(block_count) /
212 (divup<int>(block_count, numThreads()) * numThreads());
215 for (Index prev_block_count = block_count; prev_block_count > 1;) {
218 Index coarser_block_size = divup(n, prev_block_count - 1);
220 Index new_block_size = block_align(coarser_block_size);
221 eigen_assert(new_block_size >= coarser_block_size);
222 coarser_block_size = numext::mini(n, new_block_size);
224 if (coarser_block_size > max_block_size) {
228 const Index coarser_block_count = divup(n, coarser_block_size);
229 eigen_assert(coarser_block_count < prev_block_count);
230 prev_block_count = coarser_block_count;
231 const double coarser_efficiency =
232 static_cast<double>(coarser_block_count) /
233 (divup<int>(coarser_block_count, numThreads()) * numThreads());
234 if (coarser_efficiency + 0.01 >= max_efficiency) {
236 block_size = coarser_block_size;
237 block_count = coarser_block_count;
238 if (max_efficiency < coarser_efficiency) {
239 max_efficiency = coarser_efficiency;
247 Barrier barrier(static_cast<unsigned int>(block_count));
248 std::function<void(Index, Index)> handleRange;
249 handleRange = [=, &handleRange, &barrier, &f](Index first, Index last) {
250 if (last - first <= block_size) {
257 Index mid = first + divup((last - first) / 2, block_size) * block_size;
258 pool_->Schedule([=, &handleRange]() { handleRange(mid, last); });
259 pool_->Schedule([=, &handleRange]() { handleRange(first, mid); });
266 void parallelFor(Index n,
const TensorOpCost& cost,
267 std::function<
void(Index, Index)> f)
const {
268 parallelFor(n, cost,
nullptr, std::move(f));
272 ThreadPoolInterface* pool_;
279 #endif // EIGEN_CXX11_TENSOR_TENSOR_DEVICE_THREAD_POOL_H Namespace containing all symbols from the Eigen library.
Definition: AdolcForward:45