1 /* Copyright 2015 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_CORE_FRAMEWORK_RUN_HANDLER_H_ 17 #define TENSORFLOW_CORE_FRAMEWORK_RUN_HANDLER_H_ 18 19 #include "tensorflow/core/lib/core/threadpool.h" 20 #include "tensorflow/core/lib/histogram/histogram.h" 21 #include "tensorflow/core/platform/mutex.h" 22 #include "tensorflow/core/platform/thread_annotations.h" 23 #include "tensorflow/core/protobuf/config.pb.h" 24 25 namespace tensorflow { 26 27 class RunHandler; 28 29 // RunHandlerPool is a fixed size pool of pre-allocated RunHandlers 30 // that can be used for tracking inter-op work for a given Session::Run(). 31 // RunHandler(s) in the pool are initially 'inactive'. A RunHandler becomes 32 // 'active' when its unique_ptr is returned by Get() and is being used by a 33 // client. It becomes 'inactive' once more when its unique_ptr gets destroyed. 34 // 35 // Expected usage: 36 // 37 // * Create a single RunHandlerPool (say run_handler_pool_). 38 // 39 // * When a Session::Run() is invoked, obtain a handler by: 40 // auto handler = run_handler_pool_->Get(); 41 // 42 // * Use handler for scheduling all inter-op work by: 43 // handler->ScheduleInterOpClosure(closure); 44 // 45 // This class is thread safe. 46 class RunHandlerPool { 47 public: 48 explicit RunHandlerPool(int num_inter_op_threads); 49 ~RunHandlerPool(); 50 51 // Returns an inactive RunHandler from the pool. 52 // 53 // RunHandlers in RunHandlerPool are initially 'inactive'. 54 // A RunHandler becomes 'active' when its unique_ptr its returned by Get() 55 // and is being used by a client. It becomes 'inactive' once more when the 56 // unique_ptr is destroyed. 57 // 58 // Will block unless there is an inactive handler. 59 std::unique_ptr<RunHandler> Get(); 60 61 private: 62 class Impl; 63 friend class RunHandler; 64 65 std::unique_ptr<Impl> impl_; 66 }; 67 68 // RunHandler can be used to schedule inter-op closures to run on a global pool 69 // shared across all Session::Run(s). 70 // 71 // It can only be created via RunHandlerPool::Get(). 72 // 73 // This class can be used instead of directly scheduling closures on a global 74 // pool since it maintains a global view across all sessions and optimizes pool 75 // scheduling to improve (median and tail) latency. 76 // 77 // This class is thread safe. 78 class RunHandler { 79 public: 80 void ScheduleInterOpClosure(std::function<void()> fn); 81 82 ~RunHandler(); 83 84 private: 85 class Impl; 86 friend class RunHandlerPool::Impl; 87 88 explicit RunHandler(Impl* impl); 89 90 Impl* impl_; // NOT OWNED. 91 }; 92 93 } // end namespace tensorflow. 94 95 #endif // TENSORFLOW_CORE_FRAMEWORK_RUN_HANDLER_H_ 96