[MLIR] Add label to print-ir pass

[jpienaar]

Remember to clean up commit message

mlir/include/mlir/Transforms/Passes.h

 
 //===----------------------------------------------------------------------===//
 // Passes
 //===----------------------------------------------------------------------===//
 
 #define GEN_PASS_DECL_CANONICALIZER
 #define GEN_PASS_DECL_CONTROLFLOWSINK
 #define GEN_PASS_DECL_CSEPASS
+#define GEN_PASS_DECL_INLINER
 #define GEN_PASS_DECL_LOOPINVARIANTCODEMOTION
-#define GEN_PASS_DECL_STRIPDEBUGINFO
+#define GEN_PASS_DECL_PRINTIRPASS
 #define GEN_PASS_DECL_PRINTOPSTATS
-#define GEN_PASS_DECL_INLINER
+#define GEN_PASS_DECL_STRIPDEBUGINFO
 #define GEN_PASS_DECL_SCCP
 #define GEN_PASS_DECL_SYMBOLDCE
 #define GEN_PASS_DECL_SYMBOLPRIVATIZE
 #define GEN_PASS_DECL_TOPOLOGICALSORT
 #include "mlir/Transforms/Passes.h.inc"
 
 /// Creates an instance of the Canonicalizer pass, configured with default
 /// settings (which can be overridden by pass options on the command line).

 
 /// Creates a pass to perform control-flow sinking.
 std::unique_ptr<Pass> createControlFlowSinkPass();
 
 /// Creates a pass to perform common sub expression elimination.
 std::unique_ptr<Pass> createCSEPass();
 
 /// Creates a pass to print IR on the debug stream.
-std::unique_ptr<Pass> createPrintIRPass();
+std::unique_ptr<Pass> createPrintIRPass(const PrintIRPassOptions & = {});
 
 /// Creates a pass that generates IR to verify ops at runtime.
 std::unique_ptr<Pass> createGenerateRuntimeVerificationPass();
 
 /// Creates a loop invariant code motion pass that hoists loop invariant
 /// instructions out of the loop.
 std::unique_ptr<Pass> createLoopInvariantCodeMotionPass();
 

mlir/include/mlir/Transforms/Passes.td

 
 def PrintIRPass : Pass<"print-ir"> {
   let summary = "Print IR on the debug stream";
   let description = [{
     Print the entire IR on the debug stream. This is meant for debugging
     purposes to inspect the IR at a specific point in the pipeline.
   }];
   let constructor = "mlir::createPrintIRPass()";
+  let options = [
+    Option<"label", "label", "std::string", /*default=*/"", "Label">,
+  ];
 }
 
 def GenerateRuntimeVerification : Pass<"generate-runtime-verification"> {
   let summary = "Generate additional runtime op verification checks";
   let description = [{
     This pass generates op-specific runtime checks using the
     `RuntimeVerifiableOpInterface`. It can be run for debugging purposes after
     passes that are suspected to introduce faulty IR.

mlir/lib/Transforms/PrintIR.cpp

 //===- PrintIR.cpp - Pass to dump IR on debug stream ----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/Passes.h"
+#include "llvm/Support/Debug.h"
 
 namespace mlir {
 namespace {
 
 #define GEN_PASS_DEF_PRINTIRPASS
 #include "mlir/Transforms/Passes.h.inc"
 
 struct PrintIRPass : public impl::PrintIRPassBase<PrintIRPass> {
-  PrintIRPass() = default;
+  using impl::PrintIRPassBase<PrintIRPass>::PrintIRPassBase;
 
-  void runOnOperation() override { getOperation()->dump(); }
+  void runOnOperation() override {
+    llvm::dbgs() << "// -----// IR Dump";
+    if (!this->label.empty())
+      llvm::dbgs() << " " << this->label;
+    llvm::dbgs() << " //----- //\n";
+    getOperation()->dump();
+  }
 };
 
 } // namespace
 
-std::unique_ptr<Pass> createPrintIRPass() {
-  return std::make_unique<PrintIRPass>();
+std::unique_ptr<Pass> createPrintIRPass(const PrintIRPassOptions &options) {
+  return std::make_unique<PrintIRPass>(options);
 }
 
 } // namespace mlir
 No newline at end of file