[AArch64][PAC] Rework discriminator analysis in AUT and AUTPAC

Make use of post-processing the discriminator components by custom
inserter hook to eliminate duplication for DAGISel and GlobalISel and
improve cross-BB analysis for DAGISel.

Author: atrosinenko

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

@@ -1487,39 +1487,6 @@ void AArch64DAGToDAGISel::SelectTable(SDNode *N, unsigned NumVecs, unsigned Opc,
   ReplaceNode(N, CurDAG->getMachineNode(Opc, dl, VT, Ops));
 }
 
-static std::tuple<SDValue, SDValue>
-extractPtrauthBlendDiscriminators(SDValue Disc, SelectionDAG *DAG) {
-  SDLoc DL(Disc);
-  SDValue AddrDisc;
-  SDValue ConstDisc;
-
-  // If this is a blend, remember the constant and address discriminators.
-  // Otherwise, it's either a constant discriminator, or a non-blended
-  // address discriminator.
-  if (Disc->getOpcode() == ISD::INTRINSIC_WO_CHAIN &&
-      Disc->getConstantOperandVal(0) == Intrinsic::ptrauth_blend) {
-    AddrDisc = Disc->getOperand(1);
-    ConstDisc = Disc->getOperand(2);
-  } else {
-    ConstDisc = Disc;
-  }
-
-  // If the constant discriminator (either the blend RHS, or the entire
-  // discriminator value) isn't a 16-bit constant, bail out, and let the
-  // discriminator be computed separately.
-  auto *ConstDiscN = dyn_cast<ConstantSDNode>(ConstDisc);
-  if (!ConstDiscN || !isUInt<16>(ConstDiscN->getZExtValue()))
-    return std::make_tuple(DAG->getTargetConstant(0, DL, MVT::i64), Disc);
-
-  // If there's no address discriminator, use XZR directly.
-  if (!AddrDisc)
-    AddrDisc = DAG->getRegister(AArch64::XZR, MVT::i64);
-
-  return std::make_tuple(
-      DAG->getTargetConstant(ConstDiscN->getZExtValue(), DL, MVT::i64),
-      AddrDisc);
-}
-
 void AArch64DAGToDAGISel::SelectPtrauthAuth(SDNode *N) {
   SDLoc DL(N);
   // IntrinsicID is operand #0
@@ -1530,20 +1497,18 @@ void AArch64DAGToDAGISel::SelectPtrauthAuth(SDNode *N) {
   unsigned AUTKeyC = cast<ConstantSDNode>(AUTKey)->getZExtValue();
   AUTKey = CurDAG->getTargetConstant(AUTKeyC, DL, MVT::i64);
 
-  SDValue AUTAddrDisc, AUTConstDisc;
-  std::tie(AUTConstDisc, AUTAddrDisc) =
-      extractPtrauthBlendDiscriminators(AUTDisc, CurDAG);
+  SDValue Zero = CurDAG->getTargetConstant(0, DL, MVT::i64);
 
   if (!Subtarget->isX16X17Safer()) {
-    SDValue Ops[] = {Val, AUTKey, AUTConstDisc, AUTAddrDisc};
+    SDValue Ops[] = {Val, AUTKey, Zero, AUTDisc};
 
     SDNode *AUT =
         CurDAG->getMachineNode(AArch64::AUTxMxN, DL, MVT::i64, MVT::i64, Ops);
     ReplaceNode(N, AUT);
   } else {
     SDValue X16Copy = CurDAG->getCopyToReg(CurDAG->getEntryNode(), DL,
                                            AArch64::X16, Val, SDValue());
-    SDValue Ops[] = {AUTKey, AUTConstDisc, AUTAddrDisc, X16Copy.getValue(1)};
+    SDValue Ops[] = {AUTKey, Zero, AUTDisc, X16Copy.getValue(1)};
 
     SDNode *AUT = CurDAG->getMachineNode(AArch64::AUTx16x17, DL, MVT::i64, Ops);
     ReplaceNode(N, AUT);
@@ -1565,19 +1530,13 @@ void AArch64DAGToDAGISel::SelectPtrauthResign(SDNode *N) {
   AUTKey = CurDAG->getTargetConstant(AUTKeyC, DL, MVT::i64);
   PACKey = CurDAG->getTargetConstant(PACKeyC, DL, MVT::i64);
 
-  SDValue AUTAddrDisc, AUTConstDisc;
-  std::tie(AUTConstDisc, AUTAddrDisc) =
-      extractPtrauthBlendDiscriminators(AUTDisc, CurDAG);
-
-  SDValue PACAddrDisc, PACConstDisc;
-  std::tie(PACConstDisc, PACAddrDisc) =
-      extractPtrauthBlendDiscriminators(PACDisc, CurDAG);
+  SDValue Zero = CurDAG->getTargetConstant(0, DL, MVT::i64);
 
   SDValue X16Copy = CurDAG->getCopyToReg(CurDAG->getEntryNode(), DL,
                                          AArch64::X16, Val, SDValue());
 
-  SDValue Ops[] = {AUTKey,       AUTConstDisc, AUTAddrDisc,        PACKey,
-                   PACConstDisc, PACAddrDisc,  X16Copy.getValue(1)};
+  SDValue Ops[] = {
+      AUTKey, Zero, AUTDisc, PACKey, Zero, PACDisc, X16Copy.getValue(1)};
 
   SDNode *AUTPAC = CurDAG->getMachineNode(AArch64::AUTPAC, DL, MVT::i64, Ops);
   ReplaceNode(N, AUTPAC);

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -3274,10 +3274,20 @@ MachineBasicBlock *AArch64TargetLowering::EmitInstrWithCustomInserter(
   case AArch64::MOVT_TIZ_PSEUDO:
     return EmitZTInstr(MI, BB, AArch64::MOVT_TIZ, /*Op0IsDef=*/true);
 
+  case AArch64::AUT:
+    fixupBlendComponents(MI, BB, MI.getOperand(1), MI.getOperand(2),
+                         &AArch64::GPR64noipRegClass);
+    return BB;
   case AArch64::PAC:
     fixupPtrauthDiscriminator(MI, BB, MI.getOperand(3), MI.getOperand(4),
                               &AArch64::GPR64noipRegClass);
     return 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/AArch64InstrInfo.td

@@ -2137,6 +2137,7 @@ let Predicates = [HasPAuth] in {
     let Size = 32;
     let Defs = [X16,X17,NZCV];
     let Uses = [X16];
+    let usesCustomInserter = 1;
   }
 
   def AUTxMxN : Pseudo<(outs GPR64:$AuthVal, GPR64common:$Scratch),
@@ -2187,6 +2188,7 @@ let Predicates = [HasPAuth] in {
     let Size = 48;
     let Defs = [X16,X17,NZCV];
     let Uses = [X16];
+    let usesCustomInserter = 1;
   }
 
   // Materialize a signed global address, with adrp+add and PAC.

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

@@ -6725,25 +6725,15 @@ bool AArch64InstructionSelector::selectIntrinsic(MachineInstr &I,
     uint64_t PACKey = I.getOperand(5).getImm();
     Register PACDisc = I.getOperand(6).getReg();
 
-    Register AUTAddrDisc = AUTDisc;
-    uint16_t AUTConstDiscC = 0;
-    std::tie(AUTConstDiscC, AUTAddrDisc) =
-        extractPtrauthBlendDiscriminators(AUTDisc, MRI);
-
-    Register PACAddrDisc = PACDisc;
-    uint16_t PACConstDiscC = 0;
-    std::tie(PACConstDiscC, PACAddrDisc) =
-        extractPtrauthBlendDiscriminators(PACDisc, MRI);
-
     MIB.buildCopy({AArch64::X16}, {ValReg});
     MIB.buildInstr(TargetOpcode::IMPLICIT_DEF, {AArch64::X17}, {});
     MIB.buildInstr(AArch64::AUTPAC)
         .addImm(AUTKey)
-        .addImm(AUTConstDiscC)
-        .addUse(AUTAddrDisc)
+        .addImm(0)
+        .addUse(AUTDisc)
         .addImm(PACKey)
-        .addImm(PACConstDiscC)
-        .addUse(PACAddrDisc)
+        .addImm(0)
+        .addUse(PACDisc)
         .constrainAllUses(TII, TRI, RBI);
     MIB.buildCopy({DstReg}, Register(AArch64::X16));
 
@@ -6757,18 +6747,13 @@ bool AArch64InstructionSelector::selectIntrinsic(MachineInstr &I,
     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);
-
     if (STI.isX16X17Safer()) {
       MIB.buildCopy({AArch64::X16}, {ValReg});
       MIB.buildInstr(TargetOpcode::IMPLICIT_DEF, {AArch64::X17}, {});
       MIB.buildInstr(AArch64::AUTx16x17)
           .addImm(AUTKey)
-          .addImm(AUTConstDiscC)
-          .addUse(AUTAddrDisc)
+          .addImm(0)
+          .addUse(AUTDisc)
           .constrainAllUses(TII, TRI, RBI);
       MIB.buildCopy({DstReg}, Register(AArch64::X16));
     } else {
@@ -6779,8 +6764,8 @@ bool AArch64InstructionSelector::selectIntrinsic(MachineInstr &I,
           .addDef(ScratchReg)
           .addUse(ValReg)
           .addImm(AUTKey)
-          .addImm(AUTConstDiscC)
-          .addUse(AUTAddrDisc)
+          .addImm(0)
+          .addUse(AUTDisc)
           .constrainAllUses(TII, TRI, RBI);
     }