[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.
+      BuildMI(*FuncInfo.MBB, FuncInfo.InsertPt, MIMD, TII.get(AArch64::MOVKXi), ADRPReg)
+          .addReg(ADRPReg)
+          .addGlobalAddress(GV, /*Offset=*/0x100000000, AArch64II::MO_PREL | AArch64II::MO_G3)
+          .addImm(48);
+    }
+
     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);
   }