[Mlir-commits] [mlir] 0285f1c - [mlir] Don't assume non-erased DenseMap entries remain valid after erase. NFC (#199365)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon May 25 00:42:38 PDT 2026
Author: Fangrui Song
Date: 2026-05-25T00:42:33-07:00
New Revision: 0285f1c6b791fb974a71a236e329848cf2d2fbe9
URL: https://github.com/llvm/llvm-project/commit/0285f1c6b791fb974a71a236e329848cf2d2fbe9
DIFF: https://github.com/llvm/llvm-project/commit/0285f1c6b791fb974a71a236e329848cf2d2fbe9.diff
LOG: [mlir] Don't assume non-erased DenseMap entries remain valid after erase. NFC (#199365)
Like the preceding llvm/ change, fix MLIR sites that reuse an iterator
or bucket reference after erasing from the same map, in preparation for
backward-shift DenseMap deletion which relocates surviving entries.
Use DenseMap::remove_if in ThreadLocalCache::clearExpiredEntries and the
RootOrdering cycle contraction (deferring the in-cycle graph erases
until
after iteration). ThreadLocalCache::get reads the value into a local and
bufferizeOp snapshots the worklist before folding, since those erases
re-enter via a rewriter listener.
Aided by Claude Opus 4.7
Added:
Modified:
mlir/include/mlir/Support/ThreadLocalCache.h
mlir/lib/Conversion/PDLToPDLInterp/RootOrdering.cpp
mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Support/ThreadLocalCache.h b/mlir/include/mlir/Support/ThreadLocalCache.h
index 53b6d31a09555..a1f88ec609e86 100644
--- a/mlir/include/mlir/Support/ThreadLocalCache.h
+++ b/mlir/include/mlir/Support/ThreadLocalCache.h
@@ -128,11 +128,8 @@ class ThreadLocalCache {
/// Clear out any unused entries within the map. This method is not
/// thread-safe, and should only be called by the same thread as the cache.
void clearExpiredEntries() {
- for (auto it = this->begin(), e = this->end(); it != e;) {
- auto curIt = it++;
- if (!curIt->second.ptr->second)
- this->erase(curIt);
- }
+ this->remove_if(
+ [](const auto &entry) { return !entry.second.ptr->second; });
}
};
@@ -162,8 +159,14 @@ class ThreadLocalCache {
// Before returning the new instance, take the chance to clear out any used
// entries in the static map. The cache is only cleared within the same
// thread to remove the need to lock the cache itself.
+ //
+ // `threadInstance` aliases a bucket in `staticCache`; clearExpiredEntries
+ // may erase from the map and invalidate that reference. The value itself
+ // lives in heap-stable storage reached through the shared `Observer::ptr`
+ // (not in the bucket), so load it out before clearing.
+ ValueT &value = *threadInstance.ptr->first;
staticCache.clearExpiredEntries();
- return *threadInstance.ptr->first;
+ return value;
}
ValueT &operator*() { return get(); }
ValueT *operator->() { return &get(); }
diff --git a/mlir/lib/Conversion/PDLToPDLInterp/RootOrdering.cpp b/mlir/lib/Conversion/PDLToPDLInterp/RootOrdering.cpp
index 2d9c661f7df2c..b36d5a774275d 100644
--- a/mlir/lib/Conversion/PDLToPDLInterp/RootOrdering.cpp
+++ b/mlir/lib/Conversion/PDLToPDLInterp/RootOrdering.cpp
@@ -52,12 +52,11 @@ static void contract(RootOrderingGraph &graph, ArrayRef<Value> cycle,
// Now, contract the cycle, marking the actual sources and targets.
DenseMap<Value, RootOrderingEntry> repEntries;
- for (auto outer = graph.begin(), e = graph.end(); outer != e; ++outer) {
- Value target = outer->first;
+ for (auto &[target, edges] : graph) {
if (cycleSet.contains(target)) {
// Target in the cycle => edges incoming to the cycle or within the cycle.
unsigned parentDepth = parentDepths.lookup(target);
- for (const auto &inner : outer->second) {
+ for (const auto &inner : edges) {
Value source = inner.first;
// Ignore edges within the cycle.
if (cycleSet.contains(source))
@@ -81,36 +80,37 @@ static void contract(RootOrderingGraph &graph, ArrayRef<Value> cycle,
repEntries[source].cost = cost;
}
}
- // Erase the node in the cycle.
- graph.erase(outer);
+ // Defer erasing graph[target] until after the loop; backward-shift
+ // erase would otherwise invalidate the surrounding iterator.
} else {
// Target not in cycle => edges going away from or unrelated to the cycle.
- DenseMap<Value, RootOrderingEntry> &entries = outer->second;
Value bestSource;
std::pair<unsigned, unsigned> bestCost;
- auto inner = entries.begin(), innerE = entries.end();
- while (inner != innerE) {
- Value source = inner->first;
- if (cycleSet.contains(source)) {
- // Going-away edge => get its cost and erase it.
- if (!bestSource || bestCost > inner->second.cost) {
- bestSource = source;
- bestCost = inner->second.cost;
- }
- entries.erase(inner++);
- } else {
- ++inner;
+ edges.remove_if([&](const auto &inner) {
+ Value source = inner.first;
+ if (!cycleSet.contains(source))
+ return false;
+ // Going-away edge => get its cost and erase it.
+ if (!bestSource || bestCost > inner.second.cost) {
+ bestSource = source;
+ bestCost = inner.second.cost;
}
- }
+ return true;
+ });
// There were going-away edges, contract them.
if (bestSource) {
- entries[rep].cost = bestCost;
+ edges[rep].cost = bestCost;
actualSource[target] = bestSource;
}
}
}
+ // Erase all in-cycle nodes from the graph. Done after the iteration above
+ // because backward-shift erase relocates surviving entries.
+ for (Value node : cycle)
+ graph.erase(node);
+
// Store the edges to the representative.
graph[rep] = std::move(repEntries);
}
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
index 701ab52a491a8..84d5bb2d713c1 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
@@ -345,8 +345,14 @@ LogicalResult bufferization::bufferizeOp(Operation *op,
if (erasedOps.contains(op))
return success();
- // Fold all to_buffer(to_tensor(x)) pairs.
- for (Operation *op : toBufferOps) {
+ // Fold all to_buffer(to_tensor(x)) pairs. Snapshot the set first:
+ // `foldToBufferToTensorPair` can erase ops, and the rewriter listener
+ // mutates `toBufferOps` from inside that call, which would invalidate
+ // any DenseSet iterator held across it.
+ SmallVector<Operation *> toBufferOpsSnapshot = llvm::to_vector(toBufferOps);
+ for (Operation *op : toBufferOpsSnapshot) {
+ if (erasedOps.contains(op))
+ continue;
rewriter.setInsertionPoint(op);
(void)bufferization::foldToBufferToTensorPair(
rewriter, cast<ToBufferOp>(op), options);
More information about the Mlir-commits
mailing list