[llvm] [IR] Reduce memory usage in metadata use tracking (NFC) (PR #221866)
Kazu Hirata via llvm-commits
llvm-commits at lists.llvm.org
Mon Sep 21 00:34:22 PDT 2026
https://github.com/kazutakahirata updated https://github.com/llvm/llvm-project/pull/221866
>From b6e82227848034e441e2f5367a7258e4c1bb41df Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu at google.com>
Date: Wed, 16 Sep 2026 15:23:13 -0700
Subject: [PATCH 1/2] [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 preserves insertion order, eliminating the
need to sort uses in several places. For the rare nodes with more than
16 uses (less than 2% of metadata nodes), 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 16 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 16 entries, IndexMap is allocated to provide O(1)
lookups in findRef (used by hasRef, dropRef, and moveRef), and is
reset once UseMap shrinks back to 8 or fewer entries.
- In replaceAllUsesWith, UseMap is reversed before iterating over the
copied use list. Because replaceAllUsesWith visits uses in original
forward order, each callback's dropRef finds its entry at the very back
of UseMap, turning removals into O(1) pops with zero element shifting.
Benchmark results:
Base Test
File Max RSS Max RSS Delta Wall Time
-----------------------------------------------------------------
SLPVectorizer.ii 1463 MB 1327 MB -9.29% 67.31s -> 66.40s
SemaOpenMP.ii 1461 MB 1298 MB -11.17% 54.37s -> 52.28s
X86ISelLowering.ii 974 MB 931 MB -4.40% 40.26s -> 39.86s
-----------------------------------------------------------------
Total 3898 MB 3556 MB -8.77% 161.94s -> 158.54s
Geometric Mean -8.33% -2.07%
- Compilation flags: -O3 -g -fdebug-info-for-profiling -flto=thin
- Max RSS: average of 3 runs after 1 warmup run.
- Wall Time: average of 10 runs after 1 warmup run.
Assisted-by: Antigravity
---
llvm/include/llvm/IR/Metadata.h | 15 ++-
llvm/lib/IR/Metadata.cpp | 168 ++++++++++++++++++--------------
2 files changed, 108 insertions(+), 75 deletions(-)
diff --git a/llvm/include/llvm/IR/Metadata.h b/llvm/include/llvm/IR/Metadata.h
index 017fd98e704ea..dad9d170eda41 100644
--- a/llvm/include/llvm/IR/Metadata.h
+++ b/llvm/include/llvm/IR/Metadata.h
@@ -397,8 +397,17 @@ 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 = 16;
+ // Tracked uses stored in insertion order.
+ 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 +438,8 @@ class ReplaceableUses {
unsigned getNumUses() const { return UseMap.size(); }
private:
+ UseEntry *findRef(void *Ref, bool EraseFromIndex = false);
+ bool hasRef(void *Ref) const;
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 28873462748b2..0ad4df942aaf2 100644
--- a/llvm/lib/IR/Metadata.cpp
+++ b/llvm/lib/IR/Metadata.cpp
@@ -258,82 +258,104 @@ 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,
+ bool EraseFromIndex) {
+ if (!IndexMap) {
+ // Search backward for temporal locality; recently added references are
+ // often dropped or moved first.
+ auto It = find_if(reverse(UseMap),
+ [Ref](const UseEntry &U) { return U.Ref == Ref; });
+ return It != UseMap.rend() ? &*It : nullptr;
+ }
- ++NextIndex;
- assert(NextIndex != 0 && "Unexpected overflow");
+ auto It = IndexMap->find(Ref);
+ if (It == IndexMap->end())
+ return nullptr;
+ UseEntry *Entry = &UseMap[It->second];
+ if (EraseFromIndex)
+ IndexMap->erase(It);
+ return Entry;
+}
+
+bool ReplaceableUses::hasRef(void *Ref) const {
+ return const_cast<ReplaceableUses *>(this)->findRef(Ref) != nullptr;
+}
+
+void ReplaceableUses::addRef(void *Ref, OwnerTy Owner) {
+ assert(Ref && "Expected live reference");
+ assert(!hasRef(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, /*EraseFromIndex=*/true);
+ assert(Entry && "Expected to find Ref");
+ unsigned Idx = Entry - UseMap.begin();
+ UseMap.erase(Entry);
+ if (IndexMap) {
+ if (UseMap.size() <= IndexThreshold / 2) {
+ // Tear down the index map with hysteresis to avoid thrashing.
+ IndexMap.reset();
+ } else {
+ for (unsigned I = Idx, E = UseMap.size(); I != E; ++I)
+ (*IndexMap)[UseMap[I].Ref] = I;
+ }
+ }
}
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(!hasRef(New) && "Cannot move to an existing reference");
+ UseEntry *Entry = findRef(Ref, /*EraseFromIndex=*/true);
+ assert(Entry && "Expected to move a reference");
+ Entry->Ref = New;
+ if (IndexMap)
+ (*IndexMap)[New] = Entry - UseMap.begin();
// 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 +368,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 +388,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 +398,30 @@ 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;
+
+ // Reverse UseMap so that as we visit Uses in original forward order, each
+ // callback's dropRef finds its entry at the very back of UseMap, turning
+ // linear searches into O(1) pops.
+ std::reverse(UseMap.begin(), UseMap.end());
+ if (IndexMap) {
+ for (unsigned I = 0, E = UseMap.size(); I != E; ++I)
+ (*IndexMap)[UseMap[I].Ref] = I;
+ }
+
+ 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 (!hasRef(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 +432,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 +441,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 +457,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))
>From b778d30a367323b0ca75ff4f50168d430ca0b141 Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu at google.com>
Date: Mon, 21 Sep 2026 00:24:50 -0700
Subject: [PATCH 2/2] Address comments.
---
llvm/include/llvm/IR/Metadata.h | 1 -
llvm/lib/IR/Metadata.cpp | 19 +++++--------------
2 files changed, 5 insertions(+), 15 deletions(-)
diff --git a/llvm/include/llvm/IR/Metadata.h b/llvm/include/llvm/IR/Metadata.h
index dad9d170eda41..0cef2e1c6cab4 100644
--- a/llvm/include/llvm/IR/Metadata.h
+++ b/llvm/include/llvm/IR/Metadata.h
@@ -439,7 +439,6 @@ class ReplaceableUses {
private:
UseEntry *findRef(void *Ref, bool EraseFromIndex = false);
- bool hasRef(void *Ref) const;
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 0ad4df942aaf2..9c97160e98af1 100644
--- a/llvm/lib/IR/Metadata.cpp
+++ b/llvm/lib/IR/Metadata.cpp
@@ -308,13 +308,9 @@ ReplaceableUses::UseEntry *ReplaceableUses::findRef(void *Ref,
return Entry;
}
-bool ReplaceableUses::hasRef(void *Ref) const {
- return const_cast<ReplaceableUses *>(this)->findRef(Ref) != nullptr;
-}
-
void ReplaceableUses::addRef(void *Ref, OwnerTy Owner) {
assert(Ref && "Expected live reference");
- assert(!hasRef(Ref) && "Reference already tracked");
+ assert(!findRef(Ref) && "Reference already tracked");
unsigned NewIdx = UseMap.size();
UseMap.push_back({Ref, Owner});
if (IndexMap) {
@@ -333,18 +329,13 @@ void ReplaceableUses::dropRef(void *Ref) {
unsigned Idx = Entry - UseMap.begin();
UseMap.erase(Entry);
if (IndexMap) {
- if (UseMap.size() <= IndexThreshold / 2) {
- // Tear down the index map with hysteresis to avoid thrashing.
- IndexMap.reset();
- } else {
- for (unsigned I = Idx, E = UseMap.size(); I != E; ++I)
- (*IndexMap)[UseMap[I].Ref] = I;
- }
+ for (unsigned I = Idx, E = UseMap.size(); I != E; ++I)
+ (*IndexMap)[UseMap[I].Ref] = I;
}
}
void ReplaceableUses::moveRef(void *Ref, void *New, const Metadata &MD) {
- assert(!hasRef(New) && "Cannot move to an existing reference");
+ assert(!findRef(New) && "Cannot move to an existing reference");
UseEntry *Entry = findRef(Ref, /*EraseFromIndex=*/true);
assert(Entry && "Expected to move a reference");
Entry->Ref = New;
@@ -412,7 +403,7 @@ void ReplaceableUses::replaceAllUsesWith(Metadata *MD) {
for (const auto &[Ref, Owner] : Uses) {
// Check that this Ref hasn't disappeared after RAUW (when updating a
// previous Ref).
- if (!hasRef(Ref))
+ if (!findRef(Ref))
continue;
if (!Owner) {
More information about the llvm-commits
mailing list