[mlir][transform] Fix TrackingListener in regions that are isolated from above

[springerm]

Summary

When an operation is removed/replaced, the TrackingListener updates the internal transform state mapping between handles and payload IR. All handles must be updated, even the ones that are defined in a region that is beyond the most recent region that is isolated from above.

This fixes a bug, where a payload op was erased in a named sequence. Not only handles defined inside of the named region must be updated, but also all other handles such as the ones where the sequence is included.

Depends On: D153766

mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h

public:
   ArrayRef<Attribute> getParams(Value value) const;
 
   /// Returns the list of payload IR values that the given transform IR value
   /// corresponds to.
   ArrayRef<Value> getPayloadValues(Value handleValue) const;
 
   /// Populates `handles` with all handles pointing to the given Payload IR op.
   /// Returns success if such handles exist, failure otherwise.
+  /// If `includeOutOfScope` is set to "true", handles that are defined in
+  /// regions beyond the most recent isolated from above region are included.
   LogicalResult getHandlesForPayloadOp(Operation *op,
-                                       SmallVectorImpl<Value> &handles) const;
+                                       SmallVectorImpl<Value> &handles,
+                                       bool includeOutOfScope = false) const;
 
   /// Populates `handles` with all handles pointing to the given payload IR
   /// value. Returns success if such handles exist, failure otherwise.
-  LogicalResult
-  getHandlesForPayloadValue(Value payloadValue,
-                            SmallVectorImpl<Value> &handles) const;
+  /// If `includeOutOfScope` is set to "true", handles that are defined in
+  /// regions beyond the most recent isolated from above region are included.
+  LogicalResult getHandlesForPayloadValue(Value payloadValue,
+                                          SmallVectorImpl<Value> &handles,
+                                          bool includeOutOfScope = false) const;
 
   /// Applies the transformation specified by the given transform op and updates
   /// the state accordingly.
   DiagnosedSilenceableFailure applyTransform(TransformOpInterface transform);
 
   /// Records the mapping between a block argument in the transform IR and a
   /// list of operations in the payload IR. The arguments must be defined in
   /// blocks of the currently processed transform IR region, typically after a
private:
   /// argument points to the root operation in the payload IR being transformed,
   /// which may or may not contain the region with transform ops. Additional
   /// options can be provided through the trailing configuration object.
   TransformState(Region *region, Operation *payloadRoot,
                  const RaggedArray<MappedValue> &extraMappings = {},
                  const TransformOptions &options = TransformOptions());
 
   /// Returns the mappings frame for the region in which the value is defined.
-  const Mappings &getMapping(Value value) const {
-    return const_cast<TransformState *>(this)->getMapping(value);
+  /// If `allowOutOfScope` is set to "false", asserts that the value is in
+  /// scope, based on the current stack of frames.
+  const Mappings &getMapping(Value value, bool allowOutOfScope = false) const {
+    return const_cast<TransformState *>(this)->getMapping(value,
+                                                          allowOutOfScope);
   }
-  Mappings &getMapping(Value value) {
+  Mappings &getMapping(Value value, bool allowOutOfScope = false) {
     Region *region = value.getParentRegion();
     auto it = mappings.find(region);
     assert(it != mappings.end() &&
            "trying to find a mapping for a value from an unmapped region");
 #ifndef NDEBUG
+    if (!allowOutOfScope) {
       for (Region *r : llvm::reverse(llvm::make_first_range(mappings))) {
         if (r == region)
           break;
         if (r->getParentOp()->hasTrait<OpTrait::IsIsolatedFromAbove>())
-        llvm_unreachable(
-            "trying to get mapping beyond region that is isolated from above");
+          llvm_unreachable("trying to get mapping beyond region that is "
+                           "isolated from above");
+      }
     }
 #endif // NDEBUG
     return it->second;
   }
 
   /// Returns the mappings frame for the region in which the operation resides.
-  const Mappings &getMapping(Operation *operation) const {
-    return const_cast<TransformState *>(this)->getMapping(operation);
+  /// If `allowOutOfScope` is set to "false", asserts that the operation is in
+  /// scope, based on the current stack of frames.
+  const Mappings &getMapping(Operation *operation,
+                             bool allowOutOfScope = false) const {
+    return const_cast<TransformState *>(this)->getMapping(operation,
+                                                          allowOutOfScope);
   }
-  Mappings &getMapping(Operation *operation) {
+  Mappings &getMapping(Operation *operation, bool allowOutOfScope = false) {
     Region *region = operation->getParentRegion();
     auto it = mappings.find(region);
     assert(it != mappings.end() &&
            "trying to find a mapping for an operation from an unmapped region");
 #ifndef NDEBUG
+    if (!allowOutOfScope) {
       for (Region *r : llvm::reverse(llvm::make_first_range(mappings))) {
         if (r == region)
           break;
         if (r->getParentOp()->hasTrait<OpTrait::IsIsolatedFromAbove>())
-        llvm_unreachable(
-            "trying to get mapping beyond region that is isolated from above");
+          llvm_unreachable("trying to get mapping beyond region that is "
+                           "isolated from above");
+      }
     }
 #endif // NDEBUG
     return it->second;
   }
 
   /// Updates the state to include the associations between op results and the
   /// provided result of applying a transform op.
   LogicalResult updateStateFromResults(const TransformResults &results,

mlir/lib/Dialect/Transform/IR/TransformInterfaces.cpp

transform::TransformState::getPayloadValues(Value handleValue) const {
   const ValueMapping &mapping = getMapping(handleValue).values;
   auto iter = mapping.find(handleValue);
   assert(iter != mapping.end() && "cannot find mapping for value handle "
                                   "(param/operation handle provided?)");
   return iter->getSecond();
 }
 
 LogicalResult transform::TransformState::getHandlesForPayloadOp(
-    Operation *op, SmallVectorImpl<Value> &handles) const {
+    Operation *op, SmallVectorImpl<Value> &handles,
+    bool includeOutOfScope) const {
   bool found = false;
   for (const auto &[region, mapping] : llvm::reverse(mappings)) {
     auto iterator = mapping.reverse.find(op);
     if (iterator != mapping.reverse.end()) {
       llvm::append_range(handles, iterator->getSecond());
       found = true;
     }
     // Stop looking when reaching a region that is isolated from above.
-    if (region->getParentOp()->hasTrait<OpTrait::IsIsolatedFromAbove>())
+    if (!includeOutOfScope &&
+        region->getParentOp()->hasTrait<OpTrait::IsIsolatedFromAbove>())
       break;
   }
 
   return success(found);
 }
 
 LogicalResult transform::TransformState::getHandlesForPayloadValue(
-    Value payloadValue, SmallVectorImpl<Value> &handles) const {
+    Value payloadValue, SmallVectorImpl<Value> &handles,
+    bool includeOutOfScope) const {
   bool found = false;
   for (const auto &[region, mapping] : llvm::reverse(mappings)) {
     auto iterator = mapping.reverseValues.find(payloadValue);
     if (iterator != mapping.reverseValues.end()) {
       llvm::append_range(handles, iterator->getSecond());
       found = true;
     }
     // Stop looking when reaching a region that is isolated from above.
-    if (region->getParentOp()->hasTrait<OpTrait::IsIsolatedFromAbove>())
+    if (!includeOutOfScope &&
+        region->getParentOp()->hasTrait<OpTrait::IsIsolatedFromAbove>())
       break;
   }
 
   return success(found);
 }
 
 /// Given a list of MappedValues, cast them to the value kind implied by the
 /// interface of the handle type, and dispatch to one of the callbacks.

 LogicalResult
 transform::TransformState::replacePayloadOp(Operation *op,
                                             Operation *replacement) {
   // TODO: consider invalidating the handles to nested objects here.
 
 #ifndef NDEBUG
   for (Value opResult : op->getResults()) {
     SmallVector<Value> valueHandles;
-    (void)getHandlesForPayloadValue(opResult, valueHandles);
+    (void)getHandlesForPayloadValue(opResult, valueHandles,
+                                    /*includeOutOfScope=*/true);
     assert(valueHandles.empty() && "expected no mapping to old results");
   }
 #endif // NDEBUG
 
   // Drop the mapping between the op and all handles that point to it. Fail if
   // there are no handles.
   SmallVector<Value> opHandles;
-  if (failed(getHandlesForPayloadOp(op, opHandles)))
+  if (failed(getHandlesForPayloadOp(op, opHandles, /*includeOutOfScope=*/true)))
     return failure();
   for (Value handle : opHandles) {
-    Mappings &mappings = getMapping(handle);
+    Mappings &mappings = getMapping(handle, /*allowOutOfScope=*/true);
     dropMappingEntry(mappings.reverse, op, handle);
   }
 
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
   if (options.getExpensiveChecksEnabled()) {
     auto it = cachedNames.find(op);
     assert(it != cachedNames.end() && "entry not found");
     assert(it->second == op->getName() && "operation name mismatch");
#endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
   // is stored in the mapping. These nullptrs are removed after each transform.
   // Furthermore, nullptrs are not enumerated by payload op iterators. The
   // relative order of ops is preserved.
   //
   // Removing an op from the mapping would be problematic because removing an
   // element from an array invalidates iterators; merely changing the value of
   // elements does not.
   for (Value handle : opHandles) {
-    Mappings &mappings = getMapping(handle);
+    Mappings &mappings = getMapping(handle, /*allowOutOfScope=*/true);
     auto it = mappings.direct.find(handle);
     if (it == mappings.direct.end())
       continue;
 
     SmallVector<Operation *, 2> &association = it->getSecond();
     // Note that an operation may be associated with the handle more than once.
     for (Operation *&mapped : association) {
       if (mapped == op)
         mapped = replacement;
     }
 
     if (replacement) {
       mappings.reverse[replacement].push_back(handle);
     } else {
       opHandlesToCompact.insert(handle);
     }
   }
 
   return success();
 }
 
 LogicalResult
 transform::TransformState::replacePayloadValue(Value value, Value replacement) {
   SmallVector<Value> valueHandles;
-  if (failed(getHandlesForPayloadValue(value, valueHandles)))
+  if (failed(getHandlesForPayloadValue(value, valueHandles,
+                                       /*includeOutOfScope=*/true)))
     return failure();
 
   for (Value handle : valueHandles) {
-    Mappings &mappings = getMapping(handle);
+    Mappings &mappings = getMapping(handle, /*allowOutOfScope=*/true);
     dropMappingEntry(mappings.reverseValues, value, handle);
 
     // If replacing with null, that is erasing the mapping, drop the mapping
     // between the handles and the IR objects
     if (!replacement) {
       dropMappingEntry(mappings.values, handle, value);
     } else {
       auto it = mappings.values.find(handle);
if (!seen.insert(p).second) {
       return diag;
     }
   }
   return DiagnosedSilenceableFailure::success();
 }
 
 void transform::TransformState::compactOpHandles() {
   for (Value handle : opHandlesToCompact) {
-    Mappings &mappings = getMapping(handle);
+    Mappings &mappings = getMapping(handle, /*allowOutOfScope=*/true);
     llvm::erase_value(mappings.direct[handle], nullptr);
   }
   opHandlesToCompact.clear();
 }
 
 DiagnosedSilenceableFailure
 transform::TransformState::applyTransform(TransformOpInterface transform) {
   LLVM_DEBUG({
assert(op->getNumResults() == newValues.size() &&
          "invalid number of replacement values");
 
   // Replace value handles.
   for (auto [oldValue, newValue] : llvm::zip(op->getResults(), newValues))
     (void)replacePayloadValue(oldValue, newValue);
 
   // Replace op handle.
   SmallVector<Value> opHandles;
-  if (failed(getTransformState().getHandlesForPayloadOp(op, opHandles))) {
+  if (failed(getTransformState().getHandlesForPayloadOp(
+          op, opHandles, /*includeOutOfScope=*/true))) {
     // Op is not tracked.
     return;
   }
 
   // Helper function to check if the current transform op consumes any handle
   // that is mapped to `op`.
   //
   // Note: If a handle was consumed, there shouldn't be any alive users, so it

mlir/test/Dialect/Transform/test-pattern-application.mlir

 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["test.container"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   %1 = transform.structured.match ops{["test.erase_op"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   transform.test_print_remark_at_operand %1, "matched op" : !transform.any_op
   transform.apply_patterns to %0 {
     transform.apply_patterns.transform.test_patterns
   } : !transform.any_op
+  // No marker should be printed.
   transform.test_print_remark_at_operand %1, "op was deleted" : !transform.any_op
 }
 
 // -----
 
+// CHECK-LABEL: func @erase_tracked_op_in_named_sequence()
+//       CHECK:   "test.container"() ({
+//  CHECK-NEXT:   ^bb0:
+//  CHECK-NEXT:   }) : () -> ()
+module {
+  func.func @erase_tracked_op_in_named_sequence() {
+    "test.container"() ({
+      // expected-remark @below {{matched op}}
+      %0 = "test.erase_op"() {replace_with_new_op = "test.foo"} : () -> (i32)
+    }) : () -> ()
+    return
+  }
+
+  module attributes { transform.with_named_sequence } {
+    transform.named_sequence @foo(%arg0: !transform.any_op {transform.readonly}) -> () {
+      transform.apply_patterns to %arg0 {
+        transform.apply_patterns.transform.test_patterns
+      } : !transform.any_op
+      transform.yield
+    }
+
+    transform.sequence failures(propagate) {
+    ^bb1(%arg1: !transform.any_op):
+      %0 = transform.structured.match ops{["test.container"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+      %1 = transform.structured.match ops{["test.erase_op"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+      transform.test_print_remark_at_operand %1, "matched op" : !transform.any_op
+      include @foo failures(propagate) (%0) : (!transform.any_op) -> ()
+      // No marker should be printed.
+      transform.test_print_remark_at_operand %1, "op was deleted" : !transform.any_op
+    }
+  }
+}
+
+// -----
+
 // CHECK-LABEL: func @canonicalization(
 //       CHECK:   %[[c5:.*]] = arith.constant 5 : index
 //       CHECK:   return %[[c5]]
 func.func @canonicalization(%t: tensor<5xf32>) -> index {
   %c0 = arith.constant 0 : index
   // expected-remark @below {{op was replaced}}
   %dim = tensor.dim %t, %c0 : tensor<5xf32>
   return %dim : index