From cb399aa9a69e67aa99167a752d7d3d3969f09d23 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=B9=85=E6=97=A5?= <609079623@qq.com> Date: Thu, 22 May 2025 15:49:06 +0800 Subject: [PATCH] =?UTF-8?q?=E6=A0=87=E9=A2=98:=20Fix=20pac=20causes=20comp?= =?UTF-8?q?ilation=20failure=20because=20some=20phi=20branches=20are=20opt?= =?UTF-8?q?imized=20to=20constant=20function=20addresses?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 描述: 开启pac功能后，若存在一个函数指针在phi指令中被赋值，则会跟随phi指令逻辑，生成一条相似的phi指令用于生成modifier。 同时，phi指令的来源分支，也需要补充modifier的生成逻辑。 当原函数指针的phi指令来源，被优化为了常量函数地址，同时modifier的语句也会被优化掉，这就导致在生成modifier的phi指令时， 会缺少一个分支数据。在指令生成时，会处理失败。 Issue: #ICA480: 【PAC】The pac causes compilation failure because some phi branches are optimized to constant function addresses 测试: 1. 补充了用例PhiBlrBranchOptimized，用于验证phi部分分支被优化时的编译表现 2. 补充了用例PhiMultiBlr，用于验证phi指令的来源值同时应用于多个phi指令时，其中一个phi指令存在被优化异常时，pac适配处理的表现 Signed-off-by: zhangshouxu --- .../AArch64PARTS/AArch64EarlyPartsCpiPass.cpp | 57 ++++- .../AArch64PARTS/AArch64EarlyPartsCpiPass.h | 5 +- .../AArch64EarlyPartsCpiPassTest.cpp | 220 ++++++++++++++++++ 3 files changed, 280 insertions(+), 2 deletions(-) diff --git a/llvm/lib/Target/AArch64/AArch64PARTS/AArch64EarlyPartsCpiPass.cpp b/llvm/lib/Target/AArch64/AArch64PARTS/AArch64EarlyPartsCpiPass.cpp index 577e4f887c4f..6dabee2c23e5 100644 --- a/llvm/lib/Target/AArch64/AArch64PARTS/AArch64EarlyPartsCpiPass.cpp +++ b/llvm/lib/Target/AArch64/AArch64PARTS/AArch64EarlyPartsCpiPass.cpp @@ -70,6 +70,22 @@ inline bool AArch64EarlyPartsCpiPass::handleInstruction(MachineFunction &MF, Mac } auto &MI = *MIi--; for (auto &MI_indcall : IndCallVec) { + if (isIndirectAutCall(*MI_indcall)) { + continue; + } + MachineRegisterInfo *MRI = &MF.getRegInfo(); + MachineInstr *PhiMi = MRI->getVRegDef(MI_indcall->getOperand(0).getReg()); + if (PhiMi->isPHI()) { + SmallVector AutCallVec; + if (!getAllAutCalls(PhiMi, MRI, AutCallVec)) { + /* When an exception occurs in the phi instruction, insert an autia instruction and modify the register used by phi. */ + for (auto &AutCall : AutCallVec) { + insertAuthenticateInstr(MRI, AutCall, PhiMi); + ++StatAutcall; + } + continue; + } + } replaceBranchByAuthenticatedBranch(MBB, MI_indcall, MI); } // insert a copy instruction, same dest register as MI, keep subsequent instructions do not need to modified @@ -79,7 +95,46 @@ inline bool AArch64EarlyPartsCpiPass::handleInstruction(MachineFunction &MF, Mac return true; } -inline bool AArch64EarlyPartsCpiPass::isPartsAUTCALLIntrinsic(unsigned Opcode) { +inline bool AArch64EarlyPartsCpiPass::getAllAutCalls( + MachineInstr *PhiMi, MachineRegisterInfo *MRI, SmallVector &AutCallVec) const { + bool ret = true; + for (unsigned i = 1, e = PhiMi->getNumOperands(); i < e; i += 2) { + MachineOperand &MO = PhiMi->getOperand(i); + MachineInstr *CopyMi = MRI->getVRegDef(MO.getReg()); + MachineInstr *AutCall = CopyMi->isCopy() ? MRI->getVRegDef(CopyMi->getOperand(1).getReg()) : CopyMi; + if (!isPartsAUTCALLIntrinsic(AutCall->getOpcode())) { + ret = false; + continue; + } + + AutCallVec.push_back(AutCall); + } + return ret; +} + +inline void AArch64EarlyPartsCpiPass::insertAuthenticateInstr( + MachineRegisterInfo *MRI, MachineInstr *AutCall, MachineInstr *PhiMi) { + unsigned NewDestReg = MRI->createVirtualRegister(&AArch64::GPR64spRegClass); + auto BMI = BuildMI( + *AutCall->getParent(), *AutCall, AutCall->getDebugLoc(), TII->get(AArch64::AUTIA)); + BMI.addDef(NewDestReg); + auto &Op1 = AutCall->getOperand(1); + BMI.addReg(Op1.getReg(), getRegState(Op1) & ~RegState::Kill); + auto &Op2 = AutCall->getOperand(2); + BMI.addReg(Op2.getReg(), getRegState(Op2) & ~RegState::Kill); + // Correct the register referenced by the phi instruction + for (unsigned i = 1; i < PhiMi->getNumOperands(); i += 2) { + MachineOperand &MO = PhiMi->getOperand(i); + MachineInstr *CopyMi = MRI->getVRegDef(MO.getReg()); + MachineInstr *Tmp = CopyMi->isCopy() ? MRI->getVRegDef(CopyMi->getOperand(1).getReg()) : CopyMi; + if (MO.isReg() && AutCall->getOperand(0).getReg() == Tmp->getOperand(0).getReg()) { + MO.setReg(NewDestReg); + break; + } + } +} + +inline bool AArch64EarlyPartsCpiPass::isPartsAUTCALLIntrinsic(unsigned Opcode) const { switch (Opcode) { case AArch64::PARTS_AUTCALL: return true; diff --git a/llvm/lib/Target/AArch64/AArch64PARTS/AArch64EarlyPartsCpiPass.h b/llvm/lib/Target/AArch64/AArch64PARTS/AArch64EarlyPartsCpiPass.h index cab00dda41ee..29821f0a8cdf 100644 --- a/llvm/lib/Target/AArch64/AArch64PARTS/AArch64EarlyPartsCpiPass.h +++ b/llvm/lib/Target/AArch64/AArch64PARTS/AArch64EarlyPartsCpiPass.h @@ -12,6 +12,7 @@ #include "AArch64.h" #include "AArch64Subtarget.h" +#include "llvm/ADT/SmallVector.h" #include "llvm/CodeGen/MachineFunction.h" #include "llvm/CodeGen/MachineModuleInfo.h" #include "llvm/IR/Module.h" @@ -33,11 +34,13 @@ private: const AArch64Subtarget *STI = nullptr; const AArch64InstrInfo *TII = nullptr; inline bool handleInstruction(MachineFunction &MF, MachineBasicBlock &MBB, MachineBasicBlock::instr_iterator &MIi); + inline bool getAllAutCalls(MachineInstr *PhiMi, MachineRegisterInfo *MRI, SmallVector &AutCallVec) const; + inline void insertAuthenticateInstr(MachineRegisterInfo *MRI, MachineInstr *AutCall, MachineInstr *PhiMi); inline void findIndirectCallMachineInstr(MachineFunction &MF, MachineBasicBlock &MBB, MachineInstr *MIptr, SmallVector &IndCallVec); void triggerCompilationErrorOrphanAUTCALL(MachineBasicBlock &MBB); inline bool isIndirectCall(const MachineInstr &MI) const; - inline bool isPartsAUTCALLIntrinsic(unsigned Opcode); + inline bool isPartsAUTCALLIntrinsic(unsigned Opcode) const; inline const MCInstrDesc &getIndirectCallAuth(MachineInstr *MI_indcall); inline void replaceBranchByAuthenticatedBranch(MachineBasicBlock &MBB, MachineInstr *MI_indcall, MachineInstr &MI); inline void insertCOPYInstr(MachineBasicBlock &MBB, MachineInstr *MI_indcall, MachineInstr &MI); diff --git a/llvm/unittests/Target/AArch64/AArch64PARTS/AArch64EarlyPartsCpiPassTest.cpp b/llvm/unittests/Target/AArch64/AArch64PARTS/AArch64EarlyPartsCpiPassTest.cpp index b7eaa5a61a6c..dfa4fcae73c7 100644 --- a/llvm/unittests/Target/AArch64/AArch64PARTS/AArch64EarlyPartsCpiPassTest.cpp +++ b/llvm/unittests/Target/AArch64/AArch64PARTS/AArch64EarlyPartsCpiPassTest.cpp @@ -145,4 +145,224 @@ body: | bool ret = funcPass->runOnMachineFunction(*MF); ASSERT_TRUE(ret); } + +// Verify that the phi part branch address has been optimized to an address constant, +// and expect to add autia instructions to compensate for other branches. +TEST(EarlyPartsCpi, PhiBlrBranchOptimized) { + std::unique_ptr TM = createTargetMachine(); + ASSERT_TRUE(TM); + + MachineModuleInfo MMI(TM.get()); + SmallString<5000> S; + StringRef MIRString = Twine(R"MIR( +--- | + define i32 @pac_test(ptr nocapture noundef readonly %0, i64 noundef %1, ptr nocapture noundef readonly %2, ptr noundef %3) { + unreachable + } +... +--- +name: pac_test +tracksRegLiveness: true +body: | + bb.0: + liveins: $x0, $x1, $x2, $x3 + successors: %bb.1 + %59:gpr64 = COPY $x3 + %58:gpr64 = COPY $x2 + %57:gpr64 = COPY $x1 + %56:gpr64common = COPY $x0 + %62:gpr32all = COPY $wzr + %60:gpr32all = COPY %62:gpr32all + B %bb.1 + + bb.1: + ; predecessors: %bb.0 + successors: %bb.2, %bb.3 + %99:gpr32 = MOVi32imm 22545 + %100:gpr64sp = SUBREG_TO_REG 0, killed %99:gpr32, %subreg.sub_32 + %101:gpr64 = PARTS_AUTCALL %58:gpr64, killed %100:gpr64sp + ADJCALLSTACKDOWN 0, 0, implicit-def dead $sp, implicit $sp + $x0 = COPY %57:gpr64 + $x1 = COPY %59:gpr64 + BLR killed %101:gpr64 + ADJCALLSTACKUP 0, 0, implicit-def dead $sp, implicit $sp + CBNZX %57:gpr64, %bb.3 + B %bb.2 + + bb.2: + ; predecessors: %bb.1 + successors: %bb.4 + %131:gpr32 = MOVi32imm 22545 + %132:gpr64sp = SUBREG_TO_REG 0, killed %131:gpr32, %subreg.sub_32 + %133:gpr64 = PARTS_AUTCALL %58:gpr64, killed %132:gpr64sp + %31:gpr64 = COPY %133:gpr64 + B %bb.4 + + bb.3: + ; predecessors: %bb.1 + successors: %bb.4 + %125:gpr32 = MOVi32imm 22545 + %126:gpr64sp = SUBREG_TO_REG 0, killed %125:gpr32, %subreg.sub_32 + %127:gpr64common = MOVaddr target-flags(aarch64-page) @pac_test, target-flags(aarch64-pageoff, aarch64-nc) @pac_test + %32:gpr64 = COPY %127:gpr64common + ADJCALLSTACKDOWN 0, 0, implicit-def dead $sp, implicit $sp + $x0 = COPY %57:gpr64 + $x1 = COPY %59:gpr64 + BLR %127:gpr64common + ADJCALLSTACKUP 0, 0, implicit-def dead $sp, implicit $sp + B %bb.4 + + bb.4: + ; predecessors: %bb.2, %bb.3 + successors: %bb.5 + %35:gpr64 = PHI %31:gpr64, %bb.2, %32:gpr64, %bb.3 + B %bb.5 + + bb.5: + ; predecessors: %bb.4 + ADJCALLSTACKDOWN 0, 0, implicit-def dead $sp, implicit $sp + $x0 = COPY %57:gpr64 + $x1 = COPY %59:gpr64 + BLR %35:gpr64 + ADJCALLSTACKUP 0, 0, implicit-def dead $sp, implicit $sp + RET_ReallyLR implicit $w0 +)MIR").toNullTerminatedStringRef(S);; + + LLVMContext Context; + std::unique_ptr MIR; + std::unique_ptr M = parseMIR(Context, MIR, *TM, MIRString, "pac_test", MMI); + ASSERT_TRUE(M); + + Function *F = M->getFunction("pac_test"); + auto *MF = MMI.getMachineFunction(*F); + ASSERT_TRUE(MF); + std::unique_ptr funcPass(new AArch64EarlyPartsCpiPass()); + ASSERT_TRUE(funcPass); + + bool ret = funcPass->runOnMachineFunction(*MF); + ASSERT_TRUE(ret); +} + +// Verify the same function pointer for multiple phi and multiple blr jumps, +// and some branch addresses have been optimized. +TEST(EarlyPartsCpi, PhiMultiBlr) { + std::unique_ptr TM = createTargetMachine(); + ASSERT_TRUE(TM); + + MachineModuleInfo MMI(TM.get()); + SmallString<5000> S; + StringRef MIRString = Twine(R"MIR( +--- | + define i32 @pac_test(ptr nocapture noundef readonly %0, i64 noundef %1, ptr nocapture noundef readonly %2, ptr noundef %3) { + unreachable + } +... +--- +name: pac_test +tracksRegLiveness: true +body: | + bb.0: + liveins: $x0, $x1, $x2, $x3 + successors: %bb.1 + %59:gpr64 = COPY $x3 + %58:gpr64 = COPY $x2 + %57:gpr64 = COPY $x1 + %56:gpr64common = COPY $x0 + %62:gpr32all = COPY $wzr + %60:gpr32all = COPY %62:gpr32all + B %bb.1 + + bb.1: + ; predecessors: %bb.0 + successors: %bb.2, %bb.3, %bb.4 + %99:gpr32 = MOVi32imm 22545 + %100:gpr64sp = SUBREG_TO_REG 0, killed %99:gpr32, %subreg.sub_32 + %101:gpr64 = PARTS_AUTCALL %58:gpr64, killed %100:gpr64sp + ADJCALLSTACKDOWN 0, 0, implicit-def dead $sp, implicit $sp + $x0 = COPY %57:gpr64 + $x1 = COPY %59:gpr64 + BLR killed %101:gpr64 + ADJCALLSTACKUP 0, 0, implicit-def dead $sp, implicit $sp + CBNZX %59:gpr64, %bb.3 + CBNZX %57:gpr64, %bb.4 + B %bb.2 + + bb.2: + ; predecessors: %bb.1 + successors: %bb.5, %bb.6 + %131:gpr32 = MOVi32imm 22545 + %132:gpr64sp = SUBREG_TO_REG 0, killed %131:gpr32, %subreg.sub_32 + %133:gpr64 = PARTS_AUTCALL %58:gpr64, killed %132:gpr64sp + %31:gpr64 = COPY %133:gpr64 + CBNZX %31:gpr64, %bb.6 + B %bb.5 + + bb.3: + ; predecessors: %bb.1 + successors: %bb.6 + %141:gpr32 = MOVi32imm 22545 + %142:gpr64sp = SUBREG_TO_REG 0, killed %141:gpr32, %subreg.sub_32 + %143:gpr64 = PARTS_AUTCALL %58:gpr64, killed %142:gpr64sp + %41:gpr64 = COPY %143:gpr64 + B %bb.6 + + bb.4: + ; predecessors: %bb.1 + successors: %bb.5 + %125:gpr32 = MOVi32imm 22545 + %126:gpr64sp = SUBREG_TO_REG 0, killed %125:gpr32, %subreg.sub_32 + %127:gpr64common = MOVaddr target-flags(aarch64-page) @pac_test, target-flags(aarch64-pageoff, aarch64-nc) @pac_test + %32:gpr64 = COPY %127:gpr64common + ADJCALLSTACKDOWN 0, 0, implicit-def dead $sp, implicit $sp + $x0 = COPY %57:gpr64 + $x1 = COPY %59:gpr64 + BLR %127:gpr64common + ADJCALLSTACKUP 0, 0, implicit-def dead $sp, implicit $sp + B %bb.5 + + bb.5: + ; predecessors: %bb.2, %bb.4 + successors: %bb.7 + %35:gpr64 = PHI %31:gpr64, %bb.2, %32:gpr64, %bb.4 + B %bb.7 + + bb.6: + ; predecessors: %bb.2, %bb.3 + successors: %bb.8 + %40:gpr64 = PHI %31:gpr64, %bb.2, %41:gpr64, %bb.3 + B %bb.8 + + bb.7: + ; predecessors: %bb.5 + ADJCALLSTACKDOWN 0, 0, implicit-def dead $sp, implicit $sp + $x0 = COPY %57:gpr64 + $x1 = COPY %59:gpr64 + BLR %35:gpr64 + ADJCALLSTACKUP 0, 0, implicit-def dead $sp, implicit $sp + RET_ReallyLR implicit $w0 + + bb.8: + ; predecessors: %bb.6 + ADJCALLSTACKDOWN 0, 0, implicit-def dead $sp, implicit $sp + $x0 = COPY %57:gpr64 + $x1 = COPY %59:gpr64 + BLR %40:gpr64 + ADJCALLSTACKUP 0, 0, implicit-def dead $sp, implicit $sp + RET_ReallyLR implicit $w0 +)MIR").toNullTerminatedStringRef(S);; + + LLVMContext Context; + std::unique_ptr MIR; + std::unique_ptr M = parseMIR(Context, MIR, *TM, MIRString, "pac_test", MMI); + ASSERT_TRUE(M); + + Function *F = M->getFunction("pac_test"); + auto *MF = MMI.getMachineFunction(*F); + ASSERT_TRUE(MF); + std::unique_ptr funcPass(new AArch64EarlyPartsCpiPass()); + ASSERT_TRUE(funcPass); + + bool ret = funcPass->runOnMachineFunction(*MF); + ASSERT_TRUE(ret); +} } \ No newline at end of file -- Gitee