llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT--> @llvm/pr-subscribers-bolt Author: Gergely Bálint (bgergely0) <details> <summary>Changes</summary> 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. --- Full diff: https://github.com/llvm/llvm-project/pull/167329.diff 3 Files Affected: - (modified) bolt/include/bolt/Core/MCPlusBuilder.h (+13) - (modified) bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp (+75) - (modified) bolt/unittests/Core/MCPlusBuilder.cpp (+105) ``````````diff diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h index 660c1c64b06cf..4eaf444c320bf 100644 --- a/bolt/include/bolt/Core/MCPlusBuilder.h +++ b/bolt/include/bolt/Core/MCPlusBuilder.h @@ -1890,6 +1890,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 cb0a9cc0c12db..a5c88e113f726 100644 --- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp +++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp @@ -2738,6 +2738,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(); `````````` </details> https://github.com/llvm/llvm-project/pull/167329 _______________________________________________ llvm-branch-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits
