Skip to content

Commit 59abd9c

Browse files
kbeylskbeyls
committed
[GS/StackClash] Recognize mov sp, x21-like constructs originating from allocas or similar in loops.
1 parent a3c810a commit 59abd9c

File tree

3 files changed

+147
-63
lines changed

3 files changed

+147
-63
lines changed

‎bolt/lib/Passes/DataflowAnalysis.cpp

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,8 +49,14 @@ void doForAllPreds(const BinaryBasicBlock &BB,
4949
std::function<void(ProgramPoint)> Task) {
5050
MCPlusBuilder *MIB = BB.getFunction()->getBinaryContext().MIB.get();
5151
for (BinaryBasicBlock *Pred : BB.predecessors()) {
52-
if (Pred->isValid())
52+
if (Pred->isValid()) {
53+
LLVM_DEBUG({
54+
dbgs() << "\n DoForAllPreds calling Task on ("
55+
<< BB.getName() << ", " << Pred->getName()
56+
<< ")" << "\n";
57+
});
5358
Task(ProgramPoint::getLastPointAt(*Pred));
59+
}
5460
}
5561
if (!BB.isLandingPad())
5662
return;
@@ -70,8 +76,14 @@ void doForAllPreds(const BinaryBasicBlock &BB,
7076
void doForAllSuccs(const BinaryBasicBlock &BB,
7177
std::function<void(ProgramPoint)> Task) {
7278
for (BinaryBasicBlock *Succ : BB.successors())
73-
if (Succ->isValid())
79+
if (Succ->isValid()) {
80+
LLVM_DEBUG({
81+
dbgs() << "\n DoForAllSuccs calling Task on ("
82+
<< BB.getName() << ", " << Succ->getName()
83+
<< ")" << "\n";
84+
});
7485
Task(ProgramPoint::getFirstPointAt(*Succ));
86+
}
7587
}
7688

7789
void RegStatePrinter::print(raw_ostream &OS, const BitVector &State) const {

‎bolt/lib/Passes/StackClashAnalysis.cpp

Lines changed: 90 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -78,14 +78,21 @@ struct State {
7878
/// RegMaxValues stores registers that we know have a value in the
7979
/// range [0, MaxValue-1].
8080
SmallDenseMap<MCPhysReg, uint64_t, 1> RegMaxValues;
81-
/// RegsFixedOffsetFromOrigSP contains the registers that contain the value
82-
/// of the original SP at the function start, or a fixed offset from it.
81+
/// Reg2MaxOffset contains the registers that contain the value
82+
/// of SP at some point during the running function, where it's
83+
/// guaranteed that at the time the SP value was stored in the register,
84+
/// a maximum offset for any probe into the stack is a constant.
85+
/// That constant is stored in this map.
86+
///
8387
/// This is especially useful to recognize frame pointers and the fact
8488
/// that epilogues can restore stack pointers from frame pointer values.
8589
/// This is only tracked in Basic Blocks that are known to be reachable
8690
/// from an entry block. For blocks not (yet) known to be reachable from
8791
/// an entry block, the optional does not contain a value.
88-
std::optional<SmallSet<MCPhysReg, 2>> RegsFixedOffsetFromOrigSP;
92+
std::optional<SmallDenseMap<MCPhysReg, int64_t, 2>> Reg2MaxOffset;
93+
// FIXME: It seems that conceptually it does not make sense to
94+
// track wheterh the SP value is currently at a fixed offset from
95+
// the value it was at function entry.
8996
/// SPFixedOffsetFromOrig indicates whether the current SP value is
9097
/// a constant fixed offset from the SP value at the function start.
9198
std::optional<int64_t> SPFixedOffsetFromOrig;
@@ -131,16 +138,17 @@ struct State {
131138
else if (*SPFixedOffsetFromOrig != *StateIn.SPFixedOffsetFromOrig)
132139
SPFixedOffsetFromOrig.reset();
133140

134-
if (StateIn.RegsFixedOffsetFromOrigSP && RegsFixedOffsetFromOrigSP) {
135-
SmallVector<MCPhysReg, 2> RegsToRemove;
136-
for (MCPhysReg R : *RegsFixedOffsetFromOrigSP)
137-
if (!StateIn.RegsFixedOffsetFromOrigSP->contains(R))
138-
RegsToRemove.push_back(R);
139-
for (MCPhysReg R : RegsToRemove)
140-
RegsFixedOffsetFromOrigSP->erase(R);
141-
} else if (StateIn.RegsFixedOffsetFromOrigSP &&
142-
!RegsFixedOffsetFromOrigSP) {
143-
RegsFixedOffsetFromOrigSP = StateIn.RegsFixedOffsetFromOrigSP;
141+
if (StateIn.Reg2MaxOffset && Reg2MaxOffset) {
142+
SmallDenseMap<MCPhysReg, int64_t, 2> MergedMap;
143+
for (auto R2MaxOff : *Reg2MaxOffset) {
144+
const MCPhysReg R = R2MaxOff.first;
145+
if (auto SIn_R2MaxOff = StateIn.Reg2MaxOffset->find(R);
146+
SIn_R2MaxOff != StateIn.Reg2MaxOffset->end())
147+
MergedMap[R] = std::max(R2MaxOff.second, SIn_R2MaxOff->second);
148+
}
149+
Reg2MaxOffset = MergedMap;
150+
} else if (StateIn.Reg2MaxOffset && !Reg2MaxOffset) {
151+
Reg2MaxOffset = StateIn.Reg2MaxOffset;
144152
}
145153

146154
for (auto I : StateIn.LastStackGrowingInsts)
@@ -152,7 +160,7 @@ struct State {
152160
RegConstValues == RHS.RegConstValues &&
153161
RegMaxValues == RHS.RegMaxValues &&
154162
SPFixedOffsetFromOrig == RHS.SPFixedOffsetFromOrig &&
155-
RegsFixedOffsetFromOrigSP == RHS.RegsFixedOffsetFromOrigSP;
163+
Reg2MaxOffset == RHS.Reg2MaxOffset;
156164
}
157165
bool operator!=(const State &RHS) const { return !((*this) == RHS); }
158166
};
@@ -168,8 +176,8 @@ void print_reg(raw_ostream &OS, MCPhysReg Reg, const BinaryContext *BC) {
168176
}
169177
}
170178

171-
void PrintRegMap(raw_ostream &OS,
172-
const SmallDenseMap<MCPhysReg, uint64_t, 1> &M,
179+
template <class T, unsigned N>
180+
void PrintRegMap(raw_ostream &OS, const SmallDenseMap<MCPhysReg, T, N> &M,
173181
const BinaryContext *BC = nullptr) {
174182
for (auto Reg2Value : M) {
175183
print_reg(OS, Reg2Value.first, BC);
@@ -185,18 +193,15 @@ raw_ostream &print_state(raw_ostream &OS, const State &S,
185193
else
186194
OS << *(S.MaxOffsetSinceLastProbe);
187195
OS << "), RegConstValues(";
188-
PrintRegMap(OS, S.RegConstValues);
196+
PrintRegMap(OS, S.RegConstValues, BC);
189197
OS << "), RegMaxValues(";
190-
PrintRegMap(OS, S.RegMaxValues);
198+
PrintRegMap(OS, S.RegMaxValues, BC);
191199
OS << "),";
192200
OS << "SPFixedOffsetFromOrig:" << S.SPFixedOffsetFromOrig << ",";
193-
OS << "RegsFixedOffsetFromOrigSP:";
194-
if (S.RegsFixedOffsetFromOrigSP) {
201+
OS << "Reg2MaxOffset:";
202+
if (S.Reg2MaxOffset) {
195203
OS << "(";
196-
for (auto Reg : *S.RegsFixedOffsetFromOrigSP) {
197-
print_reg(OS, Reg, BC);
198-
OS << ",";
199-
}
204+
PrintRegMap(OS, *S.Reg2MaxOffset, BC);
200205
OS << ")";
201206
} else
202207
OS << "None";
@@ -243,7 +248,7 @@ class StackClashDFAnalysis
243248
State getStartingStateAtBB(const BinaryBasicBlock &BB) {
244249
State Next;
245250
if (BB.isEntryPoint())
246-
Next.RegsFixedOffsetFromOrigSP = SmallSet<MCPhysReg, 2>();
251+
Next.Reg2MaxOffset = SmallDenseMap<MCPhysReg, int64_t, 2>();
247252
return Next;
248253
}
249254

@@ -351,16 +356,29 @@ class StackClashDFAnalysis
351356
MCPhysReg FixedOffsetRegJustSet = BC.MIB->getNoRegister();
352357
if (auto OC = BC.MIB->getOffsetChange(Point, Cur.RegConstValues,
353358
Cur.RegMaxValues))
354-
if (Next.RegsFixedOffsetFromOrigSP)
355-
if (OC.FromReg == SP ||
356-
Cur.RegsFixedOffsetFromOrigSP->contains(OC.FromReg)) {
357-
Next.RegsFixedOffsetFromOrigSP->insert(OC.ToReg);
359+
if (Next.Reg2MaxOffset) {
360+
if (OC.FromReg == SP) {
361+
(*Next.Reg2MaxOffset)[OC.ToReg] = *Cur.MaxOffsetSinceLastProbe;
358362
FixedOffsetRegJustSet = OC.ToReg;
363+
} else if (auto I = Cur.Reg2MaxOffset->find(OC.FromReg);
364+
I != Cur.Reg2MaxOffset->end()) {
365+
(*Next.Reg2MaxOffset)[OC.ToReg] = (*I).second;
366+
FixedOffsetRegJustSet = OC.ToReg;
367+
}
368+
}
369+
if (Next.Reg2MaxOffset)
370+
for (const MCOperand &Operand : BC.MIB->defOperands(Point)) {
371+
if (Operand.getReg() != FixedOffsetRegJustSet) {
372+
Next.Reg2MaxOffset->erase(Operand.getReg());
373+
LLVM_DEBUG({
374+
dbgs() << " - Removed Reg " << Operand.getReg()
375+
<< " from Next.Reg2MaxOffset"
376+
<< ". On instruction:";
377+
BC.printInstruction(dbgs(), Point);
378+
dbgs() << "\n";
379+
});
359380
}
360-
if (Next.RegsFixedOffsetFromOrigSP)
361-
for (const MCOperand &Operand : BC.MIB->defOperands(Point))
362-
if (Operand.getReg() != FixedOffsetRegJustSet)
363-
Next.RegsFixedOffsetFromOrigSP->erase(Operand.getReg());
381+
}
364382

365383
if (BC.MIB->hasDefOfPhysReg(Point, SP)) {
366384
// Next, validate that we can track by how much the SP
@@ -371,37 +389,46 @@ class StackClashDFAnalysis
371389
Next.LastStackGrowingInsts.insert(MCInstInBBReference::get(&Point, BF));
372390
if (auto OC = BC.MIB->getOffsetChange(Point, Cur.RegConstValues,
373391
Cur.RegMaxValues);
374-
OC && OC.ToReg == SP && OC.FromReg == SP) {
375-
assert(OC.MaxOffsetChange);
376-
if (*OC.MaxOffsetChange < 0)
377-
Next.MaxOffsetSinceLastProbe =
378-
*Next.MaxOffsetSinceLastProbe - *OC.MaxOffsetChange;
379-
if (OC.OffsetChange && Next.SPFixedOffsetFromOrig)
380-
Next.SPFixedOffsetFromOrig =
381-
*Next.SPFixedOffsetFromOrig + *OC.OffsetChange;
382-
// FIXME: add test case for this if test.
392+
OC && OC.ToReg == SP) {
393+
if (OC.FromReg == SP) {
394+
assert(OC.MaxOffsetChange);
395+
if (*OC.MaxOffsetChange < 0)
396+
Next.MaxOffsetSinceLastProbe =
397+
*Next.MaxOffsetSinceLastProbe - *OC.MaxOffsetChange;
398+
if (OC.OffsetChange && Next.SPFixedOffsetFromOrig)
399+
Next.SPFixedOffsetFromOrig =
400+
*Next.SPFixedOffsetFromOrig + *OC.OffsetChange;
401+
// FIXME: add test case for this if test.
383402
#if 0
384403
if (IsPreIndexOffsetChange)
385404
Next.MaxOffsetSinceLastProbe =
386405
*Next.MaxOffsetSinceLastProbe - StackAccessOffset;
387406
#endif
388-
LLVM_DEBUG({
389-
dbgs() << " Found SP Offset change: ";
390-
BC.printInstruction(dbgs(), Point);
391-
dbgs() << " OffsetChange: " << OC.OffsetChange
392-
<< "; MaxOffsetChange: " << OC.MaxOffsetChange
393-
<< "; new MaxOffsetSinceLastProbe: "
394-
<< Next.MaxOffsetSinceLastProbe
395-
<< "; new SPFixedOffsetFromOrig: "
396-
<< Next.SPFixedOffsetFromOrig
397-
<< "; IsStackAccess:" << IsStackAccess
398-
<< "; StackAccessOffset: " << StackAccessOffset << "\n";
399-
});
400-
assert(!OC.IsPreIndexOffsetChange || IsStackAccess);
401-
assert(*Next.MaxOffsetSinceLastProbe >= 0);
407+
LLVM_DEBUG({
408+
dbgs() << " Found SP Offset change: ";
409+
BC.printInstruction(dbgs(), Point);
410+
dbgs() << " OffsetChange: " << OC.OffsetChange
411+
<< "; MaxOffsetChange: " << OC.MaxOffsetChange
412+
<< "; new MaxOffsetSinceLastProbe: "
413+
<< Next.MaxOffsetSinceLastProbe
414+
<< "; new SPFixedOffsetFromOrig: "
415+
<< Next.SPFixedOffsetFromOrig
416+
<< "; IsStackAccess:" << IsStackAccess
417+
<< "; StackAccessOffset: " << StackAccessOffset << "\n";
418+
});
419+
assert(!OC.IsPreIndexOffsetChange || IsStackAccess);
420+
assert(*Next.MaxOffsetSinceLastProbe >= 0);
421+
} else if (
422+
Cur.Reg2MaxOffset && Cur.Reg2MaxOffset->contains(OC.FromReg) &&
423+
OC.OffsetChange ==
424+
0 /* FIXME: also handle other offset changes if needed. */) {
425+
const int64_t MaxOffset = Cur.Reg2MaxOffset->find(OC.FromReg)->second;
426+
Next.MaxOffsetSinceLastProbe = MaxOffset;
427+
}
402428
} else {
403429
Next.MaxOffsetSinceLastProbe.reset();
404-
Next.SPFixedOffsetFromOrig.reset();
430+
Next.SPFixedOffsetFromOrig
431+
.reset(); // FIXME - should I make this the empty set?
405432
LLVM_DEBUG({
406433
dbgs() << " Found non-const SP Offset change: ";
407434
BC.printInstruction(dbgs(), Point);
@@ -480,10 +507,12 @@ void StackClashAnalysis::runOnFunction(
480507
// function into the SP, then we know this is probably fine.
481508
// Typical case is moving the frame pointer to the stack pointer
482509
// in the function epilogue.
483-
bool IsFixedRegToSPMove =
484-
OC && OC.ToReg == SP && S.RegsFixedOffsetFromOrigSP &&
485-
S.RegsFixedOffsetFromOrigSP->contains(OC.FromReg);
486-
if (!IsFixedRegToSPMove) {
510+
bool IsFixedRegToSPMove = OC && OC.ToReg == SP && S.Reg2MaxOffset &&
511+
S.Reg2MaxOffset->contains(OC.FromReg);
512+
if (IsFixedRegToSPMove) {
513+
S.MaxOffsetSinceLastProbe =
514+
S.Reg2MaxOffset->find(OC.FromReg)->second;
515+
} else {
487516
// mark to report that this may be an SP change that is not a
488517
// constant amount.
489518
LLVM_DEBUG({

‎bolt/test/gadget-scanner/gs-stackclash.s

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -153,6 +153,49 @@ f_recognize_fp_deadcode:
153153

154154
// CHECK-NOT: GS-STACKCLASH:
155155

156+
.global f_alloca_in_loop
157+
.type f_alloca_in_loop, %function
158+
f_alloca_in_loop:
159+
stp x29, x30, [sp, -32]!
160+
mov x29, sp
161+
stp x19, x20, [sp, 16]
162+
mov x19, x0
163+
sub sp, sp, #16
164+
.L6:
165+
add x1, x19, 15
166+
mov x20, sp
167+
and x2, x1, -65536
168+
and x1, x1, -16
169+
sub x2, sp, x2
170+
cmp sp, x2
171+
beq .L4
172+
.L16:
173+
sub sp, sp, #65536
174+
str xzr, [sp, 1024]
175+
cmp sp, x2
176+
bne .L16
177+
.L4:
178+
and x1, x1, 65535
179+
sub sp, sp, x1
180+
str xzr, [sp]
181+
cmp x1, 1024
182+
bcc .L5
183+
str xzr, [sp, 1024]
184+
.L5:
185+
add x0, sp, 16
186+
bl g
187+
subs x19, x19, #1
188+
mov sp, x20 // There should not be a warning here.
189+
bne .L6
190+
mov sp, x29
191+
ldp x19, x20, [sp, 16]
192+
ldp x29, x30, [sp], 32
193+
ret
194+
.size f_alloca_in_loop, .-f_alloca_in_loop
195+
196+
// CHECK-NOT: GS-STACKCLASH:
197+
198+
// TODO: also add a test case with nested alloca loops.
156199

157200
// TODO: check access with constant offset from SP through other
158201
// register:

0 commit comments

Comments
 (0)