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

[Ellis Hoag via llvm-commits]

https://github.com/ellishg updated https://github.com/llvm/llvm-project/pull/142849

>From b18989a58c471f5b5c3d31ea38b65841295d3512 Mon Sep 17 00:00:00 2001
From: Ellis Hoag <ellishoag at meta.com>
Date: Wed, 4 Jun 2025 13:31:56 -0700
Subject: [PATCH] [LOH] Don't emit str LOH when register could be clobbered

---
 llvm/lib/Target/AArch64/AArch64CollectLOH.cpp | 37 +++++++++++--------
 .../AArch64/loh-adrp-add-ldr-clobber.mir      | 37 +++++++++++++------
 2 files changed, 48 insertions(+), 26 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64CollectLOH.cpp b/llvm/lib/Target/AArch64/AArch64CollectLOH.cpp
index 64f21c4cb2297..9b6fafb529c11 100644
--- a/llvm/lib/Target/AArch64/AArch64CollectLOH.cpp
+++ b/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 traked 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 at PAGE
       // add x1, x0, @sym at 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"
diff --git a/llvm/test/CodeGen/AArch64/loh-adrp-add-ldr-clobber.mir b/llvm/test/CodeGen/AArch64/loh-adrp-add-ldr-clobber.mir
index ce2d8f02f4cc8..a1d8bf375a19b 100644
--- a/llvm/test/CodeGen/AArch64/loh-adrp-add-ldr-clobber.mir
+++ b/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
 ...