https://github.com/bgergely0 updated https://github.com/llvm/llvm-project/pull/163381
From 5b0920828b645e54ede2525406696229ca935d88 Mon Sep 17 00:00:00 2001 From: Gergely Balint <[email protected]> Date: Tue, 7 Oct 2025 14:01:47 +0000 Subject: [PATCH 1/2] [BOLT] Improve InsertNegateRAStatePass::inferUnknownStates Previous implementation used a simple heuristic. This can be improved in several ways: - If a BasicBlock has instruction both with known RAState and unknown RAState, use the known states to work out the unknown ones. - If a BasicBlock only consists of instructions with unknown RAState, use the last known RAState from its predecessors, or the first known from its successors to set the RAStates in the BasicBlock. This includes error checking: all predecessors/successors should have the same RAState. - Some BasicBlocks may only contain instructions with unknown RAState, and have no CFG neighbors. These already have incorrect unwind info. For these, we copy the last known RAState based on the layout order. Updated bolt/docs/PacRetDesign.md to reflect changes. --- bolt/docs/PacRetDesign.md | 23 +- .../bolt/Passes/InsertNegateRAStatePass.h | 34 ++- bolt/lib/Passes/InsertNegateRAStatePass.cpp | 226 ++++++++++++++++-- 3 files changed, 255 insertions(+), 28 deletions(-) diff --git a/bolt/docs/PacRetDesign.md b/bolt/docs/PacRetDesign.md index f3fe5fbd522cb..c7c76cac3a100 100644 --- a/bolt/docs/PacRetDesign.md +++ b/bolt/docs/PacRetDesign.md @@ -200,16 +200,29 @@ This pass runs after optimizations. It performns the _inverse_ of MarkRAState pa Some BOLT passes can add new Instructions. In InsertNegateRAStatePass, we have to know what RA state these have. -The current solution has the `inferUnknownStates` function to cover these, using -a fairly simple strategy: unknown states inherit the last known state. - -This will be updated to a more robust solution. - > [!important] > As issue #160989 describes, unwind info is incorrect in stubs with multiple > callers. > For this same reason, we cannot generate correct pac-specific unwind info: > the signess > of the _incorrect_ return address is meaningless. +Assignment of RAStates to newly generated instructions is done in `inferUnknownStates`. +We have three different cases to cover: + +1. If a BasicBlock has some instructions with known RA state, and some without, we + can copy the RAState of known instructions to the unknown ones. As the control + flow only changes between BasicBlocks, instructions in the same BasicBlock have the + same return address. + +2. If all instructions in a BasicBlock are unknown, we can look at all CFG neighbors + (that is predecessors/successors). The RAState should be the same as of the + neighboring blocks. Conflicting RAStates in neighbors indicate an error. Such + functions should be ignored. + +3. If a BasicBlock has no CFG neighbors, we have to copy the RAState of the previous + BasicBlock in layout order. + +If any BasicBlocks remain with unknown instructions, the function will be ignored. + ### Optimizations requiring special attention Marking states before optimizations ensure that instructions can be moved around diff --git a/bolt/include/bolt/Passes/InsertNegateRAStatePass.h b/bolt/include/bolt/Passes/InsertNegateRAStatePass.h index 836948bf5e9c0..b4b428207b657 100644 --- a/bolt/include/bolt/Passes/InsertNegateRAStatePass.h +++ b/bolt/include/bolt/Passes/InsertNegateRAStatePass.h @@ -1,4 +1,4 @@ -//===- bolt/Passes/InsertNegateRAStatePass.cpp ----------------------------===// +//===- bolt/Passes/InsertNegateRAStatePass.h ------------------------------===// // // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. // See https://llvm.org/LICENSE.txt for license information. @@ -30,9 +30,39 @@ class InsertNegateRAState : public BinaryFunctionPass { private: /// Because states are tracked as MCAnnotations on individual instructions, /// newly inserted instructions do not have a state associated with them. - /// New states are "inherited" from the last known state. void inferUnknownStates(BinaryFunction &BF); + /// Simple case: copy RAStates to unknown insts from previous inst. + /// Account for signing and authenticating insts. + void fillUnknownStateInBB(BinaryContext &BC, BinaryBasicBlock &BB); + + /// Fill unknown RAStates in BBs with no successors/predecessors. These are + /// Stubs inserted by LongJmp. As of #160989, we have to copy the RAState from + /// the previous BB in the layout, because CFIs are already incorrect here. + void fillUnknownStubs(BinaryFunction &BF); + + /// Fills unknowns RAStates of BBs with successors/predecessors. Uses + /// getRAStateByCFG to determine the RAState. Does more than one iteration if + /// needed. Reports an error, if it cannot find the RAState for all BBs with + /// predecessors/successors. + void fillUnknownBlocksInCFG(BinaryFunction &BF); + + /// For BBs which only hold instructions with unknown RAState, we check + /// CFG neighbors (successors, predecessors) of the BB. If they have different + /// RAStates, we report an inconsistency. Otherwise, we return the found + /// RAState. + std::optional<bool> getRAStateByCFG(BinaryBasicBlock &BB, BinaryFunction &BF); + /// Returns the first known RAState from \p BB, or std::nullopt if all are + /// unknown. + std::optional<bool> getFirstKnownRAState(BinaryContext &BC, + BinaryBasicBlock &BB); + + /// \p Return true if all instructions have unknown RAState. + bool isUnknownBlock(BinaryContext &BC, BinaryBasicBlock &BB); + + /// Set all instructions in \p BB to \p State. + void markUnknownBlock(BinaryContext &BC, BinaryBasicBlock &BB, bool State); + /// Support for function splitting: /// if two consecutive BBs with Signed state are going to end up in different /// functions (so are held by different FunctionFragments), we have to add a diff --git a/bolt/lib/Passes/InsertNegateRAStatePass.cpp b/bolt/lib/Passes/InsertNegateRAStatePass.cpp index 9aa7c33bd1d6e..1b1fbf7f20e9d 100644 --- a/bolt/lib/Passes/InsertNegateRAStatePass.cpp +++ b/bolt/lib/Passes/InsertNegateRAStatePass.cpp @@ -74,6 +74,20 @@ void InsertNegateRAState::runOnFunction(BinaryFunction &BF) { } } +void InsertNegateRAState::inferUnknownStates(BinaryFunction &BF) { + BinaryContext &BC = BF.getBinaryContext(); + + // Fill in missing RAStates in simple cases (inside BBs). + for (BinaryBasicBlock &BB : BF) { + fillUnknownStateInBB(BC, BB); + } + // Some stubs have no predecessors. For those, we iterate once in the layout + // order to fill their RAState. + fillUnknownStubs(BF); + + fillUnknownBlocksInCFG(BF); +} + void InsertNegateRAState::coverFunctionFragmentStart(BinaryFunction &BF, FunctionFragment &FF) { BinaryContext &BC = BF.getBinaryContext(); @@ -91,7 +105,8 @@ void InsertNegateRAState::coverFunctionFragmentStart(BinaryFunction &BF, }); // If a function is already split in the input, the first FF can also start // with Signed state. This covers that scenario as well. - auto RAState = BC.MIB->getRAState(*(*FirstNonEmpty)->begin()); + auto II = (*FirstNonEmpty)->getFirstNonPseudo(); + auto RAState = BC.MIB->getRAState(*II); if (!RAState) { BC.errs() << "BOLT-ERROR: unknown RAState after inferUnknownStates " << " in function " << BF.getPrintName() << "\n"; @@ -99,36 +114,205 @@ void InsertNegateRAState::coverFunctionFragmentStart(BinaryFunction &BF, return; } if (*RAState) - BF.addCFIInstruction(*FirstNonEmpty, (*FirstNonEmpty)->begin(), + BF.addCFIInstruction(*FirstNonEmpty, II, MCCFIInstruction::createNegateRAState(nullptr)); } -void InsertNegateRAState::inferUnknownStates(BinaryFunction &BF) { +std::optional<bool> +InsertNegateRAState::getFirstKnownRAState(BinaryContext &BC, + BinaryBasicBlock &BB) { + for (const MCInst &Inst : BB) { + if (BC.MIB->isCFI(Inst)) + continue; + auto RAStateOpt = BC.MIB->getRAState(Inst); + if (RAStateOpt) + return RAStateOpt; + } + return std::nullopt; +} + +void InsertNegateRAState::fillUnknownStateInBB(BinaryContext &BC, + BinaryBasicBlock &BB) { + + auto First = BB.getFirstNonPseudo(); + if (First == BB.end()) + return; + // If the first instruction has unknown RAState, we should copy the first + // known RAState. + auto RAStateOpt = BC.MIB->getRAState(*First); + if (!RAStateOpt) { + auto FirstRAState = getFirstKnownRAState(BC, BB); + if (!FirstRAState) + // We fill unknown BBs later. + return; + + BC.MIB->setRAState(*First, *FirstRAState); + } + + // At this point we know the RAState of the first instruction, + // so we can propagate the RAStates to all subsequent unknown instructions. + MCInst Prev = *First; + for (auto It = BB.begin() + 1; It != BB.end(); ++It) { + MCInst &Inst = *It; + if (BC.MIB->isCFI(Inst)) + continue; + + auto PrevRAState = BC.MIB->getRAState(Prev); + if (!PrevRAState) + llvm_unreachable("Previous Instruction has no RAState."); + + auto RAState = BC.MIB->getRAState(Inst); + if (!RAState) { + if (BC.MIB->isPSignOnLR(Prev)) + PrevRAState = true; + else if (BC.MIB->isPAuthOnLR(Prev)) + PrevRAState = false; + BC.MIB->setRAState(Inst, *PrevRAState); + } + Prev = Inst; + } +} + +bool InsertNegateRAState::isUnknownBlock(BinaryContext &BC, + BinaryBasicBlock &BB) { + for (const MCInst &Inst : BB) { + if (BC.MIB->isCFI(Inst)) + continue; + auto RAState = BC.MIB->getRAState(Inst); + if (RAState) + return false; + } + return true; +} + +void InsertNegateRAState::markUnknownBlock(BinaryContext &BC, + BinaryBasicBlock &BB, bool State) { + // If we call this when an Instruction has either kRASigned or kRAUnsigned + // annotation, setRASigned or setRAUnsigned would fail. + assert(isUnknownBlock(BC, BB) && + "markUnknownBlock should only be called on unknown blocks"); + for (MCInst &Inst : BB) { + if (BC.MIB->isCFI(Inst)) + continue; + BC.MIB->setRAState(Inst, State); + } +} + +std::optional<bool> InsertNegateRAState::getRAStateByCFG(BinaryBasicBlock &BB, + BinaryFunction &BF) { BinaryContext &BC = BF.getBinaryContext(); - bool FirstIter = true; - MCInst PrevInst; - for (BinaryBasicBlock &BB : BF) { - for (MCInst &Inst : BB) { - if (BC.MIB->isCFI(Inst)) + + auto checkRAState = [&](std::optional<bool> &NeighborRAState, MCInst &Inst) { + auto RAState = BC.MIB->getRAState(Inst); + if (!RAState) + return; + if (!NeighborRAState) { + NeighborRAState = *RAState; + return; + } + if (NeighborRAState != *RAState) { + BC.outs() << "BOLT-WARNING: Conflicting RAState found in function " + << BF.getPrintName() << ". Function will not be optimized.\n"; + BF.setIgnored(); + } + }; + + // Holds the first found RAState from CFG neighbors. + std::optional<bool> NeighborRAState = std::nullopt; + if (BB.pred_size() != 0) { + for (BinaryBasicBlock *PredBB : BB.predecessors()) { + // find last inst of Predecessor with known RA State. + auto LI = PredBB->getLastNonPseudo(); + if (LI == PredBB->rend()) + continue; + MCInst &LastInst = *LI; + checkRAState(NeighborRAState, LastInst); + } + } else if (BB.succ_size() != 0) { + for (BinaryBasicBlock *SuccBB : BB.successors()) { + // find first inst of Successor with known RA State. + auto FI = SuccBB->getFirstNonPseudo(); + if (FI == SuccBB->end()) continue; + MCInst &FirstInst = *FI; + checkRAState(NeighborRAState, FirstInst); + } + } else { + llvm_unreachable("Called getRAStateByCFG on a BB with no preds or succs."); + } + + return NeighborRAState; +} - auto RAState = BC.MIB->getRAState(Inst); - if (!FirstIter && !RAState) { - if (BC.MIB->isPSignOnLR(PrevInst)) - RAState = true; - else if (BC.MIB->isPAuthOnLR(PrevInst)) - RAState = false; - else { +void InsertNegateRAState::fillUnknownStubs(BinaryFunction &BF) { + BinaryContext &BC = BF.getBinaryContext(); + bool FirstIter = true; + MCInst PrevInst; + for (FunctionFragment &FF : BF.getLayout().fragments()) { + for (BinaryBasicBlock *BB : FF) { + if (!FirstIter && isUnknownBlock(BC, *BB)) { + // If we have no predecessors or successors, current BB is a Stub called + // from another BinaryFunction. As of #160989, we have to copy the + // PrevInst's RAState, because CFIs are already incorrect here. + if (BB->pred_size() == 0 && BB->succ_size() == 0) { auto PrevRAState = BC.MIB->getRAState(PrevInst); - RAState = PrevRAState ? *PrevRAState : false; + if (!PrevRAState) { + llvm_unreachable( + "Previous Instruction has no RAState in fillUnknownStubs."); + continue; + } + + if (BC.MIB->isPSignOnLR(PrevInst)) { + PrevRAState = true; + } else if (BC.MIB->isPAuthOnLR(PrevInst)) { + PrevRAState = false; + } + markUnknownBlock(BC, *BB, *PrevRAState); } - BC.MIB->setRAState(Inst, *RAState); - } else { + } + if (FirstIter) { FirstIter = false; - if (!RAState) - BC.MIB->setRAState(Inst, BF.getInitialRAState()); + if (isUnknownBlock(BC, *BB)) + markUnknownBlock(BC, *BB, false); } - PrevInst = Inst; + auto Last = BB->getLastNonPseudo(); + if (Last != BB->rend()) + PrevInst = *Last; + } + } +} +void InsertNegateRAState::fillUnknownBlocksInCFG(BinaryFunction &BF) { + BinaryContext &BC = BF.getBinaryContext(); + + auto fillUnknowns = [&](BinaryFunction &BF) -> std::pair<int, bool> { + int Unknowns = 0; + bool Updated = false; + for (BinaryBasicBlock &BB : BF) { + // Only try to iterate if the BB has either predecessors or successors. + if (isUnknownBlock(BC, BB) && + (BB.pred_size() != 0 || BB.succ_size() != 0)) { + auto RAStateOpt = getRAStateByCFG(BB, BF); + if (RAStateOpt) { + markUnknownBlock(BC, BB, *RAStateOpt); + Updated = true; + } else { + Unknowns++; + } + } + } + return std::pair<int, bool>{Unknowns, Updated}; + }; + + while (true) { + std::pair<int, bool> Iter = fillUnknowns(BF); + if (Iter.first == 0) + break; + if (!Iter.second) { + BC.errs() << "BOLT-WARNING: Could not infer RAState for " << Iter.first + << " BBs in function " << BF.getPrintName() + << ". Function will not be optimized.\n"; + BF.setIgnored(); + break; } } } From ad33b187ac3ea5893047a25b8239af840a985180 Mon Sep 17 00:00:00 2001 From: Gergely Balint <[email protected]> Date: Fri, 31 Oct 2025 13:39:01 +0000 Subject: [PATCH 2/2] [BOLT] Single-pass unittest for InsertNegateRAState This commit creates a new directory: bolt/unittests/Passes, to be used by unittests that need to register and run passes with the BinaryFunctionPassManager. An example test is created for InsertNegateRAState pass. Actual tests will be added in followup commits. --- bolt/unittests/CMakeLists.txt | 1 + bolt/unittests/Passes/CMakeLists.txt | 30 ++++ bolt/unittests/Passes/InsertNegateRAState.cpp | 164 ++++++++++++++++++ 3 files changed, 195 insertions(+) create mode 100644 bolt/unittests/Passes/CMakeLists.txt create mode 100644 bolt/unittests/Passes/InsertNegateRAState.cpp diff --git a/bolt/unittests/CMakeLists.txt b/bolt/unittests/CMakeLists.txt index 64414b83d39fe..d47ddc46b7388 100644 --- a/bolt/unittests/CMakeLists.txt +++ b/bolt/unittests/CMakeLists.txt @@ -7,3 +7,4 @@ endfunction() add_subdirectory(Core) add_subdirectory(Profile) +add_subdirectory(Passes) diff --git a/bolt/unittests/Passes/CMakeLists.txt b/bolt/unittests/Passes/CMakeLists.txt new file mode 100644 index 0000000000000..02dcd2c8417e1 --- /dev/null +++ b/bolt/unittests/Passes/CMakeLists.txt @@ -0,0 +1,30 @@ +set(LLVM_LINK_COMPONENTS + DebugInfoDWARF + Object + MC + ${LLVM_TARGETS_TO_BUILD} + ) + +add_bolt_unittest(PassTests + InsertNegateRAState.cpp + + DISABLE_LLVM_LINK_LLVM_DYLIB + ) + +target_link_libraries(PassTests + PRIVATE + LLVMBOLTCore + LLVMBOLTRewrite + LLVMBOLTPasses + LLVMBOLTProfile + LLVMBOLTUtils + ) + +foreach (tgt ${BOLT_TARGETS_TO_BUILD}) + include_directories( + ${LLVM_MAIN_SRC_DIR}/lib/Target/${tgt} + ${LLVM_BINARY_DIR}/lib/Target/${tgt} + ) + string(TOUPPER "${tgt}" upper) + target_compile_definitions(PassTests PRIVATE "${upper}_AVAILABLE") +endforeach() diff --git a/bolt/unittests/Passes/InsertNegateRAState.cpp b/bolt/unittests/Passes/InsertNegateRAState.cpp new file mode 100644 index 0000000000000..024b696eb0534 --- /dev/null +++ b/bolt/unittests/Passes/InsertNegateRAState.cpp @@ -0,0 +1,164 @@ +//===- bolt/unittest/Passes/InsertNegateRAState.cpp -----------------------===// +// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifdef AARCH64_AVAILABLE +#include "AArch64Subtarget.h" +#include "MCTargetDesc/AArch64MCTargetDesc.h" +#endif // AARCH64_AVAILABLE + +#include "bolt/Core/BinaryBasicBlock.h" +#include "bolt/Core/BinaryFunction.h" +#include "bolt/Passes/InsertNegateRAStatePass.h" +#include "bolt/Rewrite/BinaryPassManager.h" +#include "bolt/Rewrite/RewriteInstance.h" +#include "llvm/BinaryFormat/ELF.h" +#include "llvm/MC/MCDwarf.h" +#include "llvm/MC/MCInstBuilder.h" +#include "llvm/Support/TargetSelect.h" +#include "gtest/gtest.h" + +using namespace llvm; +using namespace llvm::object; +using namespace llvm::ELF; +using namespace bolt; + +namespace { +struct PassTester : public testing::TestWithParam<Triple::ArchType> { + void SetUp() override { + initalizeLLVM(); + prepareElf(); + initializeBolt(); + } + +protected: + void initalizeLLVM() { +#define BOLT_TARGET(target) \ + LLVMInitialize##target##TargetInfo(); \ + LLVMInitialize##target##TargetMC(); \ + LLVMInitialize##target##AsmParser(); \ + LLVMInitialize##target##Disassembler(); \ + LLVMInitialize##target##Target(); \ + LLVMInitialize##target##AsmPrinter(); + +#include "bolt/Core/TargetConfig.def" + } + +#define PREPARE_FUNC(name) \ + constexpr uint64_t FunctionAddress = 0x1000; \ + BinaryFunction *BF = BC->createBinaryFunction( \ + name, *TextSection, FunctionAddress, /*Size=*/0, /*SymbolSize=*/0, \ + /*Alignment=*/16); \ + /* Make sure the pass runs on the BF.*/ \ + BF->updateState(BinaryFunction::State::CFG); \ + BF->setContainedNegateRAState(); \ + /* All tests need at least one BB. */ \ + BinaryBasicBlock *BB = BF->addBasicBlock(); \ + BF->addEntryPoint(*BB); \ + BB->setCFIState(0); + + void prepareElf() { + memcpy(ElfBuf, "\177ELF", 4); + ELF64LE::Ehdr *EHdr = reinterpret_cast<typename ELF64LE::Ehdr *>(ElfBuf); + EHdr->e_ident[llvm::ELF::EI_CLASS] = llvm::ELF::ELFCLASS64; + EHdr->e_ident[llvm::ELF::EI_DATA] = llvm::ELF::ELFDATA2LSB; + EHdr->e_machine = GetParam() == Triple::aarch64 ? EM_AARCH64 : EM_X86_64; + MemoryBufferRef Source(StringRef(ElfBuf, sizeof(ElfBuf)), "ELF"); + ObjFile = cantFail(ObjectFile::createObjectFile(Source)); + } + void initializeBolt() { + Relocation::Arch = ObjFile->makeTriple().getArch(); + BC = cantFail(BinaryContext::createBinaryContext( + ObjFile->makeTriple(), std::make_shared<orc::SymbolStringPool>(), + ObjFile->getFileName(), nullptr, true, DWARFContext::create(*ObjFile), + {llvm::outs(), llvm::errs()})); + ASSERT_FALSE(!BC); + BC->initializeTarget(std::unique_ptr<MCPlusBuilder>( + createMCPlusBuilder(GetParam(), BC->MIA.get(), BC->MII.get(), + BC->MRI.get(), BC->STI.get()))); + + PassManager = std::make_unique<BinaryFunctionPassManager>(*BC); + PassManager->registerPass(std::make_unique<InsertNegateRAState>()); + + TextSection = &BC->registerOrUpdateSection( + ".text", ELF::SHT_PROGBITS, ELF::SHF_ALLOC | ELF::SHF_EXECINSTR, + /*Data=*/nullptr, /*Size=*/0, + /*Alignment=*/16); + } + + std::vector<int> findCFIOffsets(BinaryFunction &BF) { + std::vector<int> Locations; + int Idx = 0; + int InstSize = 4; // AArch64 + for (BinaryBasicBlock &BB : BF) { + for (MCInst &Inst : BB) { + if (BC->MIB->isCFI(Inst)) { + const MCCFIInstruction *CFI = BF.getCFIFor(Inst); + if (CFI->getOperation() == MCCFIInstruction::OpNegateRAState) + Locations.push_back(Idx * InstSize); + } + Idx++; + } + } + return Locations; + } + + char ElfBuf[sizeof(typename ELF64LE::Ehdr)] = {}; + std::unique_ptr<ObjectFile> ObjFile; + std::unique_ptr<BinaryContext> BC; + std::unique_ptr<BinaryFunctionPassManager> PassManager; + BinarySection *TextSection; +}; +} // namespace + +TEST_P(PassTester, ExampleTest) { + if (GetParam() != Triple::aarch64) + GTEST_SKIP(); + + ASSERT_NE(TextSection, nullptr); + + PREPARE_FUNC("ExampleFunction"); + + MCInst UnsignedInst = MCInstBuilder(AArch64::ADDSXri) + .addReg(AArch64::X0) + .addReg(AArch64::X0) + .addImm(0) + .addImm(0); + BC->MIB->setRAState(UnsignedInst, false); + BB->addInstruction(UnsignedInst); + + MCInst SignedInst = MCInstBuilder(AArch64::ADDSXri) + .addReg(AArch64::X0) + .addReg(AArch64::X0) + .addImm(1) + .addImm(0); + BC->MIB->setRAState(SignedInst, true); + BB->addInstruction(SignedInst); + + Error E = PassManager->runPasses(); + EXPECT_FALSE(E); + + /* Expected layout of BF after the pass: + + .LBB0 (3 instructions, align : 1) + Entry Point + CFI State : 0 + 00000000: adds x0, x0, #0x0 + 00000004: !CFI $0 ; OpNegateRAState + 00000004: adds x0, x0, #0x1 + CFI State: 0 + */ + auto CFILoc = findCFIOffsets(*BF); + EXPECT_EQ(CFILoc.size(), 1u); + EXPECT_EQ(CFILoc[0], 4); +} + +#ifdef AARCH64_AVAILABLE +INSTANTIATE_TEST_SUITE_P(AArch64, PassTester, + ::testing::Values(Triple::aarch64)); +#endif _______________________________________________ llvm-branch-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits
