[AArch64][PAC] Prevent immediate modifier substitution of AUT/PAC

Expand `blend(addr, imm)` in AsmPrinter, so that the discriminator
cannot be changed between blending `imm` in `addr` and using the
discriminator for signing or authenticating a pointer.

For that purpose, a new PAC pseudo instruction is introduced that has
separate operands for register and immediate modifiers. Instruction
selection is refactored along the way, so that AUT and PAC instructions
are selected by tablegen-erated patterns (always with a zero immediate
modifier at first) and a unified fixupBlendComponents function is used
by a custom inserter to fetch operands of `blend()` when possible. Aside
from unification, this helps inspecting operands of blend operation that
is located in another basic block which can be tricky for DAGISel.

This patch is somewhat related to
llvm#132857.

Author: atrosinenko

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

@@ -171,6 +171,9 @@ class AArch64AsmPrinter : public AsmPrinter {
   // Emit the sequence for AUT or AUTPAC.
   void emitPtrauthAuthResign(const MachineInstr *MI);
 
+  // Emit the sequence for PAC.
+  void emitPtrauthSign(const MachineInstr *MI);
+
   // Emit the sequence to compute the discriminator.
   //
   // ScratchReg should be x16/x17.
@@ -2173,6 +2176,30 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(const MachineInstr *MI) {
     OutStreamer->emitLabel(EndSym);
 }
 
+void AArch64AsmPrinter::emitPtrauthSign(const MachineInstr *MI) {
+  Register Val = MI->getOperand(1).getReg();
+  auto Key = (AArch64PACKey::ID)MI->getOperand(2).getImm();
+  uint64_t Disc = MI->getOperand(3).getImm();
+  Register AddrDisc = MI->getOperand(4).getReg();
+  // Can isKill() be reliably used at this point?
+
+  // Compute aut discriminator into x17
+  assert(isUInt<16>(Disc));
+  Register DiscReg = emitPtrauthDiscriminator(Disc, AddrDisc, AArch64::X17);
+  bool IsZeroDisc = DiscReg == AArch64::XZR;
+  unsigned Opc = getPACOpcodeForKey(Key, IsZeroDisc);
+
+  //  paciza x16      ; if  IsZeroDisc
+  //  pacia x16, x17  ; if !IsZeroDisc
+  MCInst PACInst;
+  PACInst.setOpcode(Opc);
+  PACInst.addOperand(MCOperand::createReg(Val));
+  PACInst.addOperand(MCOperand::createReg(Val));
+  if (!IsZeroDisc)
+    PACInst.addOperand(MCOperand::createReg(DiscReg));
+  EmitToStreamer(*OutStreamer, PACInst);
+}
+
 void AArch64AsmPrinter::emitPtrauthBranch(const MachineInstr *MI) {
   bool IsCall = MI->getOpcode() == AArch64::BLRA;
   unsigned BrTarget = MI->getOperand(0).getReg();
@@ -2867,6 +2894,10 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
     emitPtrauthAuthResign(MI);
     return;
 
+  case AArch64::PAC:
+    emitPtrauthSign(MI);
+    return;
+
   case AArch64::LOADauthptrstatic:
     LowerLOADauthptrstatic(*MI);
     return;

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

@@ -361,7 +361,6 @@ class AArch64DAGToDAGISel : public SelectionDAGISel {
 
   bool tryIndexedLoad(SDNode *N);
 
-  void SelectPtrauthAuth(SDNode *N);
   void SelectPtrauthResign(SDNode *N);
 
   bool trySelectStackSlotTagP(SDNode *N);
@@ -1520,28 +1519,6 @@ extractPtrauthBlendDiscriminators(SDValue Disc, SelectionDAG *DAG) {
       AddrDisc);
 }
 
-void AArch64DAGToDAGISel::SelectPtrauthAuth(SDNode *N) {
-  SDLoc DL(N);
-  // IntrinsicID is operand #0
-  SDValue Val = N->getOperand(1);
-  SDValue AUTKey = N->getOperand(2);
-  SDValue AUTDisc = N->getOperand(3);
-
-  unsigned AUTKeyC = cast<ConstantSDNode>(AUTKey)->getZExtValue();
-  AUTKey = CurDAG->getTargetConstant(AUTKeyC, DL, MVT::i64);
-
-  SDValue AUTAddrDisc, AUTConstDisc;
-  std::tie(AUTConstDisc, AUTAddrDisc) =
-      extractPtrauthBlendDiscriminators(AUTDisc, CurDAG);
-
-  SDValue X16Copy = CurDAG->getCopyToReg(CurDAG->getEntryNode(), DL,
-                                         AArch64::X16, Val, SDValue());
-  SDValue Ops[] = {AUTKey, AUTConstDisc, AUTAddrDisc, X16Copy.getValue(1)};
-
-  SDNode *AUT = CurDAG->getMachineNode(AArch64::AUT, DL, MVT::i64, Ops);
-  ReplaceNode(N, AUT);
-}
-
 void AArch64DAGToDAGISel::SelectPtrauthResign(SDNode *N) {
   SDLoc DL(N);
   // IntrinsicID is operand #0
@@ -5780,10 +5757,6 @@ void AArch64DAGToDAGISel::Select(SDNode *Node) {
       SelectTagP(Node);
       return;
 
-    case Intrinsic::ptrauth_auth:
-      SelectPtrauthAuth(Node);
-      return;
-
     case Intrinsic::ptrauth_resign:
       SelectPtrauthResign(Node);
       return;

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -3103,7 +3103,14 @@ void AArch64TargetLowering::fixupBlendComponents(
   Register AddrDisc = AddrDiscOp.getReg();
   int64_t IntDisc = IntDiscOp.getImm();
 
-  assert(IntDisc == 0 && "Blend components are already expanded");
+  // This function tries to expand an opaque vreg discriminator into address
+  // modifier and immediate modifier components. If the immediate modifier
+  // operand is already non-zero, no fixup needed.
+  if (IntDisc != 0) {
+    assert(AddrDisc == AArch64::XZR || AddrDisc == AArch64::NoRegister ||
+           MRI.getRegClass(AddrDisc) == AddrDiscRC);
+    return;
+  }
 
   int64_t Offset = 0;
   MachineInstr *MaybeBlend = stripAndAccumulateOffset(MRI, AddrDisc, Offset);
@@ -3130,51 +3137,70 @@ void AArch64TargetLowering::fixupBlendComponents(
 }
 
 MachineBasicBlock *
-AArch64TargetLowering::tryRewritingPAC(MachineInstr &MI,
-                                       MachineBasicBlock *BB) const {
+AArch64TargetLowering::EmitAUTxMxN(MachineInstr &MI,
+                                   MachineBasicBlock *BB) const {
+  const TargetInstrInfo *TII = Subtarget->getInstrInfo();
+  const DebugLoc &DL = MI.getDebugLoc();
+
+  Register AuthVal = MI.getOperand(0).getReg();
+  Register Val = MI.getOperand(2).getReg();
+  unsigned Key = MI.getOperand(3).getImm();
+  int64_t IntDisc = MI.getOperand(4).getImm();
+  Register AddrDisc = MI.getOperand(5).getReg();
+
+  BuildMI(*BB, MI, DL, TII->get(AArch64::COPY), AArch64::X16).addReg(Val);
+  BuildMI(*BB, MI, DL, TII->get(AArch64::AUT))
+      .addImm(Key)
+      .addImm(IntDisc)
+      .addReg(AddrDisc);
+  BuildMI(*BB, MI, DL, TII->get(AArch64::COPY), AuthVal).addReg(AArch64::X16);
+
+  MI.eraseFromParent();
+  return BB;
+}
+
+void AArch64TargetLowering::tryRewritingPAC(MachineInstr &MI,
+                                            MachineBasicBlock *BB) const {
   const TargetInstrInfo *TII = Subtarget->getInstrInfo();
   MachineRegisterInfo &MRI = MI.getMF()->getRegInfo();
   const DebugLoc &DL = MI.getDebugLoc();
 
+  Register Val = MI.getOperand(1).getReg();
+  unsigned Key = MI.getOperand(2).getImm();
+  int64_t IntDisc = MI.getOperand(3).getImm();
+  Register AddrDisc = MI.getOperand(4).getReg();
+
   // Try to find a known address-setting instruction, accumulating the offset
   // along the way. If no known pattern is found, keep everything as-is.
 
   int64_t AddrOffset = 0;
-  MachineInstr *AddrDefInstr =
-      stripAndAccumulateOffset(MRI, MI.getOperand(1).getReg(), AddrOffset);
+  MachineInstr *AddrDefInstr = stripAndAccumulateOffset(MRI, Val, AddrOffset);
   if (!AddrDefInstr)
-    return BB;
+    return;
 
   unsigned NewOpcode;
   if (AddrDefInstr->getOpcode() == AArch64::LOADgotAUTH)
     NewOpcode = AArch64::LOADgotPAC;
   else if (AddrDefInstr->getOpcode() == AArch64::MOVaddr)
     NewOpcode = AArch64::MOVaddrPAC;
   else
-    return BB; // Unknown opcode.
+    return; // Unknown opcode.
 
   MachineOperand &AddrOp = AddrDefInstr->getOperand(1);
   unsigned TargetFlags = AddrOp.getTargetFlags() & ~AArch64II::MO_PAGE;
   const GlobalValue *GV = AddrOp.getGlobal();
   AddrOffset += AddrOp.getOffset();
 
-  Register DiscReg = isPACWithZeroDisc(MI.getOpcode())
-                         ? AArch64::XZR
-                         : MI.getOperand(2).getReg();
-
-  MachineInstr *NewMI = BuildMI(*BB, MI, DL, TII->get(NewOpcode))
-                            .addGlobalAddress(GV, AddrOffset, TargetFlags)
-                            .addImm(getKeyForPACOpcode(MI.getOpcode()))
-                            .addReg(DiscReg)
-                            .addImm(0);
-  fixupBlendComponents(*NewMI, BB, NewMI->getOperand(3), NewMI->getOperand(2),
-                       &AArch64::GPR64noipRegClass);
+  BuildMI(*BB, MI, DL, TII->get(NewOpcode))
+      .addGlobalAddress(GV, AddrOffset, TargetFlags)
+      .addImm(Key)
+      .addReg(AddrDisc)
+      .addImm(IntDisc);
 
   BuildMI(*BB, MI, DL, TII->get(AArch64::COPY), MI.getOperand(0).getReg())
       .addReg(AArch64::X16);
 
   MI.removeFromParent();
-  return BB;
 }
 
 MachineBasicBlock *AArch64TargetLowering::EmitInstrWithCustomInserter(
@@ -3276,15 +3302,21 @@ MachineBasicBlock *AArch64TargetLowering::EmitInstrWithCustomInserter(
   case AArch64::MOVT_TIZ_PSEUDO:
     return EmitZTInstr(MI, BB, AArch64::MOVT_TIZ, /*Op0IsDef=*/true);
 
-  case AArch64::PACDA:
-  case AArch64::PACDB:
-  case AArch64::PACIA:
-  case AArch64::PACIB:
-  case AArch64::PACDZA:
-  case AArch64::PACDZB:
-  case AArch64::PACIZA:
-  case AArch64::PACIZB:
-    return tryRewritingPAC(MI, BB);
+  case AArch64::PAC:
+    fixupBlendComponents(MI, BB, MI.getOperand(3), MI.getOperand(4),
+                         &AArch64::GPR64noipRegClass);
+    tryRewritingPAC(MI, BB);
+    return BB;
+  case AArch64::AUTxMxN:
+    fixupBlendComponents(MI, BB, MI.getOperand(4), MI.getOperand(5),
+                         &AArch64::GPR64noipRegClass);
+    return EmitAUTxMxN(MI, BB);
+  case AArch64::AUTPAC:
+    fixupBlendComponents(MI, BB, MI.getOperand(1), MI.getOperand(2),
+                         &AArch64::GPR64noipRegClass);
+    fixupBlendComponents(MI, BB, MI.getOperand(4), MI.getOperand(5),
+                         &AArch64::GPR64noipRegClass);
+    return BB;
   }
 }
 

llvm/lib/Target/AArch64/AArch64ISelLowering.h

@@ -187,8 +187,9 @@ class AArch64TargetLowering : public TargetLowering {
                             MachineOperand &AddrDiscOp,
                             const TargetRegisterClass *AddrDiscRC) const;
 
-  MachineBasicBlock *tryRewritingPAC(MachineInstr &MI,
-                                     MachineBasicBlock *BB) const;
+  MachineBasicBlock *EmitAUTxMxN(MachineInstr &MI, MachineBasicBlock *BB) const;
+
+  void tryRewritingPAC(MachineInstr &MI, MachineBasicBlock *BB) const;
 
   MachineBasicBlock *
   EmitInstrWithCustomInserter(MachineInstr &MI,

llvm/lib/Target/AArch64/AArch64InstrInfo.td

@@ -1988,6 +1988,8 @@ def PAUTH_PROLOGUE : Pseudo<(outs), (ins), []>, Sched<[]> {
 def PAUTH_EPILOGUE : Pseudo<(outs), (ins), []>, Sched<[]>;
 }
 
+def ScratchReg : OptionalDefOperand<OtherVT, (ops GPR64), (ops (i64 zero_reg))>;
+
 // These pointer authentication instructions require armv8.3a
 let Predicates = [HasPAuth] in {
 
@@ -2018,9 +2020,7 @@ let Predicates = [HasPAuth] in {
     def DZB  : SignAuthZero<prefix_z,  0b11, !strconcat(asm, "dzb"), op>;
   }
 
-  let usesCustomInserter = true in
-  defm PAC : SignAuth<0b000, 0b010, "pac", int_ptrauth_sign>;
-
+  defm PAC : SignAuth<0b000, 0b010, "pac", null_frag>;
   defm AUT : SignAuth<0b001, 0b011, "aut", null_frag>;
 
   def XPACI : ClearAuth<0, "xpaci">;
@@ -2126,6 +2126,39 @@ let Predicates = [HasPAuth] in {
     let Uses = [X16];
   }
 
+  // ScratchReg could be used by
+  // https://github.com/llvm/llvm-project/pull/132857.
+  def AUTxMxN : Pseudo<(outs GPR64:$AuthVal, GPR64:$ScratchWb),
+                       (ins GPR64:$Val, i32imm:$Key,
+                            i64imm:$Disc, GPR64:$AddrDisc, ScratchReg:$Scratch),
+                       [], "$AuthVal = $Val, $ScratchWb = $Scratch">,
+                Sched<[WriteI, ReadI]> {
+    let isCodeGenOnly = 1;
+    let hasSideEffects = 0;
+    let mayStore = 0;
+    let mayLoad = 0;
+    let Size = 32;
+    let Defs = [NZCV];
+    let usesCustomInserter = true;
+  }
+
+  def PAC : Pseudo<(outs GPR64:$SignedVal),
+                   (ins GPR64:$Val, i32imm:$Key, i64imm:$Disc, GPR64noip:$AddrDisc),
+                   [], "$SignedVal = $Val">, Sched<[WriteI, ReadI]> {
+    let isCodeGenOnly = 1;
+    let hasSideEffects = 0;
+    let mayStore = 0;
+    let mayLoad = 0;
+    let Size = 12;
+    let Defs = [X17];
+    let usesCustomInserter = true;
+  }
+
+  def : Pat<(int_ptrauth_auth GPR64:$Val, timm:$Key, GPR64:$AddrDisc),
+            (AUTxMxN GPR64:$Val, $Key, 0, GPR64:$AddrDisc)>;
+  def : Pat<(int_ptrauth_sign GPR64:$Val, timm:$Key, GPR64noip:$AddrDisc),
+            (PAC GPR64:$Val, $Key, 0, GPR64noip:$AddrDisc)>;
+
   // AUT and re-PAC a value, using different keys/data.
   // This directly manipulates x16/x17, which are the only registers the OS
   // guarantees are safe to use for sensitive operations.
@@ -2141,6 +2174,7 @@ let Predicates = [HasPAuth] in {
     let Size = 48;
     let Defs = [X16,X17,NZCV];
     let Uses = [X16];
+    let usesCustomInserter = true;
   }
 
   // Materialize a signed global address, with adrp+add and PAC.

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

@@ -6727,30 +6727,6 @@ bool AArch64InstructionSelector::selectIntrinsic(MachineInstr &I,
     I.eraseFromParent();
     return true;
   }
-  case Intrinsic::ptrauth_auth: {
-    Register DstReg = I.getOperand(0).getReg();
-    Register ValReg = I.getOperand(2).getReg();
-    uint64_t AUTKey = I.getOperand(3).getImm();
-    Register AUTDisc = I.getOperand(4).getReg();
-
-    Register AUTAddrDisc = AUTDisc;
-    uint16_t AUTConstDiscC = 0;
-    std::tie(AUTConstDiscC, AUTAddrDisc) =
-        extractPtrauthBlendDiscriminators(AUTDisc, MRI);
-
-    MIB.buildCopy({AArch64::X16}, {ValReg});
-    MIB.buildInstr(TargetOpcode::IMPLICIT_DEF, {AArch64::X17}, {});
-    MIB.buildInstr(AArch64::AUT)
-        .addImm(AUTKey)
-        .addImm(AUTConstDiscC)
-        .addUse(AUTAddrDisc)
-        .constrainAllUses(TII, TRI, RBI);
-    MIB.buildCopy({DstReg}, Register(AArch64::X16));
-
-    RBI.constrainGenericRegister(DstReg, AArch64::GPR64RegClass, MRI);
-    I.eraseFromParent();
-    return true;
-  }
   case Intrinsic::frameaddress:
   case Intrinsic::returnaddress: {
     MachineFunction &MF = *I.getParent()->getParent();

llvm/test/CodeGen/AArch64/GlobalISel/ptrauth-constant-in-code.ll

@@ -106,7 +106,7 @@ define void @store_signed_const_local(ptr %dest) {
 ; ISEL-MIR:       body:
 ; ISEL-MIR:         %0:gpr64common = COPY $x0
 ; ISEL-MIR-NEXT:    %10:gpr64common = MOVaddr target-flags(aarch64-page) @const_table_local + 8, target-flags(aarch64-pageoff, aarch64-nc) @const_table_local + 8
-; ISEL-MIR-NEXT:    %2:gpr64common = MOVKXi %0, 1234, 48
+; ISEL-MIR-NEXT:    %2:gpr64noip = MOVKXi %0, 1234, 48
 ; ISEL-MIR-NEXT:    %15:gpr64noip = COPY %0
 ; ISEL-MIR-NEXT:    MOVaddrPAC @const_table_local + 8, 2, %15, 1234, implicit-def $x16, implicit-def $x17
 ; ISEL-MIR-NEXT:    %4:gpr64 = COPY $x16
@@ -140,7 +140,7 @@ define void @store_signed_const_got(ptr %dest) {
 ; ISEL-MIR-ELF:         %0:gpr64common = COPY $x0
 ; ISEL-MIR-ELF-NEXT:    %7:gpr64common = LOADgotAUTH target-flags(aarch64-got) @const_table_got
 ; ISEL-MIR-ELF-NEXT:    %6:gpr64common = ADDXri %7, 8, 0
-; ISEL-MIR-ELF-NEXT:    %2:gpr64common = MOVKXi %0, 1234, 48
+; ISEL-MIR-ELF-NEXT:    %2:gpr64noip = MOVKXi %0, 1234, 48
 ; ISEL-MIR-ELF-NEXT:    %12:gpr64noip = COPY %0
 ; ISEL-MIR-ELF-NEXT:    LOADgotPAC target-flags(aarch64-got) @const_table_got + 8, 2, %12, 1234, implicit-def $x16, implicit-def $x17, implicit-def $nzcv
 ; ISEL-MIR-ELF-NEXT:    %4:gpr64 = COPY $x16
@@ -179,21 +179,21 @@ define void @store_signed_arg(ptr %dest, ptr %p) {
 ; ISEL-MIR:       body:
 ; ISEL-MIR:         %0:gpr64common = COPY $x0
 ; ISEL-MIR-NEXT:    %1:gpr64common = COPY $x1
-; ISEL-MIR-NEXT:    %3:gpr64common = MOVKXi %0, 1234, 48
+; ISEL-MIR-NEXT:    %3:gpr64noip = MOVKXi %0, 1234, 48
 ; ISEL-MIR-NEXT:    %6:gpr64common = ADDXri %1, 8, 0
-; Check that no implicit defs are added to PACDA instruction.
-; ISEL-MIR-NEXT:    %8:gpr64 = PACDA %6, %3{{$}}
+; ISEL-MIR-NEXT:    %12:gpr64noip = COPY %0
+; ISEL-MIR-NEXT:    %8:gpr64 = PAC %6, 2, 1234, %12, implicit-def dead $x17
 ; ISEL-MIR-NEXT:    %10:gpr64 = COPY %8
 ; ISEL-MIR-NEXT:    STRXui %10, %0, 0 :: (store (p0) into %ir.dest)
 ; ISEL-MIR-NEXT:    RET_ReallyLR
 ;
 ; ISEL-ASM-LABEL: store_signed_arg:
 ; ISEL-ASM-NEXT:    .cfi_startproc
-; ISEL-ASM-NEXT:    mov     x8, x0
-; ISEL-ASM-NEXT:    add     x9, x1, #8
-; ISEL-ASM-NEXT:    movk    x8, #1234, lsl #48
-; ISEL-ASM-NEXT:    pacda   x9, x8
-; ISEL-ASM-NEXT:    str     x9, [x0]
+; ISEL-ASM-NEXT:    add     x8, x1, #8
+; ISEL-ASM-NEXT:    mov     x17, x0
+; ISEL-ASM-NEXT:    movk    x17, #1234, lsl #48
+; ISEL-ASM-NEXT:    pacda   x8, x17
+; ISEL-ASM-NEXT:    str     x8, [x0]
 ; ISEL-ASM-NEXT:    ret
   %dest.i = ptrtoint ptr %dest to i64
   %discr = call i64 @llvm.ptrauth.blend(i64 %dest.i, i64 1234)