1 #pragma once 2 3 #ifdef USE_PTHREADPOOL 4 5 #ifdef USE_INTERNAL_PTHREADPOOL_IMPL 6 #include <caffe2/utils/threadpool/pthreadpool.h> 7 #else 8 #include <pthreadpool.h> 9 #endif 10 11 #include <functional> 12 #include <memory> 13 #include <mutex> 14 15 namespace caffe2 { 16 17 class PThreadPool final { 18 public: 19 explicit PThreadPool(size_t thread_count); 20 ~PThreadPool() = default; 21 22 PThreadPool(const PThreadPool&) = delete; 23 PThreadPool& operator=(const PThreadPool&) = delete; 24 25 PThreadPool(PThreadPool&&) = delete; 26 PThreadPool& operator=(PThreadPool&&) = delete; 27 28 size_t get_thread_count() const; 29 void set_thread_count(size_t thread_count); 30 31 // Run, in parallel, function fn(task_id) over task_id in range [0, range). 32 // This function is blocking. All input is processed by the time it returns. 33 void run(const std::function<void(size_t)>& fn, size_t range); 34 35 private: 36 friend pthreadpool_t pthreadpool_(); 37 38 private: 39 mutable std::mutex mutex_; 40 std::unique_ptr<pthreadpool, decltype(&pthreadpool_destroy)> threadpool_; 41 }; 42 43 // Return a singleton instance of PThreadPool for ATen/TH multithreading. 44 PThreadPool* pthreadpool(); 45 46 // Exposes the underlying implementation of PThreadPool. 47 // Only for use in external libraries so as to unify threading across 48 // internal (i.e. ATen, etc.) and external (e.g. NNPACK, QNNPACK, XNNPACK) 49 // use cases. 50 pthreadpool_t pthreadpool_(); 51 52 } // namespace caffe2 53 54 #endif /* USE_PTHREADPOOL */ 55