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

[Mikhail Romanov]

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

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.

>From e0d214c8ebedd2f3fee66d7a34cdb79f23c6d7c6 Mon Sep 17 00:00:00 2001
From: Mikhail Romanov <mmromanov at ispras.ru>
Date: Thu, 2 Jul 2026 14:26:42 +0300
Subject: [PATCH] [MLIR] Fix transform.apply_patterns with apply_cse option

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.

Signed-off-by: Mikhail Romanov <mmromanov at ispras.ru>
---
 mlir/lib/Dialect/Transform/IR/TransformOps.cpp | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

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";