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