AArch64: Relax restriction on discriminator when PAuth ifunc used.

[pcc]

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.

[llvmbot]

@llvm/pr-subscribers-backend-aarch64

Author: Peter Collingbourne (pcc)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/170945.diff

3 Files Affected:

• (modified) llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp (+14-9)
• (modified) llvm/test/CodeGen/AArch64/ptrauth-irelative.ll (+13)
• (modified) llvm/test/CodeGen/AArch64/ptrauth-reloc.ll (+17)

diff --git a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
index 8ef985451a4c1..3f6be10679517 100644
--- a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
+++ b/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();
 
@@ -2508,12 +2508,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);
   }
 
@@ -2592,11 +2596,6 @@ 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(
@@ -2604,6 +2603,12 @@ AArch64AsmPrinter::lowerConstantPtrAuth(const ConstantPtrAuth &CPA) {
           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");
diff --git a/llvm/test/CodeGen/AArch64/ptrauth-irelative.ll b/llvm/test/CodeGen/AArch64/ptrauth-irelative.ll
index 6a291497d6c46..c366459a3637d 100644
--- a/llvm/test/CodeGen/AArch64/ptrauth-irelative.ll
+++ b/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:
diff --git a/llvm/test/CodeGen/AArch64/ptrauth-reloc.ll b/llvm/test/CodeGen/AArch64/ptrauth-reloc.ll
index 02c643f101913..f2d080644e93e 100644
--- a/llvm/test/CodeGen/AArch64/ptrauth-reloc.ll
+++ b/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)