Allow outlining ADRP pair as a whole

Author: pranavk

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

@@ -9547,6 +9547,21 @@ AArch64InstrInfo::getOutliningCandidateInfo(
 
   unsigned NumBytesToCreateFrame = 0;
 
+  // Avoid splitting ADRP-ADD/LDR pair into outlined functions.
+  MachineInstr &LastMI = RepeatedSequenceLocs[0].back();
+  MachineInstr &FirstMI = RepeatedSequenceLocs[0].front();
+  if (LastMI.getOpcode() == AArch64::ADRP &&
+      (LastMI.getOperand(1).getTargetFlags() & AArch64II::MO_PAGE) != 0 &&
+      (LastMI.getOperand(1).getTargetFlags() & AArch64II::MO_GOT) != 0) {
+    return std::nullopt;
+  }
+
+  if ((FirstMI.getOpcode() == AArch64::ADDXri || FirstMI.getOpcode() == AArch64::LDRXui) &&
+       (FirstMI.getOperand(2).getTargetFlags() & AArch64II::MO_PAGEOFF) != 0 &&
+       (FirstMI.getOperand(2).getTargetFlags() & AArch64II::MO_GOT) != 0) {
+    return std::nullopt;
+  }
+
   // We only allow outlining for functions having exactly matching return
   // address signing attributes, i.e., all share the same value for the
   // attribute "sign-return-address" and all share the same type of key they
@@ -10179,34 +10194,6 @@ AArch64InstrInfo::getOutliningTypeImpl(const MachineModuleInfo &MMI,
       return outliner::InstrType::Illegal;
   }
 
-  // An ADRP instruction referencing a GOT should not be outlined.
-  // This is to avoid splitting ADRP/(LDR/ADD/etc.) pair into different
-  // functions which can lead to linker ICF merging sections incorrectly.
-  if (MI.getOpcode() == AArch64::ADRP) {
-    bool IsPage = (MI.getOperand(1).getTargetFlags() & AArch64II::MO_PAGE) != 0;
-    bool IsGot = (MI.getOperand(1).getTargetFlags() & AArch64II::MO_GOT) != 0;
-    if (IsPage && IsGot)
-      return outliner::InstrType::Illegal;
-
-    // Special cases for instructions that can always be outlined, but will fail
-    // the later tests. e.g, ADRPs, which are PC-relative use LR, but can always
-    // be outlined because they don't require a *specific* value to be in LR.
-    return outliner::InstrType::Legal;
-  }
-
-  // Similarly, any user of ADRP instruction referencing a GOT should not be
-  // outlined. It's hard/costly to check exact users of ADRP. So we use check
-  // all operands and reject any that's a page offset and references a GOT.
-  const auto &F = MI.getMF()->getFunction();
-  for (const auto &MO : MI.operands()) {
-    bool IsPageOff = (MO.getTargetFlags() & AArch64II::MO_PAGEOFF) != 0;
-    bool IsGot = (MO.getTargetFlags() & AArch64II::MO_GOT) != 0;
-    if (IsPageOff && IsGot &&
-        (MI.getMF()->getTarget().getFunctionSections() || F.hasComdat() ||
-         F.hasSection() || F.getSectionPrefix()))
-      return outliner::InstrType::Illegal;
-  }
-
   // If MI is a call we might be able to outline it. We don't want to outline
   // any calls that rely on the position of items on the stack. When we outline
   // something containing a call, we have to emit a save and restore of LR in

llvm/test/CodeGen/AArch64/machine-outliner-adrp-got-split.mir

@@ -11,28 +11,16 @@
     ret i32 0
   }
 
-  define void @bar(i32 %a) #0 {
-    ret void
-  }
-
   attributes #0 = { noinline noredzone }
 ...
 ---
-# This test ensures that we do not outline ADRP / ADD pair when it's referencing 
-# a GOT entry.
+# Check that main function body doesn't split ADRP pair
 #
 # CHECK-LABEL: name: adrp_add
 # CHECK-DAG: bb.0:
-# CHECK: $x9 = ADRP target-flags(aarch64-page, aarch64-got) @x
-# CHECK: $x12 = ADDXri $x9, target-flags(aarch64-pageoff, aarch64-got) @x, 0
-
-# CHECK-DAG: bb.1
-# CHECK: $x9 = ADRP target-flags(aarch64-page, aarch64-got) @x
-# CHECK: $x12 = ADDXri $x9, target-flags(aarch64-pageoff, aarch64-got) @x, 0
-
-# CHECK-DAG: bb.2
-# CHECK: $x9 = ADRP target-flags(aarch64-page, aarch64-got) @x
-# CHECK: $x12 = ADDXri $x9, target-flags(aarch64-pageoff, aarch64-got) @x, 0
+# CHECK: BL @OUTLINED_FUNCTION_[[F0:[0-9]+]]
+# CHECK-NEXT: BL @OUTLINED_FUNCTION_[[F2:[0-9]+]]
+# CHECK-NEXT: $lr = ORRXri $xzr, 1
 name:            adrp_add
 tracksRegLiveness: true
 body:             |
@@ -74,21 +62,13 @@ body:             |
     RET undef $lr
 ...
 ---
-# This test ensures that we do not outline ADRP / LDR pair when it's referencing 
-# a GOT entry.
+# Check that main function body doesn't split ADRP pair
 #
 # CHECK-LABEL: name: adrp_ldr
 # CHECK-DAG: bb.0:
-# CHECK: $x9 = ADRP target-flags(aarch64-page, aarch64-got) @x
-# CHECK: $x12 = LDRXui $x9, target-flags(aarch64-pageoff, aarch64-got) @x
-
-# CHECK-DAG: bb.1
-# CHECK: $x9 = ADRP target-flags(aarch64-page, aarch64-got) @x
-# CHECK: $x12 = LDRXui $x9, target-flags(aarch64-pageoff, aarch64-got) @x
-
-# CHECK-DAG: bb.2
-# CHECK: $x9 = ADRP target-flags(aarch64-page, aarch64-got) @x
-# CHECK: $x12 = LDRXui $x9, target-flags(aarch64-pageoff, aarch64-got) @x
+# CHECK: BL @OUTLINED_FUNCTION_[[F0]]
+# CHECK-NEXT: BL @OUTLINED_FUNCTION_[[F1:[0-9]+]]
+# CHECK-NEXT: $lr = ORRXri $xzr, 1
 name:            adrp_ldr
 tracksRegLiveness: true
 body:             |
@@ -127,4 +107,27 @@ body:             |
     $lr = ORRXri $xzr, 1
   bb.3:
   liveins: $lr
-    RET undef $lr
+    RET undef $lr
+
+# Check that no outlined function split the ADRP pair apart
+#
+# CHECK: OUTLINED_FUNCTION_[[F0]]
+# CHECK-DAG: bb.0
+# CHECK: $w12 = ORRWri $wzr, 1
+# CHECK-NEXT: $w12 = ORRWri $wzr, 1
+# CHECK-NEXT: $w12 = ORRWri $wzr, 1
+# CHECK-NEXT: $w12 = ORRWri $wzr, 1
+# CHECK-NEXT: $w12 = ORRWri $wzr, 1
+# CHECK-NEXT: RET $lr
+
+# CHECK: OUTLINED_FUNCTION_[[F1]]
+# CHECK-DAG: bb.0
+# CHECK: $w12 = ORRWri $wzr, 1
+# CHECK-NEXT: $x9 = ADRP target-flags(aarch64-page, aarch64-got) @x
+# CHECK-NEXT: $x12 = LDRXui $x9, target-flags(aarch64-pageoff, aarch64-got) @x
+
+# CHECK: name: OUTLINED_FUNCTION_[[F2]]
+# CHECK-DAG: bb.0
+# CHECK: $w12 = ORRWri $wzr, 1
+# CHECK-NEXT: $x9 = ADRP target-flags(aarch64-page, aarch64-got) @x
+# CHECK-NEXT: $x12 = ADDXri $x9, target-flags(aarch64-pageoff, aarch64-got) @x, 0