[Mlir-commits] [mlir] 6187354 - [mlir][transform] Fix crash when consuming an op in a named sequence (#67437)

[llvmlistbot at llvm.org]

Author: Matthias Springer
Date: 2023-09-27T12:40:41+02:00
New Revision: 6187354095a64caf573379df56b6d880abe40130

URL: https://github.com/llvm/llvm-project/commit/6187354095a64caf573379df56b6d880abe40130
DIFF: https://github.com/llvm/llvm-project/commit/6187354095a64caf573379df56b6d880abe40130.diff

LOG: [mlir][transform] Fix crash when consuming an op in a named sequence (#67437)

Fix a crash when consuming an op in a named sequence, when the same op
is also mapped in the caller's mapping. Ops must be removed from *all*
mappings during the "expensive checks". Otherwise, we may have dangling
pointers in the mappings data structures, which interfere with the
expensive checks.

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h
    mlir/lib/Dialect/Transform/IR/TransformInterfaces.cpp
    mlir/test/Dialect/Transform/expensive-checks.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h b/mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h
index b45861e6190c18a..0e72a93e685e32f 100644
--- a/mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h
+++ b/mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h
@@ -617,7 +617,11 @@ class TransformState {
   /// Forgets the payload IR ops associated with the given transform IR value,
   /// as well as any association between value handles and the results of said
   /// payload IR op.
-  void forgetMapping(Value opHandle, ValueRange origOpFlatResults);
+  ///
+  /// If `allowOutOfScope` is set to "false", asserts that the handle is in
+  /// scope, based on the current stack of frames.
+  void forgetMapping(Value opHandle, ValueRange origOpFlatResults,
+                     bool allowOutOfScope = false);
 
   void forgetValueMapping(Value valueHandle,
                           ArrayRef<Operation *> payloadOperations);

diff  --git a/mlir/lib/Dialect/Transform/IR/TransformInterfaces.cpp b/mlir/lib/Dialect/Transform/IR/TransformInterfaces.cpp
index 483b0e7f7a4f998..4a9bb2dba7d660c 100644
--- a/mlir/lib/Dialect/Transform/IR/TransformInterfaces.cpp
+++ b/mlir/lib/Dialect/Transform/IR/TransformInterfaces.cpp
@@ -305,8 +305,9 @@ void dropMappingEntry(Mapping &mapping, Key key, Mapped mapped) {
 }
 
 void transform::TransformState::forgetMapping(Value opHandle,
-                                              ValueRange origOpFlatResults) {
-  Mappings &mappings = getMapping(opHandle);
+                                              ValueRange origOpFlatResults,
+                                              bool allowOutOfScope) {
+  Mappings &mappings = getMapping(opHandle, allowOutOfScope);
   for (Operation *op : mappings.direct[opHandle])
     dropMappingEntry(mappings.reverse, op, opHandle);
   mappings.direct.erase(opHandle);
@@ -1021,9 +1022,10 @@ transform::TransformState::applyTransform(TransformOpInterface transform) {
       // pre-generated error messages, so we do not need the association to
       // still be there when the invalidated handle is accessed.
       SmallVector<Value> handles;
-      (void)getHandlesForPayloadOp(op, handles);
+      (void)getHandlesForPayloadOp(op, handles, /*includeOutOfScope=*/true);
       for (Value handle : handles)
-        forgetMapping(handle, /*origOpFlatResults=*/ValueRange());
+        forgetMapping(handle, /*origOpFlatResults=*/ValueRange(),
+                      /*allowOutOfScope=*/true);
       cachedNames.erase(op);
     }
 

diff  --git a/mlir/test/Dialect/Transform/expensive-checks.mlir b/mlir/test/Dialect/Transform/expensive-checks.mlir
index fef857bcef029df..ee9a5af8055247e 100644
--- a/mlir/test/Dialect/Transform/expensive-checks.mlir
+++ b/mlir/test/Dialect/Transform/expensive-checks.mlir
@@ -410,3 +410,22 @@ transform.sequence failures(propagate) {
     transform.yield
   }
 }
+
+// -----
+
+module @named_inclusion_and_consumption attributes { transform.with_named_sequence } {
+
+  transform.named_sequence @foo(%arg0: !transform.any_op {transform.consumed}) -> () {
+    // Consuming this handle removes the mapping from the current stack frame
+    // mapping and from the caller's stack frame mapping. (If this were not
+    // be the case, the "expensive checks" caching mechanism for op names
+    // would throw an error saying that an op is mapped but not in the cache.)
+    transform.test_consume_operand %arg0 : !transform.any_op
+    transform.yield
+  }
+
+  transform.sequence failures(propagate) {
+  ^bb0(%arg0: !transform.any_op):
+    include @foo failures(propagate) (%arg0) : (!transform.any_op) -> ()
+  }
+}