Skip to content

Commit 103728d

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 cca66a2 commit 103728d

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
@@ -198,6 +198,111 @@ TEST_P(MCPlusBuilderTester, AArch64_BTI) {
198198
ASSERT_TRUE(BC->MIB->isImplicitBTIC(*II));
199199
}
200200

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

0 commit comments

Comments
 (0)