1 /* Copyright 2017 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_COMPILER_XLA_EXECUTABLE_RUN_OPTIONS_H_ 17 #define TENSORFLOW_COMPILER_XLA_EXECUTABLE_RUN_OPTIONS_H_ 18 19 // Intentionally forward declared so that ExecutableRunOptions can be linked 20 // into an XLA-compiled binary without having to link all of the pointed-to 21 // objects (e.g., for an ahead-of-time compiled CPU binary, the gpu tools don't 22 // need to be linked). 23 namespace perftools { 24 namespace gputools { 25 class Stream; 26 class Platform; 27 } 28 } 29 30 namespace tensorflow { 31 namespace thread { 32 class ThreadPool; 33 } 34 } 35 36 namespace Eigen { 37 struct ThreadPoolDevice; 38 } 39 40 namespace xla { 41 42 class DeviceMemoryAllocator; 43 class DeviceAssignment; 44 class ExecutionProfile; 45 46 // Class containing options for running a LocalExecutable. 47 class ExecutableRunOptions { 48 public: 49 // Specifies the allocator to use during execution. 50 ExecutableRunOptions& set_allocator(DeviceMemoryAllocator* allocator); 51 DeviceMemoryAllocator* allocator() const; 52 53 // If set, this is the device to run the computation on. Valid device_ordinal 54 // values are: 0 to # of devices - 1. These values are identical to the device 55 // ordinal values used by StreamExecutor. The device must be of the same type 56 // as the executable was compiled for. A value of -1 indicates this option has 57 // not been set. 58 ExecutableRunOptions& set_device_ordinal(int device_ordinal); 59 int device_ordinal() const; 60 61 // If set, this is the stream to run the computation on. The platform of the 62 // stream must match the platform the executable was built for. A value of 63 // nullptr indicates the option has not been set. 64 ExecutableRunOptions& set_stream(perftools::gputools::Stream* stream); 65 perftools::gputools::Stream* stream() const; 66 67 // Sets the thread pool on which to run parallel CPU backend 68 // computations. Does not take ownership. 69 ExecutableRunOptions& set_inter_op_thread_pool( 70 tensorflow::thread::ThreadPool* inter_op_thread_pool); 71 tensorflow::thread::ThreadPool* inter_op_thread_pool() const; 72 73 // Sets the thread pool device on which to run Eigen subcomputations. 74 // Does not take ownership. 75 ExecutableRunOptions& set_intra_op_thread_pool( 76 const Eigen::ThreadPoolDevice* intra_op_thread_pool); 77 const Eigen::ThreadPoolDevice* intra_op_thread_pool() const; 78 79 // If set, profiling information is written to 'profile'. 80 ExecutionProfile* execution_profile() const; 81 ExecutableRunOptions& set_execution_profile(ExecutionProfile* profile); 82 83 ExecutableRunOptions& set_device_assignment( 84 DeviceAssignment* device_assignment); 85 const DeviceAssignment* device_assignment() const; 86 87 private: 88 DeviceMemoryAllocator* allocator_ = nullptr; 89 int device_ordinal_ = -1; 90 DeviceAssignment* device_assignment_ = nullptr; 91 perftools::gputools::Stream* stream_ = nullptr; 92 tensorflow::thread::ThreadPool* inter_op_thread_pool_ = nullptr; 93 const Eigen::ThreadPoolDevice* intra_op_thread_pool_ = nullptr; 94 ExecutionProfile* execution_profile_ = nullptr; 95 }; 96 97 } // namespace xla 98 99 #endif // TENSORFLOW_COMPILER_XLA_EXECUTABLE_RUN_OPTIONS_H_ 100
