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_SERVICE_HLO_RUNNER_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_RUNNER_H_ 18 19 #include <memory> 20 #include <string> 21 #include <vector> 22 23 #include "tensorflow/compiler/xla/service/backend.h" 24 #include "tensorflow/compiler/xla/service/compiler.h" 25 #include "tensorflow/compiler/xla/service/hlo_computation.h" 26 #include "tensorflow/compiler/xla/service/hlo_module.h" 27 #include "tensorflow/compiler/xla/status_macros.h" 28 #include "tensorflow/compiler/xla/statusor.h" 29 #include "tensorflow/compiler/xla/types.h" 30 #include "tensorflow/compiler/xla/xla_data.pb.h" 31 #include "tensorflow/core/lib/gtl/array_slice.h" 32 #include "tensorflow/core/platform/stream_executor_no_cuda.h" 33 34 namespace xla { 35 36 // A base class for running an HloModule. This executes the given HloModule on a 37 // certain backend directly without using the client interface. HloModule can be 38 // explicitly built, or loaded from a serialization file (e.g., hlo proto 39 // file), or parsed from a hlo textual IR string. 40 class HloRunner { 41 public: 42 HloRunner(); 43 44 HloRunner(::perftools::gputools::Platform* platform); 45 46 ~HloRunner(); 47 48 // Converts an HloModule from the given hlo textual IR string (in 49 // HloModule::ToString format). 50 static StatusOr<std::unique_ptr<HloModule>> CreateModuleFromString( 51 const tensorflow::StringPiece hlo_string, 52 const DebugOptions& debug_options); 53 54 // Reads the proto file in xla.HloProto format, creates and returns the 55 // HloModule. 56 static StatusOr<std::unique_ptr<HloModule>> ReadModuleFromBinaryProtoFile( 57 const std::string& filename, const DebugOptions& debug_options); 58 static StatusOr<std::unique_ptr<HloModule>> ReadModuleFromTextProtoFile( 59 const std::string& filename, const DebugOptions& debug_options); 60 61 // Reads the hlo text dump file in HloModule::ToString format, creates and 62 // returns the HloModule. 63 static StatusOr<std::unique_ptr<HloModule>> ReadModuleFromHloTextFile( 64 const std::string& filename, const DebugOptions& debug_options); 65 66 // Executes the given module with given literals as input and returns the 67 // result as a Literal. The LiteralPtr type accepts Literal* or 68 // std::unique_ptr<Literal>. 69 // 70 // If run_hlo_passes is false, the module will be executed without Hlo 71 // optimization. 72 template <typename LiteralPtr> 73 StatusOr<std::unique_ptr<Literal>> Execute( 74 std::unique_ptr<HloModule> module, 75 const tensorflow::gtl::ArraySlice<LiteralPtr> arguments, 76 bool run_hlo_passes = true); 77 78 // If backend is not created in the constructor, creates and returns the 79 // default backend. If creation fails, crashes the program. 80 // 81 // This creates the backend lazily so it's possible to instantiate an 82 // HloRunner in a program without any backends linked in. 83 Backend& backend(); 84 85 private: 86 StatusOr<std::unique_ptr<Literal>> ExecuteInternal( 87 std::unique_ptr<HloModule> module, 88 const tensorflow::gtl::ArraySlice<Literal*> arguments, 89 bool run_hlo_passes = true); 90 91 struct EigenThreadPoolWrapper; 92 93 std::unique_ptr<EigenThreadPoolWrapper> thread_pool_wrapper_; 94 95 std::unique_ptr<Backend> backend_; 96 }; 97 98 template <typename LiteralPtr> 99 StatusOr<std::unique_ptr<Literal>> HloRunner::Execute( 100 std::unique_ptr<HloModule> module, 101 const tensorflow::gtl::ArraySlice<LiteralPtr> arguments, 102 bool run_hlo_passes) { 103 // Construct a vector of plain pointers for the arguments. 104 std::vector<Literal*> argument_pointers; 105 for (const auto& argument : arguments) { 106 argument_pointers.push_back(&*argument); 107 } 108 return ExecuteInternal(std::move(module), argument_pointers, run_hlo_passes); 109 } 110 111 } // namespace xla 112 113 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_RUNNER_H_ 114
