Allow dumping of OpToOpAdaptorPass.

[rainwoodman]

Summary

After this change, MLIR will trigger Instrumentation callbacks on OpToOpAdaptorPass itself with the module as GetOperation(). A new GetLongName() method is added to Pass to allow formatting an informative name (like OpToOpAdapterPass[[FunctionalToExecutorConversionPass]]) that is used for dumping. Note that OpToOpAdaptorPass is in mlir::detail and cannot be used in user subclasses of IR printing Config.

Before this change, MLIR dumping of OpToOpAdaptor passes are only conducted on individual container nodes (such as FuncOp). This becomes verbose when we dump DTensor MLIR rewrites of relatively large graphs from real TensorFlow models.

The user can distinguish between an OpToOpAdaptor full pass callback or a per FuncOp callback by checking if the Op is a module op. This behavior is not as straightforward as I'd like it to be (and I'd like to get some advises from the MLIR maintainers).

We have been using this patch with locally in debugging DTensor: the change reduced the number of dumps by a few thousand, which translated to a great time saving when the dumping was to a slower remote file system.

[jpienaar]

Could there be one that's the format too? (More documentative than a test really)

mlir/lib/Pass/Pass.cpp

668

Could the format in the same way as the reproducer? (Ideally we'd have something that we could just copy and paste for testing)

705

Style in MLIR is to elide trivial braces.

793

Same

mlir/lib/Pass/PassDetail.h

81

Doxygen Comment?

[rainwoodman]

• Update name style.
• Add docstring to GetLongName.

[rainwoodman]

mlir/lib/Pass/Pass.cpp

668

I've updated the unit tests with examples how the names look like. (And updated the pass's 'overall' name to OpToOpPassAdaptor.

Note that the unit test only covered the case where there is only one pass manager and one pass in the pass manager. Could you take another look if this is what you had in mind?

[mehdi_amini]

I don't get the problem this is trying to solve right now.

Before this change, MLIR dumping of OpToOpAdaptor passes are only conducted on individual container nodes (such as FuncOp). This becomes verbose when we dump DTensor MLIR rewrites of relatively large graphs from real TensorFlow models.

I don't understand what this means...

The user can distinguish between an OpToOpAdaptor full pass callback or a per FuncOp callback by checking if the Op is a module op. This behavior is not as straightforward as I'd like it to be (and I'd like to get some advises from the MLIR maintainers).

This does not make sense to me: won't there be an OpToOpAdaptor attached to the FuncOp as well?

[rainwoodman]

The way OpToOpAdaptor works is that it runs the FuncOp pass on the FuncOp
nodes one by one.

Before this change the only dumping mechanism for nesting pass is dumping
for each FuncOp pass. For a complex graph with thousands of FuncOp this
creates thousands of almost identical dumps for every single nested pass.
Serval usability issues rise from the large number of files created. This
is the problem.

The change here offers an option to dump the OpToOpAdapter pass itself,
which runs on the ModuleOp and only dumps once per nested pass. Thus
effectively user now gains the freedom to choose dump per FuncOp pass, per
ModuleOp pass, or both for a OpToOpAdapter pass.

There is some added complication in GetLongName to handle the case where an
OpToOpAdaptor is fused from multiple FuncOp passes.

[mehdi_amini]

Thanks for elaborating, I understand how it could be useful from this point of view.

There is some added complication in GetLongName to handle the case where an OpToOpAdaptor is fused from multiple FuncOp passes.

This is something that is still not clear to me: how do you get to print "for every nested pass" since the pass manager will fuse everything in a single OpToOpAdaptor?

Another question I have is how to use this: seems like this won't mesh well with the printBeforeIfEnabled, printAfterIfEnabled, etc. mechanisms used in the IRPrinter.

[wrengr]

mlir/lib/Pass/IRPrinting.cpp

67

extra space seems like a typo?

mlir/lib/Pass/Pass.cpp

665

Nit: should use single-quotes instead. Ditto for the square brackets below.

[rainwoodman]

In D152045#4487191, @mehdi_amini wrote:

Thanks for elaborating, I understand how it could be useful from this point of view.

There is some added complication in GetLongName to handle the case where an OpToOpAdaptor is fused from multiple FuncOp passes.

This is something that is still not clear to me: how do you get to print "for every nested pass" since the pass manager will fuse everything in a single OpToOpAdaptor?

Another question I have is how to use this: seems like this won't mesh well with the printBeforeIfEnabled, printAfterIfEnabled, etc. mechanisms used in the IRPrinter.

I guess this example explains both questions:

The combined pass would be dumping on "operation is mlir::ModuleOp"; the per object nested pass runs would dump on "operation is mlir::func::FuncOp" (the type used in AddNestedPass).

void printAfterIfEnabled(mlir::Pass *pass, mlir::Operation *operation,
                         PrintCallbackFn print_callback) override {
  mlir::ModuleOp module = mlir::dyn_cast<mlir::ModuleOp>(operation);
  // Only dumps the top level module and skips intermediate steps of nested
  // passes.
  if (!module) {
    return;
  }
  if (!module->hasAttr(dtensor::kDoNotLog) && !do_not_print_)
    BridgeLoggerConfig::printAfterIfEnabled(pass, operation, print_callback);
}

Even though the nested passes are fused/combined, I think it affects the concurrency more than the dumping logic. I couldn't quite find the code hooks for triggering the prints, but I remember the print hooks was either triggered (per instance, all combined passes) , or per (pass, instance) pair.

[rainwoodman]

• Use single quote.
• Remove extra space.

[mehdi_amini]

mlir/include/mlir/Pass/Pass.h

73

Instead of this, what about changing things this way:

• Add a field std::string cachedName; to OpToOpPassAdaptor
• Change getAdaptorName() to check if cachedName is empty and otherwise populate it.
• Override getName() to return cachedName.

mlir/lib/Pass/IRPrinting.cpp

66

Maybe what we need is a different hook for the instrumentations: we have runBeforePass and runAfterPass but we could also have runBeforeAdaptor and runAfterAdaptor which would offer a cleaner separation of intent here.

[rainwoodman]

mlir/lib/Pass/IRPrinting.cpp

66

If we do that then the logic for skipping the per (nested pass, FuncOp) dump and for dumping the Adaptor whole module dump will be scattered in both functions. What about lifting OpToOpPassAdaptor out of details:: such that users can do a type check inside run[Before|After]Pass?

[rriddle]

mlir/lib/Pass/IRPrinting.cpp

66

We have runBeforePipeline/runAfterPipeline, I agree with Mehdi that we should look at expanding this, like runBeforePipelineCollection/runAfterPipelineCollection. The PassAdaptor classes are internal implementation details and we should keep them hidden as much as possible.

mlir/include/mlir/Pass/Pass.h

 #define MLIR_PASS_PASS_H
 
 #include "mlir/Pass/AnalysisManager.h"
 #include "mlir/Pass/PassRegistry.h"
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/PointerIntPair.h"
 #include "llvm/ADT/Statistic.h"
 #include <optional>
+#include <string>
 
 namespace mlir {
 namespace detail {
 class OpToOpPassAdaptor;
 struct OpPassManagerImpl;
 
 /// The state for a single execution of a pass. This provides a unified
 /// interface for accessing and initializing necessary state for pass execution.
public:
   /// Returns the pass info for this pass, or null if unknown.
   const PassInfo *lookupPassInfo() const {
     return lookupPassInfo(getArgument());
   }
 
   /// Returns the derived pass name.
   virtual StringRef getName() const = 0;
 
+  /// Returns the long derived pass name. This includes the pass managers
+  /// for an adaptor.
+  virtual std::string getLongName() { return getName().str(); }
+
   /// Register dependent dialects for the current pass.
   /// A pass is expected to register the dialects it will create entities for
   /// (Operations, Types, Attributes), other than dialect that exists in the
   /// input. For example, a pass that converts from Linalg to Affine would
   /// register the Affine dialect but does not need to register Linalg.
   virtual void getDependentDialects(DialectRegistry &registry) const {}
 
   /// Return the command line argument used when registering this pass. Return

mlir/lib/Pass/IRPrinting.cpp

static void printIR(Operation *op, bool printModuleScope, raw_ostream &out,
   // Find the top-level operation.
   auto *topLevelOp = op;
   while (auto *parentOp = topLevelOp->getParentOp())
     topLevelOp = parentOp;
   topLevelOp->print(out, flags);
 }
 
 /// Instrumentation hooks.
 void IRPrinterInstrumentation::runBeforePass(Pass *pass, Operation *op) {
-  if (isa<OpToOpPassAdaptor>(pass))
-    return;
   // If the config asked to detect changes, record the current fingerprint.
   if (config->shouldPrintAfterOnlyOnChange())
     beforePassFingerPrints.try_emplace(pass, op);
 
   config->printBeforeIfEnabled(pass, op, [&](raw_ostream &out) {
-    out << "// -----// IR Dump Before " << pass->getName() << " ("
+    out << "// -----// IR Dump Before " << pass->getLongName() << " ("
         << pass->getArgument() << ")";
     printIR(op, config->shouldPrintAtModuleScope(), out,
             config->getOpPrintingFlags());
     out << "\n\n";
   });
 }
 
 void IRPrinterInstrumentation::runAfterPass(Pass *pass, Operation *op) {
-  if (isa<OpToOpPassAdaptor>(pass))
-    return;
-
   // Check to see if we are only printing on failure.
   if (config->shouldPrintAfterOnlyOnFailure())
     return;
 
   // If the config asked to detect changes, compare the current fingerprint with
   // the previous.
   if (config->shouldPrintAfterOnlyOnChange()) {
     auto fingerPrintIt = beforePassFingerPrints.find(pass);
     assert(fingerPrintIt != beforePassFingerPrints.end() &&
            "expected valid fingerprint");
     // If the fingerprints are the same, we don't print the IR.
     if (fingerPrintIt->second == OperationFingerPrint(op)) {
       beforePassFingerPrints.erase(fingerPrintIt);
       return;
     }
     beforePassFingerPrints.erase(fingerPrintIt);
   }
 
   config->printAfterIfEnabled(pass, op, [&](raw_ostream &out) {
-    out << "// -----// IR Dump After " << pass->getName() << " ("
+    out << "// -----// IR Dump After " << pass->getLongName() << " ("
         << pass->getArgument() << ")";
     printIR(op, config->shouldPrintAtModuleScope(), out,
             config->getOpPrintingFlags());
     out << "\n\n";
   });
 }
 
 void IRPrinterInstrumentation::runAfterPassFailed(Pass *pass, Operation *op) {
-  if (isa<OpToOpPassAdaptor>(pass))
-    return;
   if (config->shouldPrintAfterOnlyOnChange())
     beforePassFingerPrints.erase(pass);
 
   config->printAfterIfEnabled(pass, op, [&](raw_ostream &out) {
-    out << formatv("// -----// IR Dump After {0} Failed ({1})", pass->getName(),
-                   pass->getArgument());
+    out << formatv("// -----// IR Dump After {0} Failed ({1})",
+                   pass->getLongName(), pass->getArgument());
     printIR(op, config->shouldPrintAtModuleScope(), out, OpPrintingFlags());
     out << "\n\n";
   });
 }
 
 //===----------------------------------------------------------------------===//
 // IRPrinterConfig
 //===----------------------------------------------------------------------===//

mlir/lib/Pass/Pass.cpp

std::string OpToOpPassAdaptor::getAdaptorName() {
   llvm::raw_string_ostream os(name);
   llvm::interleaveComma(getPassManagers(), os, [&](OpPassManager &pm) {
     os << '\'' << pm.getOpAnchorName() << '\'';
   });
   os << ']';
   return os.str();
 }
 
+/// Returns the adaptor pass long name, including the nested passes.
+std::string OpToOpPassAdaptor::getLongName() {
+  std::string name = "OpToOpPassAdaptor";
+  llvm::raw_string_ostream os(name);
+  os << '(';
+  llvm::interleaveComma(getPassManagers(), os, [&](OpPassManager &pm) {
+    os << '[';
+    llvm::interleaveComma(pm.getPasses(), os,
+                          [&](Pass &pass) { os << pass.getName(); });
+    os << ']';
+  });
+  os << ')';
+  return os.str();
+}
 void OpToOpPassAdaptor::runOnOperation() {
   llvm_unreachable(
       "Unexpected call to Pass::runOnOperation() on OpToOpPassAdaptor");
 }
 
 /// Run the held pipeline over all nested operations.
 void OpToOpPassAdaptor::runOnOperation(bool verifyPasses) {
   if (getContext().isMultithreadingEnabled())
     runOnOperationAsyncImpl(verifyPasses);
   else
     runOnOperationImpl(verifyPasses);
 }
 
 /// Run this pass adaptor synchronously.
 void OpToOpPassAdaptor::runOnOperationImpl(bool verifyPasses) {
   auto am = getAnalysisManager();
   PassInstrumentation::PipelineParentInfo parentInfo = {llvm::get_threadid(),
                                                         this};
   auto *instrumentor = am.getPassInstrumentor();
+
   for (auto &region : getOperation()->getRegions()) {
     for (auto &block : region) {
       for (auto &op : block) {
         auto *mgr = findPassManagerFor(mgrs, op.getName(), *op.getContext());
         if (!mgr)
           continue;
 
         // Run the held pipeline over the current operation.
         unsigned initGeneration = mgr->impl->initializationGeneration;
         if (failed(runPipeline(*mgr, &op, am.nest(&op), verifyPasses,
                                initGeneration, instrumentor, &parentInfo)))
           return signalPassFailure();
       }
     }
   }
 }
 
 /// Utility functor that checks if the two ranges of pass managers have a size
 /// mismatch.
LogicalResult pipelineResult = runPipeline(
         pm.impl->initializationGeneration, instrumentor, &parentInfo);
 
     // Reset the active bit for this pass manager.
     activePMs[pmIndex].store(false);
     return pipelineResult;
   };
 
   // Signal a failure if any of the executors failed.
   if (failed(failableParallelForEach(context, opInfos, processFn)))
     signalPassFailure();
 }
 
 //===----------------------------------------------------------------------===//
 // PassManager
 //===----------------------------------------------------------------------===//
 
 PassManager::PassManager(MLIRContext *ctx, StringRef operationName,

mlir/lib/Pass/PassDetail.h

public:
   /// Return the async pass managers held by this parallel adaptor.
   MutableArrayRef<SmallVector<OpPassManager, 1>> getParallelPassManagers() {
     return asyncExecutors;
   }
 
   /// Returns the adaptor pass name.
   std::string getAdaptorName();
 
+  /// Returns the long pass name. For Adaptor pass, the long name includes
+  /// names of the nested passes.
+  std::string getLongName() override;
+
 private:
   /// Run this pass adaptor synchronously.
   void runOnOperationImpl(bool verifyPasses);
 
   /// Run this pass adaptor asynchronously.
   void runOnOperationAsyncImpl(bool verifyPasses);
 
   /// Run the given operation and analysis manager on a single pass.

mlir/test/Pass/dynamic-pipeline-nested.mlir

 // RUN: mlir-opt %s -pass-pipeline='builtin.module(builtin.module(test-dynamic-pipeline{op-name=inner_mod1 dynamic-pipeline=cse}))'  --mlir-disable-threading  -mlir-print-ir-before-all 2>&1 | FileCheck %s --check-prefix=NOTNESTED --check-prefix=CHECK
 // RUN: mlir-opt %s -pass-pipeline='builtin.module(builtin.module(test-dynamic-pipeline{op-name=inner_mod1 run-on-nested-operations=1 dynamic-pipeline=cse}))'  --mlir-disable-threading  -mlir-print-ir-before-all 2>&1 | FileCheck %s --check-prefix=NESTED --check-prefix=CHECK
 
 
 // Verify that we can schedule a dynamic pipeline on a nested operation
 
 func.func @f() {
   return
 }
 
 // CHECK: IR Dump Before
-// CHECK-SAME: TestDynamicPipelinePass
+// CHECK-SAME: OpToOpPassAdaptor([{anonymous}::TestDynamicPipelinePass])
+// CHECK: IR Dump Before
+// CHECK-SAME: {anonymous}::TestDynamicPipelinePass
 // CHECK-NEXT: module @inner_mod1
 module @inner_mod1 {
 // We use the mlir-print-ir-after-all dumps to check the granularity of the
 // scheduling: if we are nesting we expect to see to individual "Dump Before
 // CSE" output: one for each of the function. If we don't nest, then we expect
 // the CSE pass to run on the `inner_mod1` module directly.
 
 // CHECK: Dump Before CSE
 // NOTNESTED-NEXT: @inner_mod1
 // NESTED-NEXT: @foo
   module @foo {}
 // Only in the nested case we have a second run of the pass here.
 // NESTED: Dump Before CSE
 // NESTED-NEXT: @baz
   module @baz {}
 }

mlir/test/Pass/dynamic-pipeline.mlir

 // RUN: mlir-opt %s -pass-pipeline='builtin.module(builtin.module(test-dynamic-pipeline{op-name=inner_mod1, dynamic-pipeline=func.func(cse,canonicalize)}))'  --mlir-disable-threading  -mlir-print-ir-before-all 2>&1 | FileCheck %s --check-prefix=MOD1 --check-prefix=MOD1-ONLY --check-prefix=CHECK
 // RUN: mlir-opt %s -pass-pipeline='builtin.module(builtin.module(test-dynamic-pipeline{op-name=inner_mod2, dynamic-pipeline=func.func(cse,canonicalize)}))'  --mlir-disable-threading  -mlir-print-ir-before-all 2>&1 | FileCheck %s --check-prefix=MOD2 --check-prefix=MOD2-ONLY --check-prefix=CHECK
 // RUN: mlir-opt %s -pass-pipeline='builtin.module(builtin.module(test-dynamic-pipeline{op-name=inner_mod1,inner_mod2, dynamic-pipeline=func.func(cse,canonicalize)}))'  --mlir-disable-threading  -mlir-print-ir-before-all 2>&1 | FileCheck %s --check-prefix=MOD1 --check-prefix=MOD2 --check-prefix=CHECK
 // RUN: mlir-opt %s -pass-pipeline='builtin.module(builtin.module(test-dynamic-pipeline{dynamic-pipeline=func.func(cse,canonicalize)}))'  --mlir-disable-threading  -mlir-print-ir-before-all 2>&1 | FileCheck %s --check-prefix=MOD1 --check-prefix=MOD2 --check-prefix=CHECK
 
 
 func.func @f() {
   return
 }
 
 // CHECK: IR Dump Before
-// CHECK-SAME: TestDynamicPipelinePass
+// CHECK-SAME: OpToOpPassAdaptor([{anonymous}::TestDynamicPipelinePass])
+// CHECK: IR Dump Before
+// CHECK-SAME: {anonymous}::TestDynamicPipelinePass
 // CHECK-NEXT: module @inner_mod1
 // MOD2-ONLY: dynamic-pipeline skip op name: inner_mod1
 module @inner_mod1 {
 // MOD1: Dump Before CSE
 // MOD1-NEXT: @foo
 // MOD1: Dump Before Canonicalizer
 // MOD1-NEXT: @foo
   func.func @foo() {
     return
   }
 // MOD1: Dump Before CSE
 // MOD1-NEXT: @baz
 // MOD1: Dump Before Canonicalizer
 // MOD1-NEXT: @baz
   func.func @baz() {
     return
   }
 }
 
 // CHECK: IR Dump Before
-// CHECK-SAME: TestDynamicPipelinePass
+// CHECK-SAME: {anonymous}::TestDynamicPipelinePass
 // CHECK-NEXT: module @inner_mod2
 // MOD1-ONLY: dynamic-pipeline skip op name: inner_mod2
 module @inner_mod2 {
 // MOD2: Dump Before CSE
 // MOD2-NEXT: @foo
 // MOD2: Dump Before Canonicalizer
 // MOD2-NEXT: @foo
   func.func @foo() {
     return
   }
 }