[PAuthLR] Add support for FEAT_PAuth_LR's DWARF frame instruction

This introduces compiler and dwarfdump support for emitting
and parsing the new `DW_CFA_AARCH64_negate_ra_state_with_pc`
DWARF instruction for FEAT_PAuth_LR.

This does mean that, when using FEAT_PAuthLR, the improvements
introduced in llvm#96337 cannot be utilised. `.cfi_negate_ra_state_with_pc`
must be emitted directly after the signing instruction, and when bundled
with other CFI calls, leads to faults when running a program. There are
no changes seen when not using FEAT_PAuthLR to how the CFI Instructions
are generated.

See ARM-software/abi-aa#245 for the ABI change
that incororates FEAT_PAuthLR.

Authored-by: pratlucas <lucas.prates@arm.com>
Co-authored by: vhscampos <victor.campos@arm.com>
Co-authored by: Stylie777 <jack.styles@arm.com>

llvm/include/llvm/BinaryFormat/Dwarf.def

@@ -1238,6 +1238,7 @@ HANDLE_DW_CFA(0x16, val_expression)
 // Vendor extensions:
 HANDLE_DW_CFA_PRED(0x1d, MIPS_advance_loc8, SELECT_MIPS64)
 HANDLE_DW_CFA_PRED(0x2d, GNU_window_save, SELECT_SPARC)
+HANDLE_DW_CFA_PRED(0x2c, AARCH64_negate_ra_state_with_pc, SELECT_AARCH64)
 HANDLE_DW_CFA_PRED(0x2d, AARCH64_negate_ra_state, SELECT_AARCH64)
 HANDLE_DW_CFA_PRED(0x2e, GNU_args_size, SELECT_X86)
 // Heterogeneous Debugging Extension defined at

llvm/include/llvm/MC/MCDwarf.h

@@ -515,6 +515,7 @@ class MCCFIInstruction {
     OpRegister,
     OpWindowSave,
     OpNegateRAState,
+    OpNegateRAStateWithPC,
     OpGnuArgsSize,
     OpLabel,
   };
@@ -642,6 +643,13 @@ class MCCFIInstruction {
     return MCCFIInstruction(OpNegateRAState, L, 0, INT64_C(0), Loc);
   }
 
+  /// .cfi_negate_ra_state_with_pc AArch64 negate RA state with PC.
+  static MCCFIInstruction createNegateRAStateWithPC(MCSymbol *L,
+                                                    SMLoc Loc = {}) {
+    return MCCFIInstruction(OpNegateRAStateWithPC, L, 0, INT64_C(0), Loc);
+  }
+
+
   /// .cfi_restore says that the rule for Register is now the same as it
   /// was at the beginning of the function, after all initial instructions added
   /// by .cfi_startproc were executed.

llvm/include/llvm/MC/MCStreamer.h

@@ -1022,6 +1022,7 @@ class MCStreamer {
                                SMLoc Loc = {});
   virtual void emitCFIWindowSave(SMLoc Loc = {});
   virtual void emitCFINegateRAState(SMLoc Loc = {});
+  virtual void emitCFINegateRAStateWithPC(SMLoc Loc = {});
   virtual void emitCFILabelDirective(SMLoc Loc, StringRef Name);
 
   virtual void emitWinCFIStartProc(const MCSymbol *Symbol, SMLoc Loc = SMLoc());

llvm/lib/CodeGen/AsmPrinter/AsmPrinterDwarf.cpp

@@ -236,6 +236,9 @@ void AsmPrinter::emitCFIInstruction(const MCCFIInstruction &Inst) const {
   case MCCFIInstruction::OpNegateRAState:
     OutStreamer->emitCFINegateRAState(Loc);
     break;
+  case MCCFIInstruction::OpNegateRAStateWithPC:
+    OutStreamer->emitCFINegateRAStateWithPC(Loc);
+    break;
   case MCCFIInstruction::OpSameValue:
     OutStreamer->emitCFISameValue(Inst.getRegister(), Loc);
     break;

llvm/lib/CodeGen/CFIInstrInserter.cpp

@@ -260,6 +260,7 @@ void CFIInstrInserter::calculateOutgoingCFAInfo(MBBCFAInfo &MBBInfo) {
       case MCCFIInstruction::OpEscape:
       case MCCFIInstruction::OpWindowSave:
       case MCCFIInstruction::OpNegateRAState:
+      case MCCFIInstruction::OpNegateRAStateWithPC:
       case MCCFIInstruction::OpGnuArgsSize:
       case MCCFIInstruction::OpLabel:
         break;

llvm/lib/CodeGen/MIRParser/MILexer.cpp

@@ -238,6 +238,8 @@ static MIToken::TokenKind getIdentifierKind(StringRef Identifier) {
       .Case("window_save", MIToken::kw_cfi_window_save)
       .Case("negate_ra_sign_state",
             MIToken::kw_cfi_aarch64_negate_ra_sign_state)
+      .Case("negate_ra_sign_state_with_pc",
+            MIToken::kw_cfi_aarch64_negate_ra_sign_state_with_pc)
       .Case("blockaddress", MIToken::kw_blockaddress)
       .Case("intrinsic", MIToken::kw_intrinsic)
       .Case("target-index", MIToken::kw_target_index)

llvm/lib/CodeGen/MIRParser/MILexer.h

@@ -96,6 +96,7 @@ struct MIToken {
     kw_cfi_undefined,
     kw_cfi_window_save,
     kw_cfi_aarch64_negate_ra_sign_state,
+    kw_cfi_aarch64_negate_ra_sign_state_with_pc,
     kw_blockaddress,
     kw_intrinsic,
     kw_target_index,

llvm/lib/CodeGen/MIRParser/MIParser.cpp

@@ -2565,6 +2565,10 @@ bool MIParser::parseCFIOperand(MachineOperand &Dest) {
   case MIToken::kw_cfi_aarch64_negate_ra_sign_state:
     CFIIndex = MF.addFrameInst(MCCFIInstruction::createNegateRAState(nullptr));
     break;
+  case MIToken::kw_cfi_aarch64_negate_ra_sign_state_with_pc:
+    CFIIndex =
+        MF.addFrameInst(MCCFIInstruction::createNegateRAStateWithPC(nullptr));
+    break;
   case MIToken::kw_cfi_escape: {
     std::string Values;
     if (parseCFIEscapeValues(Values))
@@ -2920,6 +2924,7 @@ bool MIParser::parseMachineOperand(const unsigned OpCode, const unsigned OpIdx,
   case MIToken::kw_cfi_undefined:
   case MIToken::kw_cfi_window_save:
   case MIToken::kw_cfi_aarch64_negate_ra_sign_state:
+  case MIToken::kw_cfi_aarch64_negate_ra_sign_state_with_pc:
     return parseCFIOperand(Dest);
   case MIToken::kw_blockaddress:
     return parseBlockAddressOperand(Dest);

llvm/lib/CodeGen/MachineOperand.cpp

@@ -768,6 +768,10 @@ static void printCFI(raw_ostream &OS, const MCCFIInstruction &CFI,
     if (MCSymbol *Label = CFI.getLabel())
       MachineOperand::printSymbol(OS, *Label);
     break;
+  case MCCFIInstruction::OpNegateRAStateWithPC:
+    OS << "negate_ra_sign_state_with_pc ";
+    if (MCSymbol *Label = CFI.getLabel())
+      MachineOperand::printSymbol(OS, *Label);
   default:
     // TODO: Print the other CFI Operations.
     OS << "<unserializable cfi directive>";

llvm/lib/DebugInfo/DWARF/DWARFDebugFrame.cpp

@@ -288,6 +288,7 @@ Error CFIProgram::parse(DWARFDataExtractor Data, uint64_t *Offset,
     case DW_CFA_remember_state:
     case DW_CFA_restore_state:
     case DW_CFA_GNU_window_save:
+    case DW_CFA_AARCH64_negate_ra_state_with_pc:
       // No operands
       addInstruction(Opcode);
       break;
@@ -666,6 +667,28 @@ Error UnwindTable::parseRows(const CFIProgram &CFIP, UnwindRow &Row,
       }
       break;
 
+    case dwarf::DW_CFA_AARCH64_negate_ra_state_with_pc: {
+        constexpr uint32_t AArch64DWARFPAuthRaState = 34;
+        auto LRLoc = Row.getRegisterLocations().getRegisterLocation(
+            AArch64DWARFPAuthRaState);
+        if (LRLoc) {
+          if (LRLoc->getLocation() == UnwindLocation::Constant) {
+            // Toggle the constant value of bits[1:0] from 0 to 1 or 1 to 0.
+            LRLoc->setConstant(LRLoc->getConstant() ^ 0x3);
+          } else {
+            return createStringError(
+                errc::invalid_argument,
+                "%s encountered when existing rule for this register is not "
+                "a constant",
+                CFIP.callFrameString(Inst.Opcode).str().c_str());
+          }
+        } else {
+          Row.getRegisterLocations().setRegisterLocation(
+              AArch64DWARFPAuthRaState, UnwindLocation::createIsConstant(0x3));
+        }
+        break;
+      }
+
     case dwarf::DW_CFA_undefined: {
       llvm::Expected<uint64_t> RegNum = Inst.getOperandAsUnsigned(CFIP, 0);
       if (!RegNum)
@@ -847,6 +870,7 @@ CFIProgram::getOperandTypes() {
   DECLARE_OP0(DW_CFA_remember_state);
   DECLARE_OP0(DW_CFA_restore_state);
   DECLARE_OP0(DW_CFA_GNU_window_save);
+  DECLARE_OP0(DW_CFA_AARCH64_negate_ra_state_with_pc);
   DECLARE_OP1(DW_CFA_GNU_args_size, OT_Offset);
   DECLARE_OP0(DW_CFA_nop);
 

llvm/lib/MC/MCAsmStreamer.cpp

@@ -373,6 +373,7 @@ class MCAsmStreamer final : public MCStreamer {
                        SMLoc Loc) override;
   void emitCFIWindowSave(SMLoc Loc) override;
   void emitCFINegateRAState(SMLoc Loc) override;
+  void emitCFINegateRAStateWithPC(SMLoc Loc) override;
   void emitCFIReturnColumn(int64_t Register) override;
   void emitCFILabelDirective(SMLoc Loc, StringRef Name) override;
 
@@ -2145,6 +2146,12 @@ void MCAsmStreamer::emitCFINegateRAState(SMLoc Loc) {
   EmitEOL();
 }
 
+void MCAsmStreamer::emitCFINegateRAStateWithPC(SMLoc Loc) {
+  MCStreamer::emitCFINegateRAStateWithPC(Loc);
+  OS << "\t.cfi_negate_ra_state_with_pc";
+  EmitEOL();
+}
+
 void MCAsmStreamer::emitCFIReturnColumn(int64_t Register) {
   MCStreamer::emitCFIReturnColumn(Register);
   OS << "\t.cfi_return_column ";

llvm/lib/MC/MCDwarf.cpp

@@ -1381,6 +1381,10 @@ void FrameEmitterImpl::emitCFIInstruction(const MCCFIInstruction &Instr) {
     Streamer.emitInt8(dwarf::DW_CFA_AARCH64_negate_ra_state);
     return;
 
+  case MCCFIInstruction::OpNegateRAStateWithPC:
+    Streamer.emitInt8(dwarf::DW_CFA_AARCH64_negate_ra_state_with_pc);
+    return;
+
   case MCCFIInstruction::OpUndefined: {
     unsigned Reg = Instr.getRegister();
     Streamer.emitInt8(dwarf::DW_CFA_undefined);

llvm/lib/MC/MCStreamer.cpp

@@ -688,6 +688,16 @@ void MCStreamer::emitCFINegateRAState(SMLoc Loc) {
   CurFrame->Instructions.push_back(Instruction);
 }
 
+void MCStreamer::emitCFINegateRAStateWithPC(SMLoc Loc) {
+  MCSymbol *Label = emitCFILabel();
+  MCCFIInstruction Instruction =
+      MCCFIInstruction::createNegateRAStateWithPC(Label, Loc);
+  MCDwarfFrameInfo *CurFrame = getCurrentDwarfFrameInfo();
+  if (!CurFrame)
+    return;
+  CurFrame->Instructions.push_back(Instruction);
+}
+
 void MCStreamer::emitCFIReturnColumn(int64_t Register) {
   MCDwarfFrameInfo *CurFrame = getCurrentDwarfFrameInfo();
   if (!CurFrame)

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

@@ -702,7 +702,10 @@ void AArch64FrameLowering::resetCFIToInitialState(
 
   // Flip the RA sign state.
   if (MFI.shouldSignReturnAddress(MF)) {
-    CFIIndex = MF.addFrameInst(MCCFIInstruction::createNegateRAState(nullptr));
+    auto CFIInst = MFI.branchProtectionPAuthLR()
+                       ? MCCFIInstruction::createNegateRAStateWithPC(nullptr)
+                       : MCCFIInstruction::createNegateRAState(nullptr);
+    CFIIndex = MF.addFrameInst(CFIInst);
     BuildMI(MBB, InsertPt, DL, CFIDesc).addCFIIndex(CFIIndex);
   }
 

llvm/lib/Target/AArch64/AArch64PointerAuth.cpp

@@ -71,6 +71,18 @@ FunctionPass *llvm::createAArch64PointerAuthPass() {
 
 char AArch64PointerAuth::ID = 0;
 
+static void emitPACSymOffsetIntoX16(const TargetInstrInfo &TII,
+                                    MachineBasicBlock &MBB,
+                                    MachineBasicBlock::iterator I, DebugLoc DL,
+                                    MCSymbol *PACSym) {
+  BuildMI(MBB, I, DL, TII.get(AArch64::ADRP), AArch64::X16)
+      .addSym(PACSym, AArch64II::MO_PAGE);
+  BuildMI(MBB, I, DL, TII.get(AArch64::ADDXri), AArch64::X16)
+      .addReg(AArch64::X16)
+      .addSym(PACSym, AArch64II::MO_PAGEOFF | AArch64II::MO_NC)
+      .addImm(0);
+}
+
 // Where PAuthLR support is not known at compile time, it is supported using
 // PACM. PACM is in the hint space so has no effect when PAuthLR is not
 // supported by the hardware, but will alter the behaviour of PACI*SP, AUTI*SP
@@ -81,12 +93,10 @@ static void BuildPACM(const AArch64Subtarget &Subtarget, MachineBasicBlock &MBB,
   const TargetInstrInfo *TII = Subtarget.getInstrInfo();
   auto &MFnI = *MBB.getParent()->getInfo<AArch64FunctionInfo>();
 
-  // ADR X16,<address_of_PACIASP>
+  // Offset to PAC*SP using ADRP + ADD.
   if (PACSym) {
     assert(Flags == MachineInstr::FrameDestroy);
-    BuildMI(MBB, MBBI, DL, TII->get(AArch64::ADR))
-        .addReg(AArch64::X16, RegState::Define)
-        .addSym(PACSym);
+    emitPACSymOffsetIntoX16(*TII, MBB, MBBI, DL, PACSym);
   }
 
   // Only emit PACM if -mbranch-protection has +pc and the target does not
@@ -95,12 +105,49 @@ static void BuildPACM(const AArch64Subtarget &Subtarget, MachineBasicBlock &MBB,
     BuildMI(MBB, MBBI, DL, TII->get(AArch64::PACM)).setMIFlag(Flags);
 }
 
+static void emitPACCFI(const AArch64Subtarget &Subtarget,
+                       MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
+                       DebugLoc DL, MachineInstr::MIFlag Flags, bool EmitCFI) {
+  if (!EmitCFI)
+    return;
+
+  const TargetInstrInfo *TII = Subtarget.getInstrInfo();
+  auto &MF = *MBB.getParent();
+  auto &MFnI = *MF.getInfo<AArch64FunctionInfo>();
+  bool EmitAsyncCFI = MFnI.needsAsyncDwarfUnwindInfo(MF);
+
+  auto CFIInst = MFnI.branchProtectionPAuthLR()
+                     ? MCCFIInstruction::createNegateRAStateWithPC(nullptr)
+                     : MCCFIInstruction::createNegateRAState(nullptr);
+
+  // Because of PAuthLR, when using NegateRAStateWithPC, the CFI instruction cannot
+  // be bundled with other CFI instructions in the prolog, as it needs to directly
+  // follow the signing instruction. This ensures the PC value is captured incase of
+  // an error in the following the following instructions.
+  if (!EmitAsyncCFI && !(MFnI.branchProtectionPAuthLR())) {
+    // Reduce the size of the generated call frame information for synchronous
+    // CFI by bundling the new CFI instruction with others in the prolog, so
+    // that no additional DW_CFA_advance_loc is needed.
+    for (auto I = MBBI; I != MBB.end(); ++I) {
+      if (I->getOpcode() == TargetOpcode::CFI_INSTRUCTION &&
+          I->getFlag(MachineInstr::FrameSetup)) {
+        MBBI = I;
+        break;
+      }
+    }
+  }
+
+  unsigned CFIIndex = MF.addFrameInst(CFIInst);
+  BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
+      .addCFIIndex(CFIIndex)
+      .setMIFlags(Flags);
+}
+
 void AArch64PointerAuth::signLR(MachineFunction &MF,
                                 MachineBasicBlock::iterator MBBI) const {
   auto &MFnI = *MF.getInfo<AArch64FunctionInfo>();
   bool UseBKey = MFnI.shouldSignWithBKey();
   bool EmitCFI = MFnI.needsDwarfUnwindInfo(MF);
-  bool EmitAsyncCFI = MFnI.needsAsyncDwarfUnwindInfo(MF);
   bool NeedsWinCFI = MF.hasWinCFI();
 
   MachineBasicBlock &MBB = *MBBI->getParent();
@@ -128,34 +175,18 @@ void AArch64PointerAuth::signLR(MachineFunction &MF,
                                                : AArch64::PACIASPPC))
         .setMIFlag(MachineInstr::FrameSetup)
         ->setPreInstrSymbol(MF, MFnI.getSigningInstrLabel());
+    emitPACCFI(*Subtarget, MBB, MBBI, DL, MachineInstr::FrameSetup, EmitCFI);
   } else {
     BuildPACM(*Subtarget, MBB, MBBI, DL, MachineInstr::FrameSetup);
     BuildMI(MBB, MBBI, DL,
             TII->get(MFnI.shouldSignWithBKey() ? AArch64::PACIBSP
                                                : AArch64::PACIASP))
         .setMIFlag(MachineInstr::FrameSetup)
         ->setPreInstrSymbol(MF, MFnI.getSigningInstrLabel());
+    emitPACCFI(*Subtarget, MBB, MBBI, DL, MachineInstr::FrameSetup, EmitCFI);
   }
 
-  if (EmitCFI) {
-    if (!EmitAsyncCFI) {
-      // Reduce the size of the generated call frame information for synchronous
-      // CFI by bundling the new CFI instruction with others in the prolog, so
-      // that no additional DW_CFA_advance_loc is needed.
-      for (auto I = MBBI; I != MBB.end(); ++I) {
-        if (I->getOpcode() == TargetOpcode::CFI_INSTRUCTION &&
-            I->getFlag(MachineInstr::FrameSetup)) {
-          MBBI = I;
-          break;
-        }
-      }
-    }
-    unsigned CFIIndex =
-        MF.addFrameInst(MCCFIInstruction::createNegateRAState(nullptr));
-    BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
-        .addCFIIndex(CFIIndex)
-        .setMIFlags(MachineInstr::FrameSetup);
-  } else if (NeedsWinCFI) {
+  if (!EmitCFI && NeedsWinCFI) {
     BuildMI(MBB, MBBI, DL, TII->get(AArch64::SEH_PACSignLR))
         .setMIFlag(MachineInstr::FrameSetup);
   }
@@ -190,6 +221,7 @@ void AArch64PointerAuth::authenticateLR(
       !MF.getFunction().hasFnAttribute(Attribute::ShadowCallStack)) {
     if (MFnI->branchProtectionPAuthLR() && Subtarget->hasPAuthLR()) {
       assert(PACSym && "No PAC instruction to refer to");
+      emitPACSymOffsetIntoX16(*TII, MBB, MBBI, DL, PACSym);
       BuildMI(MBB, TI, DL,
               TII->get(UseBKey ? AArch64::RETABSPPCi : AArch64::RETAASPPCi))
           .addSym(PACSym)
@@ -205,24 +237,20 @@ void AArch64PointerAuth::authenticateLR(
   } else {
     if (MFnI->branchProtectionPAuthLR() && Subtarget->hasPAuthLR()) {
       assert(PACSym && "No PAC instruction to refer to");
+      emitPACSymOffsetIntoX16(*TII, MBB, MBBI, DL, PACSym);
       BuildMI(MBB, MBBI, DL,
               TII->get(UseBKey ? AArch64::AUTIBSPPCi : AArch64::AUTIASPPCi))
           .addSym(PACSym)
           .setMIFlag(MachineInstr::FrameDestroy);
+      emitPACCFI(*Subtarget, MBB, MBBI, DL, MachineInstr::FrameDestroy, EmitAsyncCFI);
     } else {
       BuildPACM(*Subtarget, MBB, MBBI, DL, MachineInstr::FrameDestroy, PACSym);
       BuildMI(MBB, MBBI, DL,
               TII->get(UseBKey ? AArch64::AUTIBSP : AArch64::AUTIASP))
           .setMIFlag(MachineInstr::FrameDestroy);
+      emitPACCFI(*Subtarget, MBB, MBBI, DL, MachineInstr::FrameDestroy, EmitAsyncCFI);
     }
 
-    if (EmitAsyncCFI) {
-      unsigned CFIIndex =
-          MF.addFrameInst(MCCFIInstruction::createNegateRAState(nullptr));
-      BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
-          .addCFIIndex(CFIIndex)
-          .setMIFlags(MachineInstr::FrameDestroy);
-    }
     if (NeedsWinCFI) {
       BuildMI(MBB, MBBI, DL, TII->get(AArch64::SEH_PACSignLR))
           .setMIFlag(MachineInstr::FrameDestroy);

llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp

@@ -195,6 +195,7 @@ class AArch64AsmParser : public MCTargetAsmParser {
   bool parseDirectiveReq(StringRef Name, SMLoc L);
   bool parseDirectiveUnreq(SMLoc L);
   bool parseDirectiveCFINegateRAState();
+  bool parseDirectiveCFINegateRAStateWithPC();
   bool parseDirectiveCFIBKeyFrame();
   bool parseDirectiveCFIMTETaggedFrame();
 
@@ -6821,6 +6822,8 @@ bool AArch64AsmParser::ParseDirective(AsmToken DirectiveID) {
     parseDirectiveInst(Loc);
   else if (IDVal == ".cfi_negate_ra_state")
     parseDirectiveCFINegateRAState();
+  else if (IDVal == ".cfi_negate_ra_state_with_pc")
+    parseDirectiveCFINegateRAStateWithPC();
   else if (IDVal == ".cfi_b_key_frame")
     parseDirectiveCFIBKeyFrame();
   else if (IDVal == ".cfi_mte_tagged_frame")
@@ -7271,6 +7274,13 @@ bool AArch64AsmParser::parseDirectiveCFINegateRAState() {
   return false;
 }
 
+bool AArch64AsmParser::parseDirectiveCFINegateRAStateWithPC() {
+  if (parseEOL())
+    return true;
+  getStreamer().emitCFINegateRAStateWithPC();
+  return false;
+}
+
 /// parseDirectiveCFIBKeyFrame
 /// ::= .cfi_b_key
 bool AArch64AsmParser::parseDirectiveCFIBKeyFrame() {