diff --git a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp index 1a2808f4d56d8..e7135da17a8d5 100644 --- a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp +++ b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp @@ -155,6 +155,13 @@ class AArch64AsmPrinter : public AsmPrinter { void emitSled(const MachineInstr &MI, SledKind Kind); + bool isPtrauthRegSafe(Register Reg) const { + if (STI->isX16X17Safer()) + return Reg == AArch64::X16 || Reg == AArch64::X17; + + return true; + } + // Emit the sequence for BRA/BLRA (authenticate + branch/call). void emitPtrauthBranch(const MachineInstr *MI); @@ -180,13 +187,15 @@ class AArch64AsmPrinter : public AsmPrinter { // Emit the sequence to compute the discriminator. // - // The returned register is either unmodified AddrDisc or ScratchReg. + // The Scratch register passed to this function must be safe, see + // isPtrauthRegSafe(Reg) function. + // + // The returned register is either ScratchReg or AddrDisc. Furthermore, it + // is safe, unless unsafe AddrDisc was passed-through unmodified. // // If the expanded pseudo is allowed to clobber AddrDisc register, setting - // MayUseAddrAsScratch may save one MOV instruction, provided the address - // is already in x16/x17 (i.e. return x16/x17 which is the *modified* AddrDisc - // register at the same time) or the OS doesn't make it safer to use x16/x17 - // (see AArch64Subtarget::isX16X17Safer()): + // MayUseAddrAsScratch may save one MOV instruction, provided + // isPtrauthRegSafe(AddrDisc) is true: // // mov x17, x16 // movk x17, #1234, lsl #48 @@ -195,7 +204,7 @@ class AArch64AsmPrinter : public AsmPrinter { // can be replaced by // // movk x16, #1234, lsl #48 - Register emitPtrauthDiscriminator(uint16_t Disc, Register AddrDisc, + Register emitPtrauthDiscriminator(uint64_t Disc, Register AddrDisc, Register ScratchReg, bool MayUseAddrAsScratch = false); @@ -1902,12 +1911,14 @@ void AArch64AsmPrinter::emitFMov0AsFMov(const MachineInstr &MI, EmitToStreamer(*OutStreamer, FMov); } -Register AArch64AsmPrinter::emitPtrauthDiscriminator(uint16_t Disc, +Register AArch64AsmPrinter::emitPtrauthDiscriminator(uint64_t Disc, Register AddrDisc, Register ScratchReg, bool MayUseAddrAsScratch) { - assert(ScratchReg == AArch64::X16 || ScratchReg == AArch64::X17 || - !STI->isX16X17Safer()); + assert(isPtrauthRegSafe(ScratchReg) && + "Safe scratch register must be provided by the caller"); + assert(isUInt<16>(Disc) && "Constant discriminator is too wide"); + // So far we've used NoRegister in pseudos. Now we need real encodings. if (AddrDisc == AArch64::NoRegister) AddrDisc = AArch64::XZR; @@ -1926,13 +1937,13 @@ Register AArch64AsmPrinter::emitPtrauthDiscriminator(uint16_t Disc, // If there are both, emit a blend into the scratch register. // Check if we can save one MOV instruction. - assert(MayUseAddrAsScratch || ScratchReg != AddrDisc); - bool AddrDiscIsSafe = AddrDisc == AArch64::X16 || AddrDisc == AArch64::X17 || - !STI->isX16X17Safer(); - if (MayUseAddrAsScratch && AddrDiscIsSafe) + if (MayUseAddrAsScratch && isPtrauthRegSafe(AddrDisc)) { ScratchReg = AddrDisc; - else + } else { emitMovXReg(ScratchReg, AddrDisc); + assert(ScratchReg != AddrDisc && + "Forbidden to clobber AddrDisc, but have to"); + } emitMOVK(ScratchReg, Disc, 48); return ScratchReg; @@ -2151,7 +2162,6 @@ void AArch64AsmPrinter::emitPtrauthAuthResign( } // Compute aut discriminator - assert(isUInt<16>(AUTDisc)); Register AUTDiscReg = emitPtrauthDiscriminator( AUTDisc, AUTAddrDisc->getReg(), Scratch, AUTAddrDisc->isKill()); bool AUTZero = AUTDiscReg == AArch64::XZR; @@ -2188,7 +2198,6 @@ void AArch64AsmPrinter::emitPtrauthAuthResign( return; // Compute pac discriminator - assert(isUInt<16>(PACDisc)); Register PACDiscReg = emitPtrauthDiscriminator(PACDisc, PACAddrDisc, Scratch); bool PACZero = PACDiscReg == AArch64::XZR; @@ -2223,7 +2232,6 @@ void AArch64AsmPrinter::emitPtrauthSign(const MachineInstr *MI) { "Neither X16 nor X17 is available as a scratch register"); // Compute pac discriminator - assert(isUInt<16>(Disc)); Register DiscReg = emitPtrauthDiscriminator( Disc, AddrDisc, ScratchReg, /*MayUseAddrAsScratch=*/AddrDiscKilled); bool IsZeroDisc = DiscReg == AArch64::XZR; @@ -2249,7 +2257,6 @@ void AArch64AsmPrinter::emitPtrauthBranch(const MachineInstr *MI) { "Invalid auth call key"); uint64_t Disc = MI->getOperand(2).getImm(); - assert(isUInt<16>(Disc)); unsigned AddrDisc = MI->getOperand(3).getReg(); @@ -2422,8 +2429,6 @@ void AArch64AsmPrinter::LowerMOVaddrPAC(const MachineInstr &MI) { const auto Key = (AArch64PACKey::ID)KeyC; const unsigned AddrDisc = MI.getOperand(2).getReg(); const uint64_t Disc = MI.getOperand(3).getImm(); - assert(isUInt<16>(Disc) && - "constant discriminator is out of range [0, 0xffff]"); const int64_t Offset = GAOp.getOffset(); GAOp.setOffset(0); @@ -2995,7 +3000,6 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) { "Invalid auth key for tail-call return"); const uint64_t Disc = MI->getOperand(3).getImm(); - assert(isUInt<16>(Disc) && "Integer discriminator is too wide"); Register AddrDisc = MI->getOperand(4).getReg();