AArch64: Add a tagged-globals backend feature.

[pcc]

Summary

This feature instructs the backend to allow locally defined global variable
addresses to contain a pointer tag in bits 56-63 that will be ignored by
the hardware (i.e. TBI), but may be used by an instrumentation pass such
as HWASAN. It works by adding a MOVK instruction to the regular ADRP/ADD
sequence that sets bits 48-63 to the corresponding bits of the global, with
the linker bounds check disabled on the ADRP instruction to prevent the tag
from causing a link failure.

This implementation of the feature omits the MOVK when loading from or storing
to a global, which is sufficient for TBI. If the same approach is extended
to MTE, assuming that 0 is not configured as a catch-all tag, we will most
likely also need the MOVK in this case in order to avoid a tag mismatch.

[ostannard]

llvm/lib/Target/AArch64/AArch64Subtarget.cpp

231

I don't really like the idea of re-using MO_NC to mean something more complex/specialised than the description in AArch64BaseInfo.h, maybe it would be better to add a new value (MO_TAGGED) to represent this?

[pcc]

llvm/lib/Target/AArch64/AArch64Subtarget.cpp

231

That would be fine with me, but we currently only have 12 bits of storage for the target flags and I didn't want to take another one unnecessarily.

[pcc]

Address review comment

llvm/lib/Target/AArch64/AArch64Subtarget.cpp

231

Done; this required fixing a silent truncation of flags above bit 8 in D65465.

[ostannard]

LGTM

llvm/lib/Target/AArch64/AArch64.td

 
 def FeatureTRBE : SubtargetFeature<"trbe", "HasTRBE",
     "true", "Enable Trace Buffer Extension">;
 
 def FeatureETE : SubtargetFeature<"ete", "HasETE",
     "true", "Enable Embedded Trace Extension",
     [FeatureTRBE]>;
 
+def FeatureTaggedGlobals : SubtargetFeature<"tagged-globals",
+    "AllowTaggedGlobals",
+    "true", "Use an instruction sequence for taking the address of a global "
+    "that allows a memory tag in the upper address bits">;
+
 //===----------------------------------------------------------------------===//
 // Architectures.
 //
 
 def HasV8_1aOps : SubtargetFeature<"v8.1a", "HasV8_1aOps", "true",
   "Support ARM v8.1a instructions", [FeatureCRC, FeatureLSE, FeatureRDM,
   FeaturePAN, FeatureLOR, FeatureVH]>;
 

llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp

bool AArch64ExpandPseudo::expandMI(MachineBasicBlock &MBB,
   case AArch64::MOVaddrTLS:
   case AArch64::MOVaddrEXT: {
     // Expand into ADRP + ADD.
     unsigned DstReg = MI.getOperand(0).getReg();
     MachineInstrBuilder MIB1 =
         BuildMI(MBB, MBBI, MI.getDebugLoc(), TII->get(AArch64::ADRP), DstReg)
             .add(MI.getOperand(1));
 
+    if (MI.getOperand(1).getTargetFlags() & AArch64II::MO_TAGGED) {
+      // MO_TAGGED on the page indicates a tagged address. Set the tag now.
+      // We do so by creating a MOVK that sets bits 48-63 of the register to
+      // (global address + 0x100000000 - PC) >> 48. This assumes that we're in
+      // the small code model so we can assume a binary size of <= 4GB, which
+      // makes the untagged PC relative offset positive. The binary must also be
+      // loaded into address range [0, 2^48). Both of these properties need to
+      // be ensured at runtime when using tagged addresses.
+      auto Tag = MI.getOperand(1);
+      Tag.setTargetFlags(AArch64II::MO_PREL | AArch64II::MO_G3);
+      Tag.setOffset(0x100000000);
+      BuildMI(MBB, MBBI, MI.getDebugLoc(), TII->get(AArch64::MOVKXi), DstReg)
+          .addReg(DstReg)
+          .add(Tag)
+          .addImm(48);
+    }
+
     MachineInstrBuilder MIB2 =
         BuildMI(MBB, MBBI, MI.getDebugLoc(), TII->get(AArch64::ADDXri))
             .add(MI.getOperand(0))
             .addReg(DstReg)
             .add(MI.getOperand(2))
             .addImm(0);
 
     transferImpOps(MI, MIB1, MIB2);

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

 ArrayRef<std::pair<unsigned, const char *>>
 AArch64InstrInfo::getSerializableBitmaskMachineOperandTargetFlags() const {
   using namespace AArch64II;
 
   static const std::pair<unsigned, const char *> TargetFlags[] = {
       {MO_COFFSTUB, "aarch64-coffstub"},
       {MO_GOT, "aarch64-got"},   {MO_NC, "aarch64-nc"},
       {MO_S, "aarch64-s"},       {MO_TLS, "aarch64-tls"},
-      {MO_DLLIMPORT, "aarch64-dllimport"}};
+      {MO_DLLIMPORT, "aarch64-dllimport"},
+      {MO_PREL, "aarch64-prel"}};
   return makeArrayRef(TargetFlags);
 }
 
 ArrayRef<std::pair<MachineMemOperand::Flags, const char *>>
 AArch64InstrInfo::getSerializableMachineMemOperandTargetFlags() const {
   static const std::pair<MachineMemOperand::Flags, const char *> TargetFlags[] =
       {{MOSuppressPair, "aarch64-suppress-pair"},
        {MOStridedAccess, "aarch64-strided-access"}};

llvm/lib/Target/AArch64/AArch64MCInstLower.cpp

case TLSModel::LocalExec:
       break;
     case TLSModel::LocalDynamic:
       RefFlags |= AArch64MCExpr::VK_DTPREL;
       break;
     case TLSModel::GeneralDynamic:
       RefFlags |= AArch64MCExpr::VK_TLSDESC;
       break;
     }
+  } else if (MO.getTargetFlags() & AArch64II::MO_PREL) {
+    RefFlags |= AArch64MCExpr::VK_PREL;
   } else {
     // No modifier means this is a generic reference, classified as absolute for
     // the cases where it matters (:abs_g0: etc).
     RefFlags |= AArch64MCExpr::VK_ABS;
   }
 
   if ((MO.getTargetFlags() & AArch64II::MO_FRAGMENT) == AArch64II::MO_PAGE)
     RefFlags |= AArch64MCExpr::VK_PAGE;

llvm/lib/Target/AArch64/AArch64Subtarget.h

protected:
   bool HasFuseCryptoEOR = false;
   bool HasFuseLiterals = false;
   bool DisableLatencySchedHeuristic = false;
   bool UseRSqrt = false;
   bool Force32BitJumpTables = false;
   bool UseEL1ForTP = false;
   bool UseEL2ForTP = false;
   bool UseEL3ForTP = false;
+  bool AllowTaggedGlobals = false;
   uint8_t MaxInterleaveFactor = 2;
   uint8_t VectorInsertExtractBaseCost = 3;
   uint16_t CacheLineSize = 0;
   uint16_t PrefetchDistance = 0;
   uint16_t MinPrefetchStride = 1;
   unsigned MaxPrefetchIterationsAhead = UINT_MAX;
   unsigned PrefFunctionAlignment = 0;
   unsigned PrefLoopAlignment = 0;

llvm/lib/Target/AArch64/AArch64Subtarget.cpp

AArch64Subtarget::ClassifyGlobalReference(const GlobalValue *GV,
 
   // The small code model's direct accesses use ADRP, which cannot
   // necessarily produce the value 0 (if the code is above 4GB).
   // Same for the tiny code model, where we have a pc relative LDR.
   if ((useSmallAddressing() || TM.getCodeModel() == CodeModel::Tiny) &&
       GV->hasExternalWeakLinkage())
     return AArch64II::MO_GOT;
 
+  // References to tagged globals are marked with MO_NC | MO_TAGGED to indicate
+  // that their nominal addresses are tagged and outside of the code model. In
+  // AArch64ExpandPseudo::expandMI we emit an additional instruction to set the
+  // tag if necessary based on MO_TAGGED.
+  if (AllowTaggedGlobals && !isa<FunctionType>(GV->getValueType()))
+    return AArch64II::MO_NC | AArch64II::MO_TAGGED;
+
   return AArch64II::MO_NO_FLAG;
 }
 
 unsigned AArch64Subtarget::classifyGlobalFunctionReference(
     const GlobalValue *GV, const TargetMachine &TM) const {
   // MachO large model always goes via a GOT, because we don't have the
   // relocations available to do anything else..
   if (TM.getCodeModel() == CodeModel::Large && isTargetMachO() &&

llvm/lib/Target/AArch64/Utils/AArch64BaseInfo.h

enum TOF {
     /// MO_DLLIMPORT - On a symbol operand, this represents that the reference
     /// to the symbol is for an import stub.  This is used for DLL import
     /// storage class indication on Windows.
     MO_DLLIMPORT = 0x80,
 
     /// MO_S - Indicates that the bits of the symbol operand represented by
     /// MO_G0 etc are signed.
     MO_S = 0x100,
+
+    /// MO_PREL - Indicates that the bits of the symbol operand represented by
+    /// MO_G0 etc are PC relative.
+    MO_PREL = 0x200,
+
+    /// MO_TAGGED - With MO_PAGE, indicates that the page includes a memory tag
+    /// in bits 56-63.
+    MO_TAGGED = 0x400,
   };
 } // end namespace AArch64II
 
 } // end namespace llvm
 
 #endif

llvm/test/CodeGen/AArch64/tagged-globals.ll

+; RUN: llc < %s | FileCheck %s
+
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64-unknown-linux-android"
+
+@global = external hidden global i32
+declare void @func()
+
+define i32* @global_addr() #0 {
+  ; CHECK: global_addr:
+  ; CHECK: adrp x0, :pg_hi21_nc:global
+  ; CHECK: movk x0, #:prel_g3:global+4294967296
+  ; CHECK: add x0, x0, :lo12:global
+  ret i32* @global
+}
+
+define i32 @global_load() #0 {
+  ; CHECK: global_load:
+  ; CHECK: adrp x8, :pg_hi21_nc:global
+  ; CHECK: ldr w0, [x8, :lo12:global]
+  %load = load i32, i32* @global
+  ret i32 %load
+}
+
+define void ()* @func_addr() #0 {
+  ; CHECK: func_addr:
+  ; CHECK: adrp x0, func
+  ; CHECK: add x0, x0, :lo12:func
+  ret void ()* @func
+}
+
+attributes #0 = { "target-features"="+tagged-globals" }