Skip to content

Commit 7ae084c

Browse files
bgergely0bgergely0
committed
[BOLT][BTI] Add MCPlusBuilder::addBTItoBBStart
This function contains most of the logic for BTI: - it takes the BasicBlock and the instruction used to jump to it. - then it checks if the first non-pseudo instruction is a sufficient landing pad for the used call. - if not, it generates the correct BTI instruction. Also introduce the isBTIVariantCoveringCall helper to simplify the logic.
1 parent c668bf8 commit 7ae084c

File tree

3 files changed

+193
-0
lines changed

3 files changed

+193
-0
lines changed

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1888,6 +1888,19 @@ class MCPlusBuilder {
18881888
llvm_unreachable("not implemented");
18891889
}
18901890

1891+
/// Checks if the indirect call / jump is accepted by the landing pad at the
1892+
/// start of the target BasicBlock.
1893+
virtual bool isBTIVariantCoveringCall(MCInst &Call, MCInst &Pad) const {
1894+
llvm_unreachable("not implemented");
1895+
return false;
1896+
}
1897+
1898+
/// Adds a BTI landing pad to the start of the BB, that matches the indirect
1899+
/// call/jump inst.
1900+
virtual void addBTItoBBStart(BinaryBasicBlock &BB, MCInst &Call) const {
1901+
llvm_unreachable("not implemented");
1902+
}
1903+
18911904
/// Store \p Target absolute address to \p RegName
18921905
virtual InstructionListType materializeAddress(const MCSymbol *Target,
18931906
MCContext *Ctx,

‎bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2782,6 +2782,81 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
27822782
Inst.addOperand(MCOperand::createImm(HintNum));
27832783
}
27842784

2785+
bool isBTIVariantCoveringCall(MCInst &Call, MCInst &Pad) const override {
2786+
assert((isIndirectCall(Call) || isIndirectBranch(Call)) &&
2787+
"Not an indirect call or branch.");
2788+
2789+
// A BLR can be accepted by a BTI c.
2790+
if (isIndirectCall(Call))
2791+
return isBTILandingPad(Pad, true, false) ||
2792+
isBTILandingPad(Pad, true, true);
2793+
2794+
// A BR can be accepted by a BTI j or BTI c (and BTI jc) IF the operand is
2795+
// x16 or x17. If the operand is not x16 or x17, it can be accepted by a BTI
2796+
// j or BTI jc (and not BTI c).
2797+
if (isIndirectBranch(Call)) {
2798+
assert(Call.getNumOperands() == 1 &&
2799+
"Indirect branch needs to have 1 operand.");
2800+
assert(Call.getOperand(0).isReg() &&
2801+
"Indirect branch does not have a register operand.");
2802+
MCPhysReg Reg = Call.getOperand(0).getReg();
2803+
if (Reg == AArch64::X16 || Reg == AArch64::X17)
2804+
return isBTILandingPad(Pad, true, false) ||
2805+
isBTILandingPad(Pad, false, true) ||
2806+
isBTILandingPad(Pad, true, true);
2807+
return isBTILandingPad(Pad, false, true) ||
2808+
isBTILandingPad(Pad, true, true);
2809+
}
2810+
return false;
2811+
}
2812+
2813+
void addBTItoBBStart(BinaryBasicBlock &BB, MCInst &Call) const override {
2814+
auto II = BB.getFirstNonPseudo();
2815+
if (II != BB.end()) {
2816+
if (isBTIVariantCoveringCall(Call, *II))
2817+
return;
2818+
// A BLR can be accepted by a BTI c.
2819+
if (isIndirectCall(Call)) {
2820+
// if we have a BTI j at the start, extend it to a BTI jc,
2821+
// otherwise insert a new BTI c.
2822+
if (isBTILandingPad(*II, false, true)) {
2823+
updateBTIVariant(*II, true, true);
2824+
} else {
2825+
MCInst BTIInst;
2826+
createBTI(BTIInst, true, false);
2827+
BB.insertInstruction(II, BTIInst);
2828+
}
2829+
}
2830+
2831+
// A BR can be accepted by a BTI j or BTI c (and BTI jc) IF the operand is
2832+
// x16 or x17. If the operand is not x16 or x17, it can be accepted by a
2833+
// BTI j or BTI jc (and not BTI c).
2834+
if (isIndirectBranch(Call)) {
2835+
assert(Call.getNumOperands() == 1 &&
2836+
"Indirect branch needs to have 1 operand.");
2837+
assert(Call.getOperand(0).isReg() &&
2838+
"Indirect branch does not have a register operand.");
2839+
MCPhysReg Reg = Call.getOperand(0).getReg();
2840+
if (Reg == AArch64::X16 || Reg == AArch64::X17) {
2841+
// Add a new BTI c
2842+
MCInst BTIInst;
2843+
createBTI(BTIInst, true, false);
2844+
BB.insertInstruction(II, BTIInst);
2845+
} else {
2846+
// If BB starts with a BTI c, extend it to BTI jc,
2847+
// otherwise insert a new BTI j.
2848+
if (isBTILandingPad(*II, true, false)) {
2849+
updateBTIVariant(*II, true, true);
2850+
} else {
2851+
MCInst BTIInst;
2852+
createBTI(BTIInst, false, true);
2853+
BB.insertInstruction(II, BTIInst);
2854+
}
2855+
}
2856+
}
2857+
}
2858+
}
2859+
27852860
InstructionListType materializeAddress(const MCSymbol *Target, MCContext *Ctx,
27862861
MCPhysReg RegName,
27872862
int64_t Addend = 0) const override {

‎bolt/unittests/Core/MCPlusBuilder.cpp‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -196,6 +196,111 @@ TEST_P(MCPlusBuilderTester, AArch64_BTI) {
196196
ASSERT_TRUE(BC->MIB->isImplicitBTIC(*II));
197197
}
198198

199+
TEST_P(MCPlusBuilderTester, AArch64_addBTItoBBStart_0) {
200+
if (GetParam() != Triple::aarch64)
201+
GTEST_SKIP();
202+
BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
203+
std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock();
204+
MCInst Inst = MCInstBuilder(AArch64::RET).addReg(AArch64::LR);
205+
BB->addInstruction(Inst);
206+
// BR x16 needs BTI c or BTI j. We prefer adding a BTI c.
207+
MCInst CallInst = MCInstBuilder(AArch64::BR).addReg(AArch64::X16);
208+
BC->MIB->addBTItoBBStart(*BB, CallInst);
209+
auto II = BB->begin();
210+
ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, true, false));
211+
}
212+
213+
TEST_P(MCPlusBuilderTester, AArch64_addBTItoBBStart_1) {
214+
if (GetParam() != Triple::aarch64)
215+
GTEST_SKIP();
216+
BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
217+
std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock();
218+
MCInst BTIc;
219+
BC->MIB->createBTI(BTIc, true, false);
220+
BB->addInstruction(BTIc);
221+
// BR x16 needs BTI c or BTI j. We have a BTI c, no change is needed.
222+
MCInst CallInst = MCInstBuilder(AArch64::BR).addReg(AArch64::X16);
223+
BC->MIB->addBTItoBBStart(*BB, CallInst);
224+
auto II = BB->begin();
225+
ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, true, false));
226+
}
227+
228+
TEST_P(MCPlusBuilderTester, AArch64_addBTItoBBStart_2) {
229+
if (GetParam() != Triple::aarch64)
230+
GTEST_SKIP();
231+
BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
232+
std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock();
233+
MCInst BTIc;
234+
BC->MIB->createBTI(BTIc, true, false);
235+
BB->addInstruction(BTIc);
236+
// BR x5 needs BTI j
237+
// we have BTI c -> extend it to BTI jc.
238+
MCInst CallInst = MCInstBuilder(AArch64::BR).addReg(AArch64::X5);
239+
BC->MIB->addBTItoBBStart(*BB, CallInst);
240+
auto II = BB->begin();
241+
ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, true, true));
242+
}
243+
244+
TEST_P(MCPlusBuilderTester, AArch64_addBTItoBBStart_3) {
245+
if (GetParam() != Triple::aarch64)
246+
GTEST_SKIP();
247+
BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
248+
std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock();
249+
MCInst Inst = MCInstBuilder(AArch64::RET).addReg(AArch64::LR);
250+
BB->addInstruction(Inst);
251+
// BR x5 needs BTI j
252+
MCInst CallInst = MCInstBuilder(AArch64::BR).addReg(AArch64::X5);
253+
BC->MIB->addBTItoBBStart(*BB, CallInst);
254+
auto II = BB->begin();
255+
ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, false, true));
256+
}
257+
258+
TEST_P(MCPlusBuilderTester, AArch64_addBTItoBBStart_4) {
259+
if (GetParam() != Triple::aarch64)
260+
GTEST_SKIP();
261+
BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
262+
std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock();
263+
MCInst Inst = MCInstBuilder(AArch64::RET).addReg(AArch64::LR);
264+
BB->addInstruction(Inst);
265+
// BLR needs BTI c, regardless of the register used.
266+
MCInst CallInst = MCInstBuilder(AArch64::BLR).addReg(AArch64::X5);
267+
BC->MIB->addBTItoBBStart(*BB, CallInst);
268+
auto II = BB->begin();
269+
ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, true, false));
270+
}
271+
272+
TEST_P(MCPlusBuilderTester, AArch64_addBTItoBBStart_5) {
273+
if (GetParam() != Triple::aarch64)
274+
GTEST_SKIP();
275+
BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
276+
std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock();
277+
MCInst BTIj;
278+
BC->MIB->createBTI(BTIj, false, true);
279+
BB->addInstruction(BTIj);
280+
// BLR needs BTI c, regardless of the register used.
281+
// We have a BTI j -> extend it to BTI jc.
282+
MCInst CallInst = MCInstBuilder(AArch64::BLR).addReg(AArch64::X5);
283+
BC->MIB->addBTItoBBStart(*BB, CallInst);
284+
auto II = BB->begin();
285+
ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, true, true));
286+
}
287+
288+
TEST_P(MCPlusBuilderTester, AArch64_addBTItoBBStart_6) {
289+
if (GetParam() != Triple::aarch64)
290+
GTEST_SKIP();
291+
BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
292+
std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock();
293+
MCInst Paciasp =
294+
MCInstBuilder(AArch64::PACIASP).addReg(AArch64::LR).addReg(AArch64::SP);
295+
BB->addInstruction(Paciasp);
296+
// PACI(AB)SP are implicit BTI c, no change needed.
297+
MCInst CallInst = MCInstBuilder(AArch64::BR).addReg(AArch64::X17);
298+
BC->MIB->addBTItoBBStart(*BB, CallInst);
299+
auto II = BB->begin();
300+
ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, true, false));
301+
ASSERT_TRUE(BC->MIB->isPSignOnLR(*II));
302+
}
303+
199304
TEST_P(MCPlusBuilderTester, AArch64_CmpJNE) {
200305
if (GetParam() != Triple::aarch64)
201306
GTEST_SKIP();

0 commit comments

Comments
 (0)