[llvm-branch-commits] [llvm] [BOLT][BTI] Add MCPlusBuilder::addBTItoBBStart (PR #167329)
Gergely Bálint via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Mon Nov 24 01:21:40 PST 2025
https://github.com/bgergely0 updated https://github.com/llvm/llvm-project/pull/167329
>From 7ae084c66064fe5defa1feee2026cd29c6ffe2b2 Mon Sep 17 00:00:00 2001
From: Gergely Balint <gergely.balint at arm.com>
Date: Wed, 17 Sep 2025 12:24:04 +0000
Subject: [PATCH] [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.
---
bolt/include/bolt/Core/MCPlusBuilder.h | 13 +++
.../Target/AArch64/AArch64MCPlusBuilder.cpp | 75 +++++++++++++
bolt/unittests/Core/MCPlusBuilder.cpp | 105 ++++++++++++++++++
3 files changed, 193 insertions(+)
diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index b4dc7b3a8b936..4f5a3f191fc2a 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -1888,6 +1888,19 @@ class MCPlusBuilder {
llvm_unreachable("not implemented");
}
+ /// Checks if the indirect call / jump is accepted by the landing pad at the
+ /// start of the target BasicBlock.
+ virtual bool isBTIVariantCoveringCall(MCInst &Call, MCInst &Pad) const {
+ llvm_unreachable("not implemented");
+ return false;
+ }
+
+ /// Adds a BTI landing pad to the start of the BB, that matches the indirect
+ /// call/jump inst.
+ virtual void addBTItoBBStart(BinaryBasicBlock &BB, MCInst &Call) const {
+ llvm_unreachable("not implemented");
+ }
+
/// Store \p Target absolute address to \p RegName
virtual InstructionListType materializeAddress(const MCSymbol *Target,
MCContext *Ctx,
diff --git a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
index b92e5b2df062a..056cd0ec8ad54 100644
--- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
+++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
@@ -2782,6 +2782,81 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
Inst.addOperand(MCOperand::createImm(HintNum));
}
+ bool isBTIVariantCoveringCall(MCInst &Call, MCInst &Pad) const override {
+ assert((isIndirectCall(Call) || isIndirectBranch(Call)) &&
+ "Not an indirect call or branch.");
+
+ // A BLR can be accepted by a BTI c.
+ if (isIndirectCall(Call))
+ return isBTILandingPad(Pad, true, false) ||
+ isBTILandingPad(Pad, true, true);
+
+ // A BR can be accepted by a BTI j or BTI c (and BTI jc) IF the operand is
+ // x16 or x17. If the operand is not x16 or x17, it can be accepted by a BTI
+ // j or BTI jc (and not BTI c).
+ if (isIndirectBranch(Call)) {
+ assert(Call.getNumOperands() == 1 &&
+ "Indirect branch needs to have 1 operand.");
+ assert(Call.getOperand(0).isReg() &&
+ "Indirect branch does not have a register operand.");
+ MCPhysReg Reg = Call.getOperand(0).getReg();
+ if (Reg == AArch64::X16 || Reg == AArch64::X17)
+ return isBTILandingPad(Pad, true, false) ||
+ isBTILandingPad(Pad, false, true) ||
+ isBTILandingPad(Pad, true, true);
+ return isBTILandingPad(Pad, false, true) ||
+ isBTILandingPad(Pad, true, true);
+ }
+ return false;
+ }
+
+ void addBTItoBBStart(BinaryBasicBlock &BB, MCInst &Call) const override {
+ auto II = BB.getFirstNonPseudo();
+ if (II != BB.end()) {
+ if (isBTIVariantCoveringCall(Call, *II))
+ return;
+ // A BLR can be accepted by a BTI c.
+ if (isIndirectCall(Call)) {
+ // if we have a BTI j at the start, extend it to a BTI jc,
+ // otherwise insert a new BTI c.
+ if (isBTILandingPad(*II, false, true)) {
+ updateBTIVariant(*II, true, true);
+ } else {
+ MCInst BTIInst;
+ createBTI(BTIInst, true, false);
+ BB.insertInstruction(II, BTIInst);
+ }
+ }
+
+ // A BR can be accepted by a BTI j or BTI c (and BTI jc) IF the operand is
+ // x16 or x17. If the operand is not x16 or x17, it can be accepted by a
+ // BTI j or BTI jc (and not BTI c).
+ if (isIndirectBranch(Call)) {
+ assert(Call.getNumOperands() == 1 &&
+ "Indirect branch needs to have 1 operand.");
+ assert(Call.getOperand(0).isReg() &&
+ "Indirect branch does not have a register operand.");
+ MCPhysReg Reg = Call.getOperand(0).getReg();
+ if (Reg == AArch64::X16 || Reg == AArch64::X17) {
+ // Add a new BTI c
+ MCInst BTIInst;
+ createBTI(BTIInst, true, false);
+ BB.insertInstruction(II, BTIInst);
+ } else {
+ // If BB starts with a BTI c, extend it to BTI jc,
+ // otherwise insert a new BTI j.
+ if (isBTILandingPad(*II, true, false)) {
+ updateBTIVariant(*II, true, true);
+ } else {
+ MCInst BTIInst;
+ createBTI(BTIInst, false, true);
+ BB.insertInstruction(II, BTIInst);
+ }
+ }
+ }
+ }
+ }
+
InstructionListType materializeAddress(const MCSymbol *Target, MCContext *Ctx,
MCPhysReg RegName,
int64_t Addend = 0) const override {
diff --git a/bolt/unittests/Core/MCPlusBuilder.cpp b/bolt/unittests/Core/MCPlusBuilder.cpp
index 02ecb87b4a5e3..e08ae09e76027 100644
--- a/bolt/unittests/Core/MCPlusBuilder.cpp
+++ b/bolt/unittests/Core/MCPlusBuilder.cpp
@@ -196,6 +196,111 @@ TEST_P(MCPlusBuilderTester, AArch64_BTI) {
ASSERT_TRUE(BC->MIB->isImplicitBTIC(*II));
}
+TEST_P(MCPlusBuilderTester, AArch64_addBTItoBBStart_0) {
+ if (GetParam() != Triple::aarch64)
+ GTEST_SKIP();
+ BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
+ std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock();
+ MCInst Inst = MCInstBuilder(AArch64::RET).addReg(AArch64::LR);
+ BB->addInstruction(Inst);
+ // BR x16 needs BTI c or BTI j. We prefer adding a BTI c.
+ MCInst CallInst = MCInstBuilder(AArch64::BR).addReg(AArch64::X16);
+ BC->MIB->addBTItoBBStart(*BB, CallInst);
+ auto II = BB->begin();
+ ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, true, false));
+}
+
+TEST_P(MCPlusBuilderTester, AArch64_addBTItoBBStart_1) {
+ if (GetParam() != Triple::aarch64)
+ GTEST_SKIP();
+ BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
+ std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock();
+ MCInst BTIc;
+ BC->MIB->createBTI(BTIc, true, false);
+ BB->addInstruction(BTIc);
+ // BR x16 needs BTI c or BTI j. We have a BTI c, no change is needed.
+ MCInst CallInst = MCInstBuilder(AArch64::BR).addReg(AArch64::X16);
+ BC->MIB->addBTItoBBStart(*BB, CallInst);
+ auto II = BB->begin();
+ ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, true, false));
+}
+
+TEST_P(MCPlusBuilderTester, AArch64_addBTItoBBStart_2) {
+ if (GetParam() != Triple::aarch64)
+ GTEST_SKIP();
+ BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
+ std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock();
+ MCInst BTIc;
+ BC->MIB->createBTI(BTIc, true, false);
+ BB->addInstruction(BTIc);
+ // BR x5 needs BTI j
+ // we have BTI c -> extend it to BTI jc.
+ MCInst CallInst = MCInstBuilder(AArch64::BR).addReg(AArch64::X5);
+ BC->MIB->addBTItoBBStart(*BB, CallInst);
+ auto II = BB->begin();
+ ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, true, true));
+}
+
+TEST_P(MCPlusBuilderTester, AArch64_addBTItoBBStart_3) {
+ if (GetParam() != Triple::aarch64)
+ GTEST_SKIP();
+ BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
+ std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock();
+ MCInst Inst = MCInstBuilder(AArch64::RET).addReg(AArch64::LR);
+ BB->addInstruction(Inst);
+ // BR x5 needs BTI j
+ MCInst CallInst = MCInstBuilder(AArch64::BR).addReg(AArch64::X5);
+ BC->MIB->addBTItoBBStart(*BB, CallInst);
+ auto II = BB->begin();
+ ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, false, true));
+}
+
+TEST_P(MCPlusBuilderTester, AArch64_addBTItoBBStart_4) {
+ if (GetParam() != Triple::aarch64)
+ GTEST_SKIP();
+ BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
+ std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock();
+ MCInst Inst = MCInstBuilder(AArch64::RET).addReg(AArch64::LR);
+ BB->addInstruction(Inst);
+ // BLR needs BTI c, regardless of the register used.
+ MCInst CallInst = MCInstBuilder(AArch64::BLR).addReg(AArch64::X5);
+ BC->MIB->addBTItoBBStart(*BB, CallInst);
+ auto II = BB->begin();
+ ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, true, false));
+}
+
+TEST_P(MCPlusBuilderTester, AArch64_addBTItoBBStart_5) {
+ if (GetParam() != Triple::aarch64)
+ GTEST_SKIP();
+ BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
+ std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock();
+ MCInst BTIj;
+ BC->MIB->createBTI(BTIj, false, true);
+ BB->addInstruction(BTIj);
+ // BLR needs BTI c, regardless of the register used.
+ // We have a BTI j -> extend it to BTI jc.
+ MCInst CallInst = MCInstBuilder(AArch64::BLR).addReg(AArch64::X5);
+ BC->MIB->addBTItoBBStart(*BB, CallInst);
+ auto II = BB->begin();
+ ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, true, true));
+}
+
+TEST_P(MCPlusBuilderTester, AArch64_addBTItoBBStart_6) {
+ if (GetParam() != Triple::aarch64)
+ GTEST_SKIP();
+ BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
+ std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock();
+ MCInst Paciasp =
+ MCInstBuilder(AArch64::PACIASP).addReg(AArch64::LR).addReg(AArch64::SP);
+ BB->addInstruction(Paciasp);
+ // PACI(AB)SP are implicit BTI c, no change needed.
+ MCInst CallInst = MCInstBuilder(AArch64::BR).addReg(AArch64::X17);
+ BC->MIB->addBTItoBBStart(*BB, CallInst);
+ auto II = BB->begin();
+ ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, true, false));
+ ASSERT_TRUE(BC->MIB->isPSignOnLR(*II));
+}
+
TEST_P(MCPlusBuilderTester, AArch64_CmpJNE) {
if (GetParam() != Triple::aarch64)
GTEST_SKIP();
More information about the llvm-branch-commits
mailing list