[llvm] [IR] Reduce memory usage in metadata use tracking (NFC) (PR #221866)
Kazu Hirata via llvm-commits
llvm-commits at lists.llvm.org
Tue Sep 29 00:42:20 PDT 2026
https://github.com/kazutakahirata updated https://github.com/llvm/llvm-project/pull/221866
>From 23abc80d2d28fca7ed1f31cfecceae7d61fc1e8c Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu at google.com>
Date: Mon, 28 Sep 2026 22:28:09 -0700
Subject: [PATCH] [IR] Reduce memory usage in metadata use tracking (NFC)
This patch reduces the memory footprint of metadata use tracking.
Currently, ReplaceableUses uses:
SmallDenseMap<void *, std::pair<OwnerTy, uint64_t>, 4>
where the uint64_t records insertion order. The problem is that
once it exceeds 2 elements, it allocates 64 buckets (~1.5 KB). In large
compilations with debug info, thousands of metadata nodes have 3 or 4
uses, leading to significant overhead.
With this patch, uses are stored directly in a SmallVector:
struct UseEntry {
void *Ref;
OwnerTy Owner;
};
SmallVector<UseEntry, 4> UseMap;
// In addRef:
UseMap.push_back({Ref, Owner});
SmallVector stores up to 4 elements inline (~87% of metadata nodes)
without heap allocation, up from 2 elements (~64% of metadata nodes)
with SmallDenseMap. It also provides a deterministic array order,
eliminating the need to sort uses by an explicit ID in several places.
For the rare nodes with more than 32 uses, an on-demand
DenseMap<void *, unsigned> (IndexMap) is lazily allocated to map each
Ref to its index in UseMap.
To support this representation efficiently:
- For nodes with up to 32 uses, IndexMap remains null and findRef scans
UseMap backwards from the end, taking advantage of temporal locality
since recently added references are typically dropped or moved first.
- Once UseMap exceeds 32 entries, IndexMap is allocated to provide O(1)
lookups in findRef, and is reset once UseMap shrinks back to 16 or
fewer entries so it cannot thrash.
- In dropRef, the removed entry is overwritten with UseMap.back() and
the back is popped, keeping removals O(1) from any position while
maintaining a deterministic order.
Benchmark results:
Base Test
File Max RSS Max RSS Delta Wall Time
-----------------------------------------------------------------
SLPVectorizer.ii 1343 MB 1207 MB -10.14% 43.48s -> 42.80s
SemaOpenMP.ii 1358 MB 1196 MB -11.95% 34.42s -> 33.71s
X86ISelLowering.ii 881 MB 838 MB -4.85% 26.23s -> 26.03s
-----------------------------------------------------------------
Total 3582 MB 3241 MB -9.52% 104.13s -> 102.53s
Geometric Mean -9.03% -1.47%
- Compilation flags: -O3 -g -fdebug-info-for-profiling -flto=thin
- Max RSS & Wall Time: average of 3 runs after 1 warmup run.
Thanks to Fangrui Song for exempting ConstantData from metadata use
tracking (#223303), suggesting the swap-and-pop removal in dropRef, and
providing the LargeUseList unit test (#227186).
Assisted-by: Antigravity
---
llvm/include/llvm/IR/Metadata.h | 15 ++-
llvm/lib/IR/Metadata.cpp | 151 +++++++++++++++--------------
llvm/unittests/IR/MetadataTest.cpp | 23 +++++
3 files changed, 114 insertions(+), 75 deletions(-)
diff --git a/llvm/include/llvm/IR/Metadata.h b/llvm/include/llvm/IR/Metadata.h
index 9347a3dd93e13..ec486863d5198 100644
--- a/llvm/include/llvm/IR/Metadata.h
+++ b/llvm/include/llvm/IR/Metadata.h
@@ -397,8 +397,18 @@ class ReplaceableUses {
using OwnerTy = MetadataTracking::OwnerTy;
private:
- uint64_t NextIndex = 0;
- SmallDenseMap<void *, std::pair<OwnerTy, uint64_t>, 4> UseMap;
+ struct UseEntry {
+ void *Ref = nullptr;
+ OwnerTy Owner = nullptr;
+ };
+
+ static constexpr unsigned IndexThreshold = 32;
+ // Tracked uses of this. Dropping one moves the last entry into its slot, so
+ // the order is deterministic but not the order they were added in.
+ SmallVector<UseEntry, 4> UseMap;
+ // Lazily allocated map from Ref to its index in UseMap for large use lists.
+ using IndexMapTy = DenseMap<void *, unsigned>;
+ std::unique_ptr<IndexMapTy> IndexMap;
protected:
~ReplaceableUses() {
@@ -429,6 +439,7 @@ class ReplaceableUses {
unsigned getNumUses() const { return UseMap.size(); }
private:
+ UseEntry *findRef(void *Ref);
void addRef(void *Ref, OwnerTy Owner);
void dropRef(void *Ref);
void moveRef(void *Ref, void *New, const Metadata &MD);
diff --git a/llvm/lib/IR/Metadata.cpp b/llvm/lib/IR/Metadata.cpp
index a947bcfd6fa96..bb07775a6f529 100644
--- a/llvm/lib/IR/Metadata.cpp
+++ b/llvm/lib/IR/Metadata.cpp
@@ -258,82 +258,96 @@ bool MetadataTracking::isReplaceable(const Metadata &MD) {
}
SmallVector<Metadata *> ReplaceableUses::getAllArgListUsers() {
- SmallVector<std::pair<OwnerTy, uint64_t> *> MDUsersWithID;
- for (auto Pair : UseMap) {
- OwnerTy Owner = Pair.second.first;
+ SmallVector<Metadata *> MDUsers;
+ for (const auto &[Ref, Owner] : UseMap) {
if (Owner.isNull())
continue;
if (!isa<Metadata *>(Owner))
continue;
Metadata *OwnerMD = cast<Metadata *>(Owner);
if (OwnerMD->getMetadataID() == Metadata::DIArgListKind)
- MDUsersWithID.push_back(&UseMap[Pair.first]);
+ MDUsers.push_back(OwnerMD);
}
- llvm::sort(MDUsersWithID, [](auto UserA, auto UserB) {
- return UserA->second < UserB->second;
- });
- SmallVector<Metadata *> MDUsers;
- for (auto *UserWithID : MDUsersWithID)
- MDUsers.push_back(cast<Metadata *>(UserWithID->first));
return MDUsers;
}
SmallVector<DbgVariableRecord *>
ReplaceableUses::getAllDbgVariableRecordUsers() {
- SmallVector<std::pair<OwnerTy, uint64_t> *> DVRUsersWithID;
- for (auto Pair : UseMap) {
- OwnerTy Owner = Pair.second.first;
- if (Owner.isNull())
- continue;
- if (!isa<DebugValueUser *>(Owner))
- continue;
- DVRUsersWithID.push_back(&UseMap[Pair.first]);
- }
// Order DbgVariableRecord users in reverse-creation order. Normal dbg.value
// users of MetadataAsValues are ordered by their UseList, i.e. reverse order
// of when they were added: we need to replicate that here. The structure of
// debug-info output depends on the ordering of intrinsics, thus we need
// to keep them consistent for comparisons sake.
- llvm::sort(DVRUsersWithID, [](auto UserA, auto UserB) {
- return UserA->second > UserB->second;
- });
SmallVector<DbgVariableRecord *> DVRUsers;
- for (auto UserWithID : DVRUsersWithID)
- DVRUsers.push_back(cast<DebugValueUser *>(UserWithID->first)->getUser());
+ for (const auto &[Ref, Owner] : reverse(UseMap)) {
+ if (Owner.isNull())
+ continue;
+ if (!isa<DebugValueUser *>(Owner))
+ continue;
+ DVRUsers.push_back(cast<DebugValueUser *>(Owner)->getUser());
+ }
return DVRUsers;
}
-void ReplaceableUses::addRef(void *Ref, OwnerTy Owner) {
- bool WasInserted =
- UseMap.insert(std::make_pair(Ref, std::make_pair(Owner, NextIndex)))
- .second;
- (void)WasInserted;
- assert(WasInserted && "Expected to add a reference");
+ReplaceableUses::UseEntry *ReplaceableUses::findRef(void *Ref) {
+ if (!IndexMap) {
+ // Search backward for temporal locality; recently added references are
+ // often dropped or moved first.
+ for (UseEntry &U : reverse(UseMap))
+ if (U.Ref == Ref)
+ return &U;
+ return nullptr;
+ }
- ++NextIndex;
- assert(NextIndex != 0 && "Unexpected overflow");
+ auto It = IndexMap->find(Ref);
+ return It == IndexMap->end() ? nullptr : &UseMap[It->second];
+}
+
+void ReplaceableUses::addRef(void *Ref, OwnerTy Owner) {
+ assert(Ref && "Expected live reference");
+ assert(!findRef(Ref) && "Reference already tracked");
+ unsigned NewIdx = UseMap.size();
+ UseMap.push_back({Ref, Owner});
+ if (IndexMap) {
+ (*IndexMap)[Ref] = NewIdx;
+ } else if (UseMap.size() > IndexThreshold) {
+ // Build the index map once UseMap grows past the threshold.
+ IndexMap = std::make_unique<IndexMapTy>();
+ for (unsigned I = 0, E = UseMap.size(); I != E; ++I)
+ (*IndexMap)[UseMap[I].Ref] = I;
+ }
}
void ReplaceableUses::dropRef(void *Ref) {
- bool WasErased = UseMap.erase(Ref);
- (void)WasErased;
- assert(WasErased && "Expected to drop a reference");
+ UseEntry *Entry = findRef(Ref);
+ assert(Entry && "Expected to find Ref");
+ if (IndexMap)
+ IndexMap->erase(Ref);
+ if (Entry != &UseMap.back()) {
+ *Entry = UseMap.back();
+ if (IndexMap)
+ (*IndexMap)[Entry->Ref] = Entry - UseMap.begin();
+ }
+ UseMap.pop_back();
+ if (IndexMap && UseMap.size() <= IndexThreshold / 2)
+ IndexMap.reset();
}
void ReplaceableUses::moveRef(void *Ref, void *New, const Metadata &MD) {
- auto I = UseMap.find(Ref);
- assert(I != UseMap.end() && "Expected to move a reference");
- auto OwnerAndIndex = I->second;
- UseMap.erase(I);
- bool WasInserted = UseMap.insert(std::make_pair(New, OwnerAndIndex)).second;
- (void)WasInserted;
- assert(WasInserted && "Expected to add a reference");
+ assert(!findRef(New) && "Cannot move to an existing reference");
+ UseEntry *Entry = findRef(Ref);
+ assert(Entry && "Expected to move a reference");
+ if (IndexMap) {
+ IndexMap->erase(Ref);
+ (*IndexMap)[New] = Entry - UseMap.begin();
+ }
+ Entry->Ref = New;
// Check that the references are direct if there's no owner.
(void)MD;
- assert((OwnerAndIndex.first || *static_cast<Metadata **>(Ref) == &MD) &&
+ assert((Entry->Owner || *static_cast<Metadata **>(Ref) == &MD) &&
"Reference without owner must be direct");
- assert((OwnerAndIndex.first || *static_cast<Metadata **>(New) == &MD) &&
+ assert((Entry->Owner || *static_cast<Metadata **>(New) == &MD) &&
"Reference without owner must be direct");
}
@@ -346,14 +360,11 @@ void ReplaceableUses::SalvageDebugInfo(const Constant &C) {
auto &Store = Context.pImpl->ValuesAsMetadata;
auto I = Store.find(&C);
ValueAsMetadata *MD = I->second;
- using UseTy =
- std::pair<void *, std::pair<MetadataTracking::OwnerTy, uint64_t>>;
// Copy out uses and update value of Constant used by debug info metadata with
- // poison below
- SmallVector<UseTy, 8> Uses(MD->UseMap.begin(), MD->UseMap.end());
+ // poison below.
+ SmallVector<UseEntry, 4> Uses = MD->UseMap;
- for (const auto &Pair : Uses) {
- MetadataTracking::OwnerTy Owner = Pair.second.first;
+ for (const auto &[Ref, Owner] : Uses) {
if (!Owner)
continue;
// Check for MetadataAsValue.
@@ -369,7 +380,7 @@ void ReplaceableUses::SalvageDebugInfo(const Constant &C) {
continue;
if (isa<DINode>(OwnerMD)) {
OwnerMD->handleChangedOperand(
- Pair.first, ValueAsMetadata::get(PoisonValue::get(C.getType())));
+ Ref, ValueAsMetadata::get(PoisonValue::get(C.getType())));
}
}
}
@@ -379,25 +390,21 @@ void ReplaceableUses::replaceAllUsesWith(Metadata *MD) {
return;
// Copy out uses since UseMap will get touched below.
- using UseTy = std::pair<void *, std::pair<OwnerTy, uint64_t>>;
- SmallVector<UseTy, 8> Uses(UseMap.begin(), UseMap.end());
- llvm::sort(Uses, [](const UseTy &L, const UseTy &R) {
- return L.second.second < R.second.second;
- });
- for (const auto &Pair : Uses) {
+ SmallVector<UseEntry, 4> Uses = UseMap;
+
+ for (const auto &[Ref, Owner] : Uses) {
// Check that this Ref hasn't disappeared after RAUW (when updating a
// previous Ref).
- if (!UseMap.count(Pair.first))
+ if (!findRef(Ref))
continue;
- OwnerTy Owner = Pair.second.first;
if (!Owner) {
// Update unowned tracking references directly.
- Metadata *&Ref = *static_cast<Metadata **>(Pair.first);
- Ref = MD;
+ Metadata *&DirectRef = *static_cast<Metadata **>(Ref);
+ dropRef(Ref);
+ DirectRef = MD;
if (MD)
- MetadataTracking::track(Ref);
- UseMap.erase(Pair.first);
+ MetadataTracking::track(DirectRef);
continue;
}
@@ -408,7 +415,7 @@ void ReplaceableUses::replaceAllUsesWith(Metadata *MD) {
}
if (auto *DVU = dyn_cast<DebugValueUser *>(Owner)) {
- DVU->handleChangedValue(Pair.first, MD);
+ DVU->handleChangedValue(Ref, MD);
continue;
}
@@ -417,7 +424,7 @@ void ReplaceableUses::replaceAllUsesWith(Metadata *MD) {
switch (OwnerMD->getMetadataID()) {
#define HANDLE_METADATA_LEAF(CLASS) \
case Metadata::CLASS##Kind: \
- cast<CLASS>(OwnerMD)->handleChangedOperand(Pair.first, MD); \
+ cast<CLASS>(OwnerMD)->handleChangedOperand(Ref, MD); \
continue;
#include "llvm/IR/Metadata.def"
default:
@@ -433,18 +440,16 @@ void ReplaceableUses::resolveAllUses(bool ResolveUsers) {
if (!ResolveUsers) {
UseMap.clear();
+ IndexMap.reset();
return;
}
- // Copy out uses since UseMap could get touched below.
- using UseTy = std::pair<void *, std::pair<OwnerTy, uint64_t>>;
- SmallVector<UseTy, 8> Uses(UseMap.begin(), UseMap.end());
- llvm::sort(Uses, [](const UseTy &L, const UseTy &R) {
- return L.second.second < R.second.second;
- });
+ // Move uses out since UseMap could get touched below.
+ SmallVector<UseEntry, 4> Uses = std::move(UseMap);
UseMap.clear();
- for (const auto &Pair : Uses) {
- auto Owner = Pair.second.first;
+ IndexMap.reset();
+ for (const auto &U : Uses) {
+ auto Owner = U.Owner;
if (!Owner)
continue;
if (!isa<Metadata *>(Owner))
diff --git a/llvm/unittests/IR/MetadataTest.cpp b/llvm/unittests/IR/MetadataTest.cpp
index e9af072ea800e..2867fc88c6cbc 100644
--- a/llvm/unittests/IR/MetadataTest.cpp
+++ b/llvm/unittests/IR/MetadataTest.cpp
@@ -5355,6 +5355,29 @@ TEST_F(TrackingMDRefTest, UpdatesOnDeletion) {
EXPECT_TRUE(!MD);
}
+// Once the use list is indexed, look each changed reference up again right
+// away: a later swap in dropRef would repair a stale index entry.
+TEST_F(TrackingMDRefTest, LargeUseList) {
+ auto Temp = MDTuple::getTemporary(Context, {});
+ TrackingMDRef Refs[41];
+ for (unsigned I = 0; I != 40; ++I)
+ Refs[I].reset(Temp.get());
+ Refs[0].reset();
+ Refs[0].reset(Temp.get());
+ Refs[0].reset();
+ Refs[40] = std::move(Refs[1]);
+ Refs[1].reset(Temp.get());
+ Refs[40].reset();
+ Refs[39].reset();
+ EXPECT_EQ(38u, Temp->getNumTemporaryUses());
+
+ auto *N = MDTuple::getDistinct(Context, {});
+ Temp->replaceAllUsesWith(N);
+ EXPECT_EQ(0u, Temp->getNumTemporaryUses());
+ for (unsigned I = 1; I != 39; ++I)
+ EXPECT_EQ(N, Refs[I].get()) << I;
+}
+
TEST(NamedMDNodeTest, Search) {
LLVMContext Context;
ConstantAsMetadata *C =
More information about the llvm-commits
mailing list