[llvm] Teach FastISel for AArch64 about tagged globals

[leonardchan]

Summary

This addresses https://github.com/llvm/llvm-project/issues/57750. For some globals, the tag wasn't propagated correctly because the necessary movk wasn't emitted sometimes.

[pcc]

Test case?

Can we extract a function out of the code in AArch64ExpandPseudoInsts.cpp?

[leonardchan]

In D138615#3948014, @pcc wrote:

Test case?

Can we extract a function out of the code in AArch64ExpandPseudoInsts.cpp?

Done. Wasn't sure if any existing headers made sense to add a function, so I made a new one.

[pcc]

LGTM

llvm/lib/Target/AArch64/AArch64FastISel.cpp

BuildMI(*FuncInfo.MBB, FuncInfo.InsertPt, MIMD,
         .addImm(AArch64::sub_32);
     return Result64;
   } else {
     // ADRP + ADDX
     BuildMI(*FuncInfo.MBB, FuncInfo.InsertPt, MIMD, TII.get(AArch64::ADRP),
             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)
         .addReg(ADRPReg)
         .addGlobalAddress(GV, 0,
                           AArch64II::MO_PAGEOFF | AArch64II::MO_NC | OpFlags)
         .addImm(0);
   }

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

+; 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"
+
[e-mail] = 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" }