[compiler-rt] b01e5b2 - [ctxprof][nfc] Refactor `__llvm_ctx_profile_start_context` (#133744)

[via llvm-commits]

Author: Mircea Trofin
Date: 2025-03-31T12:26:25-07:00
New Revision: b01e5b23dd880e9686cc4151c7d1b1737cbdd98e

URL: https://github.com/llvm/llvm-project/commit/b01e5b23dd880e9686cc4151c7d1b1737cbdd98e
DIFF: https://github.com/llvm/llvm-project/commit/b01e5b23dd880e9686cc4151c7d1b1737cbdd98e.diff

LOG: [ctxprof][nfc] Refactor `__llvm_ctx_profile_start_context` (#133744)

Most of the functionality will be reused with the auto-root detection mechanism (which is introduced subsequently in PR #133147).

Added: 
    

Modified: 
    compiler-rt/lib/ctx_profile/CtxInstrProfiling.cpp

Removed: 
    


################################################################################
diff  --git a/compiler-rt/lib/ctx_profile/CtxInstrProfiling.cpp b/compiler-rt/lib/ctx_profile/CtxInstrProfiling.cpp
index da291e0bbabdd..10a6a8c1f71e5 100644
--- a/compiler-rt/lib/ctx_profile/CtxInstrProfiling.cpp
+++ b/compiler-rt/lib/ctx_profile/CtxInstrProfiling.cpp
@@ -277,6 +277,25 @@ ContextRoot *FunctionData::getOrAllocateContextRoot() {
   return Root;
 }
 
+ContextNode *tryStartContextGivenRoot(ContextRoot *Root, GUID Guid,
+                                      uint32_t Counters, uint32_t Callsites)
+    SANITIZER_NO_THREAD_SAFETY_ANALYSIS {
+  IsUnderContext = true;
+  __sanitizer::atomic_fetch_add(&Root->TotalEntries, 1,
+                                __sanitizer::memory_order_relaxed);
+  if (!Root->FirstMemBlock) {
+    setupContext(Root, Guid, Counters, Callsites);
+  }
+  if (Root->Taken.TryLock()) {
+    __llvm_ctx_profile_current_context_root = Root;
+    onContextEnter(*Root->FirstNode);
+    return Root->FirstNode;
+  }
+  // If this thread couldn't take the lock, return scratch context.
+  __llvm_ctx_profile_current_context_root = nullptr;
+  return TheScratchContext;
+}
+
 ContextNode *getUnhandledContext(FunctionData &Data, GUID Guid,
                                  uint32_t NumCounters) {
 
@@ -369,24 +388,8 @@ ContextNode *__llvm_ctx_profile_get_context(FunctionData *Data, void *Callee,
 ContextNode *__llvm_ctx_profile_start_context(
     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);
-
-  if (!Root->FirstMemBlock) {
-    setupContext(Root, Guid, Counters, Callsites);
-  }
-  if (Root->Taken.TryLock()) {
-    __llvm_ctx_profile_current_context_root = Root;
-    onContextEnter(*Root->FirstNode);
-    return Root->FirstNode;
-  }
-  // If this thread couldn't take the lock, return scratch context.
-  __llvm_ctx_profile_current_context_root = nullptr;
-  return TheScratchContext;
+  return tryStartContextGivenRoot(FData->getOrAllocateContextRoot(), Guid,
+                                  Counters, Callsites);
 }
 
 void __llvm_ctx_profile_release_context(FunctionData *FData)