[AArch64][GlobalISel] Fold constants into G_GLOBAL_VALUE

This is pretty much just ports `performGlobalAddressCombine` from
AArch64ISelLowering. (AArch64 doesn't use the generic DAG combine for this.)

This adds a pre-legalize combine which looks for this pattern:

```
  %g = G_GLOBAL_VALUE @x
  %ptr1 = G_PTR_ADD %g, cst1
  %ptr2 = G_PTR_ADD %g, cst2
  ...
  %ptrN = G_PTR_ADD %g, cstN
```

And then, if possible, transforms it like so:

```
  %g = G_GLOBAL_VALUE @x
  %offset_g = G_PTR_ADD %g, -min(cst)
  %ptr1 = G_PTR_ADD %offset_g, cst1
  %ptr2 = G_PTR_ADD %offset_g, cst2
  ...
  %ptrN = G_PTR_ADD %offset_g, cstN
```

Where min(cst) is the smallest out of the G_PTR_ADD constants.

This means we should save at least one G_PTR_ADD.

This also updates code in the legalizer + selector which assumes that
G_GLOBAL_VALUE will never have an offset and adds/updates relevant tests.

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

Author: Jessica Paquette

llvm/lib/Target/AArch64/AArch64Combine.td

@@ -24,10 +24,20 @@ def icmp_redundant_trunc : GICombineRule<
          [{ return matchICmpRedundantTrunc(*${root}, MRI, Helper.getKnownBits(), ${matchinfo}); }]),
   (apply [{ applyICmpRedundantTrunc(*${root}, MRI, B, Observer, ${matchinfo}); }])>;
 
+// AArch64-specific offset folding for G_GLOBAL_VALUE.
+def fold_global_offset_matchdata : GIDefMatchData<"std::pair<uint64_t, uint64_t>">;
+def fold_global_offset : GICombineRule<
+  (defs root:$root, fold_global_offset_matchdata:$matchinfo),
+  (match (wip_match_opcode G_GLOBAL_VALUE):$root,
+          [{ return matchFoldGlobalOffset(*${root}, MRI, ${matchinfo}); }]),
+  (apply [{  return applyFoldGlobalOffset(*${root}, MRI, B, Observer, ${matchinfo});}])
+>;
+
 def AArch64PreLegalizerCombinerHelper: GICombinerHelper<
   "AArch64GenPreLegalizerCombinerHelper", [all_combines,
                                            fconstant_to_constant,
-                                           icmp_redundant_trunc]> {
+                                           icmp_redundant_trunc,
+                                           fold_global_offset]> {
   let DisableRuleOption = "aarch64prelegalizercombiner-disable-rule";
   let StateClass = "AArch64PreLegalizerCombinerHelperState";
   let AdditionalArguments = [];

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

@@ -5628,8 +5628,10 @@ AArch64InstructionSelector::tryFoldAddLowIntoImm(MachineInstr &RootDef,
     return None;
 
   // TODO: add heuristics like isWorthFoldingADDlow() from SelectionDAG.
-  // TODO: Need to check GV's offset % size if doing offset folding into globals.
-  assert(Adrp.getOperand(1).getOffset() == 0 && "Unexpected offset in global");
+  auto Offset = Adrp.getOperand(1).getOffset();
+  if (Offset % Size != 0)
+    return None;
+
   auto GV = Adrp.getOperand(1).getGlobal();
   if (GV->isThreadLocal())
     return None;
@@ -5643,7 +5645,7 @@ AArch64InstructionSelector::tryFoldAddLowIntoImm(MachineInstr &RootDef,
   Register AdrpReg = Adrp.getOperand(0).getReg();
   return {{[=](MachineInstrBuilder &MIB) { MIB.addUse(AdrpReg); },
            [=](MachineInstrBuilder &MIB) {
-             MIB.addGlobalAddress(GV, /* Offset */ 0,
+             MIB.addGlobalAddress(GV, Offset,
                                   OpFlags | AArch64II::MO_PAGEOFF |
                                       AArch64II::MO_NC);
            }}};

llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp

@@ -781,7 +781,8 @@ bool AArch64LegalizerInfo::legalizeSmallCMGlobalValue(
   // G_ADD_LOW instructions.
   // By splitting this here, we can optimize accesses in the small code model by
   // folding in the G_ADD_LOW into the load/store offset.
-  auto GV = MI.getOperand(1).getGlobal();
+  auto &GlobalOp = MI.getOperand(1);
+  const auto* GV = GlobalOp.getGlobal();
   if (GV->isThreadLocal())
     return true; // Don't want to modify TLS vars.
 
@@ -791,9 +792,10 @@ bool AArch64LegalizerInfo::legalizeSmallCMGlobalValue(
   if (OpFlags & AArch64II::MO_GOT)
     return true;
 
+  auto Offset = GlobalOp.getOffset();
   Register DstReg = MI.getOperand(0).getReg();
   auto ADRP = MIRBuilder.buildInstr(AArch64::ADRP, {LLT::pointer(0, 64)}, {})
-                  .addGlobalAddress(GV, 0, OpFlags | AArch64II::MO_PAGE);
+                  .addGlobalAddress(GV, Offset, OpFlags | AArch64II::MO_PAGE);
   // Set the regclass on the dest reg too.
   MRI.setRegClass(ADRP.getReg(0), &AArch64::GPR64RegClass);
 
@@ -811,6 +813,8 @@ bool AArch64LegalizerInfo::legalizeSmallCMGlobalValue(
   // 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.
   if (OpFlags & AArch64II::MO_TAGGED) {
+    assert(!Offset &&
+           "Should not have folded in an offset for a tagged global!");
     ADRP = MIRBuilder.buildInstr(AArch64::MOVKXi, {LLT::pointer(0, 64)}, {ADRP})
                .addGlobalAddress(GV, 0x100000000,
                                  AArch64II::MO_PREL | AArch64II::MO_G3)
@@ -819,7 +823,7 @@ bool AArch64LegalizerInfo::legalizeSmallCMGlobalValue(
   }
 
   MIRBuilder.buildInstr(AArch64::G_ADD_LOW, {DstReg}, {ADRP})
-      .addGlobalAddress(GV, 0,
+      .addGlobalAddress(GV, Offset,
                         OpFlags | AArch64II::MO_PAGEOFF | AArch64II::MO_NC);
   MI.eraseFromParent();
   return true;

llvm/lib/Target/AArch64/GISel/AArch64PreLegalizerCombiner.cpp

@@ -107,6 +107,116 @@ static bool applyICmpRedundantTrunc(MachineInstr &MI, MachineRegisterInfo &MRI,
   return true;
 }
 
+/// \returns true if it is possible to fold a constant into a G_GLOBAL_VALUE.
+///
+/// e.g.
+///
+/// %g = G_GLOBAL_VALUE @x -> %g = G_GLOBAL_VALUE @x + cst
+static bool matchFoldGlobalOffset(MachineInstr &MI, MachineRegisterInfo &MRI,
+                                  std::pair<uint64_t, uint64_t> &MatchInfo) {
+  assert(MI.getOpcode() == TargetOpcode::G_GLOBAL_VALUE);
+  MachineFunction &MF = *MI.getMF();
+  auto &GlobalOp = MI.getOperand(1);
+  auto *GV = GlobalOp.getGlobal();
+
+  // Don't allow anything that could represent offsets etc.
+  if (MF.getSubtarget<AArch64Subtarget>().ClassifyGlobalReference(
+          GV, MF.getTarget()) != AArch64II::MO_NO_FLAG)
+    return false;
+
+  // Look for a G_GLOBAL_VALUE only used by G_PTR_ADDs against constants:
+  //
+  //  %g = G_GLOBAL_VALUE @x
+  //  %ptr1 = G_PTR_ADD %g, cst1
+  //  %ptr2 = G_PTR_ADD %g, cst2
+  //  ...
+  //  %ptrN = G_PTR_ADD %g, cstN
+  //
+  // Identify the *smallest* constant. We want to be able to form this:
+  //
+  //  %offset_g = G_GLOBAL_VALUE @x + min_cst
+  //  %g = G_PTR_ADD %offset_g, -min_cst
+  //  %ptr1 = G_PTR_ADD %g, cst1
+  //  ...
+  Register Dst = MI.getOperand(0).getReg();
+  uint64_t MinOffset = -1ull;
+  for (auto &UseInstr : MRI.use_nodbg_instructions(Dst)) {
+    if (UseInstr.getOpcode() != TargetOpcode::G_PTR_ADD)
+      return false;
+    auto Cst =
+        getConstantVRegValWithLookThrough(UseInstr.getOperand(2).getReg(), MRI);
+    if (!Cst)
+      return false;
+    MinOffset = std::min(MinOffset, Cst->Value.getZExtValue());
+  }
+
+  // Require that the new offset is larger than the existing one to avoid
+  // infinite loops.
+  uint64_t CurrOffset = GlobalOp.getOffset();
+  uint64_t NewOffset = MinOffset + CurrOffset;
+  if (NewOffset <= CurrOffset)
+    return false;
+
+  // Check whether folding this offset is legal. It must not go out of bounds of
+  // the referenced object to avoid violating the code model, and must be
+  // smaller than 2^21 because this is the largest offset expressible in all
+  // object formats.
+  //
+  // This check also prevents us from folding negative offsets, which will end
+  // up being treated in the same way as large positive ones. They could also
+  // cause code model violations, and aren't really common enough to matter.
+  if (NewOffset >= (1 << 21))
+    return false;
+
+  Type *T = GV->getValueType();
+  if (!T->isSized() ||
+      NewOffset > GV->getParent()->getDataLayout().getTypeAllocSize(T))
+    return false;
+  MatchInfo = std::make_pair(NewOffset, MinOffset);
+  return true;
+}
+
+static bool applyFoldGlobalOffset(MachineInstr &MI, MachineRegisterInfo &MRI,
+                                  MachineIRBuilder &B,
+                                  GISelChangeObserver &Observer,
+                                  std::pair<uint64_t, uint64_t> &MatchInfo) {
+  // Change:
+  //
+  //  %g = G_GLOBAL_VALUE @x
+  //  %ptr1 = G_PTR_ADD %g, cst1
+  //  %ptr2 = G_PTR_ADD %g, cst2
+  //  ...
+  //  %ptrN = G_PTR_ADD %g, cstN
+  //
+  // To:
+  //
+  //  %offset_g = G_GLOBAL_VALUE @x + min_cst
+  //  %g = G_PTR_ADD %offset_g, -min_cst
+  //  %ptr1 = G_PTR_ADD %g, cst1
+  //  ...
+  //  %ptrN = G_PTR_ADD %g, cstN
+  //
+  // Then, the original G_PTR_ADDs should be folded later on so that they look
+  // like this:
+  //
+  //  %ptrN = G_PTR_ADD %offset_g, cstN - min_cst
+  uint64_t Offset, MinOffset;
+  std::tie(Offset, MinOffset) = MatchInfo;
+  B.setInstrAndDebugLoc(MI);
+  Observer.changingInstr(MI);
+  auto &GlobalOp = MI.getOperand(1);
+  auto *GV = GlobalOp.getGlobal();
+  GlobalOp.ChangeToGA(GV, Offset, GlobalOp.getTargetFlags());
+  Register Dst = MI.getOperand(0).getReg();
+  Register NewGVDst = MRI.cloneVirtualRegister(Dst);
+  MI.getOperand(0).setReg(NewGVDst);
+  Observer.changedInstr(MI);
+  B.buildPtrAdd(
+      Dst, NewGVDst,
+      B.buildConstant(LLT::scalar(64), -static_cast<int64_t>(MinOffset)));
+  return true;
+}
+
 class AArch64PreLegalizerCombinerHelperState {
 protected:
   CombinerHelper &Helper;