Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2021-02-14 13:33:23

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