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_CPU_CPU_COMPILER_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_COMPILER_H_ 18 19 #include <memory> 20 21 #include "absl/types/span.h" 22 #include "llvm/Target/TargetMachine.h" 23 #include "tensorflow/compiler/xla/cpu_function_runtime.h" 24 #include "tensorflow/compiler/xla/service/cpu/target_machine_features.h" 25 #include "tensorflow/compiler/xla/service/executable.h" 26 #include "tensorflow/compiler/xla/service/hlo_module.h" 27 #include "tensorflow/compiler/xla/service/llvm_compiler.h" 28 #include "tensorflow/compiler/xla/statusor.h" 29 #include "tensorflow/core/platform/macros.h" 30 #include "tensorflow/core/platform/stream_executor_no_cuda.h" 31 32 namespace xla { 33 namespace cpu { 34 35 // This class wraps the configurability options that LLVM exposes including: the 36 // target triple, the target cpu and the target features. It also includes the 37 // desired linkage name for the computation entry point. 38 class CpuAotCompilationOptions : public AotCompilationOptions { 39 public: 40 // Relocation models available for compilation. 41 enum class RelocationModel { 42 // Corresponds to the -fno-pic compiler option. 43 Static, 44 // Corresponds to the -fpic compiler option. 45 SmallPic, 46 // Corresponds to the -fPIC compiler option. 47 BigPic, 48 // Corresponds to the -fpie compiler option. 49 SmallPie, 50 // Corresponds to the -fPIE compiler option. 51 BigPie 52 }; 53 54 CpuAotCompilationOptions(string triple, string cpu_name, string features, 55 string entry_point_name, 56 RelocationModel relocation_model); 57 58 ~CpuAotCompilationOptions() override; 59 60 se::Platform::Id PlatformId() const override; 61 62 // The triple used for compilation, similar to clang's -target flag. triple()63 const string& triple() const { return triple_; } 64 // The CPU name used for compilation, similar to clang's -mcpu flag. cpu_name()65 const string& cpu_name() const { return cpu_name_; } 66 // The target features used for compilation ("+avx2", "+neon", etc). features()67 const string& features() const { return features_; } 68 // The name to be used for the compiled code's entry point. entry_point_name()69 const string& entry_point_name() const { return entry_point_name_; } 70 // The relocation model used for compilation. relocation_model()71 RelocationModel relocation_model() const { return relocation_model_; } 72 73 private: 74 const string triple_; 75 const string cpu_name_; 76 const string features_; 77 const string entry_point_name_; 78 const RelocationModel relocation_model_; 79 }; 80 81 class CpuAotCompilationResult : public AotCompilationResult { 82 public: 83 CpuAotCompilationResult( 84 ObjectFileData object_file_data, 85 std::vector<cpu_function_runtime::BufferInfo> buffer_infos, 86 int64_t result_buffer_index, 87 std::unique_ptr<HloProfilePrinterData> hlo_profile_printer_data); 88 ~CpuAotCompilationResult(); 89 hlo_profile_printer_data()90 HloProfilePrinterData* hlo_profile_printer_data() const { 91 return hlo_profile_printer_data_.get(); 92 } 93 object_file_data()94 const ObjectFileData& object_file_data() const { return object_file_data_; } buffer_infos()95 const std::vector<cpu_function_runtime::BufferInfo>& buffer_infos() const { 96 return buffer_infos_; 97 } result_buffer_index()98 int64 result_buffer_index() const { return result_buffer_index_; } 99 100 private: 101 // Contains the compiled computation: an object file. 102 const ObjectFileData object_file_data_; 103 104 // A list of BufferInfo objects describing the buffers used by the XLA 105 // computation. 106 const std::vector<cpu_function_runtime::BufferInfo> buffer_infos_; 107 108 // Contains which buffer index into |buffer_sizes| was designated to the 109 // result of the computation. This buffer should be passed into the output 110 // parameter when calling the compiled computation. 111 const int64 result_buffer_index_; 112 113 // Contains an instance of HloProfilePrinterData if HLO profiling is enabled, 114 // otherwise is nullptr. 115 std::unique_ptr<HloProfilePrinterData> hlo_profile_printer_data_; 116 }; 117 118 // CPU-targeting implementation of the XLA Compiler interface. 119 // 120 // The compiler translates XLA HLO code into LLVM IR and uses LLVM's JIT 121 // infrastructure to create an executable "blob" that can then be returned 122 // wrapped in CpuExecutable and actually invoked. 123 class CpuCompiler : public LLVMCompiler { 124 public: 125 CpuCompiler(); ~CpuCompiler()126 ~CpuCompiler() override {} 127 128 StatusOr<std::vector<std::unique_ptr<Executable>>> Compile( 129 std::unique_ptr<HloModuleGroup> module_group, 130 std::vector<std::vector<se::StreamExecutor*>> stream_execs, 131 const CompileOptions& options) override; 132 133 StatusOr<std::unique_ptr<HloModule>> RunHloPasses( 134 std::unique_ptr<HloModule> module, se::StreamExecutor* stream_exec, 135 const CompileOptions& options) override; 136 137 StatusOr< 138 std::tuple<std::unique_ptr<HloModule>, std::unique_ptr<BufferAssignment>>> 139 RunHloPassesAndBufferAssignement(std::unique_ptr<HloModule> module, 140 se::StreamExecutor* executor, bool optimize, 141 const CompileOptions& options) override; 142 143 StatusOr<std::unique_ptr<Executable>> RunBackend( 144 std::unique_ptr<HloModule> module, se::StreamExecutor* stream_exec, 145 const CompileOptions& options) override; 146 147 StatusOr<std::vector<std::unique_ptr<AotCompilationResult>>> 148 CompileAheadOfTime(std::unique_ptr<HloModuleGroup> module_group, 149 const AotCompilationOptions& options) override; 150 151 se::Platform::Id PlatformId() const override; 152 153 HloCostAnalysis::ShapeSizeFunction ShapeSizeBytesFunction() const override; 154 155 private: 156 // Initialize the LLVM target. 157 static void InitializeLLVMTarget(); 158 159 // Runs the HLO passes which are necessary for both optimizations and 160 // correctness. 161 Status RunHloPasses(HloModule* module, bool is_aot_compile, 162 llvm::TargetMachine* target_machine); 163 164 // Runs HLO passes up to and including layout assignment. 165 Status RunHloPassesThroughLayoutAssn( 166 HloModule* module, bool /*is_aot_compile*/, 167 LLVMTargetMachineFeatures* target_machine_features); 168 169 // Runs HLO passes after layout assignment. 170 Status RunHloPassesAfterLayoutAssn( 171 HloModule* module, bool is_aot_compile, 172 LLVMTargetMachineFeatures* target_machine_features); 173 174 TF_DISALLOW_COPY_AND_ASSIGN(CpuCompiler); 175 }; 176 177 } // namespace cpu 178 } // namespace xla 179 180 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_COMPILER_H_ 181