From cff00d11c5545463e74f35b522ef97c1a6ac5b96 Mon Sep 17 00:00:00 2001 From: Teresa Johnson Date: Fri, 5 Sep 2025 16:37:28 -0700 Subject: [PATCH 1/2] [MemProf] Add ambigous memprof attribute 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. --- .../include/llvm/Analysis/MemoryProfileInfo.h | 8 +++++++ llvm/lib/Analysis/MemoryProfileInfo.cpp | 22 +++++++++++++++++++ .../IPO/MemProfContextDisambiguation.cpp | 2 ++ .../Analysis/MemoryProfileInfoTest.cpp | 21 ++++++++++++------ 4 files changed, 46 insertions(+), 7 deletions(-) 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 StackIds, std::vector 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 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::identifyClones( void ModuleCallsiteContextGraph::updateAllocationCall( CallInfo &Call, AllocationType AllocType) { std::string AllocTypeString = getAllocTypeAttributeString(AllocType); + removeAnyExistingAmbiguousAttribute(cast(Call.call())); auto A = llvm::Attribute::get(Call.call()->getFunction()->getContext(), "memprof", AllocTypeString); cast(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((*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); From b75e224df019a019b9de3209afb904bc09d33c29 Mon Sep 17 00:00:00 2001 From: Teresa Johnson Date: Fri, 5 Sep 2025 18:00:01 -0700 Subject: [PATCH 2/2] Address comment --- llvm/lib/Analysis/MemoryProfileInfo.cpp | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/llvm/lib/Analysis/MemoryProfileInfo.cpp b/llvm/lib/Analysis/MemoryProfileInfo.cpp index 53fac32642f91..b5ca6b13108fe 100644 --- a/llvm/lib/Analysis/MemoryProfileInfo.cpp +++ b/llvm/lib/Analysis/MemoryProfileInfo.cpp @@ -122,10 +122,10 @@ bool llvm::memprof::hasSingleAllocType(uint8_t AllocTypes) { } void llvm::memprof::removeAnyExistingAmbiguousAttribute(CallBase *CB) { - if (CB->hasFnAttr("memprof")) { - assert(CB->getFnAttr("memprof").getValueAsString() == "ambiguous"); - CB->removeFnAttr("memprof"); - } + if (!CB->hasFnAttr("memprof")) + return; + assert(CB->getFnAttr("memprof").getValueAsString() == "ambiguous"); + CB->removeFnAttr("memprof"); } void llvm::memprof::addAmbiguousAttribute(CallBase *CB) {