Address review feedback.

- test invoke for ELF as well
- consolidate GISel/SDAG invoke test
- check key/disc validity in consistent ways
- misc. nits

Author: ahmedbougacha

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

@@ -1552,18 +1552,20 @@ void AArch64AsmPrinter::emitPtrauthBranch(const MachineInstr *MI) {
   unsigned InstsEmitted = 0;
 
   unsigned BrTarget = MI->getOperand(0).getReg();
+
   auto Key = (AArch64PACKey::ID)MI->getOperand(1).getImm();
+  assert((Key == AArch64PACKey::IA || Key == AArch64PACKey::IB) &&
+         "Invalid auth call key");
+
   uint64_t Disc = MI->getOperand(2).getImm();
+  assert(isUInt<16>(Disc));
+
   unsigned AddrDisc = MI->getOperand(3).getReg();
 
   // Compute discriminator into x17
-  assert(isUInt<16>(Disc));
   unsigned DiscReg = emitPtrauthDiscriminator(Disc, AddrDisc, InstsEmitted);
   bool IsZeroDisc = DiscReg == AArch64::XZR;
 
-  assert((Key == AArch64PACKey::IA || Key == AArch64PACKey::IB) &&
-         "Invalid auth call key");
-
   unsigned Opc;
   if (Key == AArch64PACKey::IA)
     Opc = IsZeroDisc ? AArch64::BLRAAZ : AArch64::BLRAA;
@@ -1726,8 +1728,12 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
   case AArch64::AUTH_TCRETURN:
   case AArch64::AUTH_TCRETURN_BTI: {
     const uint64_t Key = MI->getOperand(2).getImm();
-    assert(Key < 2 && "Unknown key kind for authenticating tail-call return");
+    assert((Key == AArch64PACKey::IA || Key == AArch64PACKey::IB) &&
+           "Invalid auth key for tail-call return");
+
     const uint64_t Disc = MI->getOperand(3).getImm();
+    assert(isUInt<16>(Disc) && "Integer discriminator is too wide");
+
     Register AddrDisc = MI->getOperand(4).getReg();
 
     Register ScratchReg = MI->getOperand(0).getReg() == AArch64::X16
@@ -1736,8 +1742,6 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
 
     unsigned DiscReg = AddrDisc;
     if (Disc) {
-      assert(isUInt<16>(Disc) && "Integer discriminator is too wide");
-
       if (AddrDisc != AArch64::NoRegister) {
         EmitToStreamer(*OutStreamer, MCInstBuilder(AArch64::ORRXrs)
                                          .addReg(ScratchReg)
@@ -1758,14 +1762,14 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
       DiscReg = ScratchReg;
     }
 
-    const bool isZero = DiscReg == AArch64::NoRegister;
+    const bool IsZero = DiscReg == AArch64::NoRegister;
     const unsigned Opcodes[2][2] = {{AArch64::BRAA, AArch64::BRAAZ},
                                     {AArch64::BRAB, AArch64::BRABZ}};
 
     MCInst TmpInst;
-    TmpInst.setOpcode(Opcodes[Key][isZero]);
+    TmpInst.setOpcode(Opcodes[Key][IsZero]);
     TmpInst.addOperand(MCOperand::createReg(MI->getOperand(0).getReg()));
-    if (!isZero)
+    if (!IsZero)
       TmpInst.addOperand(MCOperand::createReg(DiscReg));
     EmitToStreamer(*OutStreamer, TmpInst);
     return;

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -8468,8 +8468,9 @@ AArch64TargetLowering::LowerCall(CallLoweringInfo &CLI,
     Ops.insert(Ops.begin() + 1, GA);
   } else if (CallConv == CallingConv::ARM64EC_Thunk_X64) {
     Opc = AArch64ISD::CALL_ARM64EC_TO_X64;
-  } else if (GuardWithBTI)
+  } else if (GuardWithBTI) {
     Opc = AArch64ISD::CALL_BTI;
+  }
 
   if (IsTailCall) {
     // Each tail call may have to adjust the stack by a different amount, so
@@ -8480,9 +8481,8 @@ AArch64TargetLowering::LowerCall(CallLoweringInfo &CLI,
 
   if (CLI.PAI) {
     const uint64_t Key = CLI.PAI->Key;
-    // Authenticated calls only support IA and IB.
-    if (Key > 1)
-      report_fatal_error("Unsupported key kind for authenticating call");
+    assert((Key == AArch64PACKey::IA || Key == AArch64PACKey::IB) &&
+           "Invalid auth call key");
 
     // Split the discriminator into address/integer components.
     SDValue AddrDisc, IntDisc;

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

@@ -1021,10 +1021,11 @@ static unsigned getCallOpcode(const MachineFunction &CallerF, bool IsIndirect,
 
   if (!IsTailCall) {
     if (!PAI)
-      return IsIndirect ? getBLRCallOpcode(CallerF) : (unsigned)AArch64::BL;
+      return IsIndirect ? getBLRCallOpcode(CallerF) : AArch64::BL;
 
-    assert(IsIndirect && "authenticated direct call");
-    assert(PAI->Key == 0 || PAI->Key == 1 && "invalid ptrauth key");
+    assert(IsIndirect && "Direct call should not be authenticated");
+    assert((PAI->Key == AArch64PACKey::IA || PAI->Key == AArch64PACKey::IB) &&
+           "Invalid auth call key");
     return AArch64::BLRA;
   }
 
@@ -1036,11 +1037,12 @@ static unsigned getCallOpcode(const MachineFunction &CallerF, bool IsIndirect,
   if (FuncInfo->branchTargetEnforcement()) {
     if (PAI)
       return AArch64::AUTH_TCRETURN_BTI;
-    else if (FuncInfo->branchProtectionPAuthLR())
+    if (FuncInfo->branchProtectionPAuthLR())
       return AArch64::TCRETURNrix17;
-    else
-      return AArch64::TCRETURNrix16x17;
-  } else if (FuncInfo->branchProtectionPAuthLR())
+    return AArch64::TCRETURNrix16x17;
+  }
+
+  if (FuncInfo->branchProtectionPAuthLR())
     return AArch64::TCRETURNrinotx16;
 
   if (PAI)
@@ -1104,7 +1106,9 @@ bool AArch64CallLowering::lowerTailCall(
 
   // Authenticated tail calls always take key/discriminator arguments.
   if (Opc == AArch64::AUTH_TCRETURN || Opc == AArch64::AUTH_TCRETURN_BTI) {
-    assert(Info.PAI->Key == 0 || Info.PAI->Key == 1 && "invalid key");
+    assert((Info.PAI->Key == AArch64PACKey::IA ||
+            Info.PAI->Key == AArch64PACKey::IB) &&
+           "Invalid auth call key");
     MIB.addImm(Info.PAI->Key);
 
     Register AddrDisc = 0;
@@ -1371,7 +1375,9 @@ bool AArch64CallLowering::lowerCall(MachineIRBuilder &MIRBuilder,
   Mask = getMaskForArgs(OutArgs, Info, *TRI, MF);
 
   if (Opc == AArch64::BLRA || Opc == AArch64::BLRA_RVMARKER) {
-    assert(Info.PAI->Key == 0 || Info.PAI->Key == 1 && "invalid key");
+    assert((Info.PAI->Key == AArch64PACKey::IA ||
+            Info.PAI->Key == AArch64PACKey::IB) &&
+           "Invalid auth call key");
     MIB.addImm(Info.PAI->Key);
 
     Register AddrDisc = 0;