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 "tensorflow/compiler/xla/service/executable.h" 22 #include "tensorflow/compiler/xla/service/hlo_module.h" 23 #include "tensorflow/compiler/xla/service/llvm_compiler.h" 24 #include "tensorflow/compiler/xla/statusor.h" 25 #include "tensorflow/core/lib/gtl/array_slice.h" 26 #include "tensorflow/core/platform/macros.h" 27 #include "tensorflow/core/platform/stream_executor_no_cuda.h" 28 29 namespace xla { 30 namespace cpu { 31 32 // This class wraps the configurability options that LLVM exposes including: the 33 // target triple, the target cpu and the target features. It also includes the 34 // desired linkage name for the computation entry point. 35 class CpuAotCompilationOptions : public AotCompilationOptions { 36 public: 37 // Relocation models available for compilation. 38 enum class RelocationModel { 39 // Corresponds to the -fno-pic compiler option. 40 Static, 41 // Corresponds to the -fpic compiler option. 42 SmallPic, 43 // Corresponds to the -fPIC compiler option. 44 BigPic, 45 // Corresponds to the -fpie compiler option. 46 SmallPie, 47 // Corresponds to the -fPIE compiler option. 48 BigPie 49 }; 50 51 CpuAotCompilationOptions(string triple, string cpu_name, string features, 52 string entry_point_name, 53 RelocationModel relocation_model); 54 ~CpuAotCompilationOptions() override; 55 56 perftools::gputools::Platform::Id PlatformId() const override; 57 58 // The triple used for compilation, similar to clang's -target flag. 59 const string& triple() const { return triple_; } 60 // The CPU name used for compilation, similar to clang's -mcpu flag. 61 const string& cpu_name() const { return cpu_name_; } 62 // The target features used for compilation ("+avx2", "+neon", etc). 63 const string& features() const { return features_; } 64 // The name to be used for the compiled code's entry point. 65 const string& entry_point_name() const { return entry_point_name_; } 66 // The relocation model used for compilation. 67 RelocationModel relocation_model() const { return relocation_model_; } 68 69 private: 70 const string triple_; 71 const string cpu_name_; 72 const string features_; 73 const string entry_point_name_; 74 const RelocationModel relocation_model_; 75 }; 76 77 class CpuAotCompilationResult : public AotCompilationResult { 78 public: 79 CpuAotCompilationResult(ObjectFileData object_file_data, 80 BufferSizes buffer_sizes, int64 result_buffer_index); 81 ~CpuAotCompilationResult(); 82 83 const ObjectFileData& object_file_data() const { return object_file_data_; } 84 const BufferSizes& buffer_sizes() const { return buffer_sizes_; } 85 int64 result_buffer_index() const { return result_buffer_index_; } 86 87 private: 88 // Contains the compiled computation: an object file. 89 const ObjectFileData object_file_data_; 90 91 // The list of buffer sizes which should be allocated in order to execute the 92 // compiled computation. These buffers are used for temporary buffers used 93 // ephemerally during computation as well as the output result. 94 const BufferSizes buffer_sizes_; 95 96 // Contains which buffer index into |buffer_sizes| was designated to the 97 // result of the computation. This buffer should be passed into the output 98 // parameter when calling the compiled computation. 99 const int64 result_buffer_index_; 100 }; 101 102 // CPU-targeting implementation of the XLA Compiler interface. 103 // 104 // The compiler translates XLA HLO code into LLVM IR and uses LLVM's JIT 105 // infrastructure to create an executable "blob" that can then be returned 106 // wrapped in CpuExecutable and actually invoked. 107 class CpuCompiler : public LLVMCompiler { 108 public: 109 CpuCompiler(); 110 ~CpuCompiler() override {} 111 112 // Bring in 113 // StatusOr<std::vector<std::unique_ptr<Executable>>> Compile( 114 // std::vector<std::unique_ptr<HloModule>> modules, 115 // std::vector<std::vector<perftools::gputools::StreamExecutor*>> 116 // stream_execs) 117 using LLVMCompiler::Compile; 118 119 StatusOr<std::unique_ptr<HloModule>> RunHloPasses( 120 std::unique_ptr<HloModule> module, 121 perftools::gputools::StreamExecutor* stream_exec, 122 DeviceMemoryAllocator* device_allocator) override; 123 124 StatusOr<std::unique_ptr<Executable>> RunBackend( 125 std::unique_ptr<HloModule> module, 126 perftools::gputools::StreamExecutor* stream_exec, 127 DeviceMemoryAllocator* device_allocator) override; 128 129 StatusOr<std::vector<std::unique_ptr<AotCompilationResult>>> 130 CompileAheadOfTime(std::vector<std::unique_ptr<HloModule>> modules, 131 const AotCompilationOptions& options) override; 132 133 perftools::gputools::Platform::Id PlatformId() const override; 134 135 HloCostAnalysis::ShapeSizeFunction ShapeSizeBytesFunction() const override; 136 137 private: 138 // Initialize the LLVM target. 139 static void InitializeLLVMTarget(); 140 141 // Runs the HLO passes which are necessary for both optimizations and 142 // correctness. 143 Status RunHloPasses(HloModule* module, bool is_aot_compile); 144 145 TF_DISALLOW_COPY_AND_ASSIGN(CpuCompiler); 146 }; 147 148 } // namespace cpu 149 } // namespace xla 150 151 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_COMPILER_H_ 152
