Remove the `run` method from `OpPassManager` and `Pass` and migrate it to `OpToOpPassAdaptor`

[mehdi_amini]

Summary

This makes OpPassManager more of a "container" of passes and not responsible to drive the execution.
As such we also make it constructible publicly, which will allow to build arbitrary pipeline decoupled from the execution. We'll make use of this facility to expose "dynamic pipeline" in the future.

[mehdi_amini]

remove parent

[mehdi_amini]

remove dynamic pass pipeline

mlir/include/mlir/Pass/Pass.h

protected:
   }
 
   /// Return the MLIR context for the current function being transformed.
   MLIRContext &getContext() { return *getOperation()->getContext(); }
 
   /// The polymorphic API that runs the pass over the currently held operation.
   virtual void runOnOperation() = 0;
 
+  /// This is the equivalent to `runOnOperation` for a DynamicPass: it is
+  /// expected to return an OpPassManager containing the dynamic pipeline to
+  /// execute.
+  virtual OpPassManager *buildDynamicPipeline() {
+    llvm_unreachable("calling buildDynamicPipeline on a non-dynamic pass");
+  }
+
+  /// Return true if this is a dynamic pass: intended to be overriden only in
+  /// the DynamicPassWrapper class.
+  virtual bool isDynamic() const { return false; }
+
   /// A clone method to create a copy of this pass.
   std::unique_ptr<Pass> clone() const {
     auto newInst = clonePass();
     newInst->copyOptionValuesFrom(this);
     return newInst;
   }
 
   /// Return the current operation being transformed.
protected:
 
   /// Returns the derived pass name.
   StringRef getName() const override { return llvm::getTypeName<PassT>(); }
 
   /// A clone method to create a copy of this pass.
   std::unique_ptr<Pass> clonePass() const override {
     return std::make_unique<PassT>(*static_cast<const PassT *>(this));
   }
+
+private:
+  /// Hide from derived class, this is only available to dynamic passes.
+  OpPassManager *buildDynamicPipeline() final {
+    llvm_unreachable("calling buildDynamicPipeline on a non-dynamic pass");
+  }
+  /// Always return false: use DynamicPassWrapper for dynamic passes.
+  bool isDynamic() const final { return false; }
+};
+
+/// This class provides a CRTP wrapper around a base dynamic pass class to
+/// define several necessary utility methods.
+template <typename PassT, typename BaseT>
+class DynamicPassWrapper : public BaseT {
+public:
+  /// Support isa/dyn_cast functionality for the derived pass class.
+  static bool classof(const Pass *pass) {
+    return pass->getTypeID() == TypeID::get<PassT>();
+  }
+
+protected:
+  DynamicPassWrapper() : BaseT(TypeID::get<PassT>()) {}
+
+  /// Returns the derived pass name.
+  StringRef getName() const override { return llvm::getTypeName<PassT>(); }
+
+  /// A clone method to create a copy of this pass.
+  std::unique_ptr<Pass> clonePass() const override {
+    return std::make_unique<PassT>(*static_cast<const PassT *>(this));
+  }
+
+  /// Build and return a dynamic pipeline.
+  virtual OpPassManager *buildDynamicPipeline() = 0;
+
+private:
+  /// Hide this method from dynamic passes.
+  void runOnOperation() final {
+    llvm::report_fatal_error(
+        "Unexpected use of runOnOperation in a dynamic pass");
+  }
+
+  /// Always returns true to indicate that it is a dynamic pass.
+  bool isDynamic() const final { return true; }
 };
 
 } // end namespace mlir
 
 #endif // MLIR_PASS_PASS_H

mlir/include/mlir/Pass/PassManager.h

 class OperationName;
 class Operation;
 class Pass;
 class PassInstrumentation;
 class PassInstrumentor;
 
 namespace detail {
 struct OpPassManagerImpl;
+struct PassExecutionState;
 } // end namespace detail
 
 //===----------------------------------------------------------------------===//
 // OpPassManager
 //===----------------------------------------------------------------------===//
 
 /// This class represents a pass manager that runs passes on a specific
 /// operation type. This class is not constructed directly, but nested within
public:
   void getDependentDialects(DialectRegistry &dialects) const;
 
 private:
   /// A pointer to an internal implementation instance.
   std::unique_ptr<detail::OpPassManagerImpl> impl;
 
   /// Allow access to the constructor.
   friend class PassManager;
+  friend class Pass;
 
   /// Allow access.
   friend detail::OpPassManagerImpl;
 };
 
 //===----------------------------------------------------------------------===//
 // PassManager
 //===----------------------------------------------------------------------===//

mlir/lib/Pass/Pass.cpp

LogicalResult OpToOpPassAdaptor::run(Pass *pass, Operation *op,
                                      AnalysisManager am) {
   pass->passState.emplace(op, am);
 
   // Instrument before the pass has run.
   PassInstrumentor *pi = am.getPassInstrumentor();
   if (pi)
     pi->runBeforePass(pass, op);
 
-  // Invoke the virtual runOnOperation method.
+  // Invoke the virtual either runOnOperation method or if a dynamic pass,
+  // retried the OpPassManager and execute it.
+  if (pass->isDynamic()) {
+    OpPassManager *dynamicPipeline = pass->buildDynamicPipeline();
+    if (dynamicPipeline)
+      OpToOpPassAdaptor::runPipeline(dynamicPipeline->getPasses(), op, am);
+  } else {
     pass->runOnOperation();
+  }
 
   // Invalidate any non preserved analyses.
   am.invalidate(pass->passState->preservedAnalyses);
 
   // Instrument after the pass has run.
   bool passFailed = pass->passState->irAndPassFailed.getInt();
   if (pi) {
     if (passFailed)

mlir/test/Pass/dynamic_pipeline.mlir

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

mlir/test/lib/Transforms/TestDynamicPipeline.cpp

+//===------ TestDynamicPipeline.cpp --- dynamic pipeline test pass --------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a pass to test the dynamic pipeline feature.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Transforms/LoopUtils.h"
+#include "mlir/Transforms/Passes.h"
+
+using namespace mlir;
+
+namespace {
+
+class TestDynamicPipelinePass
+    : public DynamicPassWrapper<TestDynamicPipelinePass, OperationPass<>> {
+public:
+  TestDynamicPipelinePass(){};
+  TestDynamicPipelinePass(const TestDynamicPipelinePass &) {}
+
+  OpPassManager *buildDynamicPipeline() override {
+    llvm::errs() << "Dynamic execute '" << pipeline << "' on "
+                 << getOperation()->getName() << "\n";
+    if (pipeline.empty())
+      return nullptr;
+    auto symbolOp = dyn_cast<SymbolOpInterface>(getOperation());
+    if (!symbolOp) {
+      getOperation()->emitWarning()
+          << "Ignoring because not implementing SymbolOpInterface\n";
+      return nullptr;
+    }
+
+    auto opName = symbolOp.getName();
+    if (!opNames.empty() && !llvm::is_contained(opNames, opName)) {
+      llvm::errs() << "dynamic-pipeline skip op name: " << opName << "\n";
+      return nullptr;
+    }
+    if (!pm) {
+      pm = std::make_unique<OpPassManager>(getOperation()->getName(), false);
+      parsePassPipeline(pipeline, *pm, llvm::errs());
+    }
+    return pm.get();
+  }
+
+  std::unique_ptr<OpPassManager> pm;
+
+  Option<std::string> pipeline{
+      *this, "dynamic-pipeline",
+      llvm::cl::desc("The pipeline description that "
+                     "will run on the filtered function.")};
+  ListOption<std::string> opNames{
+      *this, "op-name", llvm::cl::MiscFlags::CommaSeparated,
+      llvm::cl::desc("List of function name to apply the pipeline to")};
+};
+} // end namespace
+
+namespace mlir {
+void registerTestDynamicPipelinePass() {
+  PassRegistration<TestDynamicPipelinePass>(
+      "test-dynamic-pipeline", "Tests the dynamic pipeline feature by applying "
+                               "a pipeline on a selected set of functions");
+}
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestAllReduceLoweringPass();
 void registerTestBufferPlacementPreparationPass();
 void registerTestCallGraphPass();
 void registerTestConstantFold();
 void registerTestConvertGPUKernelToCubinPass();
 void registerTestConvertGPUKernelToHsacoPass();
 void registerTestDominancePass();
 void registerTestDialect(DialectRegistry &);
+void registerTestDynamicPipelinePass();
 void registerTestExpandTanhPass();
 void registerTestFunc();
 void registerTestGpuMemoryPromotionPass();
 void registerTestGpuParallelLoopMappingPass();
 void registerTestInterfaces();
 void registerTestLinalgHoisting();
 void registerTestLinalgTransforms();
 void registerTestLivenessPass();

 #if MLIR_CUDA_CONVERSIONS_ENABLED
   registerTestConvertGPUKernelToCubinPass();
 #endif
 #if MLIR_ROCM_CONVERSIONS_ENABLED
   registerTestConvertGPUKernelToHsacoPass();
 #endif
   registerTestBufferPlacementPreparationPass();
   registerTestDominancePass();
+  registerTestDynamicPipelinePass();
   registerTestFunc();
   registerTestExpandTanhPass();
   registerTestGpuMemoryPromotionPass();
   registerTestInterfaces();
   registerTestLinalgHoisting();
   registerTestLinalgTransforms();
   registerTestLivenessPass();
   registerTestLoopFusion();