[mlir][Pass] Provide default implemenation for `canScheduleOn`

[springerm]

Summary

The previous API was misleading because it suggested that a pass can be restricted to a certain op type by simply setting an opName via the constructor of Pass. However, canScheduleOn should be adapted accordingly. This revision provides a default implementation for canScheduleOn, such that setting opName alone is sufficient. (Users can still override canScheduleOn to specify other constraints.)

Depends On: D153144

[springerm]

rebase

mlir/include/mlir/Pass/Pass.h

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
+  /// Indicate if this 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;
+  ///
+  /// By default, this pass can be scheduled only on ops whose name matches the
+  /// name that was specified at the constructor. If no op name was specified,
+  /// this pass can be scheduled on any op.
+  virtual bool canScheduleOn(RegisteredOperationName opName) const {
+    if (!getOpName().has_value())
+      return true;
+    return opName.getStringRef() == *getOpName();
+  }
 
   /// 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);
   }
 
protected:
   void copyOptionValuesFrom(const Pass *other);
 
 private:
   /// Out of line virtual method to ensure vtables and metadata are emitted to a
   /// single .o file.
   virtual void anchor();
 
   /// Represents a unique identifier for the pass.
-  TypeID passID;
+  const TypeID passID;
 
   /// The name of the operation that this pass operates on, or std::nullopt if
   /// this is a generic OperationPass.
-  std::optional<StringRef> opName;
+  const std::optional<StringRef> opName;
 
   /// The current execution state for the pass.
   std::optional<detail::PassExecutionState> passState;
 
   /// The set of statistics held by this pass.
   std::vector<Statistic *> statistics;
 
   /// The pass options registered to this pass instance.
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>();

 ///   - 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 that implements the given interface.
 ///
 /// Interface 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

mlir/unittests/Pass/PassManagerTest.cpp

 }
 
 namespace {
 struct InvalidPass : Pass {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InvalidPass)
 
   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;
+  void runOnOperation() override {
+    assert(getOperation()->getName().getStringRef() == "invalid_op" &&
+           "incorrect op type");
   }
 
   /// 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));
   }
 };
EXPECT_EQ(
       diagnostic->str(),
       "'invalid_op' op trying to schedule a pass on an unregistered operation");
 
   // Check that clearing the pass manager effectively removed the pass.
   pm.clear();
   result = pm.run(module.get());
   EXPECT_TRUE(succeeded(result));
 
+  // Create a pass manager that schedules arbitrary ops, but then run a pass
+  // that expects a specific op.
+  PassManager anyOpPm(&context);
+  // Adding the pass succeeds, because the we do not know yet on what op the
+  // pass will be scheduled.
+  anyOpPm.addPass(std::make_unique<InvalidPass>());
+  result = anyOpPm.run(module.get());
+  EXPECT_TRUE(failed(result));
+  ASSERT_TRUE(diagnostic.get() != nullptr);
+  EXPECT_EQ(diagnostic->str(), "'builtin.module' op trying to schedule a pass "
+                               "on an unsupported operation");
+
   // Check that adding the pass at the top-level triggers a fatal error.
   ASSERT_DEATH(pm.addPass(std::make_unique<InvalidPass>()),
                "Can't add pass 'Invalid Pass' restricted to 'invalid_op' on a "
                "PassManager intended to run on 'builtin.module', did you "
                "intend to nest?");
 }
 
 } // namespace