Skip to content

Commit fbc121c

Browse files
bgergely0bgergely0
authored
[BOLT][BTI] Add MCPlusBuilder::insertBTI (#167329)
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 isCallCoveredByBTI helper to simplify the logic.
1 parent 8af88a4 commit fbc121c

File tree

3 files changed

+204
-0
lines changed

3 files changed

+204
-0
lines changed

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1894,6 +1894,19 @@ class MCPlusBuilder {
18941894
llvm_unreachable("not implemented");
18951895
}
18961896

1897+
/// Checks if the indirect call / jump is accepted by the landing pad at the
1898+
/// start of the target BasicBlock.
1899+
virtual bool isCallCoveredByBTI(MCInst &Call, MCInst &Pad) const {
1900+
llvm_unreachable("not implemented");
1901+
return false;
1902+
}
1903+
1904+
/// Inserts a BTI landing pad to the start of the BB, that matches the
1905+
/// indirect call inst used to call the BB.
1906+
virtual void insertBTI(BinaryBasicBlock &BB, MCInst &Call) const {
1907+
llvm_unreachable("not implemented");
1908+
}
1909+
18971910
/// Store \p Target absolute address to \p RegName
18981911
virtual InstructionListType materializeAddress(const MCSymbol *Target,
18991912
MCContext *Ctx,

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

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

2809+
bool isCallCoveredByBTI(MCInst &Call, MCInst &Pad) const override {
2810+
assert((isIndirectCall(Call) || isIndirectBranch(Call)) &&
2811+
"Not an indirect call or branch.");
2812+
2813+
// A BLR can be accepted by a BTI c.
2814+
if (isIndirectCall(Call))
2815+
return isBTILandingPad(Pad, true, false) ||
2816+
isBTILandingPad(Pad, true, true);
2817+
2818+
// A BR can be accepted by a BTI j or BTI c (and BTI jc) IF the operand is
2819+
// x16 or x17. If the operand is not x16 or x17, it can be accepted by a BTI
2820+
// j or BTI jc (and not BTI c).
2821+
if (isIndirectBranch(Call)) {
2822+
assert(Call.getNumOperands() == 1 &&
2823+
"Indirect branch needs to have 1 operand.");
2824+
assert(Call.getOperand(0).isReg() &&
2825+
"Indirect branch does not have a register operand.");
2826+
MCPhysReg Reg = Call.getOperand(0).getReg();
2827+
if (Reg == AArch64::X16 || Reg == AArch64::X17)
2828+
return isBTILandingPad(Pad, true, false) ||
2829+
isBTILandingPad(Pad, false, true) ||
2830+
isBTILandingPad(Pad, true, true);
2831+
return isBTILandingPad(Pad, false, true) ||
2832+
isBTILandingPad(Pad, true, true);
2833+
}
2834+
return false;
2835+
}
2836+
2837+
void insertBTI(BinaryBasicBlock &BB, MCInst &Call) const override {
2838+
auto II = BB.getFirstNonPseudo();
2839+
// Only check the first instruction for non-empty BasicBlocks
2840+
bool Empty = (II == BB.end());
2841+
if (!Empty && isCallCoveredByBTI(Call, *II))
2842+
return;
2843+
// A BLR can be accepted by a BTI c.
2844+
if (isIndirectCall(Call)) {
2845+
// if we have a BTI j at the start, extend it to a BTI jc,
2846+
// otherwise insert a new BTI c.
2847+
if (!Empty && isBTILandingPad(*II, false, true)) {
2848+
updateBTIVariant(*II, true, true);
2849+
} else {
2850+
MCInst BTIInst;
2851+
createBTI(BTIInst, true, false);
2852+
BB.insertInstruction(II, BTIInst);
2853+
}
2854+
}
2855+
2856+
// A BR can be accepted by a BTI j or BTI c (and BTI jc) IF the operand is
2857+
// x16 or x17. If the operand is not x16 or x17, it can be accepted by a
2858+
// BTI j or BTI jc (and not BTI c).
2859+
if (isIndirectBranch(Call)) {
2860+
assert(Call.getNumOperands() == 1 &&
2861+
"Indirect branch needs to have 1 operand.");
2862+
assert(Call.getOperand(0).isReg() &&
2863+
"Indirect branch does not have a register operand.");
2864+
MCPhysReg Reg = Call.getOperand(0).getReg();
2865+
if (Reg == AArch64::X16 || Reg == AArch64::X17) {
2866+
// Add a new BTI c
2867+
MCInst BTIInst;
2868+
createBTI(BTIInst, true, false);
2869+
BB.insertInstruction(II, BTIInst);
2870+
} else {
2871+
// If BB starts with a BTI c, extend it to BTI jc,
2872+
// otherwise insert a new BTI j.
2873+
if (!Empty && isBTILandingPad(*II, true, false)) {
2874+
updateBTIVariant(*II, true, true);
2875+
} else {
2876+
MCInst BTIInst;
2877+
createBTI(BTIInst, false, true);
2878+
BB.insertInstruction(II, BTIInst);
2879+
}
2880+
}
2881+
}
2882+
}
2883+
28092884
InstructionListType materializeAddress(const MCSymbol *Target, MCContext *Ctx,
28102885
MCPhysReg RegName,
28112886
int64_t Addend = 0) const override {

‎bolt/unittests/Core/MCPlusBuilder.cpp‎

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -198,6 +198,122 @@ TEST_P(MCPlusBuilderTester, AArch64_BTI) {
198198
ASSERT_TRUE(BC->MIB->isImplicitBTIC(*II));
199199
}
200200

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

0 commit comments

Comments
 (0)