File size: 2,455 Bytes
d1d4335 |
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 |
#ifndef CAFFE2_UTILS_THREADPOOL_H_
#define CAFFE2_UTILS_THREADPOOL_H_
#include "ThreadPoolCommon.h"
#include <atomic>
#include <functional>
#include <memory>
#include <mutex>
#include <vector>
#include "c10/util/Flags.h"
#include "caffe2/core/common.h"
//
// A work-stealing threadpool loosely based off of pthreadpool
//
namespace caffe2 {
struct Task;
class WorkersPool;
constexpr size_t kCacheLineSize = 64;
// A threadpool with the given number of threads.
// NOTE: the kCacheLineSize alignment is present only for cache
// performance, and is not strictly enforced (for example, when
// the object is created on the heap). Thus, in order to avoid
// misaligned intrinsics, no SSE instructions shall be involved in
// the ThreadPool implementation.
// Note: alignas is disabled because some compilers do not deal with
// TORCH_API and alignas annotations at the same time.
class TORCH_API /*alignas(kCacheLineSize)*/ ThreadPool {
public:
static ThreadPool* createThreadPool(int numThreads);
static std::unique_ptr<ThreadPool> defaultThreadPool();
virtual ~ThreadPool() = default;
// Returns the number of threads currently in use
virtual int getNumThreads() const = 0;
virtual void setNumThreads(size_t numThreads) = 0;
// Sets the minimum work size (range) for which to invoke the
// threadpool; work sizes smaller than this will just be run on the
// main (calling) thread
void setMinWorkSize(size_t size) {
std::lock_guard<std::mutex> guard(executionMutex_);
minWorkSize_ = size;
}
size_t getMinWorkSize() const {
return minWorkSize_;
}
virtual void run(const std::function<void(int, size_t)>& fn, size_t range) = 0;
// Run an arbitrary function in a thread-safe manner accessing the Workers
// Pool
virtual void withPool(const std::function<void(WorkersPool*)>& fn) = 0;
protected:
static size_t defaultNumThreads_;
mutable std::mutex executionMutex_;
size_t minWorkSize_;
};
size_t getDefaultNumThreads();
} // namespace caffe2
C10_DECLARE_bool(caffe2_threadpool_force_inline);
// Whether or not threadpool caps apply to Android
C10_DECLARE_int(caffe2_threadpool_android_cap);
// Whether or not threadpool caps apply to iOS and MacOS
C10_DECLARE_int(caffe2_threadpool_ios_cap);
C10_DECLARE_int(caffe2_threadpool_macos_cap);
C10_DECLARE_int(pthreadpool_size);
#endif // CAFFE2_UTILS_THREADPOOL_H_
|