[mlir][Pass] Add support for an InterfacePass and pass filtering based on OperationName

[rriddle]

Summary

This commit adds a new hook Pass bool canScheduleOn(RegisteredOperationName) that
indicates if the given pass can be scheduled on operations of the given type. This makes it
easier to define constraints on generic passes without a) adding conditional checks to
the beginning of the runOnOperation, or b) defining a new pass type that forwards
from runOnOperation (after checking the invariants) to a new hook. This new hook is
used to implement an InterfacePass pass class, that represents a generic pass that
runs on operations of the given interface type.

The PassManager will also verify that passes added to a pass manager can actually be
scheduled on that pass manager, meaning that we will properly error when an Interface
is scheduled on an operation that doesn't actually implement that interface.

[rriddle]

mlir/lib/Pass/Pass.cpp

219

@mehdi_amini I went back and forth on what behavior we want here, would love to hear what you think. The patch as-is drops passes that can't be executed (with a log), but we could also emit a fatal error/diagnostic/something here instead. We can't emit an error at the time of insertion because we don't have the OperationName at that point (the dialect might not be loaded).

[mehdi_amini]

mlir/include/mlir/Pass/Pass.h

409

The doc seems copy/pasted but you forgot to update it?

mlir/lib/Pass/Pass.cpp

219

I'm concerned about a fatal_error here in the sense that this can trip on a user input. We should have a way for a client to validate a parsed pipeline with a proper error recovery.

Think about this kind of thing: https://github.com/llvm/llvm-project/blob/main/mlir/test/python/pass_manager.py#L39-L44 ; where we shouldn't just crash the process.

Can we diagnose an error and make finalizePassList returns a LogicalResult to propagate upward?

mlir/include/mlir/Pass/Pass.h

 #include "mlir/Pass/PassRegistry.h"
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/PointerIntPair.h"
 #include "llvm/ADT/Statistic.h"
 
 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.
 struct PassExecutionState {
   PassExecutionState(Operation *ir, AnalysisManager analysisManager,
                      function_ref<LogicalResult(OpPassManager &, Operation *)>
                          pipelineExecutor)
       : irAndPassFailed(ir, false), analysisManager(analysisManager),
protected:
   /// Initialize any complex state necessary for running this pass. This hook
   /// should not rely on any state accessible during the execution of a pass.
   /// For example, `getContext`/`getOperation`/`getAnalysis`/etc. should not be
   /// invoked within this hook.
   /// Returns a LogicalResult to indicate failure, in which case the pass
   /// pipeline won't execute.
   virtual LogicalResult initialize(MLIRContext *context) { return success(); }
 
+  /// Indicate if the current pass can be scheduled on the given operation type.
+  /// This is useful for generic operation passes to add restrictions on the
+  /// operations they operate on.
+  virtual bool canScheduleOn(RegisteredOperationName opName) const = 0;
+
   /// Schedule an arbitrary pass pipeline on the provided operation.
   /// This can be invoke any time in a pass to dynamic schedule more passes.
   /// The provided operation must be the current one or one nested below.
   LogicalResult runPipeline(OpPassManager &pipeline, Operation *op) {
     return passState->pipelineExecutor(pipeline, op);
   }
 
   /// A clone method to create a copy of this pass.
private:
 
   /// A pointer to the pass this pass was cloned from, if the clone was made by
   /// the pass manager for the sake of multi-threading.
   const Pass *threadingSibling = nullptr;
 
   /// Allow access to 'clone'.
   friend class OpPassManager;
 
+  /// Allow access to 'canScheduleOn'.
+  friend detail::OpPassManagerImpl;
+
   /// Allow access to 'passState'.
   friend detail::OpToOpPassAdaptor;
 
   /// Allow access to 'passOptions'.
   friend class PassInfo;
 };
 
 //===----------------------------------------------------------------------===//
protected:
   OperationPass(TypeID passID) : Pass(passID, OpT::getOperationName()) {}
   OperationPass(const OperationPass &) = default;
 
   /// Support isa/dyn_cast functionality.
   static bool classof(const Pass *pass) {
     return pass->getOpName() == OpT::getOperationName();
   }
 
+  /// Indicate if the current pass can be scheduled on the given operation type.
+  bool canScheduleOn(RegisteredOperationName opName) const final {
+    return opName.getStringRef() == getOpName();
+  }
+
   /// Return the current operation being transformed.
   OpT getOperation() { return cast<OpT>(Pass::getOperation()); }
 
   /// Query an analysis for the current operation of the specific derived
   /// operation type.
   template <typename AnalysisT>
   AnalysisT &getAnalysis() {
     return Pass::getAnalysis<AnalysisT, OpT>();

 /// Derived operation passes are expected to provide the following:
 ///   - A 'void runOnOperation()' method.
 ///   - A 'StringRef getName() const' method.
 ///   - A 'std::unique_ptr<Pass> clonePass() const' method.
 template <> class OperationPass<void> : public Pass {
 protected:
   OperationPass(TypeID passID) : Pass(passID) {}
   OperationPass(const OperationPass &) = default;
+
+  /// Indicate if the current pass can be scheduled on the given operation type.
+  /// By default, generic operation passes can be scheduled on any operation.
+  bool canScheduleOn(RegisteredOperationName opName) const override {
+    return true;
+  }
+};
+
+/// Pass to transform an operation.
+///
+/// Operation passes must not:
+///   - modify any other operations within the parent region, as other threads
+///     may be manipulating them concurrently.
+///   - modify any state within the parent operation, this includes adding
+///     additional operations.
+///
+/// Derived operation passes are expected to provide the following:
+///   - A 'void runOnOperation()' method.
+///   - A 'StringRef getName() const' method.
+///   - A 'std::unique_ptr<Pass> clonePass() const' method.
+template <typename InterfaceT>
+class InterfacePass : public OperationPass<> {
+protected:
+  using OperationPass::OperationPass;
+
+  /// Indicate if the current pass can be scheduled on the given operation type.
+  /// For an InterfacePass, this checks if the operation implements the given
+  /// interface.
+  bool canScheduleOn(RegisteredOperationName opName) const final {
+    return opName.hasInterface<InterfaceT>();
+  }
+
+  /// Return the current operation being transformed.
+  InterfaceT getOperation() { return cast<InterfaceT>(Pass::getOperation()); }
+
+  /// Query an analysis for the current operation.
+  template <typename AnalysisT>
+  AnalysisT &getAnalysis() {
+    return Pass::getAnalysis<AnalysisT, InterfaceT>();
+  }
 };
 
 /// This class provides a CRTP wrapper around a base pass class to define
 /// several necessary utility methods. This should only be used for passes that
 /// are not suitably represented using the declarative pass specification(i.e.
 /// tablegen backend).
 template <typename PassT, typename BaseT> class PassWrapper : public BaseT {
 public:

mlir/include/mlir/Pass/PassBase.td

class PassBase<string passArg, string base> {
   // A set of statistics provided by this pass.
   list<Statistic> statistics = [];
 }
 
 // This class represents an mlir::OperationPass.
 class Pass<string passArg, string operation = "">
   : PassBase<passArg, "::mlir::OperationPass<" # operation # ">">;
 
+// This class represents an mlir::InterfacePass.
+class InterfacePass<string passArg, string interface>
+  : PassBase<passArg, "::mlir::InterfacePass<" # interface # ">">;
+
 #endif // MLIR_PASS_PASSBASE

mlir/include/mlir/Pass/PassManager.h

public:
   template <typename OpT> void addNestedPass(std::unique_ptr<Pass> pass) {
     nest<OpT>().addPass(std::move(pass));
   }
 
   /// Returns the number of passes held by this manager.
   size_t size() const;
 
   /// Return the operation name that this pass manager operates on.
-  StringAttr getOpName(MLIRContext &context) const;
+  OperationName getOpName(MLIRContext &context) const;
 
   /// Return the operation name that this pass manager operates on.
   StringRef getOpName() const;
 
   /// Returns the internal implementation instance.
   detail::OpPassManagerImpl &getImpl();
 
   /// Prints out the passes of the pass manager as the textual representation

mlir/lib/Pass/Pass.cpp

 
 //===----------------------------------------------------------------------===//
 // OpPassManagerImpl
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 namespace detail {
 struct OpPassManagerImpl {
-  OpPassManagerImpl(StringAttr identifier, OpPassManager::Nesting nesting)
-      : name(identifier.str()), identifier(identifier),
+  OpPassManagerImpl(OperationName opName, OpPassManager::Nesting nesting)
+      : name(opName.getStringRef()), opName(opName),
         initializationGeneration(0), nesting(nesting) {}
   OpPassManagerImpl(StringRef name, OpPassManager::Nesting nesting)
       : name(name), initializationGeneration(0), nesting(nesting) {}
 
   /// Merge the passes of this pass manager into the one provided.
   void mergeInto(OpPassManagerImpl &rhs);
 
   /// Nest a new operation pass manager for the given operation kind under this
   /// pass manager.
   OpPassManager &nest(StringAttr nestedName);
   OpPassManager &nest(StringRef nestedName);
 
   /// Add the given pass to this pass manager. If this pass has a concrete
   /// operation type, it must be the same type as this pass manager.
   void addPass(std::unique_ptr<Pass> pass);
 
   /// Clear the list of passes in this pass manager, other options are
   /// preserved.
   void clear();
 
-  /// Coalesce adjacent AdaptorPasses into one large adaptor. This runs
-  /// recursively through the pipeline graph.
-  void coalesceAdjacentAdaptorPasses();
-
-  /// Return the operation name of this pass manager as an identifier.
-  StringAttr getOpName(MLIRContext &context) {
-    if (!identifier)
-      identifier = StringAttr::get(&context, name);
-    return *identifier;
+  /// Finalize the pass list in preparation for execution. This includes
+  /// coalescing adjacent pass managers when possible, verifying scheduled
+  /// passes, etc.
+  void finalizePassList(MLIRContext *ctx);
+
+  /// Return the operation name of this pass manager.
+  OperationName getOpName(MLIRContext &context) {
+    if (!opName)
+      opName = OperationName(name, &context);
+    return *opName;
   }
 
   /// The name of the operation that passes of this pass manager operate on.
   std::string name;
 
-  /// The cached identifier (internalized in the context) for the name of the
+  /// The cached OperationName (internalized in the context) for the name of the
   /// operation that passes of this pass manager operate on.
-  Optional<StringAttr> identifier;
+  Optional<OperationName> opName;
 
   /// The set of passes to run as part of this pass manager.
   std::vector<std::unique_ptr<Pass>> passes;
 
   /// The current initialization generation of this pass manager. This is used
   /// to indicate when a pass manager should be reinitialized.
   unsigned initializationGeneration;
 
llvm::report_fatal_error(llvm::Twine("Can't add pass '") + pass->getName() +
                              "', did you intend to nest?");
   }
 
   passes.emplace_back(std::move(pass));
 }
 
 void OpPassManagerImpl::clear() { passes.clear(); }
 
-void OpPassManagerImpl::coalesceAdjacentAdaptorPasses() {
-  // Bail out early if there are no adaptor passes.
-  if (llvm::none_of(passes, [](std::unique_ptr<Pass> &pass) {
-        return isa<OpToOpPassAdaptor>(pass.get());
-      }))
-    return;
-
+void OpPassManagerImpl::finalizePassList(MLIRContext *ctx) {
   // Walk the pass list and merge adjacent adaptors.
   OpToOpPassAdaptor *lastAdaptor = nullptr;
-  for (auto &passe : passes) {
+  for (auto &pass : passes) {
     // Check to see if this pass is an adaptor.
-    if (auto *currentAdaptor = dyn_cast<OpToOpPassAdaptor>(passe.get())) {
+    if (auto *currentAdaptor = dyn_cast<OpToOpPassAdaptor>(pass.get())) {
       // If it is the first adaptor in a possible chain, remember it and
       // continue.
       if (!lastAdaptor) {
         lastAdaptor = currentAdaptor;
         continue;
       }
 
       // Otherwise, merge into the existing adaptor and delete the current one.
       currentAdaptor->mergeInto(*lastAdaptor);
-      passe.reset();
+      pass.reset();
     } else if (lastAdaptor) {
-      // If this pass is not an adaptor, then coalesce and forget any existing
+      // If this pass is not an adaptor, then finalize and forget any existing
       // adaptor.
       for (auto &pm : lastAdaptor->getPassManagers())
-        pm.getImpl().coalesceAdjacentAdaptorPasses();
+        pm.getImpl().finalizePassList(ctx);
       lastAdaptor = nullptr;
     }
   }
 
-  // If there was an adaptor at the end of the manager, coalesce it as well.
+  // If there was an adaptor at the end of the manager, finalize it as well.
   if (lastAdaptor) {
     for (auto &pm : lastAdaptor->getPassManagers())
-      pm.getImpl().coalesceAdjacentAdaptorPasses();
+      pm.getImpl().finalizePassList(ctx);
   }
 
-  // Now that the adaptors have been merged, erase the empty slot corresponding
+  // Now that the adaptors have been merged, erase any empty slots corresponding
   // to the merged adaptors that were nulled-out in the loop above.
-  llvm::erase_if(passes, std::logical_not<std::unique_ptr<Pass>>());
+  Optional<RegisteredOperationName> opName =
+      getOpName(*ctx).getRegisteredInfo();
+  llvm::erase_if(passes, [&](const std::unique_ptr<Pass> &pass) {
+    if (!pass)
+      return true;
+    if (opName && !pass->canScheduleOn(*opName)) {
+      llvm::report_fatal_error(
+          llvm::Twine("unable to schedule pass '") + pass->getName() +
+          "' on a PassManager intended to run on '" + name + "'!");
+    }
+    return false;
+  });
 }
 
 //===----------------------------------------------------------------------===//
 // OpPassManager
 //===----------------------------------------------------------------------===//
 
 OpPassManager::OpPassManager(StringAttr name, Nesting nesting)
     : impl(new OpPassManagerImpl(name, nesting)) {}

 
 /// Returns the internal implementation instance.
 OpPassManagerImpl &OpPassManager::getImpl() { return *impl; }
 
 /// Return the operation name that this pass manager operates on.
 StringRef OpPassManager::getOpName() const { return impl->name; }
 
 /// Return the operation name that this pass manager operates on.
-StringAttr OpPassManager::getOpName(MLIRContext &context) const {
+OperationName OpPassManager::getOpName(MLIRContext &context) const {
   return impl->getOpName(context);
 }
 
 /// Prints out the given passes as the textual representation of a pipeline.
 static void printAsTextualPipeline(ArrayRef<std::unique_ptr<Pass>> passes,
                                    raw_ostream &os) {
   llvm::interleaveComma(passes, os, [&](const std::unique_ptr<Pass> &pass) {
     pass->printAsTextualPipeline(os);
auto dynamicPipelineCallback = [&](OpPassManager &pipeline,
                                      Operation *root) -> LogicalResult {
     if (!op->isAncestor(root))
       return root->emitOpError()
              << "Trying to schedule a dynamic pipeline on an "
                 "operation that isn't "
                 "nested under the current operation the pass is processing";
     assert(pipeline.getOpName() == root->getName().getStringRef());
 
-    // Before running, make sure to coalesce any adjacent pass adaptors in the
-    // pipeline.
-    pipeline.getImpl().coalesceAdjacentAdaptorPasses();
+    // Before running, finalize the passes held by the pipeline.
+    pipeline.getImpl().finalizePassList(root->getContext());
 
     // Initialize the user provided pipeline and execute the pipeline.
     if (failed(pipeline.initialize(root->getContext(), parentInitGeneration)))
       return failure();
     AnalysisManager nestedAm = root == op ? am : am.nest(root);
     return OpToOpPassAdaptor::runPipeline(pipeline.getPasses(), root, nestedAm,
                                           verifyPasses, parentInitGeneration,
                                           pi, &parentInfo);
static OpPassManager *findPassManagerFor(MutableArrayRef<OpPassManager> mgrs,
   auto *it = llvm::find_if(
       mgrs, [&](OpPassManager &mgr) { return mgr.getOpName() == name; });
   return it == mgrs.end() ? nullptr : &*it;
 }
 
 /// Find an operation pass manager that can operate on an operation of the given
 /// type, or nullptr if one does not exist.
 static OpPassManager *findPassManagerFor(MutableArrayRef<OpPassManager> mgrs,
-                                         StringAttr name,
+                                         OperationName name,
                                          MLIRContext &context) {
   auto *it = llvm::find_if(
       mgrs, [&](OpPassManager &mgr) { return mgr.getOpName(context) == name; });
   return it == mgrs.end() ? nullptr : &*it;
 }
 
 OpToOpPassAdaptor::OpToOpPassAdaptor(OpPassManager &&mgr) {
   mgrs.emplace_back(std::move(mgr));

 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().getIdentifier(),
-                                       *op.getContext());
+        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->getPasses(), &op, am.nest(&op),
                                verifyPasses, initGeneration, instrumentor,
                                &parentInfo)))
void OpToOpPassAdaptor::runOnOperationAsyncImpl(bool verifyPasses) {
   // Run a prepass over the operation to collect the nested operations to
   // execute over. This ensures that an analysis manager exists for each
   // operation, as well as providing a queue of operations to execute over.
   std::vector<std::pair<Operation *, AnalysisManager>> opAMPairs;
   for (auto &region : getOperation()->getRegions()) {
     for (auto &block : region) {
       for (auto &op : block) {
         // Add this operation iff the name matches any of the pass managers.
-        if (findPassManagerFor(mgrs, op.getName().getIdentifier(), *context))
+        if (findPassManagerFor(mgrs, op.getName(), *context))
           opAMPairs.emplace_back(&op, am.nest(&op));
       }
     }
   }
 
   // Get the current thread for this adaptor.
   PassInstrumentation::PipelineParentInfo parentInfo = {llvm::get_threadid(),
                                                         this};
   auto *instrumentor = am.getPassInstrumentor();
 
   // An atomic failure variable for the async executors.
   std::vector<std::atomic<bool>> activePMs(asyncExecutors.size());
   std::fill(activePMs.begin(), activePMs.end(), false);
   auto processFn = [&](auto &opPMPair) {
     // Find a pass manager for this operation.
     auto it = llvm::find_if(activePMs, [](std::atomic<bool> &isActive) {
       bool expectedInactive = false;
       return isActive.compare_exchange_strong(expectedInactive, true);
     });
     unsigned pmIndex = it - activePMs.begin();
 
     // Get the pass manager for this operation and execute it.
-    auto *pm =
-        findPassManagerFor(asyncExecutors[pmIndex],
-                           opPMPair.first->getName().getIdentifier(), *context);
+    auto *pm = findPassManagerFor(asyncExecutors[pmIndex],
+                                  opPMPair.first->getName(), *context);
     assert(pm && "expected valid pass manager for operation");
 
     unsigned initGeneration = pm->impl->initializationGeneration;
     LogicalResult pipelineResult =
         runPipeline(pm->getPasses(), opPMPair.first, opPMPair.second,
                     verifyPasses, initGeneration, instrumentor, &parentInfo);
 
     // Reset the active bit for this pass manager.

 
 PassManager::~PassManager() = default;
 
 void PassManager::enableVerifier(bool enabled) { verifyPasses = enabled; }
 
 /// Run the passes within this manager on the provided operation.
 LogicalResult PassManager::run(Operation *op) {
   MLIRContext *context = getContext();
-  assert(op->getName().getIdentifier() == getOpName(*context) &&
+  assert(op->getName() == getOpName(*context) &&
          "operation has a different name than the PassManager or is from a "
          "different context");
 
-  // Before running, make sure to coalesce any adjacent pass adaptors in the
-  // pipeline.
-  getImpl().coalesceAdjacentAdaptorPasses();
-
   // Register all dialects for the current pipeline.
   DialectRegistry dependentDialects;
   getDependentDialects(dependentDialects);
   context->appendDialectRegistry(dependentDialects);
   for (StringRef name : dependentDialects.getDialectNames())
     context->getOrLoadDialect(name);
 
+  // Before running, make sure to finalize the pipeline pass list.
+  getImpl().finalizePassList(context);
+
   // Initialize all of the passes within the pass manager with a new generation.
   llvm::hash_code newInitKey = context->getRegistryHash();
   if (newInitKey != initializationKey) {
     if (failed(initialize(context, impl->initializationGeneration + 1)))
       return failure();
     initializationKey = newInitKey;
   }
 

mlir/test/Pass/interface-pass.mlir

+// RUN: mlir-opt %s -verify-diagnostics -pass-pipeline='builtin.func(test-interface-pass)' -o /dev/null
+// RUN: not not mlir-opt %s -pass-pipeline='test-interface-pass' 2>&1 | FileCheck --check-prefix=INVALID_MODULE %s
+
+// Test that we run the interface pass on the function, but emits an error for the module.
+
+// INVALID_MODULE: unable to schedule pass '{{.*}}' on a PassManager intended to run on 'builtin.module'!
+
+// expected-remark@below {{Executing interface pass on operation}}
+func @main() {
+  return
+}

mlir/test/lib/Pass/TestPassManager.cpp

 struct TestFunctionPass
     : public PassWrapper<TestFunctionPass, OperationPass<FuncOp>> {
   void runOnOperation() final {}
   StringRef getArgument() const final { return "test-function-pass"; }
   StringRef getDescription() const final {
     return "Test a function pass in the pass manager";
   }
 };
+class TestInterfacePass
+    : public PassWrapper<TestInterfacePass,
+                         InterfacePass<FunctionOpInterface>> {
+  void runOnOperation() final {
+    getOperation()->emitRemark() << "Executing interface pass on operation";
+  }
+  StringRef getArgument() const final { return "test-interface-pass"; }
+  StringRef getDescription() const final {
+    return "Test an interface pass (running on FunctionOpInterface) in the "
+           "pass manager";
+  }
+};
 class TestOptionsPass
     : public PassWrapper<TestOptionsPass, OperationPass<FuncOp>> {
 public:
   struct Options : public PassPipelineOptions<Options> {
     ListOption<int> listOption{*this, "list",
                                llvm::cl::MiscFlags::CommaSeparated,
                                llvm::cl::desc("Example list option")};
     ListOption<std::string> stringListOption{

 namespace mlir {
 void registerPassManagerTestPass() {
   PassRegistration<TestOptionsPass>();
 
   PassRegistration<TestModulePass>();
 
   PassRegistration<TestFunctionPass>();
 
+  PassRegistration<TestInterfacePass>();
+
   PassRegistration<TestCrashRecoveryPass>();
   PassRegistration<TestFailurePass>();
 
   PassRegistration<TestStatisticPass>();
 
   PassPipelineRegistration<>("test-pm-nested-pipeline",
                              "Test a nested pipeline in the pass manager",
                              testNestedPipeline);

mlir/unittests/Pass/PassManagerTest.cpp

TEST(PassManagerTest, OpSpecificAnalysis) {
   }
 }
 
 namespace {
 struct InvalidPass : Pass {
   InvalidPass() : Pass(TypeID::get<InvalidPass>(), StringRef("invalid_op")) {}
   StringRef getName() const override { return "Invalid Pass"; }
   void runOnOperation() override {}
+  bool canScheduleOn(RegisteredOperationName opName) const override {
+    return true;
+  }
 
   /// A clone method to create a copy of this pass.
   std::unique_ptr<Pass> clonePass() const override {
     return std::make_unique<InvalidPass>(
         *static_cast<const InvalidPass *>(this));
   }
 };
 } // namespace