[AArch64] Fix CollectLOH creating an AdrpAdd LOH when there's a live used reg

between the two instructions.

If there's a pattern like:
$xA = ADRP foo @PAGE
[some killing use of reg Xb]
$Xb = ADDXri $Xa, 0, @PAGEOFF

CollectLOH would create an AdrpAdd LOH that resulted in the linker optimizing
this sequence into:
$xB = ADR foo
[some killing use of reg $Xb]
... and therefore clobbers the live $Xb register that was used by the
instruction in between.

This was discovered by a GlobalISel patch D78465 which broke up global variable
accesses into two pseudos, which in some cases could be moved apart.

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

lib/Target/AArch64/AArch64CollectLOH.cpp

@@ -382,7 +382,7 @@ static bool handleMiddleInst(const MachineInstr &MI, LOHInfo &DefInfo,
 
 /// Update state when seeing and ADRP instruction.
 static void handleADRP(const MachineInstr &MI, AArch64FunctionInfo &AFI,
-                       LOHInfo &Info) {
+                       LOHInfo &Info, LOHInfo *LOHInfos) {
   if (Info.LastADRP != nullptr) {
     LLVM_DEBUG(dbgs() << "Adding MCLOH_AdrpAdrp:\n"
                       << '\t' << MI << '\t' << *Info.LastADRP);
@@ -393,12 +393,24 @@ static void handleADRP(const MachineInstr &MI, AArch64FunctionInfo &AFI,
   // Produce LOH directive if possible.
   if (Info.IsCandidate) {
     switch (Info.Type) {
-    case MCLOH_AdrpAdd:
+    case MCLOH_AdrpAdd: {
+      // ADRPs and ADDs for this candidate may be split apart if using
+      // GlobalISel instead of pseudo-expanded. If that happens, the
+      // def register of the ADD may have a use in between. Adding an LOH in
+      // this case can cause the linker to rewrite the ADRP to write to that
+      // register, clobbering the use.
+      const MachineInstr *AddMI = Info.MI0;
+      int DefIdx = mapRegToGPRIndex(MI.getOperand(0).getReg());
+      int OpIdx = mapRegToGPRIndex(AddMI->getOperand(0).getReg());
+      LOHInfo DefInfo = LOHInfos[OpIdx];
+      if (DefIdx != OpIdx && (DefInfo.OneUser || DefInfo.MultiUsers))
+        break;
       LLVM_DEBUG(dbgs() << "Adding MCLOH_AdrpAdd:\n"
                         << '\t' << MI << '\t' << *Info.MI0);
       AFI.addLOHDirective(MCLOH_AdrpAdd, {&MI, Info.MI0});
       ++NumADRSimpleCandidate;
       break;
+    }
     case MCLOH_AdrpLdr:
       if (supportLoadFromLiteral(*Info.MI0)) {
         LLVM_DEBUG(dbgs() << "Adding MCLOH_AdrpLdr:\n"
@@ -545,7 +557,7 @@ bool AArch64CollectLOH::runOnMachineFunction(MachineFunction &MF) {
         const MachineOperand &Op0 = MI.getOperand(0);
         int Idx = mapRegToGPRIndex(Op0.getReg());
         if (Idx >= 0) {
-          handleADRP(MI, AFI, LOHInfos[Idx]);
+          handleADRP(MI, AFI, LOHInfos[Idx], LOHInfos);
           continue;
         }
         break;

test/CodeGen/AArch64/loh-use-between-adrp-add.mir

@@ -0,0 +1,56 @@
+# RUN: llc -o - %s -mtriple=aarch64-apple-ios -run-pass=aarch64-collect-loh -debug-only=aarch64-collect-loh 2>&1 | FileCheck %s
+# REQUIRES: asserts
+--- |
+  @rrdpb = local_unnamed_addr global i32 zeroinitializer, align 8
+
+  define internal void @test_use_between() {
+    ret void
+  }
+  define internal void @test_no_use_between() {
+    ret void
+  }
+
+...
+# CHECK-LABEL: ********** AArch64 Collect LOH **********
+# CHECK-LABEL: Looking in function test_use_between
+# Check that we don't have an AdrpAdd LOH because there's a use of the ADD defreg
+# in between the two.
+# CHECK-NOT: MCLOH_AdrpAdd
+
+# CHECK-LABEL: Looking in function test_no_use_between
+# CHECK: MCLOH_AdrpAdd
+---
+name:            test_use_between
+alignment:       4
+tracksRegLiveness: true
+liveins:
+  - { reg: '$x11', virtual-reg: '' }
+  - { reg: '$x12', virtual-reg: '' }
+body:             |
+  bb.0:
+    liveins: $x11, $x12
+    renamable $x15 = ADRP target-flags(aarch64-page) @rrdpb
+    STRXui renamable $x12, killed renamable $x11, 1 :: (store 8)
+    renamable $x11 = ADDXri killed renamable $x15, target-flags(aarch64-pageoff, aarch64-nc) @rrdpb, 0
+    STRXui renamable $x11, killed renamable $x11, 0
+    RET undef $lr
+
+...
+
+---
+name:            test_no_use_between
+alignment:       4
+tracksRegLiveness: true
+liveins:
+  - { reg: '$x11', virtual-reg: '' }
+  - { reg: '$x12', virtual-reg: '' }
+body:             |
+  bb.0:
+    liveins: $x11, $x12
+    STRXui renamable $x12, killed renamable $x11, 1 :: (store 8)
+    renamable $x15 = ADRP target-flags(aarch64-page) @rrdpb
+    renamable $x11 = ADDXri killed renamable $x15, target-flags(aarch64-pageoff, aarch64-nc) @rrdpb, 0
+    STRXui renamable $x11, killed renamable $x11, 0
+    RET undef $lr
+
+...