[clang] [llvm] [AArch64][PAC] Sign block addresses used in indirectbr. (PR #97647)

Ahmed Bougacha via cfe-commits cfe-commits at lists.llvm.org
Mon Jul 22 16:13:04 PDT 2024


https://github.com/ahmedbougacha updated https://github.com/llvm/llvm-project/pull/97647

>From 519570896c82887a5dd878fcc16f884857d4fce3 Mon Sep 17 00:00:00 2001
From: Ahmed Bougacha <ahmed at bougacha.org>
Date: Tue, 12 Mar 2024 14:40:17 -0700
Subject: [PATCH 1/5] [AArch64][PAC] Sign block addresses used in indirectbr.

Enabled in clang using:
  -fptrauth-indirect-gotos

and at the IR level using function attribute:
  "ptrauth-indirect-gotos"

Signing uses IA and a per-function integer discriminator.
The discriminator isn't ABI-visible, and is currently:
  ptrauth_string_discriminator("<function_name> blockaddress")

A sufficiently sophisticated frontend could benefit from
per-indirectbr discrimination, which would need additional
machinery, such as allowing "ptrauth" bundles on indirectbr.
For our purposes, the simple scheme above is sufficient.
---
 clang/include/clang/Basic/Features.def        |   1 +
 clang/include/clang/Basic/LangOptions.def     |   1 +
 .../include/clang/Basic/PointerAuthOptions.h  |   3 +
 clang/include/clang/Driver/Options.td         |   2 +
 clang/lib/CodeGen/CodeGenFunction.cpp         |   2 +
 clang/lib/Driver/ToolChains/Clang.cpp         |   3 +
 clang/lib/Frontend/CompilerInvocation.cpp     |   6 +-
 .../CodeGen/ptrauth-function-attributes.c     |   5 +
 llvm/docs/PointerAuth.md                      |  24 ++++
 llvm/include/llvm/CodeGen/AsmPrinter.h        |   3 +
 llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp    |   6 +-
 llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp |  35 +++++-
 llvm/lib/Target/AArch64/AArch64FastISel.cpp   |   4 +
 .../Target/AArch64/AArch64ISelLowering.cpp    |  53 ++++++++-
 llvm/lib/Target/AArch64/AArch64ISelLowering.h |   1 +
 llvm/lib/Target/AArch64/AArch64InstrInfo.td   |  18 +++
 llvm/lib/Target/AArch64/AArch64Subtarget.cpp  |  11 ++
 llvm/lib/Target/AArch64/AArch64Subtarget.h    |   9 ++
 .../GISel/AArch64InstructionSelector.cpp      |  26 +++++
 .../CodeGen/AArch64/ptrauth-indirectbr.ll     | 106 ++++++++++++++++++
 20 files changed, 308 insertions(+), 11 deletions(-)
 create mode 100644 llvm/test/CodeGen/AArch64/ptrauth-indirectbr.ll

diff --git a/clang/include/clang/Basic/Features.def b/clang/include/clang/Basic/Features.def
index 53f410d3cb4bd..cf800afe08557 100644
--- a/clang/include/clang/Basic/Features.def
+++ b/clang/include/clang/Basic/Features.def
@@ -108,6 +108,7 @@ FEATURE(ptrauth_calls, LangOpts.PointerAuthCalls)
 FEATURE(ptrauth_returns, LangOpts.PointerAuthReturns)
 FEATURE(ptrauth_vtable_pointer_address_discrimination, LangOpts.PointerAuthVTPtrAddressDiscrimination)
 FEATURE(ptrauth_vtable_pointer_type_discrimination, LangOpts.PointerAuthVTPtrTypeDiscrimination)
+FEATURE(ptrauth_indirect_gotos, LangOpts.PointerAuthIndirectGotos)
 FEATURE(ptrauth_member_function_pointer_type_discrimination, LangOpts.PointerAuthCalls)
 FEATURE(ptrauth_init_fini, LangOpts.PointerAuthInitFini)
 EXTENSION(swiftcc,
diff --git a/clang/include/clang/Basic/LangOptions.def b/clang/include/clang/Basic/LangOptions.def
index 491759e2fcdbb..bdf77a5b35208 100644
--- a/clang/include/clang/Basic/LangOptions.def
+++ b/clang/include/clang/Basic/LangOptions.def
@@ -165,6 +165,7 @@ LANGOPT(ExperimentalLibrary, 1, 0, "enable unstable and experimental library fea
 LANGOPT(PointerAuthIntrinsics, 1, 0, "pointer authentication intrinsics")
 LANGOPT(PointerAuthCalls  , 1, 0, "function pointer authentication")
 LANGOPT(PointerAuthReturns, 1, 0, "return pointer authentication")
+LANGOPT(PointerAuthIndirectGotos, 1, 0, "indirect gotos pointer authentication")
 LANGOPT(PointerAuthAuthTraps, 1, 0, "pointer authentication failure traps")
 LANGOPT(PointerAuthVTPtrAddressDiscrimination, 1, 0, "incorporate address discrimination in authenticated vtable pointers")
 LANGOPT(PointerAuthVTPtrTypeDiscrimination, 1, 0, "incorporate type discrimination in authenticated vtable pointers")
diff --git a/clang/include/clang/Basic/PointerAuthOptions.h b/clang/include/clang/Basic/PointerAuthOptions.h
index 197d63642ca6d..2711639dbe299 100644
--- a/clang/include/clang/Basic/PointerAuthOptions.h
+++ b/clang/include/clang/Basic/PointerAuthOptions.h
@@ -154,6 +154,9 @@ class PointerAuthSchema {
 };
 
 struct PointerAuthOptions {
+  /// Do indirect goto label addresses need to be authenticated?
+  bool IndirectGotos = false;
+
   /// The ABI for C function pointers.
   PointerAuthSchema FunctionPointers;
 
diff --git a/clang/include/clang/Driver/Options.td b/clang/include/clang/Driver/Options.td
index 58ca6f2bea9e4..791b7261ddbda 100644
--- a/clang/include/clang/Driver/Options.td
+++ b/clang/include/clang/Driver/Options.td
@@ -4228,6 +4228,8 @@ defm ptrauth_vtable_pointer_address_discrimination :
 defm ptrauth_vtable_pointer_type_discrimination :
   OptInCC1FFlag<"ptrauth-vtable-pointer-type-discrimination", "Enable type discrimination of vtable pointers">;
 defm ptrauth_init_fini : OptInCC1FFlag<"ptrauth-init-fini", "Enable signing of function pointers in init/fini arrays">;
+defm ptrauth_indirect_gotos : OptInCC1FFlag<"ptrauth-indirect-gotos",
+  "Enable signing and authentication of indirect goto targets">;
 }
 
 def fenable_matrix : Flag<["-"], "fenable-matrix">, Group<f_Group>,
diff --git a/clang/lib/CodeGen/CodeGenFunction.cpp b/clang/lib/CodeGen/CodeGenFunction.cpp
index 26deeca95d326..33cb07a5bc832 100644
--- a/clang/lib/CodeGen/CodeGenFunction.cpp
+++ b/clang/lib/CodeGen/CodeGenFunction.cpp
@@ -865,6 +865,8 @@ void CodeGenFunction::StartFunction(GlobalDecl GD, QualType RetTy,
   const CodeGenOptions &CodeGenOpts = CGM.getCodeGenOpts();
   if (CodeGenOpts.PointerAuth.FunctionPointers)
     Fn->addFnAttr("ptrauth-calls");
+  if (CodeGenOpts.PointerAuth.IndirectGotos)
+    Fn->addFnAttr("ptrauth-indirect-gotos");
 
   // Apply xray attributes to the function (as a string, for now)
   bool AlwaysXRayAttr = false;
diff --git a/clang/lib/Driver/ToolChains/Clang.cpp b/clang/lib/Driver/ToolChains/Clang.cpp
index aa285c39f14b4..f9dba2c9f22cd 100644
--- a/clang/lib/Driver/ToolChains/Clang.cpp
+++ b/clang/lib/Driver/ToolChains/Clang.cpp
@@ -1789,6 +1789,9 @@ void Clang::AddAArch64TargetArgs(const ArgList &Args,
       options::OPT_fno_ptrauth_vtable_pointer_type_discrimination);
   Args.addOptInFlag(CmdArgs, options::OPT_fptrauth_init_fini,
                     options::OPT_fno_ptrauth_init_fini);
+
+  Args.addOptInFlag(CmdArgs, options::OPT_fptrauth_indirect_gotos,
+                    options::OPT_fno_ptrauth_indirect_gotos);
 }
 
 void Clang::AddLoongArchTargetArgs(const ArgList &Args,
diff --git a/clang/lib/Frontend/CompilerInvocation.cpp b/clang/lib/Frontend/CompilerInvocation.cpp
index f42e28ba7e629..a64e394a7754e 100644
--- a/clang/lib/Frontend/CompilerInvocation.cpp
+++ b/clang/lib/Frontend/CompilerInvocation.cpp
@@ -1480,13 +1480,14 @@ void CompilerInvocation::setDefaultPointerAuthOptions(
     Opts.CXXVirtualFunctionPointers = Opts.CXXVirtualVariadicFunctionPointers =
         PointerAuthSchema(Key::ASIA, true, Discrimination::Decl);
   }
+  Opts.IndirectGotos = LangOpts.PointerAuthIndirectGotos;
 }
 
 static void parsePointerAuthOptions(PointerAuthOptions &Opts,
                                     const LangOptions &LangOpts,
                                     const llvm::Triple &Triple,
                                     DiagnosticsEngine &Diags) {
-  if (!LangOpts.PointerAuthCalls)
+  if (!LangOpts.PointerAuthCalls && !LangOpts.PointerAuthIndirectGotos)
     return;
 
   CompilerInvocation::setDefaultPointerAuthOptions(Opts, LangOpts, Triple);
@@ -3390,6 +3391,8 @@ static void GeneratePointerAuthArgs(const LangOptions &Opts,
     GenerateArg(Consumer, OPT_fptrauth_calls);
   if (Opts.PointerAuthReturns)
     GenerateArg(Consumer, OPT_fptrauth_returns);
+  if (Opts.PointerAuthIndirectGotos)
+    GenerateArg(Consumer, OPT_fptrauth_indirect_gotos);
   if (Opts.PointerAuthAuthTraps)
     GenerateArg(Consumer, OPT_fptrauth_auth_traps);
   if (Opts.PointerAuthVTPtrAddressDiscrimination)
@@ -3405,6 +3408,7 @@ static void ParsePointerAuthArgs(LangOptions &Opts, ArgList &Args,
   Opts.PointerAuthIntrinsics = Args.hasArg(OPT_fptrauth_intrinsics);
   Opts.PointerAuthCalls = Args.hasArg(OPT_fptrauth_calls);
   Opts.PointerAuthReturns = Args.hasArg(OPT_fptrauth_returns);
+  Opts.PointerAuthIndirectGotos = Args.hasArg(OPT_fptrauth_indirect_gotos);
   Opts.PointerAuthAuthTraps = Args.hasArg(OPT_fptrauth_auth_traps);
   Opts.PointerAuthVTPtrAddressDiscrimination =
       Args.hasArg(OPT_fptrauth_vtable_pointer_address_discrimination);
diff --git a/clang/test/CodeGen/ptrauth-function-attributes.c b/clang/test/CodeGen/ptrauth-function-attributes.c
index 7ec30498b9d35..7f93ccc7c4bce 100644
--- a/clang/test/CodeGen/ptrauth-function-attributes.c
+++ b/clang/test/CodeGen/ptrauth-function-attributes.c
@@ -4,10 +4,15 @@
 // RUN: %clang_cc1 -triple arm64-apple-ios  -fptrauth-calls   -emit-llvm %s  -o - | FileCheck %s --check-prefixes=ALL,CALLS
 // RUN: %clang_cc1 -triple aarch64-linux-gnu -fptrauth-calls  -emit-llvm %s  -o - | FileCheck %s --check-prefixes=ALL,CALLS
 
+// RUN: %clang_cc1 -triple arm64-apple-ios  -fptrauth-indirect-gotos -emit-llvm %s -o - | FileCheck %s --check-prefixes=ALL,GOTOS
+// RUN: %clang_cc1 -triple arm64e-apple-ios -fptrauth-indirect-gotos -emit-llvm %s -o - | FileCheck %s --check-prefixes=ALL,GOTOS
+
 // ALL: define {{(dso_local )?}}void @test() #0
 void test() {
 }
 
 // CALLS: attributes #0 = {{{.*}} "ptrauth-calls" {{.*}}}
 
+// GOTOS: attributes #0 = {{{.*}} "ptrauth-indirect-gotos" {{.*}}}
+
 // OFF-NOT: attributes {{.*}} "ptrauth-
diff --git a/llvm/docs/PointerAuth.md b/llvm/docs/PointerAuth.md
index cf2cc6305f130..e027c902e58e1 100644
--- a/llvm/docs/PointerAuth.md
+++ b/llvm/docs/PointerAuth.md
@@ -18,6 +18,9 @@ At the IR level, it is represented using:
 * a [set of intrinsics](#intrinsics) (to sign/authenticate pointers)
 * a [signed pointer constant](#constant) (to sign globals)
 * a [call operand bundle](#operand-bundle) (to authenticate called pointers)
+* a [set of function attributes](#function-attributes) (to describe what
+  pointers are signed and how, to control implicit codegen in the backend, as
+  well as preserve invariants in the mid-level optimizer)
 
 The current implementation leverages the
 [Armv8.3-A PAuth/Pointer Authentication Code](#armv8-3-a-pauth-pointer-authentication-code)
@@ -287,6 +290,27 @@ but with the added guarantee that `%fp_i`, `%fp_auth`, and `%fp_auth_p`
 are not stored to (and reloaded from) memory.
 
 
+### Function Attributes
+
+Some function attributes are used to describe other pointer authentication
+operations that are not otherwise explicitly expressed in IR.
+
+#### ``ptrauth-indirect-gotos``
+
+``ptrauth-indirect-gotos`` specifies that indirect gotos in this function
+should authenticate their target.  At the IR level, no other change is needed.
+When lowering [``blockaddress`` constants](https://llvm.org/docs/LangRef.html#blockaddress),
+and [``indirectbr`` instructions](https://llvm.org/docs/LangRef.html#i-indirectbr),
+this tells the backend to respectively sign and authenticate the pointers.
+
+The specific scheme isn't ABI-visible.  Currently, the AArch64 backend
+signs blockaddresses using the `ASIA` key, with an integer discriminator
+derived from the parent function's name, using the SipHash stable discriminator:
+```
+  ptrauth_string_discriminator("<function_name> blockaddress")
+```
+
+
 ## AArch64 Support
 
 AArch64 is currently the only architecture with full support of the pointer
diff --git a/llvm/include/llvm/CodeGen/AsmPrinter.h b/llvm/include/llvm/CodeGen/AsmPrinter.h
index a60dce30c4a6c..290da83cee35f 100644
--- a/llvm/include/llvm/CodeGen/AsmPrinter.h
+++ b/llvm/include/llvm/CodeGen/AsmPrinter.h
@@ -577,6 +577,9 @@ class AsmPrinter : public MachineFunctionPass {
     report_fatal_error("ptrauth constant lowering not implemented");
   }
 
+  /// Lower the specified BlockAddress to an MCExpr.
+  virtual const MCExpr *lowerBlockAddressConstant(const BlockAddress &BA);
+
   /// Return true if the basic block has exactly one predecessor and the control
   /// transfer mechanism between the predecessor and this block is a
   /// fall-through.
diff --git a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
index c52cbff689dc5..724b7cd94adee 100644
--- a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
+++ b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
@@ -3139,7 +3139,7 @@ const MCExpr *AsmPrinter::lowerConstant(const Constant *CV) {
     return MCSymbolRefExpr::create(getSymbol(GV), Ctx);
 
   if (const BlockAddress *BA = dyn_cast<BlockAddress>(CV))
-    return MCSymbolRefExpr::create(GetBlockAddressSymbol(BA), Ctx);
+    return lowerBlockAddressConstant(*BA);
 
   if (const auto *Equiv = dyn_cast<DSOLocalEquivalent>(CV))
     return getObjFileLowering().lowerDSOLocalEquivalent(Equiv, TM);
@@ -3821,6 +3821,10 @@ MCSymbol *AsmPrinter::GetBlockAddressSymbol(const BasicBlock *BB) const {
   return const_cast<AsmPrinter *>(this)->getAddrLabelSymbol(BB);
 }
 
+const MCExpr *AsmPrinter::lowerBlockAddressConstant(const BlockAddress &BA) {
+  return MCSymbolRefExpr::create(GetBlockAddressSymbol(&BA), OutContext);
+}
+
 /// GetCPISymbol - Return the symbol for the specified constant pool entry.
 MCSymbol *AsmPrinter::GetCPISymbol(unsigned CPID) const {
   if (getSubtargetInfo().getTargetTriple().isWindowsMSVCEnvironment()) {
diff --git a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
index 64d41d4147644..093208bf70cda 100644
--- a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
+++ b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
@@ -93,6 +93,8 @@ class AArch64AsmPrinter : public AsmPrinter {
 
   const MCExpr *lowerConstantPtrAuth(const ConstantPtrAuth &CPA) override;
 
+  const MCExpr *lowerBlockAddressConstant(const BlockAddress &BA) override;
+
   void emitStartOfAsmFile(Module &M) override;
   void emitJumpTableInfo() override;
   std::tuple<const MCSymbol *, uint64_t, const MCSymbol *,
@@ -128,7 +130,7 @@ class AArch64AsmPrinter : public AsmPrinter {
 
   void emitSled(const MachineInstr &MI, SledKind Kind);
 
-  // Emit the sequence for BLRA (authenticate + branch).
+  // Emit the sequence for BRA/BLRA (authenticate + branch/call).
   void emitPtrauthBranch(const MachineInstr *MI);
   // Emit the sequence to compute a discriminator into x17, or reuse AddrDisc.
   unsigned emitPtrauthDiscriminator(uint16_t Disc, unsigned AddrDisc,
@@ -1581,6 +1583,7 @@ unsigned AArch64AsmPrinter::emitPtrauthDiscriminator(uint16_t Disc,
 
 void AArch64AsmPrinter::emitPtrauthBranch(const MachineInstr *MI) {
   unsigned InstsEmitted = 0;
+  bool IsCall = MI->getOpcode() == AArch64::BLRA;
   unsigned BrTarget = MI->getOperand(0).getReg();
 
   auto Key = (AArch64PACKey::ID)MI->getOperand(1).getImm();
@@ -1597,10 +1600,17 @@ void AArch64AsmPrinter::emitPtrauthBranch(const MachineInstr *MI) {
   bool IsZeroDisc = DiscReg == AArch64::XZR;
 
   unsigned Opc;
-  if (Key == AArch64PACKey::IA)
-    Opc = IsZeroDisc ? AArch64::BLRAAZ : AArch64::BLRAA;
-  else
-    Opc = IsZeroDisc ? AArch64::BLRABZ : AArch64::BLRAB;
+  if (IsCall) {
+    if (Key == AArch64PACKey::IA)
+      Opc = IsZeroDisc ? AArch64::BLRAAZ : AArch64::BLRAA;
+    else
+      Opc = IsZeroDisc ? AArch64::BLRABZ : AArch64::BLRAB;
+  } else {
+    if (Key == AArch64PACKey::IA)
+      Opc = IsZeroDisc ? AArch64::BRAAZ : AArch64::BRAA;
+    else
+      Opc = IsZeroDisc ? AArch64::BRABZ : AArch64::BRAB;
+  }
 
   MCInst BRInst;
   BRInst.setOpcode(Opc);
@@ -1866,6 +1876,20 @@ void AArch64AsmPrinter::LowerMOVaddrPAC(const MachineInstr &MI) {
   assert(STI->getInstrInfo()->getInstSizeInBytes(MI) >= InstsEmitted * 4);
 }
 
+const MCExpr *
+AArch64AsmPrinter::lowerBlockAddressConstant(const BlockAddress &BA) {
+  const MCExpr *BAE = AsmPrinter::lowerBlockAddressConstant(BA);
+  const Function &Fn = *BA.getFunction();
+
+  if (std::optional<uint16_t> BADisc =
+          STI->getPtrAuthBlockAddressDiscriminator(Fn))
+    return AArch64AuthMCExpr::create(BAE, *BADisc, AArch64PACKey::IA,
+                                     /* HasAddressDiversity= */ false,
+                                     OutContext);
+
+  return BAE;
+}
+
 // Simple pseudo-instructions have their lowering (with expansion to real
 // instructions) auto-generated.
 #include "AArch64GenMCPseudoLowering.inc"
@@ -2010,6 +2034,7 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
     LowerMOVaddrPAC(*MI);
     return;
 
+  case AArch64::BRA:
   case AArch64::BLRA:
     emitPtrauthBranch(MI);
     return;
diff --git a/llvm/lib/Target/AArch64/AArch64FastISel.cpp b/llvm/lib/Target/AArch64/AArch64FastISel.cpp
index e3c5a143b2889..1a99a905a47d8 100644
--- a/llvm/lib/Target/AArch64/AArch64FastISel.cpp
+++ b/llvm/lib/Target/AArch64/AArch64FastISel.cpp
@@ -2516,6 +2516,10 @@ bool AArch64FastISel::selectIndirectBr(const Instruction *I) {
   if (AddrReg == 0)
     return false;
 
+  // Authenticated indirectbr is not implemented yet.
+  if (FuncInfo.MF->getFunction().hasFnAttribute("ptrauth-indirect-gotos"))
+    return false;
+
   // Emit the indirect branch.
   const MCInstrDesc &II = TII.get(AArch64::BR);
   AddrReg = constrainOperandRegClass(II, AddrReg,  II.getNumDefs());
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index e0c3cc5eddb82..8978b254f40a3 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -85,6 +85,7 @@
 #include "llvm/Support/InstructionCost.h"
 #include "llvm/Support/KnownBits.h"
 #include "llvm/Support/MathExtras.h"
+#include "llvm/Support/SipHash.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Target/TargetOptions.h"
@@ -509,6 +510,7 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
   setOperationAction(ISD::SELECT_CC, MVT::f64, Custom);
   setOperationAction(ISD::BR_JT, MVT::Other, Custom);
   setOperationAction(ISD::JumpTable, MVT::i64, Custom);
+  setOperationAction(ISD::BRIND, MVT::Other, Custom);
   setOperationAction(ISD::SETCCCARRY, MVT::i64, Custom);
 
   setOperationAction(ISD::PtrAuthGlobalAddress, MVT::i64, Custom);
@@ -6694,6 +6696,8 @@ SDValue AArch64TargetLowering::LowerOperation(SDValue Op,
     return LowerJumpTable(Op, DAG);
   case ISD::BR_JT:
     return LowerBR_JT(Op, DAG);
+  case ISD::BRIND:
+    return LowerBRIND(Op, DAG);
   case ISD::ConstantPool:
     return LowerConstantPool(Op, DAG);
   case ISD::BlockAddress:
@@ -10685,6 +10689,26 @@ SDValue AArch64TargetLowering::LowerBR_JT(SDValue Op,
   return DAG.getNode(ISD::BRIND, DL, MVT::Other, JTInfo, SDValue(Dest, 0));
 }
 
+SDValue AArch64TargetLowering::LowerBRIND(SDValue Op, SelectionDAG &DAG) const {
+  MachineFunction &MF = DAG.getMachineFunction();
+  std::optional<uint16_t> BADisc =
+      Subtarget->getPtrAuthBlockAddressDiscriminator(MF.getFunction());
+  if (!BADisc)
+    return SDValue();
+
+  SDLoc DL(Op);
+  SDValue Chain = Op.getOperand(0);
+  SDValue Dest = Op.getOperand(1);
+
+  SDValue Disc = DAG.getTargetConstant(*BADisc, DL, MVT::i64);
+  SDValue Key = DAG.getTargetConstant(AArch64PACKey::IA, DL, MVT::i32);
+  SDValue AddrDisc = DAG.getRegister(AArch64::XZR, MVT::i64);
+
+  SDNode *BrA = DAG.getMachineNode(AArch64::BRA, DL, MVT::Other,
+                                   {Dest, Key, Disc, AddrDisc, Chain});
+  return SDValue(BrA, 0);
+}
+
 SDValue AArch64TargetLowering::LowerConstantPool(SDValue Op,
                                                  SelectionDAG &DAG) const {
   ConstantPoolSDNode *CP = cast<ConstantPoolSDNode>(Op);
@@ -10704,15 +10728,36 @@ SDValue AArch64TargetLowering::LowerConstantPool(SDValue Op,
 
 SDValue AArch64TargetLowering::LowerBlockAddress(SDValue Op,
                                                SelectionDAG &DAG) const {
-  BlockAddressSDNode *BA = cast<BlockAddressSDNode>(Op);
+  BlockAddressSDNode *BAN = cast<BlockAddressSDNode>(Op);
+  const BlockAddress *BA = BAN->getBlockAddress();
+
+  if (std::optional<uint16_t> BADisc =
+          Subtarget->getPtrAuthBlockAddressDiscriminator(*BA->getFunction())) {
+    SDLoc DL(Op);
+
+    // This isn't cheap, but BRIND is rare.
+    SDValue TargetBA = DAG.getTargetBlockAddress(BA, BAN->getValueType(0));
+
+    SDValue Disc = DAG.getTargetConstant(*BADisc, DL, MVT::i64);
+
+    SDValue Key = DAG.getTargetConstant(AArch64PACKey::IA, DL, MVT::i32);
+    SDValue AddrDisc = DAG.getRegister(AArch64::XZR, MVT::i64);
+
+    SDNode *MOV =
+        DAG.getMachineNode(AArch64::MOVaddrPAC, DL, {MVT::Other, MVT::Glue},
+                           {TargetBA, Key, AddrDisc, Disc});
+    return DAG.getCopyFromReg(SDValue(MOV, 0), DL, AArch64::X16, MVT::i64,
+                              SDValue(MOV, 1));
+  }
+
   CodeModel::Model CM = getTargetMachine().getCodeModel();
   if (CM == CodeModel::Large && !Subtarget->isTargetMachO()) {
     if (!getTargetMachine().isPositionIndependent())
-      return getAddrLarge(BA, DAG);
+      return getAddrLarge(BAN, DAG);
   } else if (CM == CodeModel::Tiny) {
-    return getAddrTiny(BA, DAG);
+    return getAddrTiny(BAN, DAG);
   }
-  return getAddr(BA, DAG);
+  return getAddr(BAN, DAG);
 }
 
 SDValue AArch64TargetLowering::LowerDarwin_VASTART(SDValue Op,
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.h b/llvm/lib/Target/AArch64/AArch64ISelLowering.h
index 047c852bb01d2..69386dac318ca 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.h
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.h
@@ -1143,6 +1143,7 @@ class AArch64TargetLowering : public TargetLowering {
                          SelectionDAG &DAG) const;
   SDValue LowerJumpTable(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerBR_JT(SDValue Op, SelectionDAG &DAG) const;
+  SDValue LowerBRIND(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerConstantPool(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerBlockAddress(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerAAPCS_VASTART(SDValue Op, SelectionDAG &DAG) const;
diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.td b/llvm/lib/Target/AArch64/AArch64InstrInfo.td
index 1e06d5fdc7562..a65e5c68a24ad 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.td
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.td
@@ -1755,6 +1755,24 @@ let Predicates = [HasPAuth] in {
     let Uses = [SP];
   }
 
+  // BRA pseudo, generalized version of BRAA/BRAB/Z.
+  // This directly manipulates x16/x17, which are the only registers the OS
+  // guarantees are safe to use for sensitive operations.
+  def BRA : Pseudo<(outs), (ins GPR64noip:$Rn, i32imm:$Key, i64imm:$Disc,
+                                GPR64noip:$AddrDisc), []>, Sched<[]> {
+    let isCodeGenOnly = 1;
+    let hasNoSchedulingInfo = 1;
+    let hasSideEffects = 1;
+    let mayStore = 0;
+    let mayLoad = 0;
+    let isBranch = 1;
+    let isTerminator = 1;
+    let isBarrier = 1;
+    let isIndirectBranch = 1;
+    let Size = 12; // 4 fixed + 8 variable, to compute discriminator.
+    let Defs = [X17];
+  }
+
   let isReturn = 1, isTerminator = 1, isBarrier = 1 in {
     def RETAA   : AuthReturn<0b010, 0, "retaa">;
     def RETAB   : AuthReturn<0b010, 1, "retab">;
diff --git a/llvm/lib/Target/AArch64/AArch64Subtarget.cpp b/llvm/lib/Target/AArch64/AArch64Subtarget.cpp
index 1fad1d5ca6d7d..7275116e38f2e 100644
--- a/llvm/lib/Target/AArch64/AArch64Subtarget.cpp
+++ b/llvm/lib/Target/AArch64/AArch64Subtarget.cpp
@@ -24,6 +24,7 @@
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineScheduler.h"
 #include "llvm/IR/GlobalValue.h"
+#include "llvm/Support/SipHash.h"
 #include "llvm/TargetParser/AArch64TargetParser.h"
 
 using namespace llvm;
@@ -574,6 +575,16 @@ AArch64Subtarget::getAuthenticatedLRCheckMethod() const {
   return AArch64PAuth::AuthCheckMethod::None;
 }
 
+std::optional<uint16_t> AArch64Subtarget::getPtrAuthBlockAddressDiscriminator(
+    const Function &ParentFn) const {
+  if (!ParentFn.hasFnAttribute("ptrauth-indirect-gotos"))
+    return std::nullopt;
+  // We currently have one simple mechanism for all targets.
+  // This isn't ABI, so we can always do better in the future.
+  return getPointerAuthStableSipHash(
+      (Twine(ParentFn.getName()) + " blockaddress").str());
+}
+
 bool AArch64Subtarget::enableMachinePipeliner() const {
   return getSchedModel().hasInstrSchedModel();
 }
diff --git a/llvm/lib/Target/AArch64/AArch64Subtarget.h b/llvm/lib/Target/AArch64/AArch64Subtarget.h
index 5faba09aa67bd..172beaaaabc01 100644
--- a/llvm/lib/Target/AArch64/AArch64Subtarget.h
+++ b/llvm/lib/Target/AArch64/AArch64Subtarget.h
@@ -412,6 +412,15 @@ class AArch64Subtarget final : public AArch64GenSubtargetInfo {
   /// Choose a method of checking LR before performing a tail call.
   AArch64PAuth::AuthCheckMethod getAuthenticatedLRCheckMethod() const;
 
+  /// Compute the integer discriminator for a given BlockAddress constant, if
+  /// blockaddress signing is enabled (using function attribute
+  /// "ptrauth-indirect-gotos").
+  /// Note that this assumes the discriminator is independent of the indirect
+  /// goto branch site itself, i.e., it's the same for all BlockAddresses in
+  /// a function.
+  std::optional<uint16_t>
+  getPtrAuthBlockAddressDiscriminator(const Function &ParentFn) const;
+
   const PseudoSourceValue *getAddressCheckPSV() const {
     return AddressCheckPSV.get();
   }
diff --git a/llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp b/llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp
index 9e0860934f777..b5787b6bd8b82 100644
--- a/llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp
+++ b/llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp
@@ -2547,6 +2547,15 @@ bool AArch64InstructionSelector::select(MachineInstr &I) {
     return selectCompareBranch(I, MF, MRI);
 
   case TargetOpcode::G_BRINDIRECT: {
+    if (std::optional<uint16_t> BADisc =
+            STI.getPtrAuthBlockAddressDiscriminator(MF.getFunction())) {
+      auto MI = MIB.buildInstr(AArch64::BRA, {}, {I.getOperand(0).getReg()});
+      MI.addImm(AArch64PACKey::IA);
+      MI.addImm(*BADisc);
+      MI.addReg(/*AddrDisc=*/AArch64::XZR);
+      I.eraseFromParent();
+      return constrainSelectedInstRegOperands(*MI, TII, TRI, RBI);
+    }
     I.setDesc(TII.get(AArch64::BR));
     return constrainSelectedInstRegOperands(I, TII, TRI, RBI);
   }
@@ -3461,6 +3470,23 @@ bool AArch64InstructionSelector::select(MachineInstr &I) {
     return true;
   }
   case TargetOpcode::G_BLOCK_ADDR: {
+    Function *BAFn = I.getOperand(1).getBlockAddress()->getFunction();
+    if (std::optional<uint16_t> BADisc =
+            STI.getPtrAuthBlockAddressDiscriminator(*BAFn)) {
+      MIB.buildInstr(TargetOpcode::IMPLICIT_DEF, {AArch64::X16}, {});
+      MIB.buildInstr(TargetOpcode::IMPLICIT_DEF, {AArch64::X17}, {});
+      MIB.buildInstr(AArch64::MOVaddrPAC)
+          .addBlockAddress(I.getOperand(1).getBlockAddress())
+          .addImm(AArch64PACKey::IA)
+          .addReg(/*AddrDisc=*/AArch64::XZR)
+          .addImm(*BADisc)
+          .constrainAllUses(TII, TRI, RBI);
+      MIB.buildCopy(I.getOperand(0).getReg(), Register(AArch64::X16));
+      RBI.constrainGenericRegister(I.getOperand(0).getReg(),
+                                   AArch64::GPR64RegClass, MRI);
+      I.eraseFromParent();
+      return true;
+    }
     if (TM.getCodeModel() == CodeModel::Large && !TM.isPositionIndependent()) {
       materializeLargeCMVal(I, I.getOperand(1).getBlockAddress(), 0);
       I.eraseFromParent();
diff --git a/llvm/test/CodeGen/AArch64/ptrauth-indirectbr.ll b/llvm/test/CodeGen/AArch64/ptrauth-indirectbr.ll
new file mode 100644
index 0000000000000..db49422cf1abb
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/ptrauth-indirectbr.ll
@@ -0,0 +1,106 @@
+; RUN: llc -mtriple arm64e-apple-darwin \
+; RUN:   -asm-verbose=false -aarch64-enable-collect-loh=false \
+; RUN:   -o - %s | FileCheck %s
+
+; RUN: llc -mtriple arm64e-apple-darwin \
+; RUN:   -global-isel -global-isel-abort=1 -verify-machineinstrs \
+; RUN:   -asm-verbose=false -aarch64-enable-collect-loh=false \
+; RUN:   -o - %s | FileCheck %s
+
+; The discriminator is the same for all blockaddresses in the function.
+; ptrauth_string_discriminator("test_blockaddress blockaddress") == 52152
+
+; CHECK-LABEL: _test_blockaddress:
+; CHECK:         adrp x16, [[F1BB1ADDR:Ltmp[0-9]+]]@PAGE
+; CHECK-NEXT:    add x16, x16, [[F1BB1ADDR]]@PAGEOFF
+; CHECK-NEXT:    mov x17, #[[F1DISCVAL:52152]]
+; CHECK-NEXT:    pacia x16, x17
+; CHECK-NEXT:    mov x0, x16
+; CHECK-NEXT:    adrp x16, [[F1BB2ADDR:Ltmp[0-9]+]]@PAGE
+; CHECK-NEXT:    add x16, x16, [[F1BB2ADDR]]@PAGEOFF
+; CHECK-NEXT:    mov x17, #[[F1DISCVAL]]
+; CHECK-NEXT:    pacia x16, x17
+; CHECK-NEXT:    mov x1, x16
+; CHECK-NEXT:    bl _dummy_choose
+; CHECK-NEXT:    mov x17, #[[F1DISCVAL]]
+; CHECK-NEXT:    braa x0, x17
+; CHECK:        [[F1BB1ADDR]]:
+; CHECK-NEXT:   [[F1BB1:LBB[0-9_]+]]:
+; CHECK-NEXT:    mov w0, #1
+; CHECK:        [[F1BB2ADDR]]:
+; CHECK-NEXT:   [[F1BB2:LBB[0-9_]+]]:
+; CHECK-NEXT:    mov w0, #2
+define i32 @test_blockaddress() #0 {
+entry:
+  %tmp0 = call ptr @dummy_choose(ptr blockaddress(@test_blockaddress, %bb1), ptr blockaddress(@test_blockaddress, %bb2))
+  indirectbr ptr %tmp0, [label %bb1, label %bb2]
+
+bb1:
+  ret i32 1
+
+bb2:
+  ret i32 2
+}
+
+; Test another function to compare the discriminator.
+; ptrauth_string_discriminator("test_blockaddress_2 blockaddress") == 22012
+
+; CHECK-LABEL: _test_blockaddress_2:
+; CHECK:         adrp x16, [[F2BB1ADDR:Ltmp[0-9]+]]@PAGE
+; CHECK-NEXT:    add x16, x16, [[F2BB1ADDR]]@PAGEOFF
+; CHECK-NEXT:    mov x17, #[[F2DISCVAL:22012]]
+; CHECK-NEXT:    pacia x16, x17
+; CHECK-NEXT:    mov x0, x16
+; CHECK-NEXT:    adrp x16, [[F2BB2ADDR:Ltmp[0-9]+]]@PAGE
+; CHECK-NEXT:    add x16, x16, [[F2BB2ADDR]]@PAGEOFF
+; CHECK-NEXT:    mov x17, #[[F2DISCVAL]]
+; CHECK-NEXT:    pacia x16, x17
+; CHECK-NEXT:    mov x1, x16
+; CHECK-NEXT:    bl _dummy_choose
+; CHECK-NEXT:    mov x17, #[[F2DISCVAL]]
+; CHECK-NEXT:    braa x0, x17
+; CHECK:        [[F2BB1ADDR]]:
+; CHECK-NEXT:   [[F2BB1:LBB[0-9_]+]]:
+; CHECK-NEXT:    mov w0, #1
+; CHECK:        [[F2BB2ADDR]]:
+; CHECK-NEXT:   [[F2BB2:LBB[0-9_]+]]:
+; CHECK-NEXT:    mov w0, #2
+define i32 @test_blockaddress_2() #0 {
+entry:
+  %tmp0 = call ptr @dummy_choose(ptr blockaddress(@test_blockaddress_2, %bb1), ptr blockaddress(@test_blockaddress_2, %bb2))
+  indirectbr ptr %tmp0, [label %bb1, label %bb2]
+
+bb1:
+  ret i32 1
+
+bb2:
+  ret i32 2
+}
+
+; CHECK-LABEL: _test_blockaddress_other_function:
+; CHECK:         adrp x16, [[F1BB1ADDR]]@PAGE
+; CHECK-NEXT:    add x16, x16, [[F1BB1ADDR]]@PAGEOFF
+; CHECK-NEXT:    mov x17, #[[F1DISCVAL]]
+; CHECK-NEXT:    pacia x16, x17
+; CHECK-NEXT:    mov x0, x16
+; CHECK-NEXT:    ret
+define ptr @test_blockaddress_other_function() #0 {
+  ret ptr blockaddress(@test_blockaddress, %bb1)
+}
+
+; CHECK-LABEL: .section __DATA,__const
+; CHECK-NEXT:  .globl _test_blockaddress_array
+; CHECK-NEXT:  .p2align 4
+; CHECK-NEXT:  _test_blockaddress_array:
+; CHECK-NEXT:   .quad [[F1BB1ADDR]]@AUTH(ia,[[F1DISCVAL]]
+; CHECK-NEXT:   .quad [[F1BB2ADDR]]@AUTH(ia,[[F1DISCVAL]]
+; CHECK-NEXT:   .quad [[F2BB1ADDR]]@AUTH(ia,[[F2DISCVAL]]
+; CHECK-NEXT:   .quad [[F2BB2ADDR]]@AUTH(ia,[[F2DISCVAL]]
+ at test_blockaddress_array = constant [4 x ptr] [
+  ptr blockaddress(@test_blockaddress, %bb1), ptr blockaddress(@test_blockaddress, %bb2),
+  ptr blockaddress(@test_blockaddress_2, %bb1), ptr blockaddress(@test_blockaddress_2, %bb2)
+]
+
+declare ptr @dummy_choose(ptr, ptr)
+
+attributes #0 = { "ptrauth-indirect-gotos" nounwind }

>From c21435b2321c6b8c9104bfe58e5e16805b457adb Mon Sep 17 00:00:00 2001
From: Ahmed Bougacha <ahmed at bougacha.org>
Date: Mon, 22 Jul 2024 11:08:32 -0700
Subject: [PATCH 2/5] Error on addrlabel add/sub with -fptrauth-indirect-gotos
 in frontend.

---
 .../include/clang/Basic/DiagnosticSemaKinds.td  |  3 +++
 clang/lib/Sema/SemaExpr.cpp                     | 17 +++++++++++++++++
 clang/test/Sema/ptrauth-indirect-goto.c         | 16 ++++++++++++++++
 3 files changed, 36 insertions(+)
 create mode 100644 clang/test/Sema/ptrauth-indirect-goto.c

diff --git a/clang/include/clang/Basic/DiagnosticSemaKinds.td b/clang/include/clang/Basic/DiagnosticSemaKinds.td
index c4a4e8064fac2..b8d97a6b14fe6 100644
--- a/clang/include/clang/Basic/DiagnosticSemaKinds.td
+++ b/clang/include/clang/Basic/DiagnosticSemaKinds.td
@@ -949,6 +949,9 @@ def note_ptrauth_virtual_function_pointer_incomplete_arg_ret :
 def note_ptrauth_virtual_function_incomplete_arg_ret_type :
   Note<"%0 is incomplete">;
 
+def err_ptrauth_indirect_goto_addrlabel_arithmetic : Error<
+  "%select{subtraction|addition}0 of address-of-label expressions is not "
+  "supported with ptrauth indirect gotos">;
 
 /// main()
 // static main() is not an error in C, just in C++.
diff --git a/clang/lib/Sema/SemaExpr.cpp b/clang/lib/Sema/SemaExpr.cpp
index 1a441d99515f4..439db55668cc6 100644
--- a/clang/lib/Sema/SemaExpr.cpp
+++ b/clang/lib/Sema/SemaExpr.cpp
@@ -10909,6 +10909,14 @@ QualType Sema::CheckAdditionOperands(ExprResult &LHS, ExprResult &RHS,
   if (isObjCPointer && checkArithmeticOnObjCPointer(*this, Loc, PExp))
     return QualType();
 
+  // Arithmetic on label addresses is normally allowed, except when we add
+  // a ptrauth signature to the addresses.
+  if (isa<AddrLabelExpr>(PExp) && getLangOpts().PointerAuthIndirectGotos) {
+    Diag(Loc, diag::err_ptrauth_indirect_goto_addrlabel_arithmetic)
+        << /*addition*/ 1;
+    return QualType();
+  }
+
   // Check array bounds for pointer arithemtic
   CheckArrayAccess(PExp, IExp);
 
@@ -10983,6 +10991,15 @@ QualType Sema::CheckSubtractionOperands(ExprResult &LHS, ExprResult &RHS,
         checkArithmeticOnObjCPointer(*this, Loc, LHS.get()))
       return QualType();
 
+    // Arithmetic on label addresses is normally allowed, except when we add
+    // a ptrauth signature to the addresses.
+    if (isa<AddrLabelExpr>(LHS.get()) &&
+        getLangOpts().PointerAuthIndirectGotos) {
+      Diag(Loc, diag::err_ptrauth_indirect_goto_addrlabel_arithmetic)
+          << /*subtraction*/ 0;
+      return QualType();
+    }
+
     // The result type of a pointer-int computation is the pointer type.
     if (RHS.get()->getType()->isIntegerType()) {
       // Subtracting from a null pointer should produce a warning.
diff --git a/clang/test/Sema/ptrauth-indirect-goto.c b/clang/test/Sema/ptrauth-indirect-goto.c
new file mode 100644
index 0000000000000..54d24740c07e4
--- /dev/null
+++ b/clang/test/Sema/ptrauth-indirect-goto.c
@@ -0,0 +1,16 @@
+// RUN: %clang_cc1 -fsyntax-only -verify %s -fptrauth-indirect-gotos
+
+int f() {
+  static int diffs[] = {
+    &&l1 - &&l1, // expected-error{{subtraction of address-of-label expressions is not supported with ptrauth indirect gotos}}
+    &&l1 - &&l2 // expected-error{{subtraction of address-of-label expressions is not supported with ptrauth indirect gotos}}
+  };
+
+  int diff_32 = &&l1 - &&l2; // expected-error{{subtraction of address-of-label expressions is not supported with ptrauth indirect gotos}}
+  goto *(&&l1 + diff_32); // expected-error{{addition of address-of-label expressions is not supported with ptrauth indirect gotos}}
+
+l1:
+  return 0;
+l2:
+  return 1;
+}

>From 46274c3089e2c75110a1863118373e4a1dd470fc Mon Sep 17 00:00:00 2001
From: Ahmed Bougacha <ahmed at bougacha.org>
Date: Mon, 22 Jul 2024 11:09:14 -0700
Subject: [PATCH 3/5] Address review feedback.

- add IfEnabled to function name
- test fast-isel, elf
- various nits
---
 llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp |   5 +-
 .../Target/AArch64/AArch64ISelLowering.cpp    |   7 +-
 llvm/lib/Target/AArch64/AArch64Subtarget.cpp  |   3 +-
 llvm/lib/Target/AArch64/AArch64Subtarget.h    |   7 +-
 .../GISel/AArch64InstructionSelector.cpp      |   5 +-
 .../CodeGen/AArch64/ptrauth-indirectbr.ll     | 257 ++++++++++++------
 6 files changed, 196 insertions(+), 88 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
index 2e64a262170c9..99a60533c5412 100644
--- a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
+++ b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
@@ -1937,10 +1937,9 @@ AArch64AsmPrinter::lowerBlockAddressConstant(const BlockAddress &BA) {
   const Function &Fn = *BA.getFunction();
 
   if (std::optional<uint16_t> BADisc =
-          STI->getPtrAuthBlockAddressDiscriminator(Fn))
+          STI->getPtrAuthBlockAddressDiscriminatorIfEnabled(Fn))
     return AArch64AuthMCExpr::create(BAE, *BADisc, AArch64PACKey::IA,
-                                     /* HasAddressDiversity= */ false,
-                                     OutContext);
+                                     /*HasAddressDiversity=*/false, OutContext);
 
   return BAE;
 }
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index b4f06f2bcd085..923df31421e90 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -10754,9 +10754,9 @@ SDValue AArch64TargetLowering::LowerBR_JT(SDValue Op,
 }
 
 SDValue AArch64TargetLowering::LowerBRIND(SDValue Op, SelectionDAG &DAG) const {
-  MachineFunction &MF = DAG.getMachineFunction();
+  const MachineFunction &MF = DAG.getMachineFunction();
   std::optional<uint16_t> BADisc =
-      Subtarget->getPtrAuthBlockAddressDiscriminator(MF.getFunction());
+      Subtarget->getPtrAuthBlockAddressDiscriminatorIfEnabled(MF.getFunction());
   if (!BADisc)
     return SDValue();
 
@@ -10796,7 +10796,8 @@ SDValue AArch64TargetLowering::LowerBlockAddress(SDValue Op,
   const BlockAddress *BA = BAN->getBlockAddress();
 
   if (std::optional<uint16_t> BADisc =
-          Subtarget->getPtrAuthBlockAddressDiscriminator(*BA->getFunction())) {
+          Subtarget->getPtrAuthBlockAddressDiscriminatorIfEnabled(
+              *BA->getFunction())) {
     SDLoc DL(Op);
 
     // This isn't cheap, but BRIND is rare.
diff --git a/llvm/lib/Target/AArch64/AArch64Subtarget.cpp b/llvm/lib/Target/AArch64/AArch64Subtarget.cpp
index 7275116e38f2e..32a355fe38f1c 100644
--- a/llvm/lib/Target/AArch64/AArch64Subtarget.cpp
+++ b/llvm/lib/Target/AArch64/AArch64Subtarget.cpp
@@ -575,7 +575,8 @@ AArch64Subtarget::getAuthenticatedLRCheckMethod() const {
   return AArch64PAuth::AuthCheckMethod::None;
 }
 
-std::optional<uint16_t> AArch64Subtarget::getPtrAuthBlockAddressDiscriminator(
+std::optional<uint16_t>
+AArch64Subtarget::getPtrAuthBlockAddressDiscriminatorIfEnabled(
     const Function &ParentFn) const {
   if (!ParentFn.hasFnAttribute("ptrauth-indirect-gotos"))
     return std::nullopt;
diff --git a/llvm/lib/Target/AArch64/AArch64Subtarget.h b/llvm/lib/Target/AArch64/AArch64Subtarget.h
index 64440f64fdfd0..e585aad2f7a68 100644
--- a/llvm/lib/Target/AArch64/AArch64Subtarget.h
+++ b/llvm/lib/Target/AArch64/AArch64Subtarget.h
@@ -416,13 +416,14 @@ class AArch64Subtarget final : public AArch64GenSubtargetInfo {
   AArch64PAuth::AuthCheckMethod getAuthenticatedLRCheckMethod() const;
 
   /// Compute the integer discriminator for a given BlockAddress constant, if
-  /// blockaddress signing is enabled (using function attribute
-  /// "ptrauth-indirect-gotos").
+  /// blockaddress signing is enabled, or std::nullopt otherwise.
+  /// Blockaddress signing is controlled by the function attribute
+  /// "ptrauth-indirect-gotos" on the parent function.
   /// Note that this assumes the discriminator is independent of the indirect
   /// goto branch site itself, i.e., it's the same for all BlockAddresses in
   /// a function.
   std::optional<uint16_t>
-  getPtrAuthBlockAddressDiscriminator(const Function &ParentFn) const;
+  getPtrAuthBlockAddressDiscriminatorIfEnabled(const Function &ParentFn) const;
 
   const PseudoSourceValue *getAddressCheckPSV() const {
     return AddressCheckPSV.get();
diff --git a/llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp b/llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp
index 31b489ac5ada3..f1f5f108c0181 100644
--- a/llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp
+++ b/llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp
@@ -2552,8 +2552,9 @@ bool AArch64InstructionSelector::select(MachineInstr &I) {
     return selectCompareBranch(I, MF, MRI);
 
   case TargetOpcode::G_BRINDIRECT: {
+    const Function &Fn = MF.getFunction();
     if (std::optional<uint16_t> BADisc =
-            STI.getPtrAuthBlockAddressDiscriminator(MF.getFunction())) {
+            STI.getPtrAuthBlockAddressDiscriminatorIfEnabled(Fn)) {
       auto MI = MIB.buildInstr(AArch64::BRA, {}, {I.getOperand(0).getReg()});
       MI.addImm(AArch64PACKey::IA);
       MI.addImm(*BADisc);
@@ -3477,7 +3478,7 @@ bool AArch64InstructionSelector::select(MachineInstr &I) {
   case TargetOpcode::G_BLOCK_ADDR: {
     Function *BAFn = I.getOperand(1).getBlockAddress()->getFunction();
     if (std::optional<uint16_t> BADisc =
-            STI.getPtrAuthBlockAddressDiscriminator(*BAFn)) {
+            STI.getPtrAuthBlockAddressDiscriminatorIfEnabled(*BAFn)) {
       MIB.buildInstr(TargetOpcode::IMPLICIT_DEF, {AArch64::X16}, {});
       MIB.buildInstr(TargetOpcode::IMPLICIT_DEF, {AArch64::X17}, {});
       MIB.buildInstr(AArch64::MOVaddrPAC)
diff --git a/llvm/test/CodeGen/AArch64/ptrauth-indirectbr.ll b/llvm/test/CodeGen/AArch64/ptrauth-indirectbr.ll
index db49422cf1abb..c96d9dce6f928 100644
--- a/llvm/test/CodeGen/AArch64/ptrauth-indirectbr.ll
+++ b/llvm/test/CodeGen/AArch64/ptrauth-indirectbr.ll
@@ -1,38 +1,89 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+
+; RUN: llc -mtriple arm64e-apple-darwin \
+; RUN:   -aarch64-enable-collect-loh=false \
+; RUN:   -o - %s | FileCheck %s --check-prefix=MACHO
+
 ; RUN: llc -mtriple arm64e-apple-darwin \
-; RUN:   -asm-verbose=false -aarch64-enable-collect-loh=false \
-; RUN:   -o - %s | FileCheck %s
+; RUN:   -fast-isel \
+; RUN:   -aarch64-enable-collect-loh=false \
+; RUN:   -o - %s | FileCheck %s --check-prefix=MACHO
 
 ; RUN: llc -mtriple arm64e-apple-darwin \
 ; RUN:   -global-isel -global-isel-abort=1 -verify-machineinstrs \
-; RUN:   -asm-verbose=false -aarch64-enable-collect-loh=false \
-; RUN:   -o - %s | FileCheck %s
-
-; The discriminator is the same for all blockaddresses in the function.
-; ptrauth_string_discriminator("test_blockaddress blockaddress") == 52152
-
-; CHECK-LABEL: _test_blockaddress:
-; CHECK:         adrp x16, [[F1BB1ADDR:Ltmp[0-9]+]]@PAGE
-; CHECK-NEXT:    add x16, x16, [[F1BB1ADDR]]@PAGEOFF
-; CHECK-NEXT:    mov x17, #[[F1DISCVAL:52152]]
-; CHECK-NEXT:    pacia x16, x17
-; CHECK-NEXT:    mov x0, x16
-; CHECK-NEXT:    adrp x16, [[F1BB2ADDR:Ltmp[0-9]+]]@PAGE
-; CHECK-NEXT:    add x16, x16, [[F1BB2ADDR]]@PAGEOFF
-; CHECK-NEXT:    mov x17, #[[F1DISCVAL]]
-; CHECK-NEXT:    pacia x16, x17
-; CHECK-NEXT:    mov x1, x16
-; CHECK-NEXT:    bl _dummy_choose
-; CHECK-NEXT:    mov x17, #[[F1DISCVAL]]
-; CHECK-NEXT:    braa x0, x17
-; CHECK:        [[F1BB1ADDR]]:
-; CHECK-NEXT:   [[F1BB1:LBB[0-9_]+]]:
-; CHECK-NEXT:    mov w0, #1
-; CHECK:        [[F1BB2ADDR]]:
-; CHECK-NEXT:   [[F1BB2:LBB[0-9_]+]]:
-; CHECK-NEXT:    mov w0, #2
-define i32 @test_blockaddress() #0 {
+; RUN:   -aarch64-enable-collect-loh=false \
+; RUN:   -o - %s | FileCheck %s --check-prefix=MACHO
+
+; RUN: llc -mtriple aarch64-elf -mattr=+pauth \
+; RUN:   -o - %s | FileCheck %s --check-prefix=ELF
+
+; RUN: llc -mtriple aarch64-elf -mattr=+pauth \
+; RUN:   -fast-isel \
+; RUN:   -o - %s | FileCheck %s --check-prefix=ELF
+
+; RUN: llc -mtriple aarch64-elf -mattr=+pauth \
+; RUN:   -global-isel -global-isel-abort=1 -verify-machineinstrs \
+; RUN:   -o - %s | FileCheck %s --check-prefix=ELF
+
+;; The discriminator is the same for all blockaddresses in the function.
+;; ptrauth_string_discriminator("test_indirectbr blockaddress") == 34947
+
+define i32 @test_indirectbr() #0 {
+; MACHO-LABEL: test_indirectbr:
+; MACHO:       ; %bb.0: ; %entry
+; MACHO-NEXT:    stp x29, x30, [sp, #-16]! ; 16-byte Folded Spill
+; MACHO-NEXT:    adrp x16, Ltmp0 at PAGE
+; MACHO-NEXT:    add x16, x16, Ltmp0 at PAGEOFF
+; MACHO-NEXT:    mov x17, #34947 ; =0x8883
+; MACHO-NEXT:    pacia x16, x17
+; MACHO-NEXT:    mov x0, x16
+; MACHO-NEXT:    adrp x16, Ltmp1 at PAGE
+; MACHO-NEXT:    add x16, x16, Ltmp1 at PAGEOFF
+; MACHO-NEXT:    mov x17, #34947 ; =0x8883
+; MACHO-NEXT:    pacia x16, x17
+; MACHO-NEXT:    mov x1, x16
+; MACHO-NEXT:    bl _dummy_choose
+; MACHO-NEXT:    mov x17, #34947 ; =0x8883
+; MACHO-NEXT:    braa x0, x17
+; MACHO-NEXT:  Ltmp0: ; Block address taken
+; MACHO-NEXT:  LBB0_1: ; %common.ret
+; MACHO-NEXT:    mov w0, #1 ; =0x1
+; MACHO-NEXT:    ldp x29, x30, [sp], #16 ; 16-byte Folded Reload
+; MACHO-NEXT:    ret
+; MACHO-NEXT:  Ltmp1: ; Block address taken
+; MACHO-NEXT:  LBB0_2: ; %bb2
+; MACHO-NEXT:    mov w0, #2 ; =0x2
+; MACHO-NEXT:    ldp x29, x30, [sp], #16 ; 16-byte Folded Reload
+; MACHO-NEXT:    ret
+;
+; ELF-LABEL: test_indirectbr:
+; ELF:       // %bb.0: // %entry
+; ELF-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
+; ELF-NEXT:    adrp x16, .Ltmp0
+; ELF-NEXT:    add x16, x16, :lo12:.Ltmp0
+; ELF-NEXT:    mov x17, #34947 // =0x8883
+; ELF-NEXT:    pacia x16, x17
+; ELF-NEXT:    mov x0, x16
+; ELF-NEXT:    adrp x16, .Ltmp1
+; ELF-NEXT:    add x16, x16, :lo12:.Ltmp1
+; ELF-NEXT:    mov x17, #34947 // =0x8883
+; ELF-NEXT:    pacia x16, x17
+; ELF-NEXT:    mov x1, x16
+; ELF-NEXT:    bl dummy_choose
+; ELF-NEXT:    mov x17, #34947 // =0x8883
+; ELF-NEXT:    braa x0, x17
+; ELF-NEXT:  .Ltmp0: // Block address taken
+; ELF-NEXT:  .LBB0_1: // %common.ret
+; ELF-NEXT:    mov w0, #1 // =0x1
+; ELF-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
+; ELF-NEXT:    ret
+; ELF-NEXT:  .Ltmp1: // Block address taken
+; ELF-NEXT:  .LBB0_2: // %bb2
+; ELF-NEXT:    mov w0, #2 // =0x2
+; ELF-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
+; ELF-NEXT:    ret
 entry:
-  %tmp0 = call ptr @dummy_choose(ptr blockaddress(@test_blockaddress, %bb1), ptr blockaddress(@test_blockaddress, %bb2))
+  %tmp0 = call ptr @dummy_choose(ptr blockaddress(@test_indirectbr, %bb1), ptr blockaddress(@test_indirectbr, %bb2))
   indirectbr ptr %tmp0, [label %bb1, label %bb2]
 
 bb1:
@@ -42,32 +93,86 @@ bb2:
   ret i32 2
 }
 
-; Test another function to compare the discriminator.
-; ptrauth_string_discriminator("test_blockaddress_2 blockaddress") == 22012
-
-; CHECK-LABEL: _test_blockaddress_2:
-; CHECK:         adrp x16, [[F2BB1ADDR:Ltmp[0-9]+]]@PAGE
-; CHECK-NEXT:    add x16, x16, [[F2BB1ADDR]]@PAGEOFF
-; CHECK-NEXT:    mov x17, #[[F2DISCVAL:22012]]
-; CHECK-NEXT:    pacia x16, x17
-; CHECK-NEXT:    mov x0, x16
-; CHECK-NEXT:    adrp x16, [[F2BB2ADDR:Ltmp[0-9]+]]@PAGE
-; CHECK-NEXT:    add x16, x16, [[F2BB2ADDR]]@PAGEOFF
-; CHECK-NEXT:    mov x17, #[[F2DISCVAL]]
-; CHECK-NEXT:    pacia x16, x17
-; CHECK-NEXT:    mov x1, x16
-; CHECK-NEXT:    bl _dummy_choose
-; CHECK-NEXT:    mov x17, #[[F2DISCVAL]]
-; CHECK-NEXT:    braa x0, x17
-; CHECK:        [[F2BB1ADDR]]:
-; CHECK-NEXT:   [[F2BB1:LBB[0-9_]+]]:
-; CHECK-NEXT:    mov w0, #1
-; CHECK:        [[F2BB2ADDR]]:
-; CHECK-NEXT:   [[F2BB2:LBB[0-9_]+]]:
-; CHECK-NEXT:    mov w0, #2
-define i32 @test_blockaddress_2() #0 {
+define ptr @test_indirectbr_other_function() #0 {
+; MACHO-LABEL: test_indirectbr_other_function:
+; MACHO:       ; %bb.0:
+; MACHO-NEXT:    adrp x16, Ltmp0 at PAGE
+; MACHO-NEXT:    add x16, x16, Ltmp0 at PAGEOFF
+; MACHO-NEXT:    mov x17, #34947 ; =0x8883
+; MACHO-NEXT:    pacia x16, x17
+; MACHO-NEXT:    mov x0, x16
+; MACHO-NEXT:    ret
+;
+; ELF-LABEL: test_indirectbr_other_function:
+; ELF:       // %bb.0:
+; ELF-NEXT:    adrp x16, .Ltmp0
+; ELF-NEXT:    add x16, x16, :lo12:.Ltmp0
+; ELF-NEXT:    mov x17, #34947 // =0x8883
+; ELF-NEXT:    pacia x16, x17
+; ELF-NEXT:    mov x0, x16
+; ELF-NEXT:    ret
+  ret ptr blockaddress(@test_indirectbr, %bb1)
+}
+
+;; Test another function to compare the discriminator.
+;; ptrauth_string_discriminator("test_indirectbr_2 blockaddress") == 40224
+
+define i32 @test_indirectbr_2() #0 {
+; MACHO-LABEL: test_indirectbr_2:
+; MACHO:       ; %bb.0: ; %entry
+; MACHO-NEXT:    stp x29, x30, [sp, #-16]! ; 16-byte Folded Spill
+; MACHO-NEXT:    adrp x16, Ltmp2 at PAGE
+; MACHO-NEXT:    add x16, x16, Ltmp2 at PAGEOFF
+; MACHO-NEXT:    mov x17, #40224 ; =0x9d20
+; MACHO-NEXT:    pacia x16, x17
+; MACHO-NEXT:    mov x0, x16
+; MACHO-NEXT:    adrp x16, Ltmp3 at PAGE
+; MACHO-NEXT:    add x16, x16, Ltmp3 at PAGEOFF
+; MACHO-NEXT:    mov x17, #40224 ; =0x9d20
+; MACHO-NEXT:    pacia x16, x17
+; MACHO-NEXT:    mov x1, x16
+; MACHO-NEXT:    bl _dummy_choose
+; MACHO-NEXT:    mov x17, #40224 ; =0x9d20
+; MACHO-NEXT:    braa x0, x17
+; MACHO-NEXT:  Ltmp2: ; Block address taken
+; MACHO-NEXT:  LBB2_1: ; %common.ret
+; MACHO-NEXT:    mov w0, #1 ; =0x1
+; MACHO-NEXT:    ldp x29, x30, [sp], #16 ; 16-byte Folded Reload
+; MACHO-NEXT:    ret
+; MACHO-NEXT:  Ltmp3: ; Block address taken
+; MACHO-NEXT:  LBB2_2: ; %bb2
+; MACHO-NEXT:    mov w0, #2 ; =0x2
+; MACHO-NEXT:    ldp x29, x30, [sp], #16 ; 16-byte Folded Reload
+; MACHO-NEXT:    ret
+;
+; ELF-LABEL: test_indirectbr_2:
+; ELF:       // %bb.0: // %entry
+; ELF-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
+; ELF-NEXT:    adrp x16, .Ltmp2
+; ELF-NEXT:    add x16, x16, :lo12:.Ltmp2
+; ELF-NEXT:    mov x17, #40224 // =0x9d20
+; ELF-NEXT:    pacia x16, x17
+; ELF-NEXT:    mov x0, x16
+; ELF-NEXT:    adrp x16, .Ltmp3
+; ELF-NEXT:    add x16, x16, :lo12:.Ltmp3
+; ELF-NEXT:    mov x17, #40224 // =0x9d20
+; ELF-NEXT:    pacia x16, x17
+; ELF-NEXT:    mov x1, x16
+; ELF-NEXT:    bl dummy_choose
+; ELF-NEXT:    mov x17, #40224 // =0x9d20
+; ELF-NEXT:    braa x0, x17
+; ELF-NEXT:  .Ltmp2: // Block address taken
+; ELF-NEXT:  .LBB2_1: // %common.ret
+; ELF-NEXT:    mov w0, #1 // =0x1
+; ELF-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
+; ELF-NEXT:    ret
+; ELF-NEXT:  .Ltmp3: // Block address taken
+; ELF-NEXT:  .LBB2_2: // %bb2
+; ELF-NEXT:    mov w0, #2 // =0x2
+; ELF-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
+; ELF-NEXT:    ret
 entry:
-  %tmp0 = call ptr @dummy_choose(ptr blockaddress(@test_blockaddress_2, %bb1), ptr blockaddress(@test_blockaddress_2, %bb2))
+  %tmp0 = call ptr @dummy_choose(ptr blockaddress(@test_indirectbr_2, %bb1), ptr blockaddress(@test_indirectbr_2, %bb2))
   indirectbr ptr %tmp0, [label %bb1, label %bb2]
 
 bb1:
@@ -77,28 +182,28 @@ bb2:
   ret i32 2
 }
 
-; CHECK-LABEL: _test_blockaddress_other_function:
-; CHECK:         adrp x16, [[F1BB1ADDR]]@PAGE
-; CHECK-NEXT:    add x16, x16, [[F1BB1ADDR]]@PAGEOFF
-; CHECK-NEXT:    mov x17, #[[F1DISCVAL]]
-; CHECK-NEXT:    pacia x16, x17
-; CHECK-NEXT:    mov x0, x16
-; CHECK-NEXT:    ret
-define ptr @test_blockaddress_other_function() #0 {
-  ret ptr blockaddress(@test_blockaddress, %bb1)
-}
+; MACHO-LABEL: .section __DATA,__const
+; MACHO-NEXT:  .globl _test_indirectbr_array
+; MACHO-NEXT:  .p2align 4
+; MACHO-NEXT:  _test_indirectbr_array:
+; MACHO-NEXT:   .quad Ltmp0 at AUTH(ia,34947)
+; MACHO-NEXT:   .quad Ltmp1 at AUTH(ia,34947)
+; MACHO-NEXT:   .quad Ltmp2 at AUTH(ia,40224)
+; MACHO-NEXT:   .quad Ltmp3 at AUTH(ia,40224)
+
+; ELF-LABEL: .section .rodata,"a", at progbits
+; ELF-NEXT:  .globl test_indirectbr_array
+; ELF-NEXT:  .p2align 4, 0x0
+; ELF-NEXT:  test_indirectbr_array:
+; ELF-NEXT:   .xword .Ltmp0 at AUTH(ia,34947)
+; ELF-NEXT:   .xword .Ltmp1 at AUTH(ia,34947)
+; ELF-NEXT:   .xword .Ltmp2 at AUTH(ia,40224)
+; ELF-NEXT:   .xword .Ltmp3 at AUTH(ia,40224)
+; ELF-NEXT:   .size test_indirectbr_array, 32
 
-; CHECK-LABEL: .section __DATA,__const
-; CHECK-NEXT:  .globl _test_blockaddress_array
-; CHECK-NEXT:  .p2align 4
-; CHECK-NEXT:  _test_blockaddress_array:
-; CHECK-NEXT:   .quad [[F1BB1ADDR]]@AUTH(ia,[[F1DISCVAL]]
-; CHECK-NEXT:   .quad [[F1BB2ADDR]]@AUTH(ia,[[F1DISCVAL]]
-; CHECK-NEXT:   .quad [[F2BB1ADDR]]@AUTH(ia,[[F2DISCVAL]]
-; CHECK-NEXT:   .quad [[F2BB2ADDR]]@AUTH(ia,[[F2DISCVAL]]
- at test_blockaddress_array = constant [4 x ptr] [
-  ptr blockaddress(@test_blockaddress, %bb1), ptr blockaddress(@test_blockaddress, %bb2),
-  ptr blockaddress(@test_blockaddress_2, %bb1), ptr blockaddress(@test_blockaddress_2, %bb2)
+ at test_indirectbr_array = constant [4 x ptr] [
+  ptr blockaddress(@test_indirectbr, %bb1), ptr blockaddress(@test_indirectbr, %bb2),
+  ptr blockaddress(@test_indirectbr_2, %bb1), ptr blockaddress(@test_indirectbr_2, %bb2)
 ]
 
 declare ptr @dummy_choose(ptr, ptr)

>From 6cf9845e7e66f2afd9b465f821fb583f4ac8894b Mon Sep 17 00:00:00 2001
From: Ahmed Bougacha <ahmed at bougacha.org>
Date: Mon, 22 Jul 2024 16:10:29 -0700
Subject: [PATCH 4/5] Add (currently failing) test for BR_JT -> BRIND lowering.

---
 .../CodeGen/AArch64/ptrauth-indirectbr.ll     | 53 ++++++++++++++++---
 1 file changed, 45 insertions(+), 8 deletions(-)

diff --git a/llvm/test/CodeGen/AArch64/ptrauth-indirectbr.ll b/llvm/test/CodeGen/AArch64/ptrauth-indirectbr.ll
index c96d9dce6f928..94de1b4f949e4 100644
--- a/llvm/test/CodeGen/AArch64/ptrauth-indirectbr.ll
+++ b/llvm/test/CodeGen/AArch64/ptrauth-indirectbr.ll
@@ -2,27 +2,33 @@
 
 ; RUN: llc -mtriple arm64e-apple-darwin \
 ; RUN:   -aarch64-enable-collect-loh=false \
+; RUN:   -aarch64-min-jump-table-entries=1 -aarch64-enable-atomic-cfg-tidy=0 \
 ; RUN:   -o - %s | FileCheck %s --check-prefix=MACHO
 
 ; RUN: llc -mtriple arm64e-apple-darwin \
 ; RUN:   -fast-isel \
 ; RUN:   -aarch64-enable-collect-loh=false \
+; RUN:   -aarch64-min-jump-table-entries=1 -aarch64-enable-atomic-cfg-tidy=0 \
 ; RUN:   -o - %s | FileCheck %s --check-prefix=MACHO
 
 ; RUN: llc -mtriple arm64e-apple-darwin \
 ; RUN:   -global-isel -global-isel-abort=1 -verify-machineinstrs \
 ; RUN:   -aarch64-enable-collect-loh=false \
+; RUN:   -aarch64-min-jump-table-entries=1 -aarch64-enable-atomic-cfg-tidy=0 \
 ; RUN:   -o - %s | FileCheck %s --check-prefix=MACHO
 
 ; RUN: llc -mtriple aarch64-elf -mattr=+pauth \
+; RUN:   -aarch64-min-jump-table-entries=1 -aarch64-enable-atomic-cfg-tidy=0 \
 ; RUN:   -o - %s | FileCheck %s --check-prefix=ELF
 
 ; RUN: llc -mtriple aarch64-elf -mattr=+pauth \
 ; RUN:   -fast-isel \
+; RUN:   -aarch64-min-jump-table-entries=1 -aarch64-enable-atomic-cfg-tidy=0 \
 ; RUN:   -o - %s | FileCheck %s --check-prefix=ELF
 
 ; RUN: llc -mtriple aarch64-elf -mattr=+pauth \
 ; RUN:   -global-isel -global-isel-abort=1 -verify-machineinstrs \
+; RUN:   -aarch64-min-jump-table-entries=1 -aarch64-enable-atomic-cfg-tidy=0 \
 ; RUN:   -o - %s | FileCheck %s --check-prefix=ELF
 
 ;; The discriminator is the same for all blockaddresses in the function.
@@ -46,7 +52,7 @@ define i32 @test_indirectbr() #0 {
 ; MACHO-NEXT:    mov x17, #34947 ; =0x8883
 ; MACHO-NEXT:    braa x0, x17
 ; MACHO-NEXT:  Ltmp0: ; Block address taken
-; MACHO-NEXT:  LBB0_1: ; %common.ret
+; MACHO-NEXT:  LBB0_1: ; %bb1
 ; MACHO-NEXT:    mov w0, #1 ; =0x1
 ; MACHO-NEXT:    ldp x29, x30, [sp], #16 ; 16-byte Folded Reload
 ; MACHO-NEXT:    ret
@@ -73,7 +79,7 @@ define i32 @test_indirectbr() #0 {
 ; ELF-NEXT:    mov x17, #34947 // =0x8883
 ; ELF-NEXT:    braa x0, x17
 ; ELF-NEXT:  .Ltmp0: // Block address taken
-; ELF-NEXT:  .LBB0_1: // %common.ret
+; ELF-NEXT:  .LBB0_1: // %bb1
 ; ELF-NEXT:    mov w0, #1 // =0x1
 ; ELF-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
 ; ELF-NEXT:    ret
@@ -135,7 +141,7 @@ define i32 @test_indirectbr_2() #0 {
 ; MACHO-NEXT:    mov x17, #40224 ; =0x9d20
 ; MACHO-NEXT:    braa x0, x17
 ; MACHO-NEXT:  Ltmp2: ; Block address taken
-; MACHO-NEXT:  LBB2_1: ; %common.ret
+; MACHO-NEXT:  LBB2_1: ; %bb1
 ; MACHO-NEXT:    mov w0, #1 ; =0x1
 ; MACHO-NEXT:    ldp x29, x30, [sp], #16 ; 16-byte Folded Reload
 ; MACHO-NEXT:    ret
@@ -162,7 +168,7 @@ define i32 @test_indirectbr_2() #0 {
 ; ELF-NEXT:    mov x17, #40224 // =0x9d20
 ; ELF-NEXT:    braa x0, x17
 ; ELF-NEXT:  .Ltmp2: // Block address taken
-; ELF-NEXT:  .LBB2_1: // %common.ret
+; ELF-NEXT:  .LBB2_1: // %bb1
 ; ELF-NEXT:    mov w0, #1 // =0x1
 ; ELF-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
 ; ELF-NEXT:    ret
@@ -182,8 +188,40 @@ bb2:
   ret i32 2
 }
 
-; MACHO-LABEL: .section __DATA,__const
-; MACHO-NEXT:  .globl _test_indirectbr_array
+;; Check we don't interfere with jump-table BRIND lowering.
+
+; MACHO-LABEL: test_jumptable:
+; MACHO:        adrp x9, LJTI3_0 at PAGE
+; MACHO-NEXT:   add x9, x9, LJTI3_0 at PAGEOFF
+; MACHO-NEXT:   adr x10, LBB3_2
+; MACHO-NEXT:   ldrb w11, [x9, x8]
+; MACHO-NEXT:   add x10, x10, x11, lsl #2
+; MACHO-NEXT:   br x10
+
+; ELF-LABEL: test_jumptable:
+; ELF:        adrp x9, .LJTI3_0
+; ELF-NEXT:   add x9, x9, :lo12:.LJTI3_0
+; ELF-NEXT:   adr x10, .LBB3_2
+; ELF-NEXT:   ldrb w11, [x9, x8]
+; ELF-NEXT:   add x10, x10, x11, lsl #2
+; ELF-NEXT:   br x10
+define i32 @test_jumptable(i32 %in) #0 {
+  switch i32 %in, label %def [
+    i32 0, label %lbl1
+    i32 1, label %lbl2
+  ]
+
+def:
+  ret i32 0
+
+lbl1:
+  ret i32 1
+
+lbl2:
+  ret i32 2
+}
+
+; MACHO-LABEL: .globl _test_indirectbr_array
 ; MACHO-NEXT:  .p2align 4
 ; MACHO-NEXT:  _test_indirectbr_array:
 ; MACHO-NEXT:   .quad Ltmp0 at AUTH(ia,34947)
@@ -191,8 +229,7 @@ bb2:
 ; MACHO-NEXT:   .quad Ltmp2 at AUTH(ia,40224)
 ; MACHO-NEXT:   .quad Ltmp3 at AUTH(ia,40224)
 
-; ELF-LABEL: .section .rodata,"a", at progbits
-; ELF-NEXT:  .globl test_indirectbr_array
+; ELF-LABEL: .globl test_indirectbr_array
 ; ELF-NEXT:  .p2align 4, 0x0
 ; ELF-NEXT:  test_indirectbr_array:
 ; ELF-NEXT:   .xword .Ltmp0 at AUTH(ia,34947)

>From ecac22c1bd8075ae1a4628f1bbf2b614bea67a4e Mon Sep 17 00:00:00 2001
From: Ahmed Bougacha <ahmed at bougacha.org>
Date: Mon, 22 Jul 2024 16:12:07 -0700
Subject: [PATCH 5/5] Don't re-lower BR_JT lowering result BRIND to ptrauth
 BRA.

---
 llvm/lib/Target/AArch64/AArch64ISelLowering.cpp | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 923df31421e90..d31fa56f93f83 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -10754,6 +10754,15 @@ SDValue AArch64TargetLowering::LowerBR_JT(SDValue Op,
 }
 
 SDValue AArch64TargetLowering::LowerBRIND(SDValue Op, SelectionDAG &DAG) const {
+  SDValue Chain = Op.getOperand(0);
+  SDValue Dest = Op.getOperand(1);
+
+  // BR_JT is lowered to BRIND, but the later lowering is specific to indirectbr
+  // Skip over the jump-table BRINDs, where the destination is JumpTableDest32.
+  if (Dest->isMachineOpcode() &&
+      Dest->getMachineOpcode() == AArch64::JumpTableDest32)
+    return SDValue();
+
   const MachineFunction &MF = DAG.getMachineFunction();
   std::optional<uint16_t> BADisc =
       Subtarget->getPtrAuthBlockAddressDiscriminatorIfEnabled(MF.getFunction());
@@ -10761,8 +10770,6 @@ SDValue AArch64TargetLowering::LowerBRIND(SDValue Op, SelectionDAG &DAG) const {
     return SDValue();
 
   SDLoc DL(Op);
-  SDValue Chain = Op.getOperand(0);
-  SDValue Dest = Op.getOperand(1);
 
   SDValue Disc = DAG.getTargetConstant(*BADisc, DL, MVT::i64);
   SDValue Key = DAG.getTargetConstant(AArch64PACKey::IA, DL, MVT::i32);



More information about the cfe-commits mailing list