Automerge: [MachineOutliner] Remove LOHs from outlined candidates (#143617)

Remove Linker Optimization Hints (LOHs) from outlining candidates
instead of simply preventing outlining if LOH labels are found in the
candidate. This will improve the effectiveness of the machine outliner
when LOHs are enabled (which is the default).

In
https://discourse.llvm.org/t/loh-conflicting-with-machineoutliner/83279/1
it was observed that the machine outliner is much more effective when
LOHs are disabled. Rather than completely disabling LOH, this PR aims to
keep LOH in most places and removing them from outlined functions where
it could be illegal. Note that we are conservatively removing all LOHs
from outlined functions for simplicity, but I believe we could retain
LOHs that are in the intersection of all candidates.

It should be ok to remove these LOHs since these blocks are being
outlined anyway, which will harm performance much more than the gain
from keeping the LOHs.

Author: ellishg

llvm/include/llvm/Target/TargetMachine.h

@@ -38,6 +38,7 @@ using ModulePassManager = PassManager<Module>;
 
 class Function;
 class GlobalValue;
+class MachineInstr;
 class MachineModuleInfoWrapperPass;
 struct MachineSchedContext;
 class Mangler;
@@ -519,6 +520,15 @@ class LLVM_ABI TargetMachine {
 
   // MachineRegisterInfo callback function
   virtual void registerMachineRegisterInfoCallback(MachineFunction &MF) const {}
+
+  /// Remove all Linker Optimization Hints (LOH) associated with instructions in
+  /// \p MIs and \return the number of hints removed. This is useful in
+  /// transformations that cause these hints to be illegal, like in the machine
+  /// outliner.
+  virtual size_t clearLinkerOptimizationHints(
+      const SmallPtrSetImpl<MachineInstr *> &MIs) const {
+    return 0;
+  }
 };
 
 } // end namespace llvm

llvm/lib/CodeGen/MachineOutliner.cpp

@@ -67,6 +67,7 @@
 #include "llvm/CodeGen/MachineOptimizationRemarkEmitter.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
+#include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/DIBuilder.h"
 #include "llvm/IR/IRBuilder.h"
@@ -77,6 +78,7 @@
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/SuffixTree.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 #include <tuple>
 #include <vector>
@@ -104,6 +106,7 @@ STATISTIC(StableHashAttempts,
           "Count of hashing attempts made for outlined functions");
 STATISTIC(StableHashDropped,
           "Count of unsuccessful hashing attempts for outlined functions");
+STATISTIC(NumRemovedLOHs, "Total number of Linker Optimization Hints removed");
 
 // Set to true if the user wants the outliner to run on linkonceodr linkage
 // functions. This is false by default because the linker can dedupe linkonceodr
@@ -426,6 +429,7 @@ struct MachineOutliner : public ModulePass {
   static char ID;
 
   MachineModuleInfo *MMI = nullptr;
+  const TargetMachine *TM = nullptr;
 
   /// Set to true if the outliner should consider functions with
   /// linkonceodr linkage.
@@ -461,6 +465,7 @@ struct MachineOutliner : public ModulePass {
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<MachineModuleInfoWrapperPass>();
+    AU.addRequired<TargetPassConfig>();
     AU.addPreserved<MachineModuleInfoWrapperPass>();
     AU.addUsedIfAvailable<ImmutableModuleSummaryIndexWrapperPass>();
     AU.setPreservesAll();
@@ -1075,6 +1080,17 @@ bool MachineOutliner::outline(
                       << " B) > threshold (" << OutlinerBenefitThreshold
                       << " B)\n");
 
+    // Remove all Linker Optimization Hints from the candidates.
+    // TODO: The intersection of the LOHs from all candidates should be legal in
+    // the outlined function.
+    SmallPtrSet<MachineInstr *, 2> MIs;
+    for (Candidate &C : OF->Candidates) {
+      for (MachineInstr &MI : C)
+        MIs.insert(&MI);
+      NumRemovedLOHs += TM->clearLinkerOptimizationHints(MIs);
+      MIs.clear();
+    }
+
     // It's beneficial. Create the function and outline its sequence's
     // occurrences.
     OF->MF = createOutlinedFunction(M, *OF, Mapper, OutlinedFunctionNum);
@@ -1386,6 +1402,7 @@ bool MachineOutliner::runOnModule(Module &M) {
   initializeOutlinerMode(M);
 
   MMI = &getAnalysis<MachineModuleInfoWrapperPass>().getMMI();
+  TM = &getAnalysis<TargetPassConfig>().getTM<TargetMachine>();
 
   // Number to append to the current outlined function.
   unsigned OutlinedFunctionNum = 0;

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

@@ -9696,9 +9696,6 @@ AArch64InstrInfo::getOutliningTypeImpl(const MachineModuleInfo &MMI,
                                        MachineBasicBlock::iterator &MIT,
                                        unsigned Flags) const {
   MachineInstr &MI = *MIT;
-  MachineBasicBlock *MBB = MI.getParent();
-  MachineFunction *MF = MBB->getParent();
-  AArch64FunctionInfo *FuncInfo = MF->getInfo<AArch64FunctionInfo>();
 
   // Don't outline anything used for return address signing. The outlined
   // function will get signed later if needed
@@ -9726,10 +9723,6 @@ AArch64InstrInfo::getOutliningTypeImpl(const MachineModuleInfo &MMI,
     return outliner::InstrType::Illegal;
   }
 
-  // Don't outline LOHs.
-  if (FuncInfo->getLOHRelated().count(&MI))
-    return outliner::InstrType::Illegal;
-
   // We can only outline these if we will tail call the outlined function, or
   // fix up the CFI offsets. Currently, CFI instructions are outlined only if
   // in a tail call.

llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h

@@ -504,6 +504,20 @@ class AArch64FunctionInfo final : public MachineFunctionInfo {
     LOHRelated.insert_range(Args);
   }
 
+  size_t
+  clearLinkerOptimizationHints(const SmallPtrSetImpl<MachineInstr *> &MIs) {
+    size_t InitialSize = LOHContainerSet.size();
+    erase_if(LOHContainerSet, [&](const auto &D) {
+      return any_of(D.getArgs(), [&](auto *Arg) { return MIs.contains(Arg); });
+    });
+    // In theory there could be an LOH with one label in MIs and another label
+    // outside MIs, however we don't know if the label outside MIs is used in
+    // any other LOHs, so we can't remove them from LOHRelated. In that case, we
+    // might produce a few extra labels, but it won't break anything.
+    LOHRelated.remove_if([&](auto *MI) { return MIs.contains(MI); });
+    return InitialSize - LOHContainerSet.size();
+  };
+
   SmallVectorImpl<ForwardedRegister> &getForwardedMustTailRegParms() {
     return ForwardedMustTailRegParms;
   }

llvm/lib/Target/AArch64/AArch64TargetMachine.cpp

@@ -512,6 +512,15 @@ AArch64TargetMachine::createPostMachineScheduler(MachineSchedContext *C) const {
   return DAG;
 }
 
+size_t AArch64TargetMachine::clearLinkerOptimizationHints(
+    const SmallPtrSetImpl<MachineInstr *> &MIs) const {
+  if (MIs.empty())
+    return 0;
+  auto *MI = *MIs.begin();
+  auto *FuncInfo = MI->getMF()->getInfo<AArch64FunctionInfo>();
+  return FuncInfo->clearLinkerOptimizationHints(MIs);
+}
+
 void AArch64leTargetMachine::anchor() { }
 
 AArch64leTargetMachine::AArch64leTargetMachine(

llvm/lib/Target/AArch64/AArch64TargetMachine.h

@@ -76,6 +76,9 @@ class AArch64TargetMachine : public CodeGenTargetMachineImpl {
   ScheduleDAGInstrs *
   createPostMachineScheduler(MachineSchedContext *C) const override;
 
+  size_t clearLinkerOptimizationHints(
+      const SmallPtrSetImpl<MachineInstr *> &MIs) const override;
+
 private:
   bool isLittle;
 };

llvm/test/CodeGen/AArch64/machine-outliner-loh.ll

@@ -0,0 +1,74 @@
+; RUN: llc -verify-machineinstrs -mtriple=aarch64-apple-darwin < %s | FileCheck %s --implicit-check-not=.loh --check-prefixes=CHECK,LOH
+; RUN: llc -verify-machineinstrs -mtriple=aarch64-apple-darwin -enable-machine-outliner < %s | FileCheck %s --implicit-check-not=.loh --check-prefixes=CHECK,OUTLINE
+
+@A = global i32 0, align 4
+@B = global i32 0, align 4
+
+declare void @foo();
+declare void @bar(ptr %a);
+declare void @goo(ptr %a);
+
+; CHECK-LABEL: _a0:
+define void @a0(i32 %a) {
+
+  ; This becomes AdrpAdd when outlining is disabled, otherwise it is outlined
+  ; and there should be no LOH.
+  %addr = getelementptr inbounds i32, ptr @A, i32 0
+  %res = load i32, ptr %addr, align 4
+  ; LOH:      [[L0:Lloh.+]]:
+  ; LOH-NEXT:   adrp x19, _A@PAGE
+  ; LOH-NEXT: [[L1:Lloh.+]]:
+  ; LOH-NEXT:   add x19, x19, _A@PAGEOFF
+
+  call void @foo()
+  ; OUTLINE:      bl _OUTLINED_FUNCTION_0
+  ; OUTLINE-NEXT: mov x0, x19
+  ; OUTLINE-NEXT: bl _bar
+  call void @bar(ptr %addr)
+
+  ; This becomes AdrpAddStr.
+  %addr2 = getelementptr inbounds i32, ptr @B, i32 4
+  store i32 %res, ptr %addr2, align 4
+  ; CHECK:      [[L2:Lloh.+]]:
+  ; CHECK-NEXT:   adrp x8, _B@PAGE
+  ; CHECK-NEXT: [[L3:Lloh.+]]:
+  ; CHECK-NEXT:   add x8, x8, _B@PAGEOFF
+  ; CHECK-NEXT: [[L4:Lloh.+]]:
+  ; CHECK-NEXT:   str w20, [x8, #16]
+  ret void
+
+  ; LOH-DAG:   .loh AdrpAdd [[L0]], [[L1]]
+  ; CHECK-DAG: .loh AdrpAddStr [[L2]], [[L3]], [[L4]]
+  ; CHECK:     .cfi_endproc
+}
+
+; CHECK-LABEL: _a1:
+define i32 @a1(i32 %a) {
+
+  ; This becomes AdrpAdd when outlining is disabled, otherwise it is outlined
+  ; and there should be no LOH.
+  %addr = getelementptr inbounds i32, ptr @A, i32 0
+  %res = load i32, ptr %addr, align 4
+  ; LOH:      [[L5:Lloh.+]]:
+  ; LOH-NEXT:   adrp x19, _A@PAGE
+  ; LOH-NEXT: [[L6:Lloh.+]]:
+  ; LOH-NEXT:   add x19, x19, _A@PAGEOFF
+
+  call void @foo()
+  ; OUTLINE:      bl _OUTLINED_FUNCTION_0
+  ; OUTLINE-NEXT: mov x0, x19
+  ; OUTLINE-NEXT: bl _goo
+  call void @goo(ptr %addr)
+  ret i32 %res
+
+  ; LOH:   .loh AdrpAdd [[L5]], [[L6]]
+  ; CHECK: .cfi_endproc
+}
+
+; Note: it is not safe to add LOHs to this function as outlined functions do not
+; follow calling convention and thus x19 could be live across the call.
+; OUTLINE: _OUTLINED_FUNCTION_0:
+; OUTLINE:   adrp x19, _A@PAGE
+; OUTLINE:   add x19, x19, _A@PAGEOFF
+; OUTLINE:   ldr w20, [x19]
+; OUTLINE:   b _foo

llvm/test/CodeGen/AArch64/machine-outliner-loh.mir

@@ -0,0 +1,50 @@
+# RUN: llc %s -mtriple=aarch64-apple-ios -run-pass=aarch64-collect-loh -run-pass=machine-outliner -stats -o - 2>%t | FileCheck %s
+# RUN: FileCheck %s --input-file=%t --check-prefix=STATS
+# REQUIRES: asserts
+
+--- |
+  define void  @func0() noredzone minsize { ret void }
+
+  @g0 = external global i32
+...
+---
+# CHECK-LABEL: name: func0
+name: func0
+tracksRegLiveness: true
+body: |
+  ; CHECK-LABEL: bb.0:
+  bb.0:
+    ; CHECK: BL @OUTLINED_FUNCTION_0
+    $x10 = ADRP target-flags(aarch64-page) @g0
+    $x11 = ADDXri $x10, target-flags(aarch64-pageoff) @g0, 0
+    $w0 = ORRWri $wzr, 1
+    $w0 = ORRWri $wzr, 1
+    $w0 = ORRWri $wzr, 1
+    $w0 = ORRWri $wzr, 1
+
+  ; CHECK-LABEL: bb.1:
+  bb.1:
+    ; CHECK: BL @OUTLINED_FUNCTION_0
+    ; MCLOH_AdrpAdd is not generated because $x11 is still live. If we want to
+    ; outline these instructions, we must remove the MCLOH_AdrpAdd from bb.0.
+    $x10 = ADRP target-flags(aarch64-page) @g0
+    $x11 = ADDXri $x10, target-flags(aarch64-pageoff) @g0, 0
+    $w0 = ORRWri $wzr, 1
+    $w0 = ORRWri $wzr, 1
+    $w0 = ORRWri $wzr, 1
+    $w0 = ORRWri $wzr, 1
+    ; CHECK: $x12 = ADDXri $x11, target-flags(aarch64-pageoff) @g0, 0
+    $x12 = ADDXri $x11, target-flags(aarch64-pageoff) @g0, 0
+
+...
+
+# CHECK-LABEL: name: OUTLINED_FUNCTION_0
+# CHECK:           $x10 = ADRP target-flags(aarch64-page) @g0
+# CHECK:           $x11 = ADDXri $x10, target-flags(aarch64-pageoff) @g0, 0
+# CHECK:           $w0 = ORRWri $wzr, 1
+# CHECK:           $w0 = ORRWri $wzr, 1
+# CHECK:           $w0 = ORRWri $wzr, 1
+# CHECK:           $w0 = ORRWri $wzr, 1
+
+# STATS: 1 aarch64-collect-loh - Number of simplifiable ADRP + ADD
+# STATS: 1 machine-outliner    - Total number of Linker Optimization Hints removed