[mlir] use shared pointer to prevent vector resizes from destroying ops

[ashay-github]

Summary

The MapVector type stores key-value pairs in a vector, which, when
resized, copies the entries and destroys the old ones. This causes the
underlying operations to be deleted, subsequently causing segfaults.

This patch makes the mappings map type refer to a shared pointer
instead, so that resizing the vector doesn't call the operations'
destructors.

[ftynse]

Do you have a test case where this was triggered?

[springerm]

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

832 ↗

(On Diff #537352)

getPayloadOp returns an iterator that skips deleted ops. It sounds like there is bug somewhere. A test case would be useful, then I can debug it.

[ashay-github]

Sadly, I can reproduce the problem only with MSVC Debug builds. The following test case (extracted verbatim from test-interpreter.mlir) triggers the bug:

func.func @bar() {
  // expected-remark @below {{transform applied}}
  %0 = arith.constant 0 : i32
  return
}

transform.with_pdl_patterns {
^bb0(%arg0: !transform.any_op):
  pdl.pattern @const : benefit(1) {
    %r = pdl.types
    %0 = pdl.operation "arith.constant" -> (%r : !pdl.range<type>)
    pdl.rewrite %0 with "transform.dialect"
  }

  transform.sequence %arg0 : !transform.any_op failures(propagate) {
  ^bb1(%arg1: !transform.any_op):
    %f = pdl_match @const in %arg1 : (!transform.any_op) -> !transform.any_op
    transform.foreach %f : !transform.any_op {
    ^bb2(%arg2: !transform.any_op):
      // expected-remark @below {{1}}
      transform.test_print_number_of_associated_payload_ir_ops %arg2 : !transform.any_op
      transform.test_print_remark_at_operand %arg2, "transform applied" : !transform.any_op
    }
  }
}

If it helps, this problem surfaced after the switch from DenseMap to MapVector (https://reviews.llvm.org/D153765).

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

832 ↗

(On Diff #537352)

True, but in checking for deleted ops, it checks for nullptr, whereas here, the pointer points to 0xdddddddddddddddddd, so it doesn't get filtered out.

[springerm]

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

832 ↗

(On Diff #537352)

Do you know which C++ expression points to 0xdddddddddddddddddd?

[springerm]

I tried running the test case with ASAN and it does not detect any invalid memory accesses. Looking at the test case, there is no transform op that deletes operations. It would be interesting to see where the memory was deallocated. When there is a memory issue, ASAN shows where the memory was allocated, freed and accessed. Can you get such debug information from MSVC?

[ashay-github]

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

832 ↗

(On Diff #537352)

It's the op pointer, in the second iteration of the loop (the first iteration op pointer points to a valid operation).

[ashay-github]

I tried running the test case with ASAN and it does not detect any invalid memory accesses. Can you get such debug information from MSVC?

I see. Thanks, let me check.

[ashay-github]

Thanks much for the tip to use the sanitizer! The problem is that since
MapVector uses a vector for storing the key-value pairs, adding an entry to the
map can trigger a vector resize, which calls the destructor after copying the
mappings. The destructor ends up deleting the operation, ultimately causing
the segfault.

You can see the calls to the destructor using the following patch:

--- a/mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h
+++ b/mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h
@@ -177,6 +177,7 @@ private:
     ParamMapping params;
     ValueMapping values;
     ValueMapping reverseValues;
+    ~Mappings() { llvm::errs() << "destroying mappings\n"; }
   };

   friend LogicalResult applyTransforms(Operation *, TransformOpInterface,
@@ -285,7 +286,9 @@ public:
     /// transform IR region and payload IR objects.
     RegionScope(TransformState &state, Region &region)
         : state(state), region(&region) {
-      auto res = state.mappings.insert(std::make_pair(&region, Mappings()));
+      llvm::errs() << "about to add mappings\n";
+      auto res = state.mappings.insert(std::make_pair(&region, std::make_shared<Mappings>()));
+      llvm::errs() << "finished adding mappings\n";
       assert(res.second && "the region scope is already present");
       (void)res;
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS

If you now run `mlir-opt --test-transform-dialect-interpreter
--verify-diagnostics` on the test case I pasted a few hours ago, you'll see
messages like here:

...
about to add mappings
destroying mappings
destroying mappings
destroying mappings
finished adding mappings
...

I wasn't able to get a move constructor working for the Mappings struct, so I
modified the map definition to use a shared pointer to the struct. Is there a
better way to resolve the problem?

[springerm]

The destructor ends up deleting the operation, ultimately causing
the segfault.

This does not make sense to me yet. Where is the operation being deleted? Mappings is just a collection of DenseMaps and they do not own any operations.

[ashay-github]

Here's the slightly truncated output from running mlir-opt with ASAN. When Mappings is destroyed, it looks like it ends up deleting the operation as well. Let me know if you'd prefer to see the full ASAN output.

=================================================================
==3936==ERROR: AddressSanitizer: heap-use-after-free on address 0x1269b44a5680 at pc 0x7ff762b7b7cd bp 0x0076331811d0 sp 0x0076331811d8
READ of size 8 at 0x1269b44a5680 thread T0
    #0 0x7ff762b7b7cc in llvm::filter_iterator_base<class mlir::Operation *const *, class `public: __cdecl mlir::transform::TransformState::getPayloadOps(class mlir::Value) const'::`2'::<lambda_1>, struct std::bidirectional_iterator_tag>::findNextValid(void) llvm-project\llvm\include\llvm\ADT\STLExtras.h:468
    #1 0x7ff762b76bfc in llvm::filter_iterator_base<class mlir::Operation *const *, class `public: __cdecl mlir::transform::TransformState::getPayloadOps(class mlir::Value) const'::`2'::<lambda_1>, struct std::bidirectional_iterator_tag>::operator++(void) llvm-project\llvm\include\llvm\ADT\STLExtras.h:488
    #2 0x7ff76c066161 in mlir::transform::ForeachOp::apply(class mlir::transform::TransformRewriter &, class mlir::transform::TransformResults &, class mlir::transform::TransformState &) llvm-project\mlir\lib\Dialect\Transform\IR\TransformOps.cpp:831
...

0x1269b44a5680 is located 512 bytes inside of 2560-byte region [0x1269b44a5480,0x1269b44a5e80)
freed by thread T0 here:
    #0 0x7ff770de4601 in operator delete(void *, unsigned __int64, enum std::align_val_t) D:\a\_work\1\s\src\vctools\asan\llvm\compiler-rt\lib\asan\asan_win_delete_scalar_size_align_thunk.cpp:42
    #1 0x7ff76254c146 in llvm::deallocate_buffer(void *, unsigned __int64, unsigned __int64) llvm-project\llvm\lib\Support\MemAlloc.cpp:25
    #2 0x7ff7657de1b2 in llvm::DenseMap<class mlir::Value, class llvm::SmallVector<class mlir::Operation *, 2>, struct llvm::DenseMapInfo<class mlir::Value, void>, struct llvm::detail::DenseMapPair<class mlir::Value, class llvm::SmallVector<class mlir::Operation *, 2>>>::~DenseMap<class mlir::Value, class llvm::SmallVector<class mlir::Operation *, 2>, struct llvm::DenseMapInfo<class mlir::Value, void>, struct llvm::detail::DenseMapPair<class mlir::Value, class llvm::SmallVector<class mlir::Operation *, 2>>>(void) llvm-project\llvm\include\llvm\ADT\DenseMap.h:782
    #3 0x7ff7657de87b in mlir::transform::TransformState::Mappings::~Mappings(void) (llvm-project\build\bin\mlir-opt.exe+0x1435de87b)
...
    #13 0x7ff7657dda81 in mlir::transform::TransformState::RegionScope::RegionScope(class mlir::transform::TransformState &, class mlir::Region &) llvm-project\mlir\include\mlir\Dialect\Transform\IR\TransformInterfaces.h:288
    #14 0x7ff7657ed5e6 in mlir::transform::TransformState::make_region_scope(class mlir::Region &) llvm-project\mlir\include\mlir\Dialect\Transform\IR\TransformInterfaces.h:839
    #15 0x7ff76c066301 in mlir::transform::ForeachOp::apply(class mlir::transform::TransformRewriter &, class mlir::transform::TransformResults &, class mlir::transform::TransformState &) llvm-project\mlir\lib\Dialect\Transform\IR\TransformOps.cpp:832
...

[springerm]

Yes, the full ASAN output would be useful. I think this is not deleting the operation itself, but the data structure of Operation * that getPayloadOps is iterating over. That data structure is part of the Mappings. The realloc will copy+remove that data structure. I think I'm starting to see the issue here....

[springerm]

I was thinking of ways to avoid the shared_ptr, but this may actually be the simplest way to fix the issue here.

[ashay-github]

I think this is not deleting the operation itself, but the data structure of Operation * that getPayloadOps is iterating over.

I see, yeah, that makes sense. Here's the full ASAN output: https://pastebin.com/raw/SZautzim.

[springerm]

Actually can you try if a unique ptr works? Or is that what you meant when you said you couldn’t make the move semantics work?

[ashay-github]

Using unique pointer worked! Earlier, I was just trying to add a move
constructor to the Mappings struct, but since I ran into more problems, I had
used a shared pointer instead.

[springerm]

Thanks for fixing this. This would be huge pain for the next person who runs into it.

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

441

You can return *it->second.get() and return a Mappings & from this function (and all the other functions), then this change will be really small.

687–690

Let's explain in the comment why a pointer indirection is necessary: The mappings are being iterated over while nested transform ops are applied and a reallocation inside the MapVector must not invalidate the iterators.

[ashay-github]

Updated to make getMapping() return a reference instead of a pointer.

Thanks for your help! The next time something like this happens, I should be
able to use ASAN to submit better bug reports.

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

public:
     /// entering this scope.
     ~RegionScope();
 
   private:
     /// 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.insert(std::make_pair(&region, Mappings()));
+      auto res = state.mappings.insert(
+          std::make_pair(&region, std::make_unique<Mappings>()));
       assert(res.second && "the region scope is already present");
       (void)res;
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
       state.regionStack.push_back(&region);
 #endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
     }
 
     /// Back-reference to the transform state.
if (!allowOutOfScope) {
         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;
+    return *it->second.get();
   }
 
   /// Returns the mappings frame for the region in which the operation resides.
   /// 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,
if (!allowOutOfScope) {
         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;
+    return *it->second.get();
   }
 
   /// 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);
 
   /// Returns a list of all ops that the given transform IR value corresponds to
LogicalResult checkAndRecordHandleInvalidationImpl(
       transform::TransformOpInterface transform,
       transform::TransformState::InvalidatedHandleMap &newlyInvalidated) const;
 
   /// Remove all nullptrs from op handles that were added by `replacePayloadOp`.
   void compactOpHandles();
 
   /// 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;
+  /// We use a pointer to the Mappings struct so that reallocations inside
+  /// MapVector don't invalidate iterators when we apply nested transform ops
+  /// while also iterating over the mappings.
+  llvm::MapVector<Region *, std::unique_ptr<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

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.insert(std::make_pair(region, Mappings()));
+  auto result =
+      mappings.insert(std::make_pair(region, std::make_unique<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 {
   return iter->getSecond();
 }
 
 LogicalResult transform::TransformState::getHandlesForPayloadOp(
     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()) {
+    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 (!includeOutOfScope &&
         region->getParentOp()->hasTrait<OpTrait::IsIsolatedFromAbove>())
       break;
   }
 
   return success(found);
 }
 
 LogicalResult transform::TransformState::getHandlesForPayloadValue(
     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()) {
+    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 (!includeOutOfScope &&
         region->getParentOp()->hasTrait<OpTrait::IsIsolatedFromAbove>())
       break;
   }
void transform::TransformState::recordOpHandleInvalidation(
   // for handles associated with each operation and value.
   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 (const auto &[payloadOp, otherHandles] : mapping->reverse) {
       for (Value otherHandle : otherHandles)
         recordOpHandleInvalidationOne(handle, potentialAncestors, payloadOp,
                                       otherHandle, throughValue,
                                       newlyInvalidated);
     }
     // Go over all value handle mappings and mark as invalidated any handle
     // pointing to any result of the payload op associated with the given handle
     // or any op nested in them. Similarly invalidate handles to argument of
     // blocks belonging to any region of any payload op associated with the
     // given handle or any op nested in them.
-    for (const auto &[payloadValue, valueHandles] : mapping.reverseValues) {
+    for (const auto &[payloadValue, valueHandles] : mapping->reverseValues) {
       for (Value valueHandle : valueHandles)
         recordValueHandleInvalidationByOpHandleOne(handle, potentialAncestors,
                                                    payloadValue, valueHandle,
                                                    newlyInvalidated);
     }
   }
 }
 
for (OpOperand &operand : transform->getOpOperands()) {
         FULL_LDBG("--not a TransformHandle -> SKIP AND DROP ON THE FLOOR\n");
       }
     }
 
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
     // Cache Operation* -> OperationName mappings. These will be checked after
     // the transform has been applied to detect incorrect memory side effects
     // and missing op tracking.
-    for (Mappings &mapping : llvm::make_second_range(mappings)) {
-      for (Operation *op : llvm::make_first_range(mapping.reverse)) {
+    for (std::unique_ptr<Mappings> &mapping :
+         llvm::make_second_range(mappings)) {
+      for (Operation *op : llvm::make_first_range(mapping->reverse)) {
         auto insertion = cachedNames.insert({op, op->getName()});
         if (!insertion.second) {
           if (insertion.first->second != op->getName()) {
             // Operation is already in the cache, but with a different name.
             DiagnosedDefiniteFailure diag =
                 emitDefiniteFailure(transform->getLoc())
                 << "expensive checks failure: operation mismatch, expected "
                 << insertion.first->second;
for (Operation *op : consumedPayloadOps) {
       SmallVector<Value> handles;
       (void)getHandlesForPayloadOp(op, handles);
       for (Value handle : handles)
         forgetMapping(handle, /*origOpFlatResults=*/ValueRange());
       cachedNames.erase(op);
     }
 
     // Check cached operation names.
-    for (Mappings &mapping : llvm::make_second_range(mappings)) {
-      for (Operation *op : llvm::make_first_range(mapping.reverse)) {
+    for (std::unique_ptr<Mappings> &mapping :
+         llvm::make_second_range(mappings)) {
+      for (Operation *op : llvm::make_first_range(mapping->reverse)) {
         // Make sure that the name of the op has not changed. If it has changed,
         // the op was removed and a new op was allocated at the same memory
         // location. This means that we are missing op tracking somewhere.
         auto cacheIt = cachedNames.find(op);
         if (cacheIt == cachedNames.end()) {
           DiagnosedDefiniteFailure diag =
               emitDefiniteFailure(transform->getLoc())
               << "expensive checks failure: operation not found in cache";
for (Operation &op : block) {
       }
     }
   }
 
 #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));
+      llvm::to_vector(llvm::make_first_range(state.mappings[region]->reverse));
 #endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
 
   state.mappings.erase(region);
 
 #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.