[Mlir-commits] [mlir] [mlir][transform] Reject non-transform ops in transform.foreach body (PR #236566)

[Kevin Jay Chang]

https://github.com/kevinjaychang created https://github.com/llvm/llvm-project/pull/236566

ForeachOp::getEffects cast every body op to TransformOpInterface, so a non-transform op in the body (e.g. arith.constant) hit an assertion during verification. Use dyn_cast in getEffects, and reject such ops in ForeachOp::verify with the same diagnostic transform.sequence emits.

Fixes #234637

>From 46d9f86643c69c7a5f63714d85beb816dbebdbb0 Mon Sep 17 00:00:00 2001
From: Kevin <kevinjaychang2003 at gmail.com>
Date: Wed, 28 Oct 2026 13:42:32 -0700
Subject: [PATCH] [mlir][transform] Reject non-transform ops in
 transform.foreach body

ForeachOp::getEffects cast every body op to TransformOpInterface, so a
non-transform op in the body (e.g. arith.constant) hit an assertion
during verification. Use dyn_cast in getEffects, and reject such ops in
ForeachOp::verify with the same diagnostic transform.sequence emits.

Fixes #234637
---
 .../lib/Dialect/Transform/IR/TransformOps.cpp | 20 +++++++++++++++----
 mlir/test/Dialect/Transform/ops-invalid.mlir  | 17 +++++++++++++++-
 2 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
index e40f3b7885c273..ccc086c6e93a10 100644
--- a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
+++ b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
@@ -1677,8 +1677,8 @@ void transform::ForeachOp::getEffects(
        llvm::zip(getTargetsMutable(), getBody().front().getArguments())) {
     BlockArgument blockArgument = blockArg;
     if (any_of(getBody().front().without_terminator(), [&](Operation &op) {
-          return isHandleConsumed(blockArgument,
-                                  cast<TransformOpInterface>(&op));
+          auto transformOp = dyn_cast<TransformOpInterface>(&op);
+          return transformOp && isHandleConsumed(blockArgument, transformOp);
         })) {
       consumesHandle(target, effects);
     } else {
@@ -1687,11 +1687,13 @@ void transform::ForeachOp::getEffects(
   }
 
   if (any_of(getBody().front().without_terminator(), [&](Operation &op) {
-        return doesModifyPayload(cast<TransformOpInterface>(&op));
+        auto transformOp = dyn_cast<TransformOpInterface>(&op);
+        return transformOp && doesModifyPayload(transformOp);
       })) {
     modifiesPayload(effects);
   } else if (any_of(getBody().front().without_terminator(), [&](Operation &op) {
-               return doesReadPayload(cast<TransformOpInterface>(&op));
+               auto transformOp = dyn_cast<TransformOpInterface>(&op);
+               return transformOp && doesReadPayload(transformOp);
              })) {
     onlyReadsPayload(effects);
   }
@@ -1754,6 +1756,16 @@ LogicalResult transform::ForeachOp::verify() {
                          "operands to have the same op/value/param type");
   }
 
+  for (Operation &child : getBody().front().without_terminator()) {
+    if (!isa<TransformOpInterface>(child)) {
+      InFlightDiagnostic diag =
+          emitOpError()
+          << "expected children ops to implement TransformOpInterface";
+      diag.attachNote(child.getLoc()) << "op without interface";
+      return diag;
+    }
+  }
+
   return success();
 }
 
diff --git a/mlir/test/Dialect/Transform/ops-invalid.mlir b/mlir/test/Dialect/Transform/ops-invalid.mlir
index 585736d30269e7..4f9eba16cd32e8 100644
--- a/mlir/test/Dialect/Transform/ops-invalid.mlir
+++ b/mlir/test/Dialect/Transform/ops-invalid.mlir
@@ -339,6 +339,21 @@ transform.sequence failures(propagate) {
 
 // -----
 
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg0: !transform.any_op) {
+    // expected-error @below {{expected children ops to implement TransformOpInterface}}
+    transform.foreach %arg0 : !transform.any_op {
+    ^bb0(%arg1: !transform.any_op):
+      // expected-note @below {{op without interface}}
+      %c0 = arith.constant 0 : i32
+      transform.yield
+    }
+    transform.yield
+  }
+}
+
+// -----
+
 transform.sequence failures(suppress) {
 ^bb0(%arg0: !transform.any_op):
   // expected-error @below {{TransformOpInterface requires memory effects on operands to be specified}}
@@ -496,7 +511,7 @@ module {
 module attributes { transform.with_named_sequence} {
   transform.sequence failures(suppress) {
   ^bb0(%arg0: !transform.any_op):
-    // expected-error @below {{op symbol's parent must have the SymbolTable trai}}
+    // expected-error @below {{op symbol's parent must have the SymbolTable trait}}
     transform.named_sequence @nested() {
       transform.yield
     }