[mlir] transform dialect: don't crash in verifiers

[ftynse]

Fix crashes in the verifier of transform.with_named_sequence attribute attached to a symbol table operation caused by it constructing a call graph inside the symbol table. The call graph construction assumes calls and callables, such as functions or named sequences, have been verified, but it is not yet the case when the attribute verifier on the (parent) symbol table operation runs. Trigger such verification manually before constructing the call graph. This adds redundancy in verification, but there is currently no mechanism to change the order of verificaiton. In performance-critical scenarios, verification can be disabled altogether.

Remove unnecessary verfificaton from transform::IncludeOp::getEffects. It was introduced along with the op definition as the op used to inspect the body of callee, which assumed the body existed, to identify handle consumption behavior. This was later evolved to having explicit argument attributes on the callee, which handles the absence of such attributes gracefully without the need for verification, but the verification was never removed. It would have been causing infinite recursion if kept in place.

Fixes #159646.
Fixes #159734.
Fixes #159736.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Oleksandr "Alex" Zinenko (ftynse)

Changes

Fix crashes in the verifier of transform.with_named_sequence attribute attached to a symbol table operation caused by it constructing a call graph inside the symbol table. The call graph construction assumes calls and callables, such as functions or named sequences, have been verified, but it is not yet the case when the attribute verifier on the (parent) symbol table operation runs. Trigger such verification manually before constructing the call graph. This adds redundancy in verification, but there is currently no mechanism to change the order of verificaiton. In performance-critical scenarios, verification can be disabled altogether.

Remove unnecessary verfificaton from transform::IncludeOp::getEffects. It was introduced along with the op definition as the op used to inspect the body of callee, which assumed the body existed, to identify handle consumption behavior. This was later evolved to having explicit argument attributes on the callee, which handles the absence of such attributes gracefully without the need for verification, but the verification was never removed. It would have been causing infinite recursion if kept in place.

Fixes #159646.
Fixes #159734.

---

Full diff: https://github.com/llvm/llvm-project/pull/161098.diff

3 Files Affected:

• (modified) mlir/lib/Dialect/Transform/IR/TransformDialect.cpp (+15)
• (modified) mlir/lib/Dialect/Transform/IR/TransformOps.cpp (+1-7)
• (modified) mlir/test/Dialect/Transform/ops-invalid.mlir (+32)

diff --git a/mlir/lib/Dialect/Transform/IR/TransformDialect.cpp b/mlir/lib/Dialect/Transform/IR/TransformDialect.cpp
index a500228d68c77..45cef9c162c70 100644
--- a/mlir/lib/Dialect/Transform/IR/TransformDialect.cpp
+++ b/mlir/lib/Dialect/Transform/IR/TransformDialect.cpp
@@ -13,6 +13,7 @@
 #include "mlir/Dialect/Transform/IR/Utils.h"
 #include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"
 #include "mlir/IR/DialectImplementation.h"
+#include "mlir/IR/Verifier.h"
 #include "llvm/ADT/SCCIterator.h"
 #include "llvm/ADT/TypeSwitch.h"
 
@@ -140,6 +141,20 @@ LogicalResult transform::TransformDialect::verifyOperationAttribute(
                                         "operations with symbol tables";
     }
 
+    // Pre-verify calls and callables because call graph construction below
+    // assumes they are valid, but this verifier runs before verifying the
+    // nested operations.
+    WalkResult walkResult = op->walk([](Operation *nested) {
+      if (!isa<CallableOpInterface, CallOpInterface>(nested))
+        return WalkResult::advance();
+
+      if (failed(verify(nested, /*verifyRecursively=*/false)))
+        return WalkResult::interrupt();
+      return WalkResult::advance();
+    });
+    if (walkResult.wasInterrupted())
+      return failure();
+
     const mlir::CallGraph callgraph(op);
     for (auto scc = llvm::scc_begin(&callgraph); !scc.isAtEnd(); ++scc) {
       if (!scc.hasCycle())
diff --git a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
index 48df1a0ba12c9..d5781a3e456b1 100644
--- a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
+++ b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
@@ -2098,17 +2098,11 @@ void transform::IncludeOp::getEffects(
       getOperation(), getTarget());
   if (!callee)
     return defaultEffects();
-  DiagnosedSilenceableFailure earlyVerifierResult =
-      verifyNamedSequenceOp(callee, /*emitWarnings=*/false);
-  if (!earlyVerifierResult.succeeded()) {
-    (void)earlyVerifierResult.silence();
-    return defaultEffects();
-  }
 
   for (unsigned i = 0, e = getNumOperands(); i < e; ++i) {
     if (callee.getArgAttr(i, TransformDialect::kArgConsumedAttrName))
       consumesHandle(getOperation()->getOpOperand(i), effects);
-    else
+    else if (callee.getArgAttr(i, TransformDialect::kArgReadOnlyAttrName))
       onlyReadsHandle(getOperation()->getOpOperand(i), effects);
   }
 }
diff --git a/mlir/test/Dialect/Transform/ops-invalid.mlir b/mlir/test/Dialect/Transform/ops-invalid.mlir
index 71a260f1196e9..0e34a426bbfc8 100644
--- a/mlir/test/Dialect/Transform/ops-invalid.mlir
+++ b/mlir/test/Dialect/Transform/ops-invalid.mlir
@@ -369,6 +369,7 @@ module attributes { transform.with_named_sequence } {
   // expected-error @below {{recursion not allowed in named sequences}}
   transform.named_sequence @self_recursion() -> () {
     transform.include @self_recursion failures(suppress) () : () -> ()
+    transform.yield
   }
 }
 
@@ -908,3 +909,34 @@ module attributes { transform.with_named_sequence } {
     transform.yield
   }
 }
+
+// -----
+
+module attributes { transform.with_named_sequence } {
+  transform.named_sequence @__transform_main(%arg0: !transform.any_op) -> () {
+    // Intentionally malformed func with no region. This shouldn't crash the
+    // verifier of `with_named_sequence` that runs before we get to the
+    // function.
+    // expected-error @below {{requires one region}}
+    "func.func"() : () -> ()
+    transform.yield
+  }
+}
+
+// -----
+
+module attributes { transform.with_named_sequence } {
+  transform.named_sequence @__transform_main(%arg0: !transform.any_op) -> () {
+    // Intentionally malformed call with a region. This shouldn't crash the
+    // verifier of `with_named_sequence` that runs before we get to the call.
+    // expected-error @below {{requires zero regions}}
+    "func.call"() <{
+      function_type = () -> (),
+      sym_name = "lambda_function"
+    }> ({
+    ^bb0:
+      "func.return"() : () -> ()
+    }) : () -> ()
+    transform.yield
+  }
+}