Skip to content

Commit fbd3fc2

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 7c2404b commit fbd3fc2

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
@@ -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 isBTIVariantCoveringCall(MCInst &Call, MCInst &Pad) const {
1900+
llvm_unreachable("not implemented");
1901+
return false;
1902+
}
1903+
1904+
/// Adds a BTI landing pad to the start of the BB, that matches the indirect
1905+
/// call/jump inst.
1906+
virtual void addBTItoBBStart(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
@@ -2808,6 +2808,81 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
28082808
Inst.addOperand(MCOperand::createImm(HintNum));
28092809
}
28102810

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