[mlir][Transform] Relax the applicability of transform.foreach_match …

[nicolasvasilache]

…to also take into account the op itself

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-linalg

Author: Nicolas Vasilache (nicolasvasilache)

Changes

…to also take into account the op itself

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/Transform/IR/TransformOps.cpp (+2-6)
• (modified) mlir/test/Dialect/Linalg/match-ops-interpreter.mlir (+1)

diff --git a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
index 8db77b6059dd2e3..d39a1847e643cc7 100644
--- a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
+++ b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
@@ -850,10 +850,6 @@ transform::ForeachMatchOp::apply(transform::TransformRewriter &rewriter,
 
   for (Operation *root : state.getPayloadOps(getRoot())) {
     WalkResult walkResult = root->walk([&](Operation *op) {
-      // Skip over the root op itself so we don't invalidate it.
-      if (op == root)
-        return WalkResult::advance();
-
       DEBUG_MATCHER({
         DBGS_MATCHER() << "matching ";
         op->print(llvm::dbgs(),
@@ -1556,10 +1552,10 @@ DiagnosedSilenceableFailure transform::MatchOperationEmptyOp::matchOperation(
     ::std::optional<::mlir::Operation *> maybeCurrent,
     transform::TransformResults &results, transform::TransformState &state) {
   if (!maybeCurrent.has_value()) {
-    DBGS_MATCHER() << "MatchOperationEmptyOp success\n";
+    DEBUG_MATCHER({DBGS_MATCHER() << "MatchOperationEmptyOp success\n";});
     return DiagnosedSilenceableFailure::success();
   }
-  DBGS_MATCHER() << "MatchOperationEmptyOp failure\n";
+  DEBUG_MATCHER({ DBGS_MATCHER() << "MatchOperationEmptyOp failure\n"; });
   return emitSilenceableError() << "operation is not empty";
 }
 
diff --git a/mlir/test/Dialect/Linalg/match-ops-interpreter.mlir b/mlir/test/Dialect/Linalg/match-ops-interpreter.mlir
index 9489aadac843d7b..05bfe294754d7c6 100644
--- a/mlir/test/Dialect/Linalg/match-ops-interpreter.mlir
+++ b/mlir/test/Dialect/Linalg/match-ops-interpreter.mlir
@@ -106,6 +106,7 @@ module attributes { transform.with_named_sequence } {
     transform.yield
   }
 
+  // expected-remark @below {{other}}
   func.func @payload() attributes { transform.target_tag = "start_here" } {
     // expected-remark @below {{other}}
     %D = arith.constant dense<1.0> : tensor<2x4xf32>

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[ftynse]

What's your use case? Not matching the root is load bearing as the rewrite part can erase it and chaos will ensue when we try to associate the result with the erased operation in line 895 below.

[nicolasvasilache]

What's your use case? Not matching the root is load bearing as the rewrite part can erase it and chaos will ensue when we try to associate the result with the erased operation in line 895 below.

Added a restrict_root unit attr like we discussed offline.
It is a pretty common idiom to apply transforms on a func, check some property of the func in the matcher and then apply actions on the func itself, modifying only its content.