Skip to content

Commit 49d6fe2

Browse files
kbeylskbeyls
committed
[BOLT] Fix memory leak in dataflowanalysis
Do this by making sure that the underlying BumpPtr allocator can be freed if the DataflowAnalysis is only instantiated for a "short" period of time, i.e., does get destructed after analyzing one BinaryFunction and before analyzing the next BinaryFunction.
1 parent 696422e commit 49d6fe2

20 files changed

+174
-97
lines changed

‎bolt/include/bolt/Core/MCPlusBuilder.h

Lines changed: 49 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -351,34 +351,71 @@ class MCPlusBuilder {
351351
return nullptr;
352352
}
353353

354-
/// Initialize a new annotation allocator and return its id
355-
AllocatorIdTy initializeNewAnnotationAllocator() {
354+
private:
355+
std::vector<AllocatorIdTy> AvailablePrivateAllocIds;
356+
std::vector<AllocatorIdTy> AvailableSharedAllocIds;
357+
llvm::sys::RWMutex AllocIdMgmtMutex;
358+
359+
AllocatorIdTy _initializeNewAnnotationAllocator() {
356360
AnnotationAllocators.emplace(MaxAllocatorId, AnnotationAllocator());
361+
assert(MaxAllocatorId < std::numeric_limits<AllocatorIdTy>::max());
357362
return MaxAllocatorId++;
358363
}
359364

365+
public:
366+
AllocatorIdTy getPrivateAllocatorId() {
367+
std::unique_lock<llvm::sys::RWMutex> Lock(AllocIdMgmtMutex);
368+
if (AvailablePrivateAllocIds.size() > 0) {
369+
AllocatorIdTy AllocId = AvailablePrivateAllocIds.back();
370+
AvailablePrivateAllocIds.pop_back();
371+
AnnotationAllocators.emplace(AllocId, AnnotationAllocator());
372+
return AllocId;
373+
}
374+
AllocatorIdTy AllocId = _initializeNewAnnotationAllocator();
375+
return AllocId;
376+
}
377+
378+
std::vector<AllocatorIdTy> getSharedAllocIds(unsigned nrAllocs) {
379+
std::unique_lock<llvm::sys::RWMutex> Lock(AllocIdMgmtMutex);
380+
for (unsigned I = AvailableSharedAllocIds.size(); I <= nrAllocs; ++I) {
381+
// allocate new shared alloc Ids
382+
AllocatorIdTy AllocId = _initializeNewAnnotationAllocator();
383+
AvailableSharedAllocIds.push_back(AllocId);
384+
}
385+
std::vector<AllocatorIdTy> Result;
386+
for (unsigned I = 0; I < nrAllocs; ++I)
387+
Result.push_back(AvailableSharedAllocIds[I]);
388+
return Result;
389+
}
390+
391+
void freePrivateAllocatorId(AllocatorIdTy AllocId) {
392+
std::unique_lock<llvm::sys::RWMutex> Lock(AllocIdMgmtMutex);
393+
assert(AnnotationAllocators.count(AllocId) != 0);
394+
AnnotationAllocator &Allocator = AnnotationAllocators.find(AllocId)->second;
395+
for (MCPlus::MCAnnotation *Annotation : Allocator.AnnotationPool)
396+
Annotation->~MCAnnotation();
397+
398+
Allocator.AnnotationPool
399+
.clear(); // FIXME: wouldn't this call the destructor anyway?
400+
Allocator.ValueAllocator.Reset();
401+
AnnotationAllocators.erase(AllocId);
402+
AvailablePrivateAllocIds.push_back(AllocId);
403+
}
404+
360405
/// Return the annotation allocator of a given id
361406
AnnotationAllocator &getAnnotationAllocator(AllocatorIdTy AllocatorId) {
407+
std::shared_lock<llvm::sys::RWMutex> Lock(AllocIdMgmtMutex);
362408
assert(AnnotationAllocators.count(AllocatorId) &&
363409
"allocator not initialized");
364410
return AnnotationAllocators.find(AllocatorId)->second;
365411
}
366412

367413
// Check if an annotation allocator with the given id exists
368414
bool checkAllocatorExists(AllocatorIdTy AllocatorId) {
415+
std::shared_lock<llvm::sys::RWMutex> Lock(AllocIdMgmtMutex);
369416
return AnnotationAllocators.count(AllocatorId);
370417
}
371418

372-
/// Free the values allocator within the annotation allocator
373-
void freeValuesAllocator(AllocatorIdTy AllocatorId) {
374-
AnnotationAllocator &Allocator = getAnnotationAllocator(AllocatorId);
375-
for (MCPlus::MCAnnotation *Annotation : Allocator.AnnotationPool)
376-
Annotation->~MCAnnotation();
377-
378-
Allocator.AnnotationPool.clear();
379-
Allocator.ValueAllocator.Reset();
380-
}
381-
382419
virtual ~MCPlusBuilder() { freeAnnotations(); }
383420

384421
/// Free all memory allocated for annotations

‎bolt/include/bolt/Passes/DataflowAnalysis.h

Lines changed: 40 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -163,6 +163,8 @@ class DataflowAnalysis {
163163
/// The id of the annotation allocator to be used
164164
MCPlusBuilder::AllocatorIdTy AllocatorId = 0;
165165

166+
bool IsPrivateAllocatorId;
167+
166168
/// Tracks the state at basic block start (end) if direction of the dataflow
167169
/// is forward (backward).
168170
std::unordered_map<const BinaryBasicBlock *, StateTy> StateAtBBEntry;
@@ -243,9 +245,6 @@ class DataflowAnalysis {
243245
}
244246

245247
public:
246-
/// Return the allocator id
247-
unsigned getAllocatorId() { return AllocatorId; }
248-
249248
/// If the direction of the dataflow is forward, operates on the last
250249
/// instruction of all predecessors when performing an iteration of the
251250
/// dataflow equation for the start of this BB. If backwards, operates on
@@ -259,11 +258,40 @@ class DataflowAnalysis {
259258

260259
/// We need the current binary context and the function that will be processed
261260
/// in this dataflow analysis.
262-
DataflowAnalysis(BinaryFunction &BF,
263-
MCPlusBuilder::AllocatorIdTy AllocatorId = 0)
264-
: BC(BF.getBinaryContext()), Func(BF), AllocatorId(AllocatorId) {}
265-
266-
virtual ~DataflowAnalysis() { cleanAnnotations(); }
261+
///
262+
/// There are 2 constructors:
263+
/// * One version takes an AllocId, where it is assumed that:
264+
/// * If multiple threads are used, each different thread will get a different
265+
/// AllocId.
266+
/// * AllocIds may be reused by other functions/objects/DataflowAnalyses that
267+
/// take an AllocId of where to allocate MCAnnotations.
268+
/// As a result, the MCAnnotations this dataflow analysis allocates will
269+
/// not get freed before the end of the program run.
270+
/// * The second version takes a bool argument UsePrivateAllocators, which
271+
/// must be set to true.
272+
/// * When constructed this way, the dataflowanalysis will construct its own
273+
/// allocator, only used for MCAnnotations allocated by this DataFlowAnalysis.
274+
/// * When the DataflowAnalysis is destructed, that allocator gets freed,
275+
/// so that the memory used by this DataflowAnalysis gets completely freed.
276+
/// * Note however, that there cannot be a huge number of allocators all live
277+
/// at the same time. Therefore, only use this function when you're going
278+
/// to destruct this DataFlowAnalysis after analyzing one BinaryFunction,
279+
/// before creating a new DataFlowAnalysis to analyze another BinaryFunction.
280+
DataflowAnalysis(BinaryFunction &BF, const bool UsePrivateAllocators)
281+
: BC(BF.getBinaryContext()), Func(BF),
282+
AllocatorId(BC.MIB->getPrivateAllocatorId()),
283+
IsPrivateAllocatorId(true) {
284+
assert(UsePrivateAllocators == true);
285+
}
286+
DataflowAnalysis(BinaryFunction &BF, MCPlusBuilder::AllocatorIdTy AllocId)
287+
: BC(BF.getBinaryContext()), Func(BF), AllocatorId(AllocId),
288+
IsPrivateAllocatorId(false) {}
289+
290+
virtual ~DataflowAnalysis() {
291+
cleanAnnotations();
292+
if (IsPrivateAllocatorId)
293+
BC.MIB->freePrivateAllocatorId(AllocatorId);
294+
}
267295

268296
/// Track the state at basic block start (end) if direction of the dataflow
269297
/// is forward (backward).
@@ -529,8 +557,11 @@ class InstrsDataflowAnalysis
529557
return count(*Expressions[PointIdx], Expr);
530558
}
531559

560+
InstrsDataflowAnalysis(BinaryFunction &BF, const bool UsePrivateAllocators)
561+
: DataflowAnalysis<Derived, BitVector, Backward, StatePrinterTy>(
562+
BF, UsePrivateAllocators) {}
532563
InstrsDataflowAnalysis(BinaryFunction &BF,
533-
MCPlusBuilder::AllocatorIdTy AllocId = 0)
564+
MCPlusBuilder::AllocatorIdTy AllocId)
534565
: DataflowAnalysis<Derived, BitVector, Backward, StatePrinterTy>(
535566
BF, AllocId) {}
536567
virtual ~InstrsDataflowAnalysis() {}

‎bolt/include/bolt/Passes/DataflowInfoManager.h

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -44,16 +44,11 @@ class DataflowInfoManager {
4444
std::unique_ptr<std::unordered_map<const MCInst *, BinaryBasicBlock *>>
4545
InsnToBB;
4646

47-
// Id of the allocator to be used for annotations added by any of the managed
48-
// analysis
49-
MCPlusBuilder::AllocatorIdTy AllocatorId;
50-
5147
public:
5248
DataflowInfoManager(BinaryFunction &BF, const RegAnalysis *RA,
53-
const FrameAnalysis *FA,
54-
MCPlusBuilder::AllocatorIdTy AllocId = 0)
55-
: RA(RA), FA(FA), BC(BF.getBinaryContext()), BF(BF),
56-
AllocatorId(AllocId){};
49+
const FrameAnalysis *FA)
50+
: RA(RA), FA(FA), BC(BF.getBinaryContext()), BF(BF)
51+
{};
5752

5853
/// Helper function to fetch the parent BB associated with a program point
5954
/// If PP is a BB itself, then return itself (cast to a BinaryBasicBlock)

‎bolt/include/bolt/Passes/DominatorAnalysis.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,9 @@ class DominatorAnalysis
3232
Backward>;
3333

3434
public:
35+
DominatorAnalysis(BinaryFunction &BF)
36+
: InstrsDataflowAnalysis<DominatorAnalysis<Backward>, Backward>(BF) {
37+
}
3538
DominatorAnalysis(BinaryFunction &BF, MCPlusBuilder::AllocatorIdTy AllocId)
3639
: InstrsDataflowAnalysis<DominatorAnalysis<Backward>, Backward>(BF,
3740
AllocId) {

‎bolt/include/bolt/Passes/FrameAnalysis.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -208,7 +208,7 @@ class FrameAnalysis {
208208
/// Get or create an SPT object and run the analysis
209209
StackPointerTracking &getSPT(BinaryFunction &BF) {
210210
if (!SPTMap.count(&BF)) {
211-
SPTMap.emplace(&BF, std::make_unique<StackPointerTracking>(BF));
211+
SPTMap.emplace(&BF, std::make_unique<StackPointerTracking>(BF, true));
212212
auto Iter = SPTMap.find(&BF);
213213
assert(Iter != SPTMap.end() && "item should exist");
214214
Iter->second->run();

‎bolt/include/bolt/Passes/LivenessAnalysis.h

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,8 @@ class LivenessAnalysis : public DataflowAnalysis<LivenessAnalysis, BitVector,
3131

3232
public:
3333
LivenessAnalysis(const RegAnalysis &RA, BinaryFunction &BF,
34-
MCPlusBuilder::AllocatorIdTy AllocId)
35-
: Parent(BF, AllocId), RA(RA),
34+
const bool UsePrivateAllocators)
35+
: Parent(BF, UsePrivateAllocators), RA(RA),
3636
NumRegs(BF.getBinaryContext().MRI->getNumRegs()) {}
3737
virtual ~LivenessAnalysis();
3838

‎bolt/include/bolt/Passes/ReachingDefOrUse.h

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -34,9 +34,10 @@ class ReachingDefOrUse
3434

3535
public:
3636
ReachingDefOrUse(const RegAnalysis &RA, BinaryFunction &BF,
37-
std::optional<MCPhysReg> TrackingReg = std::nullopt,
38-
MCPlusBuilder::AllocatorIdTy AllocId = 0)
39-
: InstrsDataflowAnalysis<ReachingDefOrUse<Def>, !Def>(BF, AllocId),
37+
const bool UsePrivateAllocators,
38+
std::optional<MCPhysReg> TrackingReg = std::nullopt)
39+
: InstrsDataflowAnalysis<ReachingDefOrUse<Def>, !Def>(
40+
BF, UsePrivateAllocators),
4041
RA(RA), TrackingReg(TrackingReg) {}
4142
virtual ~ReachingDefOrUse() {}
4243

‎bolt/include/bolt/Passes/ReachingInsns.h

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,9 @@ class ReachingInsns
2626
friend class DataflowAnalysis<ReachingInsns<Backward>, BitVector, Backward>;
2727

2828
public:
29-
ReachingInsns(BinaryFunction &BF, MCPlusBuilder::AllocatorIdTy AllocId = 0)
30-
: InstrsDataflowAnalysis<ReachingInsns, Backward>(BF, AllocId) {}
29+
ReachingInsns(BinaryFunction &BF, const bool UsePrivateAllocators)
30+
: InstrsDataflowAnalysis<ReachingInsns, Backward>(BF,
31+
UsePrivateAllocators) {}
3132
virtual ~ReachingInsns() {}
3233

3334
bool isInLoop(const BinaryBasicBlock &BB) {

‎bolt/include/bolt/Passes/StackAllocationAnalysis.h

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,9 @@ class StackAllocationAnalysis
3131

3232
public:
3333
StackAllocationAnalysis(BinaryFunction &BF, StackPointerTracking &SPT,
34-
MCPlusBuilder::AllocatorIdTy AllocId)
35-
: InstrsDataflowAnalysis<StackAllocationAnalysis, false>(BF, AllocId),
34+
const bool UsePrivateAllocators)
35+
: InstrsDataflowAnalysis<StackAllocationAnalysis, false>(
36+
BF, UsePrivateAllocators),
3637
SPT(SPT) {}
3738
virtual ~StackAllocationAnalysis() {}
3839

‎bolt/include/bolt/Passes/StackAvailableExpressions.h

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,8 @@ class StackAvailableExpressions
2828

2929
public:
3030
StackAvailableExpressions(const RegAnalysis &RA, const FrameAnalysis &FA,
31-
BinaryFunction &BF);
31+
BinaryFunction &BF,
32+
const bool UsePrivateAllocators);
3233
virtual ~StackAvailableExpressions() {}
3334

3435
void run() { InstrsDataflowAnalysis<StackAvailableExpressions>::run(); }

‎bolt/include/bolt/Passes/StackPointerTracking.h

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -175,9 +175,12 @@ class StackPointerTrackingBase
175175
}
176176

177177
public:
178+
StackPointerTrackingBase(BinaryFunction &BF, const bool UsePrivateAllocators)
179+
: DataflowAnalysis<Derived, std::pair<int, int>>(BF,
180+
UsePrivateAllocators) {}
178181
StackPointerTrackingBase(BinaryFunction &BF,
179-
MCPlusBuilder::AllocatorIdTy AllocatorId = 0)
180-
: DataflowAnalysis<Derived, std::pair<int, int>>(BF, AllocatorId) {}
182+
MCPlusBuilder::AllocatorIdTy AllocId)
183+
: DataflowAnalysis<Derived, std::pair<int, int>>(BF, AllocId) {}
181184

182185
virtual ~StackPointerTrackingBase() {}
183186

@@ -192,8 +195,9 @@ class StackPointerTracking
192195
friend class DataflowAnalysis<StackPointerTracking, std::pair<int, int>>;
193196

194197
public:
198+
StackPointerTracking(BinaryFunction &BF, const bool UsePrivateAllocators);
195199
StackPointerTracking(BinaryFunction &BF,
196-
MCPlusBuilder::AllocatorIdTy AllocatorId = 0);
200+
MCPlusBuilder::AllocatorIdTy AllocId);
197201
virtual ~StackPointerTracking() {}
198202

199203
void run() { StackPointerTrackingBase<StackPointerTracking>::run(); }

‎bolt/include/bolt/Passes/StackReachingUses.h

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,8 @@ class StackReachingUses
2828

2929
public:
3030
StackReachingUses(const FrameAnalysis &FA, BinaryFunction &BF,
31-
MCPlusBuilder::AllocatorIdTy AllocId = 0)
32-
: InstrsDataflowAnalysis(BF, AllocId), FA(FA) {}
31+
const bool UsePrivateAllocators)
32+
: InstrsDataflowAnalysis(BF, UsePrivateAllocators), FA(FA) {}
3333
virtual ~StackReachingUses() {}
3434

3535
/// Return true if the stack position written by the store in \p StoreFIE was

‎bolt/lib/Core/ParallelUtilities.cpp

Lines changed: 15 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -187,20 +187,11 @@ void runOnEachFunctionWithUniqueAllocId(
187187
LLVM_DEBUG(T.stopTimer());
188188
};
189189

190-
unsigned AllocId = 1;
191-
auto EnsureAllocatorExists = [&BC](unsigned AllocId) {
192-
if (!BC.MIB->checkAllocatorExists(AllocId)) {
193-
MCPlusBuilder::AllocatorIdTy Id =
194-
BC.MIB->initializeNewAnnotationAllocator();
195-
(void)Id;
196-
assert(AllocId == Id && "unexpected allocator id created");
197-
}
198-
};
199-
200190
if (opts::NoThreads || ForceSequential) {
201-
EnsureAllocatorExists(AllocId);
191+
std::vector<MCPlusBuilder::AllocatorIdTy> AvailableSharedIds =
192+
BC.MIB->getSharedAllocIds(1);
202193
runBlock(BC.getBinaryFunctions().begin(), BC.getBinaryFunctions().end(),
203-
AllocId);
194+
AvailableSharedIds[0]);
204195
return;
205196
}
206197
// This lock is used to postpone task execution
@@ -213,7 +204,8 @@ void runOnEachFunctionWithUniqueAllocId(
213204
TotalCost > BlocksCount ? TotalCost / BlocksCount : 1;
214205

215206
// Divide work into blocks of equal cost
216-
ThreadPool &Pool = getThreadPool();
207+
std::vector<std::map<uint64_t, BinaryFunction>::iterator>
208+
BinaryFunctionBlocks;
217209
auto BlockBegin = BC.getBinaryFunctions().begin();
218210
unsigned CurrentCost = 0;
219211
for (auto It = BC.getBinaryFunctions().begin();
@@ -222,17 +214,21 @@ void runOnEachFunctionWithUniqueAllocId(
222214
CurrentCost += computeCostFor(BF, SkipPredicate, SchedPolicy);
223215

224216
if (CurrentCost >= BlockCost) {
225-
EnsureAllocatorExists(AllocId);
226-
Pool.async(runBlock, BlockBegin, std::next(It), AllocId);
227-
AllocId++;
228-
BlockBegin = std::next(It);
217+
BinaryFunctionBlocks.push_back(BlockBegin);
229218
CurrentCost = 0;
230219
}
231220
}
221+
BinaryFunctionBlocks.push_back(BlockBegin);
222+
BinaryFunctionBlocks.push_back(BC.getBinaryFunctions().end());
223+
224+
std::vector<MCPlusBuilder::AllocatorIdTy> AvailableSharedIds =
225+
BC.MIB->getSharedAllocIds(BinaryFunctionBlocks.size()-1);
232226

233-
EnsureAllocatorExists(AllocId);
227+
ThreadPool &Pool = getThreadPool();
228+
for (unsigned I = 0; I < BinaryFunctionBlocks.size() - 1; ++I)
229+
Pool.async(runBlock, BinaryFunctionBlocks[I], BinaryFunctionBlocks[I + 1],
230+
AvailableSharedIds[I]);
234231

235-
Pool.async(runBlock, BlockBegin, BC.getBinaryFunctions().end(), AllocId);
236232
Lock.unlock();
237233
Pool.wait();
238234
}

0 commit comments

Comments
 (0)