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 #include "tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/dump_ir_pass.h" 17 18 #include "llvm/IR/Module.h" 19 #include "llvm/Support/FileSystem.h" 20 #include "llvm/Support/raw_ostream.h" 21 #include "tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/utils.h" 22 #include "tensorflow/compiler/xla/types.h" 23 #include "tensorflow/core/lib/core/stringpiece.h" 24 #include "tensorflow/core/lib/io/path.h" 25 #include "tensorflow/core/lib/strings/stringprintf.h" 26 #include "tensorflow/core/platform/logging.h" 27 28 namespace xla { 29 namespace gpu { 30 31 // Pass which dumps the IR of a module into a file. 32 // 33 // Because it is implemented as a FunctionPass (IR is dumped 34 // function-by-function) rather than as a ModulePass the resulting IR is not 35 // valid (missing metadata, for example) but is still useful for inspection. 36 // The pass needs to be a FunctionPass rather than a ModulePass because 37 // inserting ModulePasses is disruptive to LLVM's pass manager. For sequential 38 // FunctionPasses (also SCC passes, etc) the pass manager executes the passes 39 // sequentially on each function (SCC, etc). Inserting a ModulePass between 40 // FunctionPasses acts as a barrier forcing the FunctionPasses to execute fully 41 // across all functions prior to advancing to the next pass. For some reason 42 // this results in different generated code resulting in an undesirable 43 // Heisenberg effect when dumping the IR. 44 class DumpIrPass : public llvm::FunctionPass { 45 public: 46 explicit DumpIrPass(const string &output_filename) 47 : llvm::FunctionPass(id_), output_filename_(output_filename) {} 48 49 bool doInitialization(llvm::Module &M) override { 50 out_.reset(new llvm::raw_fd_ostream(llvm::StringRef(output_filename_), ec_, 51 llvm::sys::fs::F_None)); 52 if (ec_) { 53 LOG(FATAL) << "Unable to open " << output_filename_ 54 << " to dump LLVM IR: " << ec_.message(); 55 } 56 return false; 57 } 58 59 bool runOnFunction(llvm::Function &Function) override { 60 Function.print(*out_); 61 return false; 62 } 63 64 void getAnalysisUsage(llvm::AnalysisUsage &AU) const override { 65 AU.setPreservesAll(); 66 } 67 68 bool doFinalization(llvm::Module &M) override { 69 out_->close(); 70 return false; 71 } 72 73 private: 74 static char id_; 75 string output_filename_; 76 std::error_code ec_; 77 std::unique_ptr<llvm::raw_fd_ostream> out_; 78 }; 79 80 char DumpIrPass::id_ = 0; 81 82 void IrDumpingPassManager::run(llvm::Module &module) { 83 for (int i = 0; i < passes_.size(); ++i) { 84 llvm::Pass *P = passes_[i]; 85 if (dump_ir_) { 86 const llvm::PassInfo *PI = 87 llvm::PassRegistry::getPassRegistry()->getPassInfo(P->getPassID()); 88 const string basename = ReplaceFilenameExtension( 89 tensorflow::io::Basename(input_filename_), 90 tensorflow::strings::Printf( 91 "pass-%02d.before.%s.ll", i, 92 (PI == nullptr ? "unknown" : PI->getPassArgument().data()))); 93 llvm::legacy::PassManager::add( 94 new DumpIrPass(tensorflow::io::JoinPath(output_dir_, basename))); 95 } 96 llvm::legacy::PassManager::add(P); 97 } 98 passes_.clear(); 99 llvm::legacy::PassManager::run(module); 100 } 101 102 } // namespace gpu 103 } // namespace xla 104
