Address review feedback.

- AArch64ExpandPseudos: generalize & use createCall
- AArch64InstrInfo: describe x16/x17 usage in pseudo
- AArch64InstrInfo: group SDNodes with others
- test call for ELF as well
- GlobalISel: rename PointerAuthInfo, reorder fields
- various nits, auto, const

Author: ahmedbougacha

llvm/include/llvm/CodeGen/GlobalISel/CallLowering.h

@@ -99,9 +99,9 @@ class CallLowering {
     ArgInfo() = default;
   };
 
-  struct PointerAuthInfo {
-    Register Discriminator;
+  struct PtrAuthInfo {
     uint64_t Key;
+    Register Discriminator;
   };
 
   struct CallLoweringInfo {
@@ -130,7 +130,8 @@ class CallLowering {
 
     MDNode *KnownCallees = nullptr;
 
-    std::optional<PointerAuthInfo> PAI;
+    /// The auth-call information in the "ptrauth" bundle, if present.
+    std::optional<PtrAuthInfo> PAI;
 
     /// True if the call must be tail call optimized.
     bool IsMustTailCall = false;
@@ -594,8 +595,7 @@ class CallLowering {
   bool lowerCall(MachineIRBuilder &MIRBuilder, const CallBase &Call,
                  ArrayRef<Register> ResRegs,
                  ArrayRef<ArrayRef<Register>> ArgRegs, Register SwiftErrorVReg,
-                 std::optional<PointerAuthInfo> PAI,
-                 Register ConvergenceCtrlToken,
+                 std::optional<PtrAuthInfo> PAI, Register ConvergenceCtrlToken,
                  std::function<unsigned()> GetCalleeReg) const;
 
   /// For targets which want to use big-endian can enable it with

llvm/lib/CodeGen/GlobalISel/CallLowering.cpp

@@ -92,7 +92,7 @@ bool CallLowering::lowerCall(MachineIRBuilder &MIRBuilder, const CallBase &CB,
                              ArrayRef<Register> ResRegs,
                              ArrayRef<ArrayRef<Register>> ArgRegs,
                              Register SwiftErrorVReg,
-                             std::optional<PointerAuthInfo> PAI,
+                             std::optional<PtrAuthInfo> PAI,
                              Register ConvergenceCtrlToken,
                              std::function<unsigned()> GetCalleeReg) const {
   CallLoweringInfo Info;

llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp

@@ -2642,18 +2642,18 @@ bool IRTranslator::translateCallBase(const CallBase &CB,
     }
   }
 
-  std::optional<CallLowering::PointerAuthInfo> PAI;
+  std::optional<CallLowering::PtrAuthInfo> PAI;
   if (CB.countOperandBundlesOfType(LLVMContext::OB_ptrauth)) {
     // Functions should never be ptrauth-called directly.
     assert(!CB.getCalledFunction() && "invalid direct ptrauth call");
 
     auto PAB = CB.getOperandBundle("ptrauth");
-    Value *Key = PAB->Inputs[0];
-    Value *Discriminator = PAB->Inputs[1];
+    const Value *Key = PAB->Inputs[0];
+    const Value *Discriminator = PAB->Inputs[1];
 
     Register DiscReg = getOrCreateVReg(*Discriminator);
-    PAI = CallLowering::PointerAuthInfo{DiscReg,
-                                        cast<ConstantInt>(Key)->getZExtValue()};
+    PAI = CallLowering::PtrAuthInfo{cast<ConstantInt>(Key)->getZExtValue(),
+                                    DiscReg};
   }
 
   Register ConvergenceCtrlToken = 0;

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

@@ -9225,12 +9225,12 @@ void SelectionDAGBuilder::visitCall(const CallInst &I) {
 void SelectionDAGBuilder::LowerCallSiteWithPtrAuthBundle(
     const CallBase &CB, const BasicBlock *EHPadBB) {
   auto PAB = CB.getOperandBundle("ptrauth");
-  auto *CalleeV = CB.getCalledOperand();
+  const Value *CalleeV = CB.getCalledOperand();
 
   // Gather the call ptrauth data from the operand bundle:
   //   [ i32 <key>, i64 <discriminator> ]
-  auto *Key = cast<ConstantInt>(PAB->Inputs[0]);
-  Value *Discriminator = PAB->Inputs[1];
+  const auto *Key = cast<ConstantInt>(PAB->Inputs[0]);
+  const Value *Discriminator = PAB->Inputs[1];
 
   assert(Key->getType()->isIntegerTy(32) && "Invalid ptrauth key");
   assert(Discriminator->getType()->isIntegerTy(64) &&

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

@@ -1550,7 +1550,6 @@ unsigned AArch64AsmPrinter::emitPtrauthDiscriminator(uint16_t Disc,
 
 void AArch64AsmPrinter::emitPtrauthBranch(const MachineInstr *MI) {
   unsigned InstsEmitted = 0;
-
   unsigned BrTarget = MI->getOperand(0).getReg();
 
   auto Key = (AArch64PACKey::ID)MI->getOperand(1).getImm();

llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp

@@ -774,26 +774,24 @@ bool AArch64ExpandPseudo::expandSVESpillFill(MachineBasicBlock &MBB,
   return true;
 }
 
-// Create a call to CallTarget, copying over all the operands from *MBBI,
-// starting at the regmask.
-static MachineInstr *createCall(MachineBasicBlock &MBB,
-                                MachineBasicBlock::iterator MBBI,
-                                const AArch64InstrInfo *TII,
-                                MachineOperand &CallTarget,
-                                unsigned RegMaskStartIdx) {
-  unsigned Opc = CallTarget.isGlobal() ? AArch64::BL : AArch64::BLR;
-  MachineInstr *Call =
-      BuildMI(MBB, MBBI, MBBI->getDebugLoc(), TII->get(Opc)).getInstr();
-
-  assert((CallTarget.isGlobal() || CallTarget.isReg()) &&
-         "invalid operand for regular call");
-  Call->addOperand(CallTarget);
+// Create a call with the passed opcode and explicit operands, copying over all
+// the implicit operands from *MBBI, starting at the regmask.
+static MachineInstr *createCallWithOps(MachineBasicBlock &MBB,
+                                       MachineBasicBlock::iterator MBBI,
+                                       const AArch64InstrInfo *TII,
+                                       unsigned Opcode,
+                                       ArrayRef<MachineOperand> ExplicitOps,
+                                       unsigned RegMaskStartIdx) {
+  // Build the MI, with explicit operands first (including the call target).
+  MachineInstr *Call = BuildMI(MBB, MBBI, MBBI->getDebugLoc(), TII->get(Opcode))
+                           .add(ExplicitOps)
+                           .getInstr();
 
   // Register arguments are added during ISel, but cannot be added as explicit
   // operands of the branch as it expects to be B <target> which is only one
   // operand. Instead they are implicit operands used by the branch.
   while (!MBBI->getOperand(RegMaskStartIdx).isRegMask()) {
-    auto MOP = MBBI->getOperand(RegMaskStartIdx);
+    MachineOperand &MOP = MBBI->getOperand(RegMaskStartIdx);
     assert(MOP.isReg() && "can only add register operands");
     Call->addOperand(MachineOperand::CreateReg(
         MOP.getReg(), /*Def=*/false, /*Implicit=*/true, /*isKill=*/false,
@@ -807,6 +805,20 @@ static MachineInstr *createCall(MachineBasicBlock &MBB,
   return Call;
 }
 
+// Create a call to CallTarget, copying over all the operands from *MBBI,
+// starting at the regmask.
+static MachineInstr *createCall(MachineBasicBlock &MBB,
+                                MachineBasicBlock::iterator MBBI,
+                                const AArch64InstrInfo *TII,
+                                MachineOperand &CallTarget,
+                                unsigned RegMaskStartIdx) {
+  unsigned Opc = CallTarget.isGlobal() ? AArch64::BL : AArch64::BLR;
+
+  assert((CallTarget.isGlobal() || CallTarget.isReg()) &&
+         "invalid operand for regular call");
+  return createCallWithOps(MBB, MBBI, TII, Opc, CallTarget, RegMaskStartIdx);
+}
+
 bool AArch64ExpandPseudo::expandCALL_RVMARKER(
     MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI) {
   // Expand CALL_RVMARKER pseudo to:
@@ -822,33 +834,19 @@ bool AArch64ExpandPseudo::expandCALL_RVMARKER(
 
   if (MI.getOpcode() == AArch64::BLRA_RVMARKER) {
     // Pointer auth call.
+    MachineOperand &CallTarget = MI.getOperand(1);
     MachineOperand &Key = MI.getOperand(2);
-    assert((Key.getImm() == 0 || Key.getImm() == 1) &&
-           "invalid key for ptrauth call");
     MachineOperand &IntDisc = MI.getOperand(3);
     MachineOperand &AddrDisc = MI.getOperand(4);
 
-    OriginalCall = BuildMI(MBB, MBBI, MI.getDebugLoc(), TII->get(AArch64::BLRA))
-                       .getInstr();
-    OriginalCall->addOperand(MI.getOperand(1));
-    OriginalCall->addOperand(Key);
-    OriginalCall->addOperand(IntDisc);
-    OriginalCall->addOperand(AddrDisc);
-
-    unsigned RegMaskStartIdx = 5;
-    // Skip register arguments. Those are added during ISel, but are not
-    // needed for the concrete branch.
-    while (!MI.getOperand(RegMaskStartIdx).isRegMask()) {
-      auto MOP = MI.getOperand(RegMaskStartIdx);
-      assert(MOP.isReg() && "can only add register operands");
-      OriginalCall->addOperand(MachineOperand::CreateReg(
-          MOP.getReg(), /*Def=*/false, /*Implicit=*/true, /*isKill=*/false,
-          /*isDead=*/false, /*isUndef=*/MOP.isUndef()));
-      RegMaskStartIdx++;
-    }
-    for (const MachineOperand &MO :
-         llvm::drop_begin(MI.operands(), RegMaskStartIdx))
-      OriginalCall->addOperand(MO);
+    assert((Key.getImm() == AArch64PACKey::IA ||
+            Key.getImm() == AArch64PACKey::IB) &&
+           "Invalid auth call key");
+
+    MachineOperand Ops[] = {CallTarget, Key, IntDisc, AddrDisc};
+
+    OriginalCall = createCallWithOps(MBB, MBBI, TII, AArch64::BLRA, Ops,
+                                     /*RegMaskStartIdx=*/5);
   } else {
     assert(MI.getOpcode() == AArch64::BLR_RVMARKER && "unknown rvmarker MI");
     OriginalCall = createCall(MBB, MBBI, TII, MI.getOperand(1),

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -349,7 +349,7 @@ extractPtrauthBlendDiscriminators(SDValue Disc, SelectionDAG *DAG) {
   // 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);
+  const auto *ConstDiscN = dyn_cast<ConstantSDNode>(ConstDisc);
   if (!ConstDiscN || !isUInt<16>(ConstDiscN->getZExtValue()))
     return std::make_tuple(DAG->getTargetConstant(0, DL, MVT::i64), Disc);
 

llvm/lib/Target/AArch64/AArch64InstrInfo.td

@@ -662,6 +662,30 @@ def AArch64call_arm64ec_to_x64 : SDNode<"AArch64ISD::CALL_ARM64EC_TO_X64",
                                       [SDNPHasChain, SDNPOptInGlue, SDNPOutGlue,
                                        SDNPVariadic]>;
 
+def AArch64authcall : SDNode<"AArch64ISD::AUTH_CALL",
+                             SDTypeProfile<0, -1, [SDTCisPtrTy<0>,
+                                                   SDTCisVT<1, i32>,
+                                                   SDTCisVT<2, i64>,
+                                                   SDTCisVT<3, i64>]>,
+                             [SDNPHasChain, SDNPOptInGlue, SDNPOutGlue,
+                              SDNPVariadic]>;
+
+def AArch64authtcret: SDNode<"AArch64ISD::AUTH_TC_RETURN",
+                             SDTypeProfile<0, 5, [SDTCisPtrTy<0>,
+                                                  SDTCisVT<2, i32>,
+                                                  SDTCisVT<3, i64>,
+                                                  SDTCisVT<4, i64>]>,
+                             [SDNPHasChain,  SDNPOptInGlue, SDNPVariadic]>;
+
+def AArch64authcall_rvmarker : SDNode<"AArch64ISD::AUTH_CALL_RVMARKER",
+                                 SDTypeProfile<0, -1, [SDTCisPtrTy<0>,
+                                                       SDTCisPtrTy<1>,
+                                                       SDTCisVT<2, i32>,
+                                                       SDTCisVT<3, i64>,
+                                                       SDTCisVT<4, i64>]>,
+                                 [SDNPHasChain, SDNPOptInGlue, SDNPOutGlue,
+                                  SDNPVariadic]>;
+
 def AArch64brcond        : SDNode<"AArch64ISD::BRCOND", SDT_AArch64Brcond,
                                 [SDNPHasChain]>;
 def AArch64cbz           : SDNode<"AArch64ISD::CBZ", SDT_AArch64cbz,
@@ -1564,30 +1588,6 @@ let Predicates = [HasComplxNum, HasNEON] in {
                          (v4f32 (bitconvert (v2i64 (AArch64duplane64 (v2i64 V128:$Rm), VectorIndexD:$idx))))>;
 }
 
-def AArch64authcall : SDNode<"AArch64ISD::AUTH_CALL",
-                             SDTypeProfile<0, -1, [SDTCisPtrTy<0>,
-                                                   SDTCisVT<1, i32>,
-                                                   SDTCisVT<2, i64>,
-                                                   SDTCisVT<3, i64>]>,
-                             [SDNPHasChain, SDNPOptInGlue, SDNPOutGlue,
-                              SDNPVariadic]>;
-
-def AArch64authtcret: SDNode<"AArch64ISD::AUTH_TC_RETURN",
-                             SDTypeProfile<0, 5, [SDTCisPtrTy<0>,
-                                                  SDTCisVT<2, i32>,
-                                                  SDTCisVT<3, i64>,
-                                                  SDTCisVT<4, i64>]>,
-                             [SDNPHasChain,  SDNPOptInGlue, SDNPVariadic]>;
-
-def AArch64authcall_rvmarker : SDNode<"AArch64ISD::AUTH_CALL_RVMARKER",
-                                 SDTypeProfile<0, -1, [SDTCisPtrTy<0>,
-                                                       SDTCisPtrTy<1>,
-                                                       SDTCisVT<2, i32>,
-                                                       SDTCisVT<3, i64>,
-                                                       SDTCisVT<4, i64>]>,
-                                 [SDNPHasChain, SDNPOptInGlue, SDNPOutGlue,
-                                  SDNPVariadic]>;
-
 // v8.3a Pointer Authentication
 // These instructions inhabit part of the hint space and so can be used for
 // armv8 targets. Keeping the old HINT mnemonic when compiling without PA is
@@ -1716,9 +1716,12 @@ let Predicates = [HasPAuth] in {
     def BLRABZ  : AuthOneOperand<0b001, 1, "blrabz">;
   }
 
-  // BLRA pseudo, generalized version of BLRAA/BLRAB/Z.
-  // This directly manipulates x16/x17, which are the only registers the OS
-  // guarantees are safe to use for sensitive operations.
+  // BLRA pseudo, a generalized version of BLRAA/BLRAB/Z.
+  // This directly manipulates x16/x17 to materialize the discriminator.
+  // x16/x17 are generally used as the safe registers for sensitive ptrauth
+  // operations (such as raw address manipulation or discriminator
+  // materialization here), in part because they're handled in a safer way by
+  // the kernel, notably on Darwin.
   def BLRA : Pseudo<(outs), (ins GPR64noip:$Rn, i32imm:$Key, i64imm:$Disc,
                                  GPR64noip:$AddrDisc),
                     [(AArch64authcall GPR64noip:$Rn, timm:$Key, timm:$Disc,

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

@@ -1015,7 +1015,7 @@ bool AArch64CallLowering::isEligibleForTailCallOptimization(
 
 static unsigned getCallOpcode(const MachineFunction &CallerF, bool IsIndirect,
                               bool IsTailCall,
-                              std::optional<CallLowering::PointerAuthInfo> &PAI,
+                              std::optional<CallLowering::PtrAuthInfo> &PAI,
                               MachineRegisterInfo &MRI) {
   const AArch64FunctionInfo *FuncInfo = CallerF.getInfo<AArch64FunctionInfo>();