correctly check uses of pattern descriptor transform ops

[ftynse]

In the transform dialect use-after-free chcker pass, account for pattern descriptor operations that intentionally have no declared side effects. They are not destroying any handles.

Closes #118761.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Oleksandr "Alex" Zinenko (ftynse)

Changes

In the transform dialect use-after-free chcker pass, account for pattern descriptor operations that intentionally have no declared side effects. They are not destroying any handles.

Closes #118761.

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/Transform/Transforms/CheckUses.cpp (+2)
• (modified) mlir/test/Dialect/Transform/check-use-after-free.mlir (+14)

diff --git a/mlir/lib/Dialect/Transform/Transforms/CheckUses.cpp b/mlir/lib/Dialect/Transform/Transforms/CheckUses.cpp
index e6db819b51c22c..bfe1d9682177d8 100644
--- a/mlir/lib/Dialect/Transform/Transforms/CheckUses.cpp
+++ b/mlir/lib/Dialect/Transform/Transforms/CheckUses.cpp
@@ -338,6 +338,8 @@ class TransformOpMemFreeAnalysis {
   void collectFreedValues(Operation *root) {
     SmallVector<MemoryEffects::EffectInstance> instances;
     root->walk([&](Operation *child) {
+      if (isa<transform::PatternDescriptorOpInterface>(child))
+        return;
       // TODO: extend this to conservatively handle operations with undeclared
       // side effects as maybe freeing the operands.
       auto iface = cast<MemoryEffectOpInterface>(child);
diff --git a/mlir/test/Dialect/Transform/check-use-after-free.mlir b/mlir/test/Dialect/Transform/check-use-after-free.mlir
index 742c5dcc78f608..0fe8b5da173558 100644
--- a/mlir/test/Dialect/Transform/check-use-after-free.mlir
+++ b/mlir/test/Dialect/Transform/check-use-after-free.mlir
@@ -177,3 +177,17 @@ transform.sequence failures(propagate) {
   ^bb0(%arg1: !transform.any_op):
   }
 }
+
+// -----
+
+// This should not crash.
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+    transform.apply_patterns to %0 {
+      transform.apply_patterns.memref.extract_address_computations
+    } : !transform.any_op
+    transform.yield
+  }
+}

[martin-luecke]

LGTM!