TensorFlow/interface/NoThreadPool.h

0001 /*
0002  * Custom TensorFlow thread pool implementation that does no threading at all,
0003  * but schedules all tasks in the caller thread.
0004  * Based on TensorFlow 2.1.
0005  * For more info, see https://gitlab.cern.ch/mrieger/CMSSW-DNN.
0006  *
0007  * Author: Marcel Rieger
0008  */
0009
0010 #ifndef PHYSICSTOOLS_TENSORFLOW_NOTHREADPOOL_H
0011 #define PHYSICSTOOLS_TENSORFLOW_NOTHREADPOOL_H
0012
0013 #include "FWCore/Utilities/interface/thread_safety_macros.h"
0014
0015 #include "tensorflow/core/lib/core/threadpool.h"
0016 #include "tensorflow/core/lib/core/threadpool_options.h"
0017
0018 namespace tensorflow {
0019
0020   class NoThreadPool : public tensorflow::thread::ThreadPoolInterface {
0021   public:
0022     static NoThreadPool& instance() {
0023       CMS_THREAD_SAFE static NoThreadPool pool;
0024       return pool;
0025     }
0026
0027     explicit NoThreadPool() : numScheduleCalled_(0) {}
0028
0029     void Schedule(std::function<void()> fn) override {
0030       numScheduleCalled_ += 1;
0031       fn();
0032     }
0033
0034     void ScheduleWithHint(std::function<void()> fn, int start, int end) override { Schedule(fn); }
0035
0036     void Cancel() override {}
0037
0038     int NumThreads() const override { return 1; }
0039
0040     int CurrentThreadId() const override { return -1; }
0041
0042     int GetNumScheduleCalled() { return numScheduleCalled_; }
0043
0044   private:
0045     std::atomic<int> numScheduleCalled_;
0046   };
0047
0048 }  // namespace tensorflow
0049
0050 #endif  // PHYSICSTOOLS_TENSORFLOW_NOTHREADPOOL_H