AArch64: Relax restriction on discriminator when PAuth ifunc used.

When a PAuth ifunc is being used, we can represent any discriminator that
we want in the code and don't need to be restricted to 16 bits. For now we
only need this capability for address discriminated ptrauth expressions,
so keep the restriction in place for other discriminators.

Reviewers: atrosinenko, fmayer

Reviewed By: fmayer

Pull Request: #170945

Author: pcc

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

@@ -220,7 +220,7 @@ class AArch64AsmPrinter : public AsmPrinter {
   void LowerLOADgotAUTH(const MachineInstr &MI);
 
   const MCExpr *emitPAuthRelocationAsIRelative(
-      const MCExpr *Target, uint16_t Disc, AArch64PACKey::ID KeyID,
+      const MCExpr *Target, uint64_t Disc, AArch64PACKey::ID KeyID,
       bool HasAddressDiversity, bool IsDSOLocal, const MCExpr *DSExpr);
 
   /// tblgen'erated driver function for lowering simple MI->MC
@@ -2461,7 +2461,7 @@ static bool targetSupportsIRelativeRelocation(const Triple &TT) {
 // ret
 // .popsection
 const MCExpr *AArch64AsmPrinter::emitPAuthRelocationAsIRelative(
-    const MCExpr *Target, uint16_t Disc, AArch64PACKey::ID KeyID,
+    const MCExpr *Target, uint64_t Disc, AArch64PACKey::ID KeyID,
     bool HasAddressDiversity, bool IsDSOLocal, const MCExpr *DSExpr) {
   const Triple &TT = TM.getTargetTriple();
 
@@ -2509,12 +2509,16 @@ const MCExpr *AArch64AsmPrinter::emitPAuthRelocationAsIRelative(
   if (HasAddressDiversity) {
     auto *PlacePlusDisc = MCBinaryExpr::createAdd(
         MCSymbolRefExpr::create(Place, OutStreamer->getContext()),
-        MCConstantExpr::create(static_cast<int16_t>(Disc),
-                               OutStreamer->getContext()),
+        MCConstantExpr::create(Disc, OutStreamer->getContext()),
         OutStreamer->getContext());
     emitAddress(*OutStreamer, AArch64::X1, PlacePlusDisc, /*IsDSOLocal=*/true,
                 *STI);
   } else {
+    if (!isUInt<16>(Disc)) {
+      OutContext.reportError(SMLoc(), "AArch64 PAC Discriminator '" +
+                                          Twine(Disc) +
+                                          "' out of range [0, 0xFFFF]");
+    }
     emitMOVZ(AArch64::X1, Disc, 0);
   }
 
@@ -2593,18 +2597,19 @@ AArch64AsmPrinter::lowerConstantPtrAuth(const ConstantPtrAuth &CPA) {
   }
 
   uint64_t Disc = CPA.getDiscriminator()->getZExtValue();
-  if (!isUInt<16>(Disc)) {
-    CPA.getContext().emitError("AArch64 PAC Discriminator '" + Twine(Disc) +
-                               "' out of range [0, 0xFFFF]");
-    Disc = 0;
-  }
 
   // Check if we need to represent this with an IRELATIVE and emit it if so.
   if (auto *IFuncSym = emitPAuthRelocationAsIRelative(
           Sym, Disc, AArch64PACKey::ID(KeyID), CPA.hasAddressDiscriminator(),
           BaseGVB && BaseGVB->isDSOLocal(), DSExpr))
     return IFuncSym;
 
+  if (!isUInt<16>(Disc)) {
+    CPA.getContext().emitError("AArch64 PAC Discriminator '" + Twine(Disc) +
+                               "' out of range [0, 0xFFFF]");
+    Disc = 0;
+  }
+
   if (DSExpr)
     report_fatal_error("deactivation symbols unsupported in constant "
                        "expressions on this target");

llvm/test/CodeGen/AArch64/ptrauth-irelative.ll

@@ -67,6 +67,19 @@
 ; CHECK-NEXT: .xword [[FUNC]]@FUNCINIT
 @disc = constant ptr ptrauth (ptr @dsolocal, i32 2, i64 0, ptr @disc), align 8
 
+; CHECK: disc65536:
+; CHECK-NEXT: [[PLACE:.*]]:
+; CHECK-NEXT: .section .text.startup
+; CHECK-NEXT: [[FUNC:.*]]:
+; CHECK-NEXT: adrp x0, dsolocal
+; CHECK-NEXT: add x0, x0, :lo12:dsolocal
+; CHECK-NEXT: adrp x1, [[PLACE]]+65536
+; CHECK-NEXT: add x1, x1, :lo12:[[PLACE]]+65536
+; CHECK-NEXT: b __emupac_pacda
+; CHECK-NEXT: .section .rodata
+; CHECK-NEXT: .xword [[FUNC]]@FUNCINIT
+@disc65536 = constant ptr ptrauth (ptr @dsolocal, i32 2, i64 65536, ptr @disc), align 8
+
 @global = external global i8
 
 ; CHECK: globalref:

llvm/test/CodeGen/AArch64/ptrauth-reloc.ll

@@ -175,3 +175,20 @@
 
 @g = external global i32
 @g.ref.ia.65536 = constant ptr ptrauth (ptr @g, i32 0, i64 65536)
+
+;--- err-disc-elf.ll
+
+; RUN: not llc < err-disc-elf.ll -mtriple aarch64-elf -mattr=+pauth 2>&1 \
+; RUN:   | FileCheck %s --check-prefix=CHECK-ERR-DISC-ELF
+; RUN: not llc < err-disc-elf.ll -mtriple aarch64-elf -mattr=+pauth \
+; RUN:   -global-isel -verify-machineinstrs -global-isel-abort=1 2>&1 \
+; RUN:   | FileCheck %s --check-prefix=CHECK-ERR-DISC-ELF
+
+@g = external global i32
+
+@ds = external global i8
+; CHECK-ERR-DISC-ELF-NOT: error: AArch64 PAC Discriminator '65537' out of range [0, 0xFFFF]
+@g.ref.da.65537 = constant ptr ptrauth (ptr @g, i32 2, i64 65537, ptr @g.ref.da.65537, ptr @ds)
+
+; CHECK-ERR-DISC-ELF: error: AArch64 PAC Discriminator '65538' out of range [0, 0xFFFF]
+@g.ref.da.65538 = constant ptr ptrauth (ptr @g, i32 2, i64 65538, ptr null, ptr @ds)