[AArch64] Fix CollectLOH creating an AdrpAdd LOH when there's a live used reg between the two instructions.

[aemerson]

Summary

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.

[paquette]

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

27–30

Will this patch prevent us from making the LOH if we have something like this?

$add_def = ... stuff ...
$adrp_def = adrp  ...
$add_def = add ...

If I understand correctly, it looks like the patch doesn't check if the use is actually between the adrp and add?

[aemerson]

Add a test case for a use before the ADRP.

[aemerson]

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

27–30

This change checks if the add def register has a use at the point of the ADRP, as it goes bottom up. Therefore if the use is before the ADRP, then OneUse won't be true because we haven't hit it yet. And having a def doesn't make sense because the def would be overwritten by the ADD.

[paquette]

LGTM

[loladiro]

I believe this fix is incomplete and also needs to apply adrp/add/ldr sequences (and maybe others). We encountered this as a miscompilation while porting to Apple Silicon (https://github.com/JuliaLang/julia/issues/39820).

[aemerson]

In D80834#2586675, @loladiro wrote:

I believe this fix is incomplete and also needs to apply adrp/add/ldr sequences (and maybe others). We encountered this as a miscompilation while porting to Apple Silicon (https://github.com/JuliaLang/julia/issues/39820).

I'll take a look. Are you by any chance enabling GlobalISel at all? I assume this is with optimizations enabled?

[loladiro]

I don't believe GlobalIsel is enabled. This is in debug mode, so
optimizations are enabled, but minimal. Thanks!

llvm/lib/Target/AArch64/AArch64CollectLOH.cpp

if (OpInfo.Type == MCLOH_AdrpAddStr && OpInfo.MI1 == nullptr) {
       return true;
     }
   }
   return false;
 }
 
 /// 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);
     AFI.addLOHDirective(MCLOH_AdrpAdrp, {&MI, Info.LastADRP});
     ++NumADRPSimpleCandidate;
   }
 
   // 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"
                           << '\t' << MI << '\t' << *Info.MI0);
         AFI.addLOHDirective(MCLOH_AdrpLdr, {&MI, Info.MI0});
         ++NumADRPToLDR;
       }
       break;
for (const MachineInstr &MI :
               handleMiddleInst(MI, LOHInfos[DefIdx], LOHInfos[OpIdx]))
             continue;
         }
         break;
       case AArch64::ADRP:
         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;
       }
       handleNormalInst(MI, LOHInfos);
     }
   }
 
   // Return "no change": The pass only collects information.
   return false;
 }
 
 FunctionPass *llvm::createAArch64CollectLOHPass() {
   return new AArch64CollectLOH();
 }

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

+# 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
+
+...