Skip to content

[ctxprof] Make ContextRoot an implementation detail #131416

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Mar 19, 2025
Merged

[ctxprof] Make ContextRoot an implementation detail #131416

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 22 additions & 3 deletions compiler-rt/lib/ctx_profile/CtxInstrProfiling.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -336,10 +336,28 @@ void setupContext(ContextRoot *Root, GUID Guid, uint32_t NumCounters,
AllContextRoots.PushBack(Root);
}

ContextRoot *FunctionData::getOrAllocateContextRoot() {
auto *Root = CtxRoot;
if (Root)
return Root;
__sanitizer::GenericScopedLock<__sanitizer::StaticSpinMutex> L(&Mutex);
Root = CtxRoot;
if (!Root) {
Root = new (__sanitizer::InternalAlloc(sizeof(ContextRoot))) ContextRoot();
CtxRoot = Root;
}

assert(Root);
return Root;
}

ContextNode *__llvm_ctx_profile_start_context(
ContextRoot *Root, GUID Guid, uint32_t Counters,
FunctionData *FData, GUID Guid, uint32_t Counters,
uint32_t Callsites) SANITIZER_NO_THREAD_SAFETY_ANALYSIS {
IsUnderContext = true;

auto *Root = FData->getOrAllocateContextRoot();

__sanitizer::atomic_fetch_add(&Root->TotalEntries, 1,
__sanitizer::memory_order_relaxed);

Expand All @@ -356,12 +374,13 @@ ContextNode *__llvm_ctx_profile_start_context(
return TheScratchContext;
}

void __llvm_ctx_profile_release_context(ContextRoot *Root)
void __llvm_ctx_profile_release_context(FunctionData *FData)
SANITIZER_NO_THREAD_SAFETY_ANALYSIS {
IsUnderContext = false;
if (__llvm_ctx_profile_current_context_root) {
__llvm_ctx_profile_current_context_root = nullptr;
Root->Taken.Unlock();
assert(FData->CtxRoot);
FData->CtxRoot->Taken.Unlock();
}
}

Expand Down
31 changes: 16 additions & 15 deletions compiler-rt/lib/ctx_profile/CtxInstrProfiling.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,9 +73,8 @@ static_assert(alignof(Arena) == ExpectedAlignment);
// sure the inlined vectors are appropriately aligned.
static_assert(alignof(ContextNode) == ExpectedAlignment);

/// ContextRoots are allocated by LLVM for entrypoints. LLVM is only concerned
/// with allocating and zero-initializing the global value (as in, GlobalValue)
/// for it.
/// ContextRoots hold memory and the start of the contextual profile tree for a
/// root function.
struct ContextRoot {
ContextNode *FirstNode = nullptr;
Arena *FirstMemBlock = nullptr;
Expand All @@ -84,7 +83,6 @@ struct ContextRoot {
// Count the number of entries - regardless if we could take the `Taken` mutex
::__sanitizer::atomic_uint64_t TotalEntries = {};

// This is init-ed by the static zero initializer in LLVM.
// Taken is used to ensure only one thread traverses the contextual graph -
// either to read it or to write it. On server side, the same entrypoint will
// be entered by numerous threads, but over time, the profile aggregated by
Expand All @@ -109,12 +107,7 @@ struct ContextRoot {
// or with more concurrent collections (==more memory) and less collection
// time. Note that concurrent collection does happen for different
// entrypoints, regardless.
::__sanitizer::StaticSpinMutex Taken;

// If (unlikely) StaticSpinMutex internals change, we need to modify the LLVM
// instrumentation lowering side because it is responsible for allocating and
// zero-initializing ContextRoots.
static_assert(sizeof(Taken) == 1);
::__sanitizer::SpinMutex Taken;
};

// This is allocated and zero-initialized by the compiler, the in-place
Expand All @@ -139,8 +132,16 @@ struct FunctionData {
FunctionData() { Mutex.Init(); }

FunctionData *Next = nullptr;
ContextRoot *volatile CtxRoot = nullptr;
ContextNode *volatile FlatCtx = nullptr;

ContextRoot *getOrAllocateContextRoot();

::__sanitizer::StaticSpinMutex Mutex;
// If (unlikely) StaticSpinMutex internals change, we need to modify the LLVM
// instrumentation lowering side because it is responsible for allocating and
// zero-initializing ContextRoots.
static_assert(sizeof(Mutex) == 1);
};

/// This API is exposed for testing. See the APIs below about the contract with
Expand Down Expand Up @@ -172,17 +173,17 @@ extern __thread __ctx_profile::ContextRoot

/// called by LLVM in the entry BB of a "entry point" function. The returned
/// pointer may be "tainted" - its LSB set to 1 - to indicate it's scratch.
ContextNode *__llvm_ctx_profile_start_context(__ctx_profile::ContextRoot *Root,
GUID Guid, uint32_t Counters,
uint32_t Callsites);
ContextNode *
__llvm_ctx_profile_start_context(__ctx_profile::FunctionData *FData, GUID Guid,
uint32_t Counters, uint32_t Callsites);

/// paired with __llvm_ctx_profile_start_context, and called at the exit of the
/// entry point function.
void __llvm_ctx_profile_release_context(__ctx_profile::ContextRoot *Root);
void __llvm_ctx_profile_release_context(__ctx_profile::FunctionData *FData);

/// called for any other function than entry points, in the entry BB of such
/// function. Same consideration about LSB of returned value as .._start_context
ContextNode *__llvm_ctx_profile_get_context(__ctx_profile::FunctionData *Data,
ContextNode *__llvm_ctx_profile_get_context(__ctx_profile::FunctionData *FData,
void *Callee, GUID Guid,
uint32_t NumCounters,
uint32_t NumCallsites);
Expand Down
30 changes: 16 additions & 14 deletions compiler-rt/lib/ctx_profile/tests/CtxInstrProfilingTest.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,11 +5,11 @@
using namespace __ctx_profile;

class ContextTest : public ::testing::Test {
void SetUp() override { memset(&Root, 0, sizeof(ContextRoot)); }
void SetUp() override { Root.getOrAllocateContextRoot(); }
void TearDown() override { __llvm_ctx_profile_free(); }

public:
ContextRoot Root;
FunctionData Root;
};

TEST(ArenaTest, ZeroInit) {
Expand Down Expand Up @@ -43,19 +43,20 @@ TEST_F(ContextTest, Basic) {
__llvm_ctx_profile_start_collection();
auto *Ctx = __llvm_ctx_profile_start_context(&Root, 1, 10, 4);
ASSERT_NE(Ctx, nullptr);
EXPECT_NE(Root.CurrentMem, nullptr);
EXPECT_EQ(Root.FirstMemBlock, Root.CurrentMem);
auto &CtxRoot = *Root.CtxRoot;
EXPECT_NE(CtxRoot.CurrentMem, nullptr);
EXPECT_EQ(CtxRoot.FirstMemBlock, CtxRoot.CurrentMem);
EXPECT_EQ(Ctx->size(), sizeof(ContextNode) + 10 * sizeof(uint64_t) +
4 * sizeof(ContextNode *));
EXPECT_EQ(Ctx->counters_size(), 10U);
EXPECT_EQ(Ctx->callsites_size(), 4U);
EXPECT_EQ(__llvm_ctx_profile_current_context_root, &Root);
Root.Taken.CheckLocked();
EXPECT_FALSE(Root.Taken.TryLock());
EXPECT_EQ(__llvm_ctx_profile_current_context_root, &CtxRoot);
CtxRoot.Taken.CheckLocked();
EXPECT_FALSE(CtxRoot.Taken.TryLock());
__llvm_ctx_profile_release_context(&Root);
EXPECT_EQ(__llvm_ctx_profile_current_context_root, nullptr);
EXPECT_TRUE(Root.Taken.TryLock());
Root.Taken.Unlock();
EXPECT_TRUE(CtxRoot.Taken.TryLock());
CtxRoot.Taken.Unlock();
}

TEST_F(ContextTest, Callsite) {
Expand Down Expand Up @@ -172,7 +173,8 @@ TEST_F(ContextTest, NeedMoreMemory) {
int FakeCalleeAddress = 0;
const bool IsScratch = isScratch(Ctx);
EXPECT_FALSE(IsScratch);
const auto *CurrentMem = Root.CurrentMem;
auto &CtxRoot = *Root.CtxRoot;
const auto *CurrentMem = CtxRoot.CurrentMem;
__llvm_ctx_profile_expected_callee[0] = &FakeCalleeAddress;
__llvm_ctx_profile_callsite[0] = &Ctx->subContexts()[2];
FunctionData FData;
Expand All @@ -181,8 +183,8 @@ TEST_F(ContextTest, NeedMoreMemory) {
__llvm_ctx_profile_get_context(&FData, &FakeCalleeAddress, 3, 1 << 20, 1);
EXPECT_EQ(FData.FlatCtx, nullptr);
EXPECT_EQ(Ctx->subContexts()[2], Subctx);
EXPECT_NE(CurrentMem, Root.CurrentMem);
EXPECT_NE(Root.CurrentMem, nullptr);
EXPECT_NE(CurrentMem, CtxRoot.CurrentMem);
EXPECT_NE(CtxRoot.CurrentMem, nullptr);
}

TEST_F(ContextTest, ConcurrentRootCollection) {
Expand Down Expand Up @@ -277,7 +279,7 @@ TEST_F(ContextTest, Dump) {
void endFlatSection() override { ++ExitedFlatCount; }
};

TestProfileWriter W(&Root, 1);
TestProfileWriter W(Root.CtxRoot, 1);
EXPECT_FALSE(W.State);
__llvm_ctx_profile_fetch(W);
EXPECT_TRUE(W.State);
Expand All @@ -289,7 +291,7 @@ TEST_F(ContextTest, Dump) {
(void)Flat;
EXPECT_NE(FData.FlatCtx, nullptr);
FData.FlatCtx->counters()[0] = 15U;
TestProfileWriter W2(&Root, 0);
TestProfileWriter W2(Root.CtxRoot, 0);
EXPECT_FALSE(W2.State);
__llvm_ctx_profile_fetch(W2);
EXPECT_TRUE(W2.State);
Expand Down
50 changes: 19 additions & 31 deletions llvm/lib/Transforms/Instrumentation/PGOCtxProfLowering.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,10 +53,9 @@ class CtxInstrumentationLowerer final {
Module &M;
ModuleAnalysisManager &MAM;
Type *ContextNodeTy = nullptr;
Type *ContextRootTy = nullptr;
Type *FunctionDataTy = nullptr;

DenseMap<const Function *, Constant *> ContextRootMap;
DenseSet<const Function *> ContextRootSet;
Function *StartCtx = nullptr;
Function *GetCtx = nullptr;
Function *ReleaseCtx = nullptr;
Expand Down Expand Up @@ -114,18 +113,10 @@ CtxInstrumentationLowerer::CtxInstrumentationLowerer(Module &M,
auto *I32Ty = Type::getInt32Ty(M.getContext());
auto *I64Ty = Type::getInt64Ty(M.getContext());

// The ContextRoot type
ContextRootTy =
StructType::get(M.getContext(), {
PointerTy, /*FirstNode*/
PointerTy, /*FirstMemBlock*/
PointerTy, /*CurrentMem*/
I64Ty, /*TotalEntries*/
SanitizerMutexType, /*Taken*/
});
FunctionDataTy =
StructType::get(M.getContext(), {
PointerTy, /*Next*/
PointerTy, /*CtxRoot*/
PointerTy, /*FlatCtx*/
SanitizerMutexType, /*Mutex*/
});
Expand All @@ -144,10 +135,7 @@ CtxInstrumentationLowerer::CtxInstrumentationLowerer(Module &M,
if (const auto *F = M.getFunction(Fname)) {
if (F->isDeclaration())
continue;
auto *G = M.getOrInsertGlobal(Fname + "_ctx_root", ContextRootTy);
cast<GlobalVariable>(G)->setInitializer(
Constant::getNullValue(ContextRootTy));
ContextRootMap.insert(std::make_pair(F, G));
ContextRootSet.insert(F);
for (const auto &BB : *F)
for (const auto &I : BB)
if (const auto *CB = dyn_cast<CallBase>(&I))
Expand All @@ -165,7 +153,7 @@ CtxInstrumentationLowerer::CtxInstrumentationLowerer(Module &M,
M.getOrInsertFunction(
CompilerRtAPINames::StartCtx,
FunctionType::get(PointerTy,
{PointerTy, /*ContextRoot*/
{PointerTy, /*FunctionData*/
I64Ty, /*Guid*/ I32Ty,
/*NumCounters*/ I32Ty /*NumCallsites*/},
false))
Expand All @@ -184,7 +172,7 @@ CtxInstrumentationLowerer::CtxInstrumentationLowerer(Module &M,
M.getOrInsertFunction(CompilerRtAPINames::ReleaseCtx,
FunctionType::get(Type::getVoidTy(M.getContext()),
{
PointerTy, /*ContextRoot*/
PointerTy, /*FunctionData*/
},
false))
.getCallee());
Expand Down Expand Up @@ -224,7 +212,7 @@ bool CtxInstrumentationLowerer::lowerFunction(Function &F) {
Value *RealContext = nullptr;

StructType *ThisContextType = nullptr;
Value *TheRootContext = nullptr;
Value *TheRootFuctionData = nullptr;
Value *ExpectedCalleeTLSAddr = nullptr;
Value *CallsiteInfoTLSAddr = nullptr;

Expand All @@ -246,23 +234,23 @@ bool CtxInstrumentationLowerer::lowerFunction(Function &F) {
ArrayType::get(Builder.getPtrTy(), NumCallsites)});
// Figure out which way we obtain the context object for this function -
// if it's an entrypoint, then we call StartCtx, otherwise GetCtx. In the
// former case, we also set TheRootContext since we need to release it
// former case, we also set TheRootFuctionData since we need to release it
// at the end (plus it can be used to know if we have an entrypoint or a
// regular function)
auto Iter = ContextRootMap.find(&F);
if (Iter != ContextRootMap.end()) {
TheRootContext = Iter->second;
// Don't set a name, they end up taking a lot of space and we don't need
// them.
auto *FData = new GlobalVariable(M, FunctionDataTy, false,
GlobalVariable::InternalLinkage,
Constant::getNullValue(FunctionDataTy));

if (ContextRootSet.contains(&F)) {
Context = Builder.CreateCall(
StartCtx, {TheRootContext, Guid, Builder.getInt32(NumCounters),
StartCtx, {FData, Guid, Builder.getInt32(NumCounters),
Builder.getInt32(NumCallsites)});
TheRootFuctionData = FData;
ORE.emit(
[&] { return OptimizationRemark(DEBUG_TYPE, "Entrypoint", &F); });
} else {
// Make up a compact name, these names end up taking up a lot of space
// in the binary.
auto *FData = new GlobalVariable(
M, FunctionDataTy, false, GlobalVariable::InternalLinkage,
Constant::getNullValue(FunctionDataTy));
Context = Builder.CreateCall(GetCtx, {FData, &F, Guid,
Builder.getInt32(NumCounters),
Builder.getInt32(NumCallsites)});
Expand Down Expand Up @@ -347,10 +335,10 @@ bool CtxInstrumentationLowerer::lowerFunction(Function &F) {
break;
}
I.eraseFromParent();
} else if (TheRootContext && isa<ReturnInst>(I)) {
} else if (TheRootFuctionData && isa<ReturnInst>(I)) {
// Remember to release the context if we are an entrypoint.
IRBuilder<> Builder(&I);
Builder.CreateCall(ReleaseCtx, {TheRootContext});
Builder.CreateCall(ReleaseCtx, {TheRootFuctionData});
ContextWasReleased = true;
}
}
Expand All @@ -359,7 +347,7 @@ bool CtxInstrumentationLowerer::lowerFunction(Function &F) {
// to disallow this, (so this then stays as an error), another is to detect
// that and then do a wrapper or disallow the tail call. This only affects
// instrumentation, when we want to detect the call graph.
if (TheRootContext && !ContextWasReleased)
if (TheRootFuctionData && !ContextWasReleased)
F.getContext().emitError(
"[ctx_prof] An entrypoint was instrumented but it has no `ret` "
"instructions above which to release the context: " +
Expand Down
Loading