diff --git a/llvm/lib/Target/AArch64/AArch64PARTS/AArch64EarlyPartsCpiPass.cpp b/llvm/lib/Target/AArch64/AArch64PARTS/AArch64EarlyPartsCpiPass.cpp index 577e4f887c4f5a432535e45c55b8c72848ae20b6..6dabee2c23e5e31f5bac91ed4d0130a1f5b96cee 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 cab00dda41ee0b4f7ec6a711f02a1f3737740f35..29821f0a8cdf3aedaa70f104310c9fad316dfa2f 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 b7eaa5a61a6cd51e5b7d2e42cd23ab7016447072..dfa4fcae73c7ac2f72c475650544832b42105286 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