[Mlir-commits] [mlir] [MLIR] Fix transform.apply_patterns with apply_cse option (PR #212818)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Mikhail Romanov (Mmi257ia)

<details>
<summary>Changes</summary>

This patch fixes a crash in `transform.apply_patterns` when `apply_cse` is enabled and the target operation does not have the `IsolatedFromAbove` trait.

For non-`IsolatedFromAbove` targets, `transform.apply_patterns` cannot apply patterns directly to the target operation. Instead, it collects the operations within the target into a worklist and passes that worklist to `applyOpPatternsGreedily`. When `apply_cse` is enabled, pattern application and CSE are performed repeatedly until a fixpoint is reached.

Previously, the worklist was collected only once before entering the fixpoint loop. However, pattern application may erase operations from the IR, leaving stale pointers in the cached worklist. Reusing this worklist in subsequent iterations may therefore result in a crash.

This patch rebuilds the worklist at the beginning of each iteration, ensuring that it always reflects the current IR state.

A representative reproducer is shown below:
```llvm
module attributes {transform.with_named_sequence} {
  func.func @<!-- -->test(%A: tensor<128x256xf32>,
                  %B: tensor<256x512xf32>,
                  %C: tensor<128x512xf32>) -> tensor<128x512xf32> {
    %res = linalg.matmul
        ins(%A, %B : tensor<128x256xf32>, tensor<256x512xf32>)
        outs(%C : tensor<128x512xf32>) -> tensor<128x512xf32>
    return %res : tensor<128x512xf32>
  }

  transform.named_sequence @<!-- -->__transform_main(%root: !transform.any_op) {
    %matmul = transform.structured.match
        ops{["linalg.matmul"]}
        in %root
        : (!transform.any_op) -> !transform.any_op

    %tiled, %loops:3 =
        transform.structured.tile_using_for %matmul
            tile_sizes [4, 4, 4]
        : (!transform.any_op)
          -> (!transform.any_op,
              !transform.any_op,
              !transform.any_op,
              !transform.any_op)

    transform.structured.vectorize %tiled : !transform.any_op

    transform.apply_patterns to %loops#<!-- -->0 {
      transform.apply_patterns.canonicalization
    } {apply_cse} : !transform.any_op

    transform.yield
  }
}
```

Here, `transform.apply_patterns` is applied to the `scf.for` loop produced by tiling, which is not `IsolatedFromAbove`. Running `mlir-opt --transform-interpreter` on this example crashes before this change and completes successfully with this patch.

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


1 Files Affected:

- (modified) mlir/lib/Dialect/Transform/IR/TransformOps.cpp (+6-5) 


``````````diff
diff --git a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
index 9a9119e572c70..2f1648a0b0c30 100644
--- a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
+++ b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
@@ -390,11 +390,6 @@ DiagnosedSilenceableFailure transform::ApplyPatternsOp::applyToOne(
   // Non-isolated case: gather the ops manually because the op-list
   // GreedyPatternRewriteDriver overload only performs a single iteration and
   // does not simplify regions. CSE is driven externally to reach a fixpoint.
-  SmallVector<Operation *> ops;
-  target->walk([&](Operation *nestedOp) {
-    if (target != nestedOp)
-      ops.push_back(nestedOp);
-  });
 
   // One or two iterations should be sufficient. Stop iterating after a certain
   // threshold to make debugging easier.
@@ -402,6 +397,12 @@ DiagnosedSilenceableFailure transform::ApplyPatternsOp::applyToOne(
   int64_t iteration = 0;
   bool cseChanged = false;
   do {
+    SmallVector<Operation *> ops;
+    target->walk([&](Operation *nestedOp) {
+      if (target != nestedOp)
+        ops.push_back(nestedOp);
+    });
+
     if (failed(applyOpPatternsGreedily(ops, frozenPatterns, config))) {
       return emitSilenceableFailure(target)
              << "greedy pattern application failed";

``````````

</details>


https://github.com/llvm/llvm-project/pull/212818