[llvm] [BOLT] Fix runOnEachFunctionWithUniqueAllocId (PR #90039)

[Kristof Beyls via llvm-commits]

https://github.com/kbeyls updated https://github.com/llvm/llvm-project/pull/90039

>From 7ce310a7a76b713cd6531935b878a2d282794a70 Mon Sep 17 00:00:00 2001
From: Kristof Beyls <kristof.beyls at arm.com>
Date: Thu, 25 Apr 2024 10:25:31 +0200
Subject: [PATCH 1/3] [BOLT] Fix runOnEachFunctionWithUniqueAllocId

When runOnEachFunctionWithUniqueAllocId is invoked with
ForceSequential=true, then the current implementation runs the function
with AllocId==0, which is the Id for the shared, non-unique, default
AnnotationAllocator.

However, the documentation for runOnEachFunctionWithUniqueAllocId
states:
/// Perform the work on each BinaryFunction except those that are rejected
/// by SkipPredicate, and create a unique annotation allocator for each
/// task. This should be used whenever the work function creates annotations to
/// allow thread-safe annotation creation.

Therefore, even when ForceSequential==true, a unique AllocId should be
used, i.e. different from 0.

In the current upstream BOLT this is presumably not depended on, but it
is needed to reduce memory usage for analyses that use a lot of
memory/annotations. Examples are the pac-ret and stack-clash analyses
that currently have prototype implementations as described in
https://discourse.llvm.org/t/rfc-bolt-based-binary-analysis-tool-to-verify-correctness-of-security-hardening/78148
These analyses use the DataFlowAnalysis framework to sometimes store
quite a lot of information on each MCInst. They run in parallel on each
function. When the dataflow analysis is finished, the annotations on
each MCInst can be removed, hugely saving on memory consumption.
The only annotations that need to remain are those that indicate some
unexpected properties somewhere in the binary.

Fixing this bug enables implementing the deletion of the memory used by
those huge number of DataFlowAnalysis annotations (by invoking
BC.MIB->freeValuesAllocator(AllocatorId)), even when run with
--no-threads. Without this bug fixed, the invocation of
BC.MIB->freeValuesAllocator(AllocatorId) results in also the memory
for all other annotations to be deleted, as AllocatorId is 0.
---
 bolt/lib/Core/ParallelUtilities.cpp | 23 +++++++++++++++--------
 1 file changed, 15 insertions(+), 8 deletions(-)

diff --git a/bolt/lib/Core/ParallelUtilities.cpp b/bolt/lib/Core/ParallelUtilities.cpp
index 5f5e96e0e7881c..dfdb23b4cd2268 100644
--- a/bolt/lib/Core/ParallelUtilities.cpp
+++ b/bolt/lib/Core/ParallelUtilities.cpp
@@ -164,6 +164,15 @@ void runOnEachFunction(BinaryContext &BC, SchedulingPolicy SchedPolicy,
   Pool.wait();
 }
 
+static void EnsureAllocatorExists(BinaryContext& BC, unsigned AllocId) {
+    if (!BC.MIB->checkAllocatorExists(AllocId)) {
+    MCPlusBuilder::AllocatorIdTy Id =
+        BC.MIB->initializeNewAnnotationAllocator();
+    (void)Id;
+    assert(AllocId == Id && "unexpected allocator id created");
+  }
+}
+
 void runOnEachFunctionWithUniqueAllocId(
     BinaryContext &BC, SchedulingPolicy SchedPolicy,
     WorkFuncWithAllocTy WorkFunction, PredicateTy SkipPredicate,
@@ -188,8 +197,12 @@ void runOnEachFunctionWithUniqueAllocId(
     LLVM_DEBUG(T.stopTimer());
   };
 
+  unsigned AllocId = 1;
+
   if (opts::NoThreads || ForceSequential) {
-    runBlock(BC.getBinaryFunctions().begin(), BC.getBinaryFunctions().end(), 0);
+    EnsureAllocatorExists(BC, AllocId);
+    runBlock(BC.getBinaryFunctions().begin(), BC.getBinaryFunctions().end(),
+             AllocId);
     return;
   }
   // This lock is used to postpone task execution
@@ -205,19 +218,13 @@ void runOnEachFunctionWithUniqueAllocId(
   ThreadPoolInterface &Pool = getThreadPool();
   auto BlockBegin = BC.getBinaryFunctions().begin();
   unsigned CurrentCost = 0;
-  unsigned AllocId = 1;
   for (auto It = BC.getBinaryFunctions().begin();
        It != BC.getBinaryFunctions().end(); ++It) {
     BinaryFunction &BF = It->second;
     CurrentCost += computeCostFor(BF, SkipPredicate, SchedPolicy);
 
     if (CurrentCost >= BlockCost) {
-      if (!BC.MIB->checkAllocatorExists(AllocId)) {
-        MCPlusBuilder::AllocatorIdTy Id =
-            BC.MIB->initializeNewAnnotationAllocator();
-        (void)Id;
-        assert(AllocId == Id && "unexpected allocator id created");
-      }
+      EnsureAllocatorExists(BC, AllocId);
       Pool.async(runBlock, BlockBegin, std::next(It), AllocId);
       AllocId++;
       BlockBegin = std::next(It);

>From 47f2dbf08989c54a9bf5aa16b43fefd239a5c71f Mon Sep 17 00:00:00 2001
From: Kristof Beyls <kristof.beyls at arm.com>
Date: Thu, 25 Apr 2024 13:54:18 +0200
Subject: [PATCH 2/3] Fix formatting

---
 bolt/lib/Core/ParallelUtilities.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/bolt/lib/Core/ParallelUtilities.cpp b/bolt/lib/Core/ParallelUtilities.cpp
index dfdb23b4cd2268..8e5f2f0f1f23d3 100644
--- a/bolt/lib/Core/ParallelUtilities.cpp
+++ b/bolt/lib/Core/ParallelUtilities.cpp
@@ -164,8 +164,8 @@ void runOnEachFunction(BinaryContext &BC, SchedulingPolicy SchedPolicy,
   Pool.wait();
 }
 
-static void EnsureAllocatorExists(BinaryContext& BC, unsigned AllocId) {
-    if (!BC.MIB->checkAllocatorExists(AllocId)) {
+static void EnsureAllocatorExists(BinaryContext &BC, unsigned AllocId) {
+  if (!BC.MIB->checkAllocatorExists(AllocId)) {
     MCPlusBuilder::AllocatorIdTy Id =
         BC.MIB->initializeNewAnnotationAllocator();
     (void)Id;

>From 34dcca4c84d78af08cc0ae2f737c2a875c15a815 Mon Sep 17 00:00:00 2001
From: Kristof Beyls <kristof.beyls at gmail.com>
Date: Fri, 3 May 2024 11:51:23 +0200
Subject: [PATCH 3/3] Update bolt/lib/Core/ParallelUtilities.cpp

Co-authored-by: Maksim Panchenko <maks at meta.com>
---
 bolt/lib/Core/ParallelUtilities.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/bolt/lib/Core/ParallelUtilities.cpp b/bolt/lib/Core/ParallelUtilities.cpp
index 8e5f2f0f1f23d3..3473d820b68db4 100644
--- a/bolt/lib/Core/ParallelUtilities.cpp
+++ b/bolt/lib/Core/ParallelUtilities.cpp
@@ -164,7 +164,7 @@ void runOnEachFunction(BinaryContext &BC, SchedulingPolicy SchedPolicy,
   Pool.wait();
 }
 
-static void EnsureAllocatorExists(BinaryContext &BC, unsigned AllocId) {
+static void ensureAllocatorExists(BinaryContext &BC, unsigned AllocId) {
   if (!BC.MIB->checkAllocatorExists(AllocId)) {
     MCPlusBuilder::AllocatorIdTy Id =
         BC.MIB->initializeNewAnnotationAllocator();