| 793 | } |
| 794 | |
| 795 | DirectSession::DirectSession(const SessionOptions& options, |
| 796 | const DeviceMgr* device_mgr, |
| 797 | DirectSessionFactory* factory, |
| 798 | const std::vector<unsigned>& visible_cpus) |
| 799 | : options_(options), |
| 800 | device_mgr_(device_mgr), |
| 801 | factory_(factory), |
| 802 | cancellation_manager_(new CancellationManager()), |
| 803 | operation_timeout_in_ms_(options_.config.operation_timeout_in_ms()), |
| 804 | use_multi_stream_(false), multi_stream_num_(0), |
| 805 | multi_stream_shared_rmgr_(nullptr) { |
| 806 | const int thread_pool_size = |
| 807 | options_.config.session_inter_op_thread_pool_size(); |
| 808 | if (thread_pool_size > 0) { |
| 809 | for (int i = 0; i < thread_pool_size; ++i) { |
| 810 | thread::ThreadPool* pool = nullptr; |
| 811 | bool owned = false; |
| 812 | init_error_.Update(NewThreadPoolFromThreadPoolOptions( |
| 813 | options_, options_.config.session_inter_op_thread_pool(i), i, &pool, |
| 814 | &owned)); |
| 815 | thread_pools_.emplace_back(pool, owned); |
| 816 | } |
| 817 | } else if (options_.config.use_per_session_threads()) { |
| 818 | thread_pools_.emplace_back(NewThreadPoolFromSessionOptions(options_), |
| 819 | true /* owned */); |
| 820 | } else { |
| 821 | thread_pools_.emplace_back(GlobalThreadPool(options), false /* owned */); |
| 822 | // Run locally if environment value of TF_NUM_INTEROP_THREADS is negative |
| 823 | // and config.inter_op_parallelism_threads is unspecified or negative. |
| 824 | static const int env_num_threads = NumInterOpThreadsFromEnvironment(); |
| 825 | if (options_.config.inter_op_parallelism_threads() < 0 || |
| 826 | (options_.config.inter_op_parallelism_threads() == 0 && |
| 827 | env_num_threads < 0)) { |
| 828 | run_in_caller_thread_ = true; |
| 829 | } |
| 830 | MemoryPlannerFactory::GetMemoryPlanner()->SetThreadPool(GlobalThreadPool(options)); |
| 831 | } |
| 832 | |
| 833 | const int stage_subgraph_thread_pool_size = |
| 834 | options_.config.session_stage_subgraph_thread_pool_size(); |
| 835 | for (int i = 0; i < stage_subgraph_thread_pool_size; ++i) { |
| 836 | std::pair<thread::ThreadPoolInterface*, thread::ThreadPoolInterface*> |
| 837 | pool_interface; |
| 838 | NewStageSubGraphThreadPoolFromStageSubGraphThreadPoolOptions( |
| 839 | options_, options_.config.session_stage_subgraph_thread_pool(i), i, |
| 840 | pool_interface); |
| 841 | stage_subgraph_thread_pools_.emplace_back(pool_interface); |
| 842 | } |
| 843 | |
| 844 | bool use_cost_model_executor = false; |
| 845 | bool use_inline_executor = false; |
| 846 | Status s = |
| 847 | ReadBoolFromEnvVar("USE_COST_MODEL_EXECUTOR", false, &use_cost_model_executor); |
| 848 | if (!s.ok()) { |
| 849 | LOG(FATAL) << s.error_message(); |
| 850 | } |
| 851 | s = ReadBoolFromEnvVar("USE_INLINE_EXECUTOR", false, &use_inline_executor); |
| 852 | if (!s.ok()) { |
nothing calls this directly
no test coverage detected