[LOH] Don't emit AdrpAddStr when register could be clobbered

[ellishg]

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

llvm-project/llvm/lib/Target/AArch64/AArch64CollectLOH.cpp

Lines 419 to 431 in 50c5704

	// There is a possibility that the linker may try to rewrite:
	// adrp x0, @sym@PAGE
	// add x1, x0, @sym@PAGEOFF
	// [x0 = some other def]
	// ldr x2, [x1]
	// ...into...
	// adrp x0, @sym
	// nop
	// [x0 = some other def]
	// ldr x2, [x0]
	// ...if the offset to the symbol won't fit within a literal load.
	// This causes the load to use the result of the adrp, which in this
	// case has already been clobbered.

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

llvm-project/lld/MachO/Arch/ARM64.cpp

Lines 699 to 702 in 50c5704

	case LOH_ARM64_ADRP_ADD_STR:
	case LOH_ARM64_ADRP_LDR_GOT_STR:
	// TODO: Implement these
	break;

[llvmbot]

@llvm/pr-subscribers-backend-aarch64

Author: Ellis Hoag (ellishg)

Changes

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

llvm-project/llvm/lib/Target/AArch64/AArch64CollectLOH.cpp

Lines 419 to 431 in 50c5704

	// There is a possibility that the linker may try to rewrite:
	// adrp x0, @sym@PAGE
	// add x1, x0, @sym@PAGEOFF
	// [x0 = some other def]
	// ldr x2, [x1]
	// ...into...
	// adrp x0, @sym
	// nop
	// [x0 = some other def]
	// ldr x2, [x0]
	// ...if the offset to the symbol won't fit within a literal load.
	// This causes the load to use the result of the adrp, which in this
	// case has already been clobbered.

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

llvm-project/lld/MachO/Arch/ARM64.cpp

Lines 699 to 702 in 50c5704

	case LOH_ARM64_ADRP_ADD_STR:
	case LOH_ARM64_ADRP_LDR_GOT_STR:
	// TODO: Implement these
	break;

---

Full diff: https://github.com/llvm/llvm-project/pull/142849.diff

2 Files Affected:

• (modified) llvm/lib/Target/AArch64/AArch64CollectLOH.cpp (+22-15)
• (modified) llvm/test/CodeGen/AArch64/loh-adrp-add-ldr-clobber.mir (+26-11)

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@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"
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
 ...

[kyulee-com]

This check already exists. I'm curious why this check is not needed for AdrpAddLdr. Can we conservatively integrate a null-check for both MI0 and MI1 inside areInstructionsConsecutive(), return false if either one is null?

[ellishg]

I was also confused about this. The reason we need it for AdrpAddStr and not AdrpAddLdr is because AdrpAddLdr starts out as a AdrpLdr and then is converted to a AdrpAddLdr by setting MI1. There is no AdrpStr, so MI1 is set to null until it finds the add instruction.

llvm-project/llvm/lib/Target/AArch64/AArch64CollectLOH.cpp

Lines 297 to 308 in 2c0a226

	// Start new LOHInfo if applicable.
	if (isCandidateLoad(MI)) {
	Info.Type = MCLOH_AdrpLdr;
	Info.IsCandidate = true;
	Info.MI0 = &MI;
	// Note that even this is AdrpLdr now, we can switch to a Ldr variant
	// later.
	} else if (isCandidateStore(MI, MO)) {
	Info.Type = MCLOH_AdrpAddStr;
	Info.IsCandidate = true;
	Info.MI0 = &MI;
	Info.MI1 = nullptr;

[kyulee-com]

Thanks for the point! The code change looks good to me.

[kyulee-com]

LGTM

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.