[mlir][Pass] Check supported op types before running a pass

[springerm]

Summary

Add extra error checking to prevent passes from being run on unsupported ops through the pass manager infrastructure.

Depends On: D153143

[mehdi_amini]

Can this be tested without using the transform dialect?

[springerm]

add unit test

[springerm]

In D153144#4428652, @mehdi_amini wrote:

Can this be tested without using the transform dialect?

I added a unit test. This issue can only be triggered from C++. (I cannot be triggered from test/Pass/pipeline-invalid.mlir.)

[rkayaith]

In D153144#4432401, @springerm wrote:

In D153144#4428652, @mehdi_amini wrote:

Can this be tested without using the transform dialect?

I added a unit test. This issue can only be triggered from C++. (I cannot be triggered from test/Pass/pipeline-invalid.mlir.)

Something like this should work:

// RUN: mlir-opt %s --pass-pipeline='any(test-function-pass)'
module {}

mlir/lib/Pass/Pass.cpp

441–443

fyi there's a bunch of similar (but slightly different) checks spread out in other places as well, and they all fail in slightly different ways:

https://github.com/llvm/llvm-project/blob/b7a86d03cb49229634419be4ae44df00fadd990d/mlir/lib/Pass/Pass.cpp#L804-L807

https://github.com/llvm/llvm-project/blob/b7a86d03cb49229634419be4ae44df00fadd990d/mlir/lib/Pass/PassRegistry.cpp#L48-L51

https://github.com/llvm/llvm-project/blob/b7a86d03cb49229634419be4ae44df00fadd990d/mlir/lib/Pass/Pass.cpp#L200-L206

https://github.com/llvm/llvm-project/blob/b7a86d03cb49229634419be4ae44df00fadd990d/mlir/lib/Pass/Pass.cpp#L266-L271

just pointing them out in case you see a way to unify anything

[springerm]

In D153144#4433658, @rkayaith wrote:

Something like this should work:

// RUN: mlir-opt %s --pass-pipeline='any(test-function-pass)'
module {}

I tried this but it does not trigger the error. I spent some time looking into this yesterday and I think I found the reason why but I can't remember now of the top of my head...

mlir/lib/Pass/Pass.cpp

441–443

These checks are all slightly different. E.g., the error is reported at a different time (adding pass vs. running pipeline) or a different error handler is used. It may be possible to simply this, but there is probably some benefit in keeping some redundant checks, so that we can get error messages earlier (when adding a pass and not when running the pipeline later).

Another thing that I tried was putting the entire "op type check" logic into canScheduleOn. This was not possible because there is no MLIRContext available in various parts of the pass infrastructure. (A StringRef is not enough because it cannot be queried for implemented interfaces, which is needed for InterfacePass.)

[springerm]

In D153144#4434069, @springerm wrote:

In D153144#4433658, @rkayaith wrote:

Something like this should work:

// RUN: mlir-opt %s --pass-pipeline='any(test-function-pass)'
module {}

I tried this but it does not trigger the error. I spent some time looking into this yesterday and I think I found the reason why but I can't remember now of the top of my head...

Actually, this is slightly different can *does* trigger the error! Thx!

[springerm]

add pipeline-invalid.mlir test

mlir/lib/Pass/Pass.cpp

LogicalResult OpToOpPassAdaptor::run(Pass *pass, Operation *op,
                                      unsigned parentInitGeneration) {
   std::optional<RegisteredOperationName> opInfo = op->getRegisteredInfo();
   if (!opInfo)
     return op->emitOpError()
            << "trying to schedule a pass on an unregistered operation";
   if (!opInfo->hasTrait<OpTrait::IsIsolatedFromAbove>())
     return op->emitOpError() << "trying to schedule a pass on an operation not "
                                 "marked as 'IsolatedFromAbove'";
+  if (!pass->canScheduleOn(*op->getName().getRegisteredInfo()))
+    return op->emitOpError()
+           << "trying to schedule a pass on an unsupported operation";
 
   // Initialize the pass state with a callback for the pass to dynamically
   // execute a pipeline on the currently visited operation.
   PassInstrumentor *pi = am.getPassInstrumentor();
   PassInstrumentation::PipelineParentInfo parentInfo = {llvm::get_threadid(),
                                                         pass};
   auto dynamicPipelineCallback = [&](OpPassManager &pipeline,
                                      Operation *root) -> LogicalResult {

mlir/test/Dialect/Transform/test-pass-application.mlir

func.func @valid_pass_option() {
   return
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %1 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   transform.apply_registered_pass "canonicalize" to %1 {options = "top-down=false"} : (!transform.any_op) -> !transform.any_op
 }
+
+// -----
+
+module {
+  // expected-error @below {{trying to schedule a pass on an unsupported operation}}
+  // expected-note @below {{target op}}
+  func.func @invalid_target_op_type() {
+    return
+  }
+
+  transform.sequence failures(propagate) {
+  ^bb1(%arg1: !transform.any_op):
+    %1 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+
+    // func-bufferize can be applied only to ModuleOps.
+    // expected-error @below {{pass pipeline failed}}
+    transform.apply_registered_pass "func-bufferize" to %1 : (!transform.any_op) -> !transform.any_op
+  }
+}

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 {}
+  void runOnOperation() override {
+    assert(getOperation()->getName().getStringRef() == "invalid_op" &&
+           "incorrect op type");
+  }
   bool canScheduleOn(RegisteredOperationName opName) const override {
-    return true;
+    return opName.getStringRef() == "invalid_op";
   }
 
   /// 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