[MemProf] Add ambigous memprof attribute #157204

teresajohnson · 2025-09-05T23:51:59Z

To help track allocations that we matched with memprof profiles but
for which we weren't able to disambiguate the different hotness
contexts, apply an "ambiguous" memprof attribute to all allocations with
matched profiles. These will be replaced if we can identify a single
hotness type, possibly after cloning.

Eventually we plan to translate this to a special hotness hint on the
allocation call.

To help track allocations that we matched with memprof profiles but for which we weren't able to disambiguate the different hotness contexts, apply an "ambiguous" memprof attribute to all allocations with matched profiles. These will be replaced if we can identify a single hotness type, possibly after cloning. Eventually we plan to translate this to a special hotness hint on the allocation call.

llvmbot · 2025-09-05T23:52:30Z

@llvm/pr-subscribers-llvm-analysis

@llvm/pr-subscribers-llvm-transforms

Author: Teresa Johnson (teresajohnson)

Changes

To help track allocations that we matched with memprof profiles but
for which we weren't able to disambiguate the different hotness
contexts, apply an "ambiguous" memprof attribute to all allocations with
matched profiles. These will be replaced if we can identify a single
hotness type, possibly after cloning.

Eventually we plan to translate this to a special hotness hint on the
allocation call.

Full diff: https://github.com/llvm/llvm-project/pull/157204.diff

4 Files Affected:

(modified) llvm/include/llvm/Analysis/MemoryProfileInfo.h (+8)
(modified) llvm/lib/Analysis/MemoryProfileInfo.cpp (+22)
(modified) llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp (+2)
(modified) llvm/unittests/Analysis/MemoryProfileInfoTest.cpp (+14-7)

diff --git a/llvm/include/llvm/Analysis/MemoryProfileInfo.h b/llvm/include/llvm/Analysis/MemoryProfileInfo.h
index 571caf95f275d..be690a49767b4 100644
--- a/llvm/include/llvm/Analysis/MemoryProfileInfo.h
+++ b/llvm/include/llvm/Analysis/MemoryProfileInfo.h
@@ -59,6 +59,14 @@ LLVM_ABI std::string getAllocTypeAttributeString(AllocationType Type);
 /// True if the AllocTypes bitmask contains just a single type.
 LLVM_ABI bool hasSingleAllocType(uint8_t AllocTypes);
 
+/// Removes any existing "ambiguous" memprof attribute. Called before we apply a
+/// specific allocation type such as "cold", "notcold", or "hot".
+LLVM_ABI void removeAnyExistingAmbiguousAttribute(CallBase *CB);
+
+/// Adds an "ambiguous" memprof attribute to call with a matched allocation
+/// profile but that we haven't yet been able to disambiguate.
+LLVM_ABI void addAmbiguousAttribute(CallBase *CB);
+
 /// Class to build a trie of call stack contexts for a particular profiled
 /// allocation call, along with their associated allocation types.
 /// The allocation will be at the root of the trie, which is then used to
diff --git a/llvm/lib/Analysis/MemoryProfileInfo.cpp b/llvm/lib/Analysis/MemoryProfileInfo.cpp
index b3c8a7d4563b7..53fac32642f91 100644
--- a/llvm/lib/Analysis/MemoryProfileInfo.cpp
+++ b/llvm/lib/Analysis/MemoryProfileInfo.cpp
@@ -121,6 +121,24 @@ bool llvm::memprof::hasSingleAllocType(uint8_t AllocTypes) {
   return NumAllocTypes == 1;
 }
 
+void llvm::memprof::removeAnyExistingAmbiguousAttribute(CallBase *CB) {
+  if (CB->hasFnAttr("memprof")) {
+    assert(CB->getFnAttr("memprof").getValueAsString() == "ambiguous");
+    CB->removeFnAttr("memprof");
+  }
+}
+
+void llvm::memprof::addAmbiguousAttribute(CallBase *CB) {
+  // We may have an existing ambiguous attribute if we are reanalyzing
+  // after inlining.
+  if (CB->hasFnAttr("memprof")) {
+    assert(CB->getFnAttr("memprof").getValueAsString() == "ambiguous");
+  } else {
+    auto A = llvm::Attribute::get(CB->getContext(), "memprof", "ambiguous");
+    CB->addFnAttr(A);
+  }
+}
+
 void CallStackTrie::addCallStack(
     AllocationType AllocType, ArrayRef<uint64_t> StackIds,
     std::vector<ContextTotalSize> ContextSizeInfo) {
@@ -466,6 +484,9 @@ void CallStackTrie::addSingleAllocTypeAttribute(CallBase *CI, AllocationType AT,
                                                 StringRef Descriptor) {
   auto AllocTypeString = getAllocTypeAttributeString(AT);
   auto A = llvm::Attribute::get(CI->getContext(), "memprof", AllocTypeString);
+  // After inlining we may be able to convert an existing ambiguous allocation
+  // to an unambiguous one.
+  removeAnyExistingAmbiguousAttribute(CI);
   CI->addFnAttr(A);
   if (MemProfReportHintedSizes) {
     std::vector<ContextTotalSize> ContextSizeInfo;
@@ -525,6 +546,7 @@ bool CallStackTrie::buildAndAttachMIBMetadata(CallBase *CI) {
     assert(MIBCallStack.size() == 1 &&
            "Should only be left with Alloc's location in stack");
     CI->setMetadata(LLVMContext::MD_memprof, MDNode::get(Ctx, MIBNodes));
+    addAmbiguousAttribute(CI);
     return true;
   }
   // If there exists corner case that CallStackTrie has one chain to leaf
diff --git a/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp b/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
index b8c99f1f33891..7f9693169af0c 100644
--- a/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
+++ b/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
@@ -3965,6 +3965,7 @@ void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::identifyClones(
 void ModuleCallsiteContextGraph::updateAllocationCall(
     CallInfo &Call, AllocationType AllocType) {
   std::string AllocTypeString = getAllocTypeAttributeString(AllocType);
+  removeAnyExistingAmbiguousAttribute(cast<CallBase>(Call.call()));
   auto A = llvm::Attribute::get(Call.call()->getFunction()->getContext(),
                                 "memprof", AllocTypeString);
   cast<CallBase>(Call.call())->addFnAttr(A);
@@ -5501,6 +5502,7 @@ bool MemProfContextDisambiguation::applyImport(Module &M) {
               // clone J-1 (J==0 is the original clone and does not have a VMaps
               // entry).
               CBClone = cast<CallBase>((*VMaps[J - 1])[CB]);
+            removeAnyExistingAmbiguousAttribute(CBClone);
             CBClone->addFnAttr(A);
             ORE.emit(OptimizationRemark(DEBUG_TYPE, "MemprofAttribute", CBClone)
                      << ore::NV("AllocationCall", CBClone) << " in clone "
diff --git a/llvm/unittests/Analysis/MemoryProfileInfoTest.cpp b/llvm/unittests/Analysis/MemoryProfileInfoTest.cpp
index aa4d712cde09e..8c4fd8bb449d5 100644
--- a/llvm/unittests/Analysis/MemoryProfileInfoTest.cpp
+++ b/llvm/unittests/Analysis/MemoryProfileInfoTest.cpp
@@ -228,7 +228,8 @@ declare dso_local noalias noundef i8* @malloc(i64 noundef)
   CallBase *Call = findCall(*Func, "call");
   Trie.buildAndAttachMIBMetadata(Call);
 
-  EXPECT_FALSE(Call->hasFnAttr("memprof"));
+  EXPECT_TRUE(Call->hasFnAttr("memprof"));
+  EXPECT_EQ(Call->getFnAttr("memprof").getValueAsString(), "ambiguous");
   EXPECT_TRUE(Call->hasMetadata(LLVMContext::MD_memprof));
   MDNode *MemProfMD = Call->getMetadata(LLVMContext::MD_memprof);
   ASSERT_EQ(MemProfMD->getNumOperands(), 2u);
@@ -277,7 +278,8 @@ declare dso_local noalias noundef i8* @malloc(i64 noundef)
   CallBase *Call = findCall(*Func, "call");
   Trie.buildAndAttachMIBMetadata(Call);
 
-  EXPECT_FALSE(Call->hasFnAttr("memprof"));
+  EXPECT_TRUE(Call->hasFnAttr("memprof"));
+  EXPECT_EQ(Call->getFnAttr("memprof").getValueAsString(), "ambiguous");
   EXPECT_TRUE(Call->hasMetadata(LLVMContext::MD_memprof));
   MDNode *MemProfMD = Call->getMetadata(LLVMContext::MD_memprof);
   ASSERT_EQ(MemProfMD->getNumOperands(), 2u);
@@ -331,7 +333,8 @@ declare dso_local noalias noundef i8* @malloc(i64 noundef)
   CallBase *Call = findCall(*Func, "call");
   Trie.buildAndAttachMIBMetadata(Call);
 
-  EXPECT_FALSE(Call->hasFnAttr("memprof"));
+  EXPECT_TRUE(Call->hasFnAttr("memprof"));
+  EXPECT_EQ(Call->getFnAttr("memprof").getValueAsString(), "ambiguous");
   EXPECT_TRUE(Call->hasMetadata(LLVMContext::MD_memprof));
   MDNode *MemProfMD = Call->getMetadata(LLVMContext::MD_memprof);
   ASSERT_EQ(MemProfMD->getNumOperands(), 2u);
@@ -390,7 +393,8 @@ declare dso_local noalias noundef i8* @malloc(i64 noundef)
   CallBase *Call = findCall(*Func, "call");
   Trie.buildAndAttachMIBMetadata(Call);
 
-  EXPECT_FALSE(Call->hasFnAttr("memprof"));
+  EXPECT_TRUE(Call->hasFnAttr("memprof"));
+  EXPECT_EQ(Call->getFnAttr("memprof").getValueAsString(), "ambiguous");
   EXPECT_TRUE(Call->hasMetadata(LLVMContext::MD_memprof));
   MDNode *MemProfMD = Call->getMetadata(LLVMContext::MD_memprof);
   ASSERT_EQ(MemProfMD->getNumOperands(), 2u);
@@ -461,7 +465,8 @@ declare dso_local noalias noundef i8* @malloc(i64 noundef)
   ASSERT_NE(Call, nullptr);
   Trie.buildAndAttachMIBMetadata(Call);
 
-  EXPECT_FALSE(Call->hasFnAttr("memprof"));
+  EXPECT_TRUE(Call->hasFnAttr("memprof"));
+  EXPECT_EQ(Call->getFnAttr("memprof").getValueAsString(), "ambiguous");
   EXPECT_TRUE(Call->hasMetadata(LLVMContext::MD_memprof));
   MDNode *MemProfMD = Call->getMetadata(LLVMContext::MD_memprof);
   EXPECT_THAT(MemProfMD, MemprofMetadataEquals(ExpectedVals));
@@ -534,7 +539,8 @@ declare dso_local noalias noundef i8* @malloc(i64 noundef)
   // Restore original option value.
   MemProfKeepAllNotColdContexts = OrigMemProfKeepAllNotColdContexts;
 
-  EXPECT_FALSE(Call->hasFnAttr("memprof"));
+  EXPECT_TRUE(Call->hasFnAttr("memprof"));
+  EXPECT_EQ(Call->getFnAttr("memprof").getValueAsString(), "ambiguous");
   EXPECT_TRUE(Call->hasMetadata(LLVMContext::MD_memprof));
   MDNode *MemProfMD = Call->getMetadata(LLVMContext::MD_memprof);
   EXPECT_THAT(MemProfMD, MemprofMetadataEquals(ExpectedVals));
@@ -662,7 +668,8 @@ declare dso_local noalias noundef i8* @malloc(i64 noundef)
   // The hot allocations will be converted to NotCold and pruned as they
   // are unnecessary to determine how to clone the cold allocation.
 
-  EXPECT_FALSE(Call->hasFnAttr("memprof"));
+  EXPECT_TRUE(Call->hasFnAttr("memprof"));
+  EXPECT_EQ(Call->getFnAttr("memprof").getValueAsString(), "ambiguous");
   EXPECT_TRUE(Call->hasMetadata(LLVMContext::MD_memprof));
   MemProfMD = Call->getMetadata(LLVMContext::MD_memprof);
   ASSERT_EQ(MemProfMD->getNumOperands(), 2u);

snehasish

lgtm

snehasish · 2025-09-06T00:06:46Z

llvm/lib/Analysis/MemoryProfileInfo.cpp

 }

+void llvm::memprof::removeAnyExistingAmbiguousAttribute(CallBase *CB) {
+  if (CB->hasFnAttr("memprof")) {


nit: if(!CB->hasFnAttr ....) return

to avoid nesting.

snehasish · 2025-09-06T00:12:22Z

llvm/unittests/Analysis/MemoryProfileInfoTest.cpp

Is it worth testing that the attribute is only added to the allocation calls and not other callinsts which are part of the allocation context?

I don't think so - we are only adding these new attributes to the same calls where we currently add memprof attributes. This unittest in particular doesn't have any non-alloc calls in any case - it's only checking what we add to the allocations.

This reverts commit cf44f19.

Reverts #157204 This caused issues in ThinLTO binaries because of the checking here, that didn't expect allocations needing cloning to have memprof metadata: https://github.com/llvm/llvm-project/blob/9133fc8cb04f8e45c9b46de85a8de99bf01e55c7/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp#L5572-L5582 I need to move the assert into the if check and guard by that condition. And add a more thorough test.

Reverts llvm/llvm-project#157204 This caused issues in ThinLTO binaries because of the checking here, that didn't expect allocations needing cloning to have memprof metadata: https://github.com/llvm/llvm-project/blob/9133fc8cb04f8e45c9b46de85a8de99bf01e55c7/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp#L5572-L5582 I need to move the assert into the if check and guard by that condition. And add a more thorough test.

Reverts llvm#157204 This caused issues in ThinLTO binaries because of the checking here, that didn't expect allocations needing cloning to have memprof metadata: https://github.com/llvm/llvm-project/blob/9133fc8cb04f8e45c9b46de85a8de99bf01e55c7/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp#L5572-L5582 I need to move the assert into the if check and guard by that condition. And add a more thorough test.

Reapply llvm#157204 with fix and a new test for the issue it caused (the test change provoked the assert that was converted to an if condition). Also, make the application of this new attribute under an (on by default) flag, so that it can be more easily disabled if needed. Add test for the new flag.

llvmbot added llvm:analysis Includes value tracking, cost tables and constant folding llvm:transforms labels Sep 5, 2025

teresajohnson requested review from snehasish and kazutakahirata September 5, 2025 23:52

snehasish approved these changes Sep 6, 2025

View reviewed changes

Address comment

b75e224

teresajohnson merged commit cf44f19 into llvm:main Sep 6, 2025
9 checks passed

teresajohnson added a commit that referenced this pull request Oct 2, 2025

Revert "[MemProf] Add ambigous memprof attribute (#157204)"

e1ae245

This reverts commit cf44f19.

teresajohnson mentioned this pull request Oct 2, 2025

Revert "[MemProf] Add ambigous memprof attribute" #161717

Merged

teresajohnson mentioned this pull request Oct 3, 2025

Reapply "[MemProf] Add ambigous memprof attribute" (#161717) #161918

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[MemProf] Add ambigous memprof attribute #157204

[MemProf] Add ambigous memprof attribute #157204

Uh oh!

teresajohnson commented Sep 5, 2025

Uh oh!

llvmbot commented Sep 5, 2025 •

edited

Loading

Uh oh!

snehasish left a comment

Uh oh!

snehasish Sep 6, 2025

Uh oh!

teresajohnson Sep 6, 2025

Uh oh!

snehasish Sep 6, 2025

Uh oh!

teresajohnson Sep 6, 2025

Uh oh!

Uh oh!

Uh oh!

[MemProf] Add ambigous memprof attribute #157204

[MemProf] Add ambigous memprof attribute #157204

Uh oh!

Conversation

teresajohnson commented Sep 5, 2025

Uh oh!

llvmbot commented Sep 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

snehasish left a comment

Choose a reason for hiding this comment

Uh oh!

snehasish Sep 6, 2025

Choose a reason for hiding this comment

Uh oh!

teresajohnson Sep 6, 2025

Choose a reason for hiding this comment

Uh oh!

snehasish Sep 6, 2025

Choose a reason for hiding this comment

Uh oh!

teresajohnson Sep 6, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

llvmbot commented Sep 5, 2025 •

edited

Loading