[LOH] Don't emit AdrpAddStr when register could be clobbered (#142849)

b783aa8
added a check to ensure an `AdrpAddLdr` LOH isn't created when there is
an instruction between the `add` and `ldr`


https://github.com/llvm/llvm-project/blob/50c5704dc000cc0af41a511aa44db03233edf0af/llvm/lib/Target/AArch64/AArch64CollectLOH.cpp#L419-L431

We need a similar check for `AdrpAddStr`. Although this technically
isn't implemented in LLD, it could be in the future.


https://github.com/llvm/llvm-project/blob/50c5704dc000cc0af41a511aa44db03233edf0af/lld/MachO/Arch/ARM64.cpp#L699-L702

Author: ellishg

llvm/lib/Target/AArch64/AArch64CollectLOH.cpp

@@ -247,6 +247,17 @@ static bool supportLoadFromLiteral(const MachineInstr &MI) {
   }
 }
 
+/// Returns \p true if there are no non-debug instructions between \p First and
+/// \p Second
+static bool areInstructionsConsecutive(const MachineInstr *First,
+                                       const MachineInstr *Second) {
+  auto It = First->getIterator();
+  auto EndIt = First->getParent()->instr_end();
+  if (It == EndIt)
+    return false;
+  return next_nodbg(It, EndIt) == Second->getIterator();
+}
+
 /// Number of GPR registers tracked by mapRegToGPRIndex()
 static const unsigned N_GPR_REGS = 31;
 /// Map register number to index from 0-30.
@@ -415,7 +426,7 @@ static void handleADRP(const MachineInstr &MI, AArch64FunctionInfo &AFI,
         ++NumADRPToLDR;
       }
       break;
-    case MCLOH_AdrpAddLdr: {
+    case MCLOH_AdrpAddLdr:
       // There is a possibility that the linker may try to rewrite:
       // adrp x0, @sym@PAGE
       // add x1, x0, @sym@PAGEOFF
@@ -432,28 +443,24 @@ static void handleADRP(const MachineInstr &MI, AArch64FunctionInfo &AFI,
       // FIXME: Implement proper liveness tracking for all registers. For now,
       // don't emit the LOH if there are any instructions between the add and
       // the ldr.
-      MachineInstr *AddMI = const_cast<MachineInstr *>(Info.MI1);
-      const MachineInstr *LdrMI = Info.MI0;
-      auto AddIt = MachineBasicBlock::iterator(AddMI);
-      auto EndIt = AddMI->getParent()->end();
-      if (AddMI->getIterator() == EndIt || LdrMI != &*next_nodbg(AddIt, EndIt))
+      if (!areInstructionsConsecutive(Info.MI1, Info.MI0))
         break;
-
       LLVM_DEBUG(dbgs() << "Adding MCLOH_AdrpAddLdr:\n"
                         << '\t' << MI << '\t' << *Info.MI1 << '\t'
                         << *Info.MI0);
       AFI.addLOHDirective(MCLOH_AdrpAddLdr, {&MI, Info.MI1, Info.MI0});
       ++NumADDToLDR;
       break;
-    }
     case MCLOH_AdrpAddStr:
-      if (Info.MI1 != nullptr) {
-        LLVM_DEBUG(dbgs() << "Adding MCLOH_AdrpAddStr:\n"
-                          << '\t' << MI << '\t' << *Info.MI1 << '\t'
-                          << *Info.MI0);
-        AFI.addLOHDirective(MCLOH_AdrpAddStr, {&MI, Info.MI1, Info.MI0});
-        ++NumADDToSTR;
-      }
+      if (!Info.MI1)
+        break;
+      if (!areInstructionsConsecutive(Info.MI1, Info.MI0))
+        break;
+      LLVM_DEBUG(dbgs() << "Adding MCLOH_AdrpAddStr:\n"
+                        << '\t' << MI << '\t' << *Info.MI1 << '\t'
+                        << *Info.MI0);
+      AFI.addLOHDirective(MCLOH_AdrpAddStr, {&MI, Info.MI1, Info.MI0});
+      ++NumADDToSTR;
       break;
     case MCLOH_AdrpLdrGotLdr:
       LLVM_DEBUG(dbgs() << "Adding MCLOH_AdrpLdrGotLdr:\n"

llvm/test/CodeGen/AArch64/loh-adrp-add-ldr-clobber.mir

@@ -1,16 +1,34 @@
-# RUN: llc -o /dev/null %s -mtriple=aarch64-apple-ios -run-pass=aarch64-collect-loh -debug-only=aarch64-collect-loh 2>&1 | FileCheck %s
+# RUN: llc -o /dev/null %s -mtriple=aarch64-apple-ios -run-pass=aarch64-collect-loh -debug-only=aarch64-collect-loh 2>&1 | FileCheck %s --implicit-check-not=MCLOH_
 # REQUIRES: asserts
+
+# Check that we don't emit LOHs when there is a clobbering def of x8.
 --- |
   @sym2 = local_unnamed_addr global [10000000 x i32] zeroinitializer, align 8
   @sym = local_unnamed_addr global i32 zeroinitializer, align 8
 
-  define i32 @main() {
-    ret i32 0
-  }
+  define i32 @adrp_add_ldr() { ret i32 0 }
+  define i32 @adrp_add_str() { ret i32 0 }
+...
+
+---
+name:            adrp_add_ldr
+alignment:       4
+tracksRegLiveness: true
+liveins:
+  - { reg: '$x21', virtual-reg: '' }
+body:             |
+  bb.0:
+    liveins: $x21
+    renamable $x8 = ADRP target-flags(aarch64-page) @sym
+    renamable $x9 = ADDXri killed renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @sym, 0
+    renamable $x8 = ADDXri killed renamable $x21, 1, 0
+    $x9 = LDRXui $x9, 0
 
+    RET undef $lr
 ...
+
 ---
-name:            main
+name:            adrp_add_str
 alignment:       4
 tracksRegLiveness: true
 liveins:
@@ -19,13 +37,10 @@ liveins:
 body:             |
   bb.0:
     liveins: $x21, $x22
-    ; Check we don't emit an loh here because there's a clobbering def of x8 before the ldr.
-    ; CHECK-LABEL: main
-    ; CHECK-NOT: MCLOH_AdrpAddLdr
     renamable $x8 = ADRP target-flags(aarch64-page) @sym
     renamable $x9 = ADDXri killed renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @sym, 0
-    renamable $x8 = ADDXri killed renamable $x22, 1, 0
-    $x9 = LDRXui $x9, 0
-    RET undef $lr
+    renamable $x8 = ADDXri killed renamable $x21, 1, 0
+    STRXui $x22, $x9, 0
 
+    RET undef $lr
 ...