Add a pass to dump a given operation

[tpopp]

Summary

For use while debugging pipelines and wanting some intermediate stages.

[tpopp]

I often add a simple dump command inside of a pass at the beginning or end when I'm quickly trying to gather some data and don't want to dump everything. I considered making a tiny rfc regarding this to see if others prefer another way, but want to at least first run it by a couple of folks. I'm sure there's several ways to do the same thing, but this would be my preferred way of debugging a pipeline not presented in a test using opt and before my morning coffee.

[tpopp]

NFC: cleanup some ugliness

[jpienaar]

And this is a more fine grained version of mlir-print-ir-after ? What is missing there? Triggering only particular pipeline, particular part of pipeline?

[tpopp]

In D139843#3988794, @jpienaar wrote:

And this is a more fine grained version of mlir-print-ir-after ? What is missing there? Triggering only particular pipeline, particular part of pipeline?

This allows, for example, printing after an (unfortunately) unnamed pass, adding print statements when looking at the C++ code without first having to locate the opt version of the pass' name, a single pipeline with printing after several milestone passes. The overall thought is that this lends itself better to simple C++ based debugging when one is not starting with a textual opt pipeline. It does not add much in the way of new functionality.

[mehdi_amini]

LGTM, but please wait for converging with Jacques on this!

And this is a more fine grained version of mlir-print-ir-after ? What is missing there? Triggering only particular pipeline, particular part of pipeline?

While mlir-print-ir-after is a good debugging tool, it does not compose well when "using MLIR as a library". That is someone building their pipeline programmatically and wanting some control on the debuggability / traceability of the flow would need something more custom that "pass name".
As an example, what if you wanted to get some print to a file after "phase1" of something like "iree_compiler" where "phase1" (whatever that means) ends with running canonicalization?
You can't really ask mlir-print-ir-after=canonicalize right?
Also, the whole PassManagerOptions::addPrinterInstrumentation() logic hard-codes llvm::errs() right now. Of course one can re-implement their own instrumentation, but that becomes kind of "not using mlir-print-ir-after" as well.

The only thing I could wonder here would be how to make it convenient to have this as an instrumentation instead of a pass so that is stays decoupled from the pass-pipeline dump, the time statistics, etc.

printing after an (unfortunately) unnamed pass

Uh, seems like a bug that I wouldn't want to justify a feature to work around this.

mlir/include/mlir/Transforms/Passes.td

145

print-op ?

mlir/test/IR/print-ops.mlir

1

Can we write the test so that we don't need this flag?

[silvas]

Nice, this is definitely useful. Agree that "print op" is better than plural "print ops".

I'm not super familiar with the concept of "instrumentation" in this context, but even on the opt command line it is still useful to be able to drop this pass in somewhere specific as just a pass. I think that if you are printing out stuff then there is no point in measuring compilation time -- all the context switches, cache invalidation, etc. make it unrepresentative of the overall running time even if you were to somehow "factor out" the time between starting/ending printing.

[mehdi_amini]

In D139843#3991660, @silvas wrote:

I'm not super familiar with the concept of "instrumentation" in this context,

See https://mlir.llvm.org/docs/PassManagement/#pass-instrumentation ; basically it is conceptually important from the infra point of view to keep the separation between the actual pass-pipeline from the thing we do around it, hence the numerous hooks there that stays out of the "pipeline configuration" itself.

but even on the opt command line it is still useful to be able to drop this pass in somewhere specific as just a pass.

--mlir-print-ir-after=pass-name is the recommended way for this right now, it'll have the same effect as this pass (except in cases where you invoke multiple times the same pass I believe, as I mentioned before).

mlir/include/mlir/Transforms/Passes.h

 //===----------------------------------------------------------------------===//
 
 #define GEN_PASS_DECL_CANONICALIZER
 #define GEN_PASS_DECL_CONTROLFLOWSINK
 #define GEN_PASS_DECL_CSEPASS
 #define GEN_PASS_DECL_LOOPINVARIANTCODEMOTION
 #define GEN_PASS_DECL_STRIPDEBUGINFO
 #define GEN_PASS_DECL_PRINTOPSTATS
+#define GEN_PASS_DECL_PRINTOPS
 #define GEN_PASS_DECL_INLINER
 #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

 
 /// Creates a loop invariant code motion pass that hoists loop invariant
 /// instructions out of the loop.
 std::unique_ptr<Pass> createLoopInvariantCodeMotionPass();
 
 /// Creates a pass to strip debug information from a function.
 std::unique_ptr<Pass> createStripDebugInfoPass();
 
+/// Creates a pass which prints the op the pass is run on.
+std::unique_ptr<Pass> createPrintOpsPass(raw_ostream &os = llvm::errs());
+
 /// Creates a pass which prints the list of ops and the number of occurrences in
 /// the module.
 std::unique_ptr<Pass> createPrintOpStatsPass(raw_ostream &os = llvm::errs());
 
 /// Creates a pass which prints the list of ops and the number of occurrences in
 /// the module with the output format option.
 std::unique_ptr<Pass> createPrintOpStatsPass(raw_ostream &os, bool printAsJSON);
 

mlir/include/mlir/Transforms/Passes.td

def LocationSnapshot : Pass<"snapshot-op-locations"> {
   ];
 }
 
 def LoopInvariantCodeMotion : Pass<"loop-invariant-code-motion"> {
   let summary = "Hoist loop invariant instructions outside of the loop";
   let constructor = "mlir::createLoopInvariantCodeMotionPass()";
 }
 
+def PrintOps : Pass<"print-ops"> {
+  let summary = "Print operations";
+  let constructor = "mlir::createPrintOpsPass()";
+}
+
 def PrintOpStats : Pass<"print-op-stats"> {
   let summary = "Print statistics of operations";
   let constructor = "mlir::createPrintOpStatsPass()";
   let options = [
     Option<"printAsJSON", "json", "bool", /*default=*/"false",
            "print the stats as JSON">
   ];
 }

mlir/lib/Transforms/CMakeLists.txt

 add_subdirectory(Utils)
 
 add_mlir_library(MLIRTransforms
   Canonicalizer.cpp
   ControlFlowSink.cpp
   CSE.cpp
   Inliner.cpp
   LocationSnapshot.cpp
   LoopInvariantCodeMotion.cpp
   OpStats.cpp
+  PrintOps.cpp
   SCCP.cpp
   StripDebugInfo.cpp
   SymbolDCE.cpp
   SymbolPrivatize.cpp
   TopologicalSort.cpp
   ViewOpGraph.cpp
 
   ADDITIONAL_HEADER_DIRS

mlir/lib/Transforms/PrintOps.cpp

+//===- PrintOps.cpp - Prints operation for debugging----- -----------------===//
+//
+// 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/Transforms/Passes.h"
+
+namespace mlir {
+#define GEN_PASS_DEF_PRINTOPS
+#include "mlir/Transforms/Passes.h.inc"
+} // namespace mlir
+
+using namespace mlir;
+
+namespace {
+struct PrintOpsPass : public impl::PrintOpsBase<PrintOpsPass> {
+  explicit PrintOpsPass(raw_ostream &os) : os(os) {}
+
+  // Prints the operation the pass is run on.
+  void runOnOperation() override;
+
+private:
+  raw_ostream &os;
+};
+} // namespace
+
+void PrintOpsPass::runOnOperation() { getOperation()->print(os); }
+
+std::unique_ptr<Pass> mlir::createPrintOpsPass(raw_ostream &os) {
+  return std::make_unique<PrintOpsPass>(os);
+}

mlir/test/IR/print-ops.mlir

+// RUN: mlir-opt -allow-unregistered-dialect -print-ops %s -o=/dev/null 2>&1 | FileCheck %s
+
+func.func @main(tensor<4xf32>, tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) {
+^bb0(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>):
+  %0 = arith.addf %arg0, %arg1 : tensor<4xf32>
+  %unregistered = "unregistered_op"(%0, %arg1) : (tensor<4xf32>,tensor<4xf32>)-> tensor<4xf32>
+  return %0, %unregistered : tensor<4xf32>, tensor<4xf32>
+}
+
+// CHECK: func.func @main(%[[ARG0:.*]]: tensor<4xf32>, %[[ARG1:.*]]: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) {
+// CHECK:   %[[ADD:.*]] = arith.addf %[[ARG0]], %[[ARG1]] : tensor<4xf32>
+// CHECK:   %[[UNREGISTERED:.*]] = "unregistered_op"(%[[ADD]], %[[ARG1]]) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
+// CHECK:   return %[[ADD]], %[[UNREGISTERED]] : tensor<4xf32>, tensor<4xf32>
+// CHECK: }