[mlir][transform][NFC] Store all Mappings in region stack

[springerm]

Summary

Do not swap the Mappings when entering a region that is isolated from above. Simply push another Mappings struct to the stack and prevent invalid accesses during lookups.

Depends On: D153764

[ashay-github]

@springerm this patch seems to have introduced a use-after-free error. Specifically, in the ForeachOp::apply() method, ops inserted in setPayloadOps() seem to be erased before the next call to getPayloadOps(). My hunch is that the pointers are invalidated as a result of applying the transformation, since copying all pointers first seems to make the problem go away. Does the following patch seem like a reasonable fix? Are there other places where we should copy the pointers first before using them?

--- a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
+++ b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
@@ -828,7 +828,9 @@ transform::ForeachOp::apply(transform::TransformRewriter &rewriter,
                             transform::TransformState &state) {
   SmallVector<SmallVector<Operation *>> resultOps(getNumResults(), {});

-  for (Operation *op : state.getPayloadOps(getTarget())) {
+  // Copy pointers to guard against invalidation by the transform.
+  SmallVector<Operation *> payloadOps{state.getPayloadOps(getTarget())};
+  for (Operation *op : payloadOps) {
     auto scope = state.make_region_scope(getBody());
     if (failed(state.mapBlockArguments(getIterationVariable(), {op})))
       return DiagnosedSilenceableFailure::definiteFailure();

[ashay-github]

Correction: this patch didn't _introduce_ the problem, it just made it apparent.

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

#endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
   class RegionScope;
 
   /// Creates a new region scope for the given region. The region is expected to
   /// be nested in the currently processed region.
   // Implementation note: this method is inline but implemented outside of the
   // class body to comply with visibility and full-declaration requirements.
   inline RegionScope make_region_scope(Region &region);
 
-  /// Creates a new region scope for the given isolated-from-above region.
-  /// Unlike the non-isolated counterpart, there is no nesting expectation.
-  // Implementation note: this method is inline but implemented outside of the
-  // class body to comply with visibility and full-declaration requirements
-  inline RegionScope make_isolated_region_scope(Region &region);
-
   /// A RAII object maintaining a "stack frame" for a transform IR region. When
   /// applying a transform IR operation that contains a region, the caller is
   /// expected to create a RegionScope before applying the ops contained in the
   /// region. This ensures that the mappings between values defined in the
   /// transform IR region and payload IR operations are cleared when the region
   /// processing ends; such values cannot be accessed outside the region.
   class RegionScope {
   public:
     /// Forgets the mapping from or to values defined in the associated
     /// transform IR region, and restores the mapping that existed before
     /// entering this scope.
     ~RegionScope();
 
   private:
-    /// Tag structure for differentiating the constructor for isolated regions.
-    struct Isolated {};
-
     /// Creates a new scope for mappings between values defined in the given
     /// transform IR region and payload IR objects.
     RegionScope(TransformState &state, Region &region)
         : state(state), region(&region) {
-      auto res = state.mappings.try_emplace(this->region);
+      auto res = state.mappings.insert(std::make_pair(&region, Mappings()));
       assert(res.second && "the region scope is already present");
       (void)res;
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
-      assert(((state.regionStack.size() == 1 && !state.regionStack.back()) ||
-              state.regionStack.back()->isProperAncestor(&region)) &&
-             "scope started at a non-nested region");
       state.regionStack.push_back(&region);
 #endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
     }
 
-    /// Creates a new scope for mappings between values defined in the given
-    /// isolated-from-above transform IR region and payload IR objects.
-    RegionScope(TransformState &state, Region &region, Isolated)
-        : state(state), region(&region) {
-      // Store the previous mapping stack locally.
-      storedMappings = llvm::SmallDenseMap<Region *, Mappings>();
-      storedMappings->swap(state.mappings);
-      state.mappings.try_emplace(this->region);
-#if LLVM_ENABLE_ABI_BREAKING_CHECKS
-      state.regionStack.push_back(this->region);
-#endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
-    }
-
     /// Back-reference to the transform state.
     TransformState &state;
 
     /// The region this scope is associated with.
     Region *region;
 
-    /// Local copy of the mappings that existed before entering the current
-    /// region. Used only when the current region is isolated so we don't
-    /// accidentally look up the values defined outside the isolated region.
-    std::optional<llvm::SmallDenseMap<Region *, Mappings>> storedMappings =
-        std::nullopt;
-
     friend RegionScope TransformState::make_region_scope(Region &);
-    friend RegionScope TransformState::make_isolated_region_scope(Region &);
   };
   friend class RegionScope;
 
   /// Base class for TransformState extensions that allow TransformState to
   /// contain user-specified information in the state object. Clients are
   /// expected to derive this class, add the desired fields, and make the
   /// derived class compatible with the MLIR TypeID mechanism:
   ///
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);
   }
   Mappings &getMapping(Value value) {
-    auto it = mappings.find(value.getParentRegion());
+    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
+    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");
+    }
+#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);
   }
   Mappings &getMapping(Operation *operation) {
-    auto it = mappings.find(operation->getParentRegion());
+    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
+    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");
+    }
+#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,
                                        ResultRange opResults);
 
#endif // NDEBUG
   /// early.
   LogicalResult checkAndRecordHandleInvalidationImpl(
       transform::TransformOpInterface transform,
       transform::TransformState::InvalidatedHandleMap &newlyInvalidated) const;
 
   /// Remove all nullptrs from op handles that were added by `replacePayloadOp`.
   void compactOpHandles();
 
-  /// The mappings between transform IR values and payload IR ops, aggregated by
-  /// the region in which the transform IR values are defined.
-  llvm::SmallDenseMap<Region *, Mappings> mappings;
+  /// A stack of mappings between transform IR values and payload IR ops,
+  /// aggregated by the region in which the transform IR values are defined.
+  llvm::MapVector<Region *, Mappings> mappings;
 
   /// Op handles may be temporarily mapped to nullptr to avoid invalidating
   /// payload op iterators. This set contains all op handles with nullptrs.
   /// These handles are "compacted" (i.e., nullptrs removed) at the end of each
   /// transform.
   DenseSet<Value> opHandlesToCompact;
 
   /// Extensions attached to the TransformState, identified by the TypeID of

 
 /// Creates a RAII object the lifetime of which corresponds to the new mapping
 /// for transform IR values defined in the given region. Values defined in
 /// surrounding regions remain accessible.
 TransformState::RegionScope TransformState::make_region_scope(Region &region) {
   return RegionScope(*this, region);
 }
 
-/// Creates a RAII object the lifetime of which corresponds to the new mapping
-/// for transform IR values defined in the given isolated-from-above region.
-/// Values defined in surrounding regions cannot be accessed.
-TransformState::RegionScope
-TransformState::make_isolated_region_scope(Region &region) {
-  return RegionScope(*this, region, RegionScope::Isolated());
-}
-
 /// A listener that updates a TransformState based on IR modifications. This
 /// listener can be used during a greedy pattern rewrite to keep the transform
 /// state up-to-date.
 class TrackingListener : public RewriterBase::Listener,
                          public TransformState::Extension {
 public:
   /// Create a new TrackingListener for usage in the specified transform op.
   TrackingListener(TransformState &state, TransformOpInterface op);

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

transform::TransformState::TransformState(
     Region *region, Operation *payloadRoot,
     const RaggedArray<MappedValue> &extraMappings,
     const TransformOptions &options)
     : topLevel(payloadRoot), options(options) {
   topLevelMappedValues.reserve(extraMappings.size());
   for (ArrayRef<MappedValue> mapping : extraMappings)
     topLevelMappedValues.push_back(mapping);
 
-  auto result = mappings.try_emplace(region);
+  auto result = mappings.insert(std::make_pair(region, Mappings()));
   assert(result.second && "the region scope is already present");
   (void)result;
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
   regionStack.push_back(region);
 #endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
 }
 
 Operation *transform::TransformState::getTopLevel() const { return topLevel; }
transform::TransformState::getPayloadValues(Value handleValue) const {
   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 {
   bool found = false;
-  for (const Mappings &mapping : llvm::make_second_range(mappings)) {
+  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>())
+      break;
   }
 
   return success(found);
 }
 
 LogicalResult transform::TransformState::getHandlesForPayloadValue(
     Value payloadValue, SmallVectorImpl<Value> &handles) const {
   bool found = false;
-  for (const Mappings &mapping : llvm::make_second_range(mappings)) {
+  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>())
+      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.
 static DiagnosedSilenceableFailure dispatchMappedValues(
void transform::TransformState::recordOpHandleInvalidation(
 
   // Iterate over the mapping and invalidate aliasing handles. This is quite
   // expensive and only necessary for error reporting in case of transform
   // dialect misuse with dangling handles. Iteration over the handles is based
   // on the assumption that the number of handles is significantly less than the
   // number of IR objects (operations and values). Alternatively, we could walk
   // the IR nested in each payload op associated with the given handle and look
   // for handles associated with each operation and value.
-  for (const transform::TransformState::Mappings &mapping :
-       llvm::make_second_range(mappings)) {
+  for (const auto &[region, mapping] : llvm::reverse(mappings)) {
+    // Stop lookup when reaching a region that is isolated from above.
+    if (region->getParentOp()->hasTrait<OpTrait::IsIsolatedFromAbove>())
+      break;
     // Go over all op handle mappings and mark as invalidated any handle
     // pointing to any of the payload ops associated with the given handle or
     // any op nested in them.
     for (const auto &[payloadOp, otherHandles] : mapping.reverse) {
       for (Value otherHandle : otherHandles)
         recordOpHandleInvalidationOne(handle, potentialAncestors, payloadOp,
                                       otherHandle, throughValue,
                                       newlyInvalidated);

 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
   // Remember pointers to payload ops referenced by the handles going out of
   // scope.
   SmallVector<Operation *> referencedOps =
       llvm::to_vector(llvm::make_first_range(state.mappings[region].reverse));
 #endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
 
   state.mappings.erase(region);
-  if (storedMappings.has_value())
-    state.mappings.swap(*storedMappings);
 
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
   // If the last handle to a payload op has gone out of scope, we no longer
   // need to store the cached name. Pointers may get reused, leading to
   // incorrect associations in the cache.
   for (Operation *op : referencedOps) {
     SmallVector<Value> handles;
     if (succeeded(state.getHandlesForPayloadOp(op, handles)))

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

 /// content to the debug output stream) and returns failure. If any of the
 /// matchers produces a definite failure, reports it and returns failure. If all
 /// matchers in the block succeed, populates `mappings` with the payload
 /// entities associated with the block terminator operands.
 static DiagnosedSilenceableFailure
 matchBlock(Block &block, Operation *op, transform::TransformState &state,
            SmallVectorImpl<SmallVector<transform::MappedValue>> &mappings) {
   assert(block.getParent() && "cannot match using a detached block");
-  auto matchScope = state.make_isolated_region_scope(*block.getParent());
+  auto matchScope = state.make_region_scope(*block.getParent());
   if (failed(state.mapBlockArgument(block.getArgument(0), {op})))
     return DiagnosedSilenceableFailure::definiteFailure();
 
   for (Operation &match : block.without_terminator()) {
     if (!isa<transform::MatchOpInterface>(match)) {
       return emitDefiniteFailure(match.getLoc())
              << "expected operations in the match part to "
                 "implement MatchOpInterface";
WalkResult walkResult = root->walk([&](Operation *op) {
         if (diag.isDefiniteFailure())
           return WalkResult::interrupt();
         if (diag.isSilenceableFailure()) {
           DEBUG_MATCHER(DBGS_MATCHER() << "matcher " << matcher.getName()
                                        << " failed: " << diag.getMessage());
           continue;
         }
 
-        auto scope = state.make_isolated_region_scope(action.getFunctionBody());
+        auto scope = state.make_region_scope(action.getFunctionBody());
         for (auto &&[arg, map] : llvm::zip_equal(
                  action.getFunctionBody().front().getArguments(), mappings)) {
           if (failed(state.mapBlockArgument(arg, map)))
             return WalkResult::interrupt();
         }
 
         for (Operation &transform :
              action.getFunctionBody().front().without_terminator()) {
transform::IncludeOp::apply(transform::TransformRewriter &rewriter,
   assert(callee && "unverified reference to unknown symbol");
 
   if (callee.isExternal())
     return emitDefiniteFailure() << "unresolved external named sequence";
 
   // Map operands to block arguments.
   SmallVector<SmallVector<MappedValue>> mappings;
   detail::prepareValueMappings(mappings, getOperands(), state);
-  auto scope = state.make_isolated_region_scope(callee.getBody());
+  auto scope = state.make_region_scope(callee.getBody());
   for (auto &&[arg, map] :
        llvm::zip_equal(callee.getBody().front().getArguments(), mappings)) {
     if (failed(state.mapBlockArgument(arg, map)))
       return DiagnosedSilenceableFailure::definiteFailure();
   }
 
   DiagnosedSilenceableFailure result = applySequenceBlock(
       callee.getBody().front(), getFailurePropagationMode(), state, results);