From 64b728128df355b9838eb7f6c19082628fb30dbe Mon Sep 17 00:00:00 2001 From: Alexandros Lamprineas Date: Thu, 2 Apr 2026 11:59:59 +0100 Subject: [PATCH] [BOLT][AArch64] Add minimal support for liveness analysis. (#183298) In this patch I am adding the missing target hooks required for the liveness analysis to run on AArch64. These are - getFlagsReg() - getRegsUsedAsParams() - getDefaultLiveOut() - getGPRegs() - isCleanRegXOR() I am also introducing the following API in LivenessAnalysis - BitVector getLiveIn/Out(const MCInst &) - MCPhysReg scavengeRegFromState(BitVector &) My intention is to allow the LongJmp pass scavenge usable registers when injecting code. --- bolt/include/bolt/Core/MCPlusBuilder.h | 7 +- bolt/include/bolt/Passes/LivenessAnalysis.h | 32 ++- .../Target/AArch64/AArch64MCPlusBuilder.cpp | 119 +++++++++++ bolt/lib/Target/X86/X86MCPlusBuilder.cpp | 6 + bolt/unittests/Core/MCPlusBuilder.cpp | 60 ++++++ bolt/unittests/Passes/CMakeLists.txt | 1 + bolt/unittests/Passes/LivenessAnalysis.cpp | 188 ++++++++++++++++++ 7 files changed, 403 insertions(+), 10 deletions(-) create mode 100644 bolt/unittests/Passes/LivenessAnalysis.cpp diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h index a672d7a45689..0f6076688b65 100644 --- a/bolt/include/bolt/Core/MCPlusBuilder.h +++ b/bolt/include/bolt/Core/MCPlusBuilder.h @@ -1595,7 +1595,7 @@ public: llvm_unreachable("not implemented"); } - /// Similar to getDefaultDefIn + /// Registers which may contain a meaningful value after a function returns. virtual void getDefaultLiveOut(BitVector &Regs) const { llvm_unreachable("not implemented"); } @@ -1605,6 +1605,11 @@ public: llvm_unreachable("not implemented"); } + /// Remove non scavengeable special registers from \p Regs + virtual void removeNonScavengeableRegs(BitVector &Regs) const { + llvm_unreachable("not implemented"); + } + /// Change \p Regs with a bitmask with all general purpose regs that can be /// encoded without extra prefix bytes. For x86 only. virtual void getClassicGPRegs(BitVector &Regs) const { diff --git a/bolt/include/bolt/Passes/LivenessAnalysis.h b/bolt/include/bolt/Passes/LivenessAnalysis.h index f4faa1dc34ec..6fcb66957f02 100644 --- a/bolt/include/bolt/Passes/LivenessAnalysis.h +++ b/bolt/include/bolt/Passes/LivenessAnalysis.h @@ -36,6 +36,16 @@ public: NumRegs(BF.getBinaryContext().MRI->getNumRegs()) {} virtual ~LivenessAnalysis(); + // Return the state before the execution of an Instruction. + BitVector getLiveIn(const MCInst &Inst) const { + return *this->getStateAt(Inst); + } + + // Return the state after the execution of an Instruction. + BitVector getLiveOut(const MCInst &Inst) const { + return *this->getStateBefore(Inst); + } + bool isAlive(ProgramPoint PP, MCPhysReg Reg) const { const BitVector &BV = *this->getStateAt(PP); const BitVector &RegAliases = BC.MIB->getAliases(Reg); @@ -46,18 +56,22 @@ public: // Return a usable general-purpose reg after point P. Return 0 if no reg is // available. - MCPhysReg scavengeRegAfter(ProgramPoint P) { + MCPhysReg scavengeRegAfter(ProgramPoint P) const { BitVector BV = *this->getStateAt(P); - BV.flip(); + return scavengeRegFromState(BV); + } + + // Return a usable general-purpose reg given a liveness state. Return 0 if + // no reg is available. + MCPhysReg scavengeRegFromState(BitVector &LiveRegs) const { BitVector GPRegs(NumRegs, false); this->BC.MIB->getGPRegs(GPRegs, /*IncludeAlias=*/false); - // Ignore the register used for frame pointer even if it is not alive (it - // may be used by CFI which is not represented in our dataflow). - BitVector FP = BC.MIB->getAliases(BC.MIB->getFramePointer()); - FP.flip(); - BV &= GPRegs; - BV &= FP; - int Reg = BV.find_first(); + LiveRegs.flip(); + LiveRegs &= GPRegs; + // Ignore target-specific special registers even if they are dead + // (they may be used by CFI which is not represented in our dataflow). + BC.MIB->removeNonScavengeableRegs(LiveRegs); + int Reg = LiveRegs.find_first(); return Reg != -1 ? Reg : 0; } diff --git a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp index 3955ff378be4..b4c08cf1f615 100644 --- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp +++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp @@ -158,6 +158,7 @@ public: MCPhysReg getStackPointer() const override { return AArch64::SP; } MCPhysReg getFramePointer() const override { return AArch64::FP; } + MCPhysReg getFlagsReg() const override { return AArch64::NZCV; } bool isBreakpoint(const MCInst &Inst) const override { return Inst.getOpcode() == AArch64::BRK; @@ -1238,6 +1239,19 @@ public: return true; } + BitVector getRegsUsedAsParams() const override { + BitVector Regs = BitVector(RegInfo->getNumRegs(), false); + Regs |= getAliases(AArch64::X0); + Regs |= getAliases(AArch64::X1); + Regs |= getAliases(AArch64::X2); + Regs |= getAliases(AArch64::X3); + Regs |= getAliases(AArch64::X4); + Regs |= getAliases(AArch64::X5); + Regs |= getAliases(AArch64::X6); + Regs |= getAliases(AArch64::X7); + return Regs; + } + void getCalleeSavedRegs(BitVector &Regs) const override { Regs |= getAliases(AArch64::X18); Regs |= getAliases(AArch64::X19); @@ -1254,6 +1268,88 @@ public: Regs |= getAliases(AArch64::FP); } + void getDefaultLiveOut(BitVector &Regs) const override { + // According to the AArch64 ABI the return registers are X0 to X7, + // which happen to be the same as the parameter registers. + Regs |= getRegsUsedAsParams(); + } + + void getGPRegs(BitVector &Regs, bool IncludeAlias = true) const override { + if (IncludeAlias) { + Regs |= getAliases(AArch64::X0); + Regs |= getAliases(AArch64::X1); + Regs |= getAliases(AArch64::X2); + Regs |= getAliases(AArch64::X3); + Regs |= getAliases(AArch64::X4); + Regs |= getAliases(AArch64::X5); + Regs |= getAliases(AArch64::X6); + Regs |= getAliases(AArch64::X7); + Regs |= getAliases(AArch64::X8); + Regs |= getAliases(AArch64::X9); + Regs |= getAliases(AArch64::X10); + Regs |= getAliases(AArch64::X11); + Regs |= getAliases(AArch64::X12); + Regs |= getAliases(AArch64::X13); + Regs |= getAliases(AArch64::X14); + Regs |= getAliases(AArch64::X15); + Regs |= getAliases(AArch64::X16); + Regs |= getAliases(AArch64::X17); + Regs |= getAliases(AArch64::X18); + Regs |= getAliases(AArch64::X19); + Regs |= getAliases(AArch64::X20); + Regs |= getAliases(AArch64::X21); + Regs |= getAliases(AArch64::X22); + Regs |= getAliases(AArch64::X23); + Regs |= getAliases(AArch64::X24); + Regs |= getAliases(AArch64::X25); + Regs |= getAliases(AArch64::X26); + Regs |= getAliases(AArch64::X27); + Regs |= getAliases(AArch64::X28); + Regs |= getAliases(AArch64::LR); + Regs |= getAliases(AArch64::FP); + return; + } + Regs.set(AArch64::X0); + Regs.set(AArch64::X1); + Regs.set(AArch64::X2); + Regs.set(AArch64::X3); + Regs.set(AArch64::X4); + Regs.set(AArch64::X5); + Regs.set(AArch64::X6); + Regs.set(AArch64::X7); + Regs.set(AArch64::X8); + Regs.set(AArch64::X9); + Regs.set(AArch64::X10); + Regs.set(AArch64::X11); + Regs.set(AArch64::X12); + Regs.set(AArch64::X13); + Regs.set(AArch64::X14); + Regs.set(AArch64::X15); + Regs.set(AArch64::X16); + Regs.set(AArch64::X17); + Regs.set(AArch64::X18); + Regs.set(AArch64::X19); + Regs.set(AArch64::X20); + Regs.set(AArch64::X21); + Regs.set(AArch64::X22); + Regs.set(AArch64::X23); + Regs.set(AArch64::X24); + Regs.set(AArch64::X25); + Regs.set(AArch64::X26); + Regs.set(AArch64::X27); + Regs.set(AArch64::X28); + Regs.set(AArch64::LR); + Regs.set(AArch64::FP); + } + + void removeNonScavengeableRegs(BitVector &Regs) const override { + BitVector ExclusionMask = getAliases(AArch64::LR); + ExclusionMask |= getAliases(AArch64::FP); + ExclusionMask |= getAliases(AArch64::X18); // platform register + ExclusionMask.flip(); + Regs &= ExclusionMask; + } + const MCExpr *getTargetExprFor(MCInst &Inst, const MCExpr *Expr, MCContext &Ctx, uint32_t RelType) const override { @@ -2412,6 +2508,29 @@ public: isAArch64ExclusiveStore(Inst); } + bool isCleanRegXOR(const MCInst &Inst) const override { + switch (Inst.getOpcode()) { + case AArch64::EORXrs: + case AArch64::EORWrs: + return Inst.getOperand(1).getReg() == Inst.getOperand(2).getReg() && + Inst.getOperand(3).getImm() == 0; + case AArch64::ORRXrs: + return Inst.getOperand(1).getReg() == AArch64::XZR && + Inst.getOperand(2).getReg() == AArch64::XZR && + Inst.getOperand(3).getImm() == 0; + case AArch64::ORRWrs: + return Inst.getOperand(1).getReg() == AArch64::WZR && + Inst.getOperand(2).getReg() == AArch64::WZR && + Inst.getOperand(3).getImm() == 0; + case AArch64::MOVZXi: + case AArch64::MOVZWi: + return Inst.getOperand(1).isImm() && Inst.getOperand(1).getImm() == 0 && + Inst.getOperand(2).getImm() == 0; + default: + return false; + } + } + bool isStoreToStack(const MCInst &Inst) const { if (!mayStore(Inst)) return false; diff --git a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp index 51e7d27f18a0..aa1a70e40669 100644 --- a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp +++ b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp @@ -827,6 +827,12 @@ public: Regs.set(X86::R15); } + void removeNonScavengeableRegs(BitVector &Regs) const override { + BitVector FP = getAliases(X86::RBP); + FP.flip(); + Regs &= FP; + } + void getClassicGPRegs(BitVector &Regs) const override { Regs |= getAliases(X86::RAX); Regs |= getAliases(X86::RBX); diff --git a/bolt/unittests/Core/MCPlusBuilder.cpp b/bolt/unittests/Core/MCPlusBuilder.cpp index d76dac04c5bb..7073b41f524f 100644 --- a/bolt/unittests/Core/MCPlusBuilder.cpp +++ b/bolt/unittests/Core/MCPlusBuilder.cpp @@ -8,6 +8,7 @@ #ifdef AARCH64_AVAILABLE #include "AArch64Subtarget.h" +#include "MCTargetDesc/AArch64MCAsmInfo.h" #include "MCTargetDesc/AArch64MCTargetDesc.h" #endif // AARCH64_AVAILABLE @@ -625,6 +626,65 @@ TEST_P(MCPlusBuilderTester, AArch64_Psign_Pauth_variants) { ASSERT_TRUE(BC->MIB->isPAuthAndRet(Retab)); } +TEST_P(MCPlusBuilderTester, AArch64_isCleanRegXOR) { + if (GetParam() != Triple::aarch64) + GTEST_SKIP(); + + BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true); + BinaryBasicBlock *BB = BF->addBasicBlock(); + + // eor x0, x0, x0 + MCInst EORXrs = MCInstBuilder(AArch64::EORXrs) + .addReg(AArch64::X0) + .addReg(AArch64::X0) + .addReg(AArch64::X0) + .addImm(0); + ASSERT_TRUE(BC->MIB->isCleanRegXOR(EORXrs)); + + // eor w0, w0, w0 + MCInst EORWrs = MCInstBuilder(AArch64::EORWrs) + .addReg(AArch64::W0) + .addReg(AArch64::W0) + .addReg(AArch64::W0) + .addImm(0); + ASSERT_TRUE(BC->MIB->isCleanRegXOR(EORWrs)); + + // mov x0, xzr + MCInst ORRXrs = MCInstBuilder(AArch64::ORRXrs) + .addReg(AArch64::X0) + .addReg(AArch64::XZR) + .addReg(AArch64::XZR) + .addImm(0); + ASSERT_TRUE(BC->MIB->isCleanRegXOR(ORRXrs)); + + // mov w0, wzr + MCInst ORRWrs = MCInstBuilder(AArch64::ORRWrs) + .addReg(AArch64::W0) + .addReg(AArch64::WZR) + .addReg(AArch64::WZR) + .addImm(0); + ASSERT_TRUE(BC->MIB->isCleanRegXOR(ORRWrs)); + + // mov x0, #0 + MCInst MOVZXi = + MCInstBuilder(AArch64::MOVZXi).addReg(AArch64::X0).addImm(0).addImm(0); + ASSERT_TRUE(BC->MIB->isCleanRegXOR(MOVZXi)); + + // mov w0, #0 + MCInst MOVZWi = + MCInstBuilder(AArch64::MOVZWi).addReg(AArch64::W0).addImm(0).addImm(0); + ASSERT_TRUE(BC->MIB->isCleanRegXOR(MOVZWi)); + + // movz x0, #:abs_g3:symbol + MCInst MOVZXiWithExpr = + MCInstBuilder(AArch64::MOVZXi) + .addReg(AArch64::X0) + .addExpr(MCSpecifierExpr::create(BB->getLabel(), AArch64::S_ABS_G3, + *BC->Ctx.get())) + .addImm(48); + ASSERT_FALSE(BC->MIB->isCleanRegXOR(MOVZXiWithExpr)); +} + #endif // AARCH64_AVAILABLE #ifdef X86_AVAILABLE diff --git a/bolt/unittests/Passes/CMakeLists.txt b/bolt/unittests/Passes/CMakeLists.txt index 17ae8802f695..bd373cd0006f 100644 --- a/bolt/unittests/Passes/CMakeLists.txt +++ b/bolt/unittests/Passes/CMakeLists.txt @@ -6,6 +6,7 @@ set(LLVM_LINK_COMPONENTS ) add_bolt_unittest(PassTests + LivenessAnalysis.cpp PointerAuthCFIFixup.cpp DISABLE_LLVM_LINK_LLVM_DYLIB diff --git a/bolt/unittests/Passes/LivenessAnalysis.cpp b/bolt/unittests/Passes/LivenessAnalysis.cpp new file mode 100644 index 000000000000..7ddf683a855e --- /dev/null +++ b/bolt/unittests/Passes/LivenessAnalysis.cpp @@ -0,0 +1,188 @@ +//===- bolt/unittest/Passes/LivenessAnalysis.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/Core/BinaryFunctionCallGraph.h" +#include "bolt/Passes/DataflowInfoManager.h" +#include "bolt/Passes/RegAnalysis.h" +#include "bolt/Rewrite/RewriteInstance.h" +#include "llvm/BinaryFormat/ELF.h" +#include "llvm/DebugInfo/DWARF/DWARFContext.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 opts { +extern cl::opt AssumeABI; +} // namespace opts + +namespace { +struct LivenessAnalysisTester + : public testing::TestWithParam { + 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" + } + + void prepareElf() { + memcpy(ElfBuf, "\177ELF", 4); + ELF64LE::Ehdr *EHdr = reinterpret_cast(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(), + ObjFile->getFileName(), nullptr, true, DWARFContext::create(*ObjFile), + {llvm::outs(), llvm::errs()})); + ASSERT_FALSE(!BC); + BC->initializeTarget(std::unique_ptr( + createMCPlusBuilder(GetParam(), BC->MIA.get(), BC->MII.get(), + BC->MRI.get(), BC->STI.get()))); + } + + char ElfBuf[sizeof(typename ELF64LE::Ehdr)] = {}; + std::unique_ptr ObjFile; + std::unique_ptr BC; +}; +} // namespace + +#ifdef AARCH64_AVAILABLE + +INSTANTIATE_TEST_SUITE_P(AArch64, LivenessAnalysisTester, + ::testing::Values(Triple::aarch64)); + +TEST_P(LivenessAnalysisTester, AArch64_scavengeRegFromState) { + if (GetParam() != Triple::aarch64) + GTEST_SKIP(); + + opts::AssumeABI = true; + BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true); + BinaryBasicBlock *EntryBB = BF->addBasicBlock(); + BinaryBasicBlock *FallThroughBB = BF->addBasicBlock(); + BinaryBasicBlock *TargetBB = BF->addBasicBlock(); + BF->addEntryPoint(*EntryBB); + EntryBB->addSuccessor(FallThroughBB); + EntryBB->addSuccessor(TargetBB); + FallThroughBB->addSuccessor(TargetBB); + EntryBB->setCFIState(0); + FallThroughBB->setCFIState(0); + TargetBB->setCFIState(0); + + // mov x8, #1 + MCInst MOVZXi = + MCInstBuilder(AArch64::MOVZXi).addReg(AArch64::X8).addImm(1).addImm(0); + // cbgt x0, #0, target + MCInst CBGTXri = MCInstBuilder(AArch64::CBGTXri) + .addReg(AArch64::X0) + .addImm(0) + .addExpr(MCSymbolRefExpr::create(TargetBB->getLabel(), + *BC->Ctx.get())); + // add x0, x8, #1 + MCInst ADDXri = MCInstBuilder(AArch64::ADDXri) + .addReg(AArch64::X0) + .addReg(AArch64::X8) + .addImm(1) + .addImm(0); + // ret + MCInst RET = MCInstBuilder(AArch64::RET).addReg(AArch64::LR); + + EntryBB->addInstruction(MOVZXi); + EntryBB->addInstruction(CBGTXri); + FallThroughBB->addInstruction(ADDXri); + TargetBB->addInstruction(RET); + + BinaryFunctionCallGraph CG(buildCallGraph(*BC)); + RegAnalysis RA(*BC, &BC->getBinaryFunctions(), &CG); + DataflowInfoManager Info(*BF, &RA, nullptr); + + auto II = EntryBB->begin(); + + // Test that parameter registers are LiveIn. + BitVector ParamRegs = BC->MIB->getRegsUsedAsParams(); + ASSERT_TRUE(ParamRegs.subsetOf(Info.getLivenessAnalysis().getLiveIn(*II))); + + BitVector LiveIn, LiveOut; + // mov x8, #1 -> LiveIn = {x0}, LiveOut = {x0, x8} + LiveIn = Info.getLivenessAnalysis().getLiveIn(*II); + LiveOut = Info.getLivenessAnalysis().getLiveOut(*II); + ASSERT_TRUE(LiveIn.test(AArch64::X0)); + ASSERT_FALSE(LiveIn.test(AArch64::X8)); + ASSERT_TRUE(LiveOut.test(AArch64::X0)); + ASSERT_TRUE(LiveOut.test(AArch64::X8)); + ASSERT_EQ(Info.getLivenessAnalysis().scavengeRegFromState(LiveOut), + AArch64::X9); + II++; + // cbgt x0, #0, target -> LiveIn = {x0, x8}, LiveOut = {x0, x8} + LiveIn = Info.getLivenessAnalysis().getLiveIn(*II); + LiveOut = Info.getLivenessAnalysis().getLiveOut(*II); + ASSERT_TRUE(LiveIn.test(AArch64::X0)); + ASSERT_TRUE(LiveIn.test(AArch64::X8)); + ASSERT_TRUE(LiveOut.test(AArch64::X0)); + ASSERT_TRUE(LiveOut.test(AArch64::X8)); + ASSERT_EQ(Info.getLivenessAnalysis().scavengeRegFromState(LiveOut), + AArch64::X9); + II = FallThroughBB->begin(); + // add x0, x8, #1 -> LiveIn = {x8}, LiveOut = {x0} + LiveIn = Info.getLivenessAnalysis().getLiveIn(*II); + LiveOut = Info.getLivenessAnalysis().getLiveOut(*II); + ASSERT_FALSE(LiveIn.test(AArch64::X0)); + ASSERT_TRUE(LiveIn.test(AArch64::X8)); + ASSERT_TRUE(LiveOut.test(AArch64::X0)); + ASSERT_FALSE(LiveOut.test(AArch64::X8)); + ASSERT_EQ(Info.getLivenessAnalysis().scavengeRegFromState(LiveOut), + AArch64::X8); + II = TargetBB->begin(); + // ret -> LiveIn = {x0}, LiveOut = {x0} + LiveIn = Info.getLivenessAnalysis().getLiveIn(*II); + LiveOut = Info.getLivenessAnalysis().getLiveOut(*II); + ASSERT_TRUE(LiveIn.test(AArch64::X0)); + ASSERT_FALSE(LiveIn.test(AArch64::X8)); + ASSERT_TRUE(LiveOut.test(AArch64::X0)); + ASSERT_FALSE(LiveOut.test(AArch64::X8)); + ASSERT_EQ(Info.getLivenessAnalysis().scavengeRegFromState(LiveOut), + AArch64::X8); + + // Test that return registers are LiveOut. + BitVector DefaultLiveOutRegs; + BC->MIB->getDefaultLiveOut(DefaultLiveOutRegs); + ASSERT_TRUE( + DefaultLiveOutRegs.subsetOf(Info.getLivenessAnalysis().getLiveOut(*II))); +} + +#endif // AARCH64_AVAILABLE