[llvm] [AArch64][PAC] Introduce AArch64::PAC pseudo instruction (PR #146488)

[Anatoly Trosinenko via llvm-commits]

================
@@ -3079,6 +3079,76 @@ AArch64TargetLowering::EmitGetSMESaveSize(MachineInstr &MI,
   return BB;
 }
 
+// Helper function to find the instruction that defined a virtual register.
+// If unable to find such instruction, returns nullptr.
+static MachineInstr *stripVRegCopies(const MachineRegisterInfo &MRI,
+                                     Register Reg) {
+  while (Reg.isVirtual()) {
+    MachineInstr *DefMI = MRI.getVRegDef(Reg);
+    assert(DefMI && "Virtual register definition not found");
+    unsigned Opcode = DefMI->getOpcode();
+
+    if (Opcode == AArch64::COPY) {
+      Reg = DefMI->getOperand(1).getReg();
+      // Vreg is defined by copying from physreg.
+      if (Reg.isPhysical())
+        return DefMI;
+      continue;
+    }
+    if (Opcode == AArch64::SUBREG_TO_REG) {
+      Reg = DefMI->getOperand(2).getReg();
+      continue;
+    }
+
+    return DefMI;
+  }
+  return nullptr;
+}
+
+void AArch64TargetLowering::fixupPtrauthDiscriminator(
+    MachineInstr &MI, MachineBasicBlock *BB, MachineOperand &IntDiscOp,
+    MachineOperand &AddrDiscOp, const TargetRegisterClass *AddrDiscRC) const {
+  const TargetInstrInfo *TII = Subtarget->getInstrInfo();
+  MachineRegisterInfo &MRI = MI.getMF()->getRegInfo();
+  const DebugLoc &DL = MI.getDebugLoc();
+
+  Register AddrDisc = AddrDiscOp.getReg();
+  int64_t IntDisc = IntDiscOp.getImm();
+
+  assert(IntDisc == 0 && "Blend components are already expanded");
+
+  MachineInstr *MaybeBlend = stripVRegCopies(MRI, AddrDisc);
+
+  if (MaybeBlend) {
+    // Detect blend(addr, imm) which is lowered as "MOVK addr, #imm, #48".
+    // Alternatively, recognize small immediate modifier passed via VReg.
+    if (MaybeBlend->getOpcode() == AArch64::MOVKXi &&
+        MaybeBlend->getOperand(3).getImm() == 48) {
+      AddrDisc = MaybeBlend->getOperand(1).getReg();
+      IntDisc = MaybeBlend->getOperand(2).getImm();
+    } else if (MaybeBlend->getOpcode() == AArch64::MOVi32imm &&
----------------
atrosinenko wrote:

Thank you for pointing this out, I rechecked which instructions can be used to materialize an immediate constant and it turned out it is possible for `MOVi64imm` to be used even for small `i64` constants when optimization is disabled (see [patterns in `AArch64InstrInfo.td`](https://github.com/llvm/llvm-project/blob/82451d0b1341a9b6c01eaa5d27088ff9f3287853/llvm/lib/Target/AArch64/AArch64InstrInfo.td#L2535)). Furthermore, I realized it is necessary to check if integer constant is an immediate operand (as opposed to some relocation referring a global symbol, for example). For that reason, I rewritten the chain of "if"-"else if" into a switch and added more tests in 33b61f5939813244d6ea812e821945bf86854a99.

https://github.com/llvm/llvm-project/pull/146488