Adding ability to customize the printer of the LLVMIRLoweringPass

Author: Renaud-K

flang/include/flang/Optimizer/CodeGen/CodeGen.h

@@ -12,6 +12,8 @@
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassRegistry.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Support/raw_ostream.h"
 #include <memory>
 
 namespace fir {
@@ -36,9 +38,13 @@ std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> createFirTargetRewritePass(
 /// Convert FIR to the LLVM IR dialect
 std::unique_ptr<mlir::Pass> createFIRToLLVMPass();
 
+using LLVMIRLoweringPrinter =
+    std::function<void(llvm::Module &, llvm::raw_ostream &)>;
 /// Convert the LLVM IR dialect to LLVM-IR proper
-std::unique_ptr<mlir::Pass>
-createLLVMDialectToLLVMPass(llvm::raw_ostream &output);
+std::unique_ptr<mlir::Pass> createLLVMDialectToLLVMPass(
+    llvm::raw_ostream &output,
+    LLVMIRLoweringPrinter printer =
+        [](llvm::Module &M, llvm::raw_ostream &out) { M.print(out, nullptr); });
 
 // declarative passes
 #define GEN_PASS_REGISTRATION

flang/lib/Optimizer/CodeGen/CodeGen.cpp

@@ -2926,7 +2926,9 @@ class FIRToLLVMLowering : public fir::FIRToLLVMLoweringBase<FIRToLLVMLowering> {
 struct LLVMIRLoweringPass
     : public mlir::PassWrapper<LLVMIRLoweringPass,
                                mlir::OperationPass<mlir::ModuleOp>> {
-  LLVMIRLoweringPass(raw_ostream &output) : output{output} {}
+  using Printer = fir::LLVMIRLoweringPrinter;
+  LLVMIRLoweringPass(raw_ostream &output, Printer p)
+      : output{output}, printer{p} {}
 
   mlir::ModuleOp getModule() { return getOperation(); }
 
@@ -2936,7 +2938,7 @@ struct LLVMIRLoweringPass
     LLVMContext llvmCtx;
     if (auto llvmModule = mlir::translateModuleToLLVMIR(
             getModule(), llvmCtx, optName ? *optName : "FIRModule")) {
-      llvmModule->print(output, nullptr);
+      printer(*llvmModule, output);
       return;
     }
 
@@ -2946,6 +2948,7 @@ struct LLVMIRLoweringPass
 
 private:
   raw_ostream &output;
+  Printer printer;
 };
 
 } // namespace
@@ -2955,6 +2958,7 @@ std::unique_ptr<mlir::Pass> fir::createFIRToLLVMPass() {
 }
 
 std::unique_ptr<mlir::Pass>
-fir::createLLVMDialectToLLVMPass(raw_ostream &output) {
-  return std::make_unique<LLVMIRLoweringPass>(output);
+fir::createLLVMDialectToLLVMPass(raw_ostream &output,
+                                 fir::LLVMIRLoweringPrinter printer) {
+  return std::make_unique<LLVMIRLoweringPass>(output, printer);
 }