[llvm] Teach FastISel for AArch64 about tagged globals

This addresses llvm#57750. For
some globals, the tag wasn't propagated correctly because the necessary
movk wasn't emitted sometimes.

Differential Revision: https://reviews.llvm.org/D138615

Author: PiJoules

llvm/lib/Target/AArch64/AArch64FastISel.cpp

@@ -496,6 +496,30 @@ unsigned AArch64FastISel::materializeGV(const GlobalValue *GV) {
             ADRPReg)
         .addGlobalAddress(GV, 0, AArch64II::MO_PAGE | OpFlags);
 
+    if (OpFlags & 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.
+      //
+      // TODO: There is duplicate logic in AArch64ExpandPseudoInsts.cpp that
+      // also uses BuildMI for making an ADRP (+ MOVK) + ADD, but the operands
+      // are not exactly 1:1 with FastISel so we cannot easily abstract this
+      // out. At some point, it would be nice to find a way to not have this
+      // duplciate code.
+      unsigned DstReg = createResultReg(&AArch64::GPR64spRegClass);
+      BuildMI(*FuncInfo.MBB, FuncInfo.InsertPt, MIMD, TII.get(AArch64::MOVKXi),
+              DstReg)
+          .addReg(ADRPReg)
+          .addGlobalAddress(GV, /*Offset=*/0x100000000,
+                            AArch64II::MO_PREL | AArch64II::MO_G3)
+          .addImm(48);
+      ADRPReg = DstReg;
+    }
+
     ResultReg = createResultReg(&AArch64::GPR64spRegClass);
     BuildMI(*FuncInfo.MBB, FuncInfo.InsertPt, MIMD, TII.get(AArch64::ADDXri),
             ResultReg)

llvm/test/CodeGen/AArch64/arm64-fast-isel-tag.ll

@@ -0,0 +1,23 @@
+; RUN: llc < %s -fast-isel -relocation-model=pic | 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-gnu"
+
+@glob.hwasan = private constant i64 0
+
+;; The constant here is 0x2F << 56. This effectively makes the alias a tagged version of the original global.
+@glob = private alias i64, inttoptr (i64 add (i64 ptrtoint (ptr @glob.hwasan to i64), i64 3386706919782612992) to ptr)
+
+; CHECK-LABEL: func
+define void @func() #0 {
+entry:
+  ; CHECK:      adrp    [[REG:x[0-9]+]], :pg_hi21_nc:.Lglob
+  ; CHECK-NEXT: movk    [[REG]], #:prel_g3:.Lglob+4294967296
+  ; CHECK-NEXT: add     x0, [[REG]], :lo12:.Lglob
+  call void @extern_func(ptr @glob)
+  ret void
+}
+
+declare void @extern_func(ptr)
+
+attributes #0 = { "target-features"="+tagged-globals" }