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_COMPILER_FUNCTOR_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_COMPILER_FUNCTOR_H_ 18 19 #include "llvm/IR/LegacyPassManager.h" 20 #include "llvm/IR/Module.h" 21 #include "llvm/Object/ObjectFile.h" 22 #include "llvm/Target/TargetMachine.h" 23 #include "tensorflow/compiler/xla/service/llvm_compiler.h" 24 #include "tensorflow/core/platform/logging.h" 25 26 namespace xla { 27 namespace cpu { 28 29 // Functor class for compiling an LLVM module down to an object file. For use by 30 // Orc JIT compile layer. 31 class CompilerFunctor { 32 public: 33 explicit CompilerFunctor( 34 llvm::TargetMachine* target_machine, int opt_level, 35 bool optimize_for_size, bool enable_fast_math, 36 bool disable_expensive_passes, 37 LLVMCompiler::ModuleHook pre_optimization_hook = nullptr, 38 LLVMCompiler::ModuleHook post_optimization_hook = nullptr, 39 std::function<void(const llvm::object::ObjectFile&)> post_codegen_hook = 40 nullptr) target_machine_(target_machine)41 : target_machine_(target_machine), 42 opt_level_(opt_level), 43 optimize_for_size_(optimize_for_size), 44 enable_fast_math_(enable_fast_math), 45 disable_expensive_passes_(disable_expensive_passes), 46 pre_optimization_hook_(std::move(pre_optimization_hook)), 47 post_optimization_hook_(std::move(post_optimization_hook)), 48 post_codegen_hook_(std::move(post_codegen_hook)) {} 49 50 // Compile a Module to an ObjectFile. 51 std::unique_ptr<llvm::MemoryBuffer> operator()( 52 llvm::Module& module) const; // NOLINT 53 54 private: 55 // Populates the given pass manager with TargetLibraryInfo and 56 // TargetTransformInfo passes. 57 void AddTargetInfoPasses(llvm::legacy::PassManagerBase* passes) const; 58 59 // Populates the given pass managers based on the optimization level. 60 void AddOptimizationPasses(llvm::legacy::PassManagerBase* module_passes, 61 llvm::legacy::FunctionPassManager* function_passes, 62 unsigned opt_level, unsigned size_level) const; 63 64 llvm::TargetMachine* target_machine_; 65 const unsigned opt_level_; 66 const bool optimize_for_size_; 67 const bool enable_fast_math_; 68 const bool disable_expensive_passes_; 69 LLVMCompiler::ModuleHook pre_optimization_hook_; 70 LLVMCompiler::ModuleHook post_optimization_hook_; 71 std::function<void(const llvm::object::ObjectFile&)> post_codegen_hook_; 72 }; 73 74 } // namespace cpu 75 } // namespace xla 76 77 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_COMPILER_FUNCTOR_H_ 78