[LOH] Remove hints when outlining

Author: ellishg

llvm/include/llvm/CodeGen/TargetInstrInfo.h

@@ -2184,6 +2184,12 @@ class LLVM_ABI TargetInstrInfo : public MCInstrInfo {
                                        MachineBasicBlock::iterator &MIT,
                                        unsigned Flags) const;
 
+  /// Remove all Linker Optimization Hints associated with instructions in
+  // \p MIs and \return the number of hints removed.
+  virtual size_t clearLOHs(const SmallPtrSetImpl<MachineInstr *> &MIs) const {
+    return 0;
+  }
+
   /// Optional target hook that returns true if \p MBB is safe to outline from,
   /// and returns any target-specific information in \p Flags.
   virtual bool isMBBSafeToOutlineFrom(MachineBasicBlock &MBB,

llvm/lib/CodeGen/MachineOutliner.cpp

@@ -104,6 +104,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
@@ -1075,6 +1076,23 @@ bool MachineOutliner::outline(
                       << " B) > threshold (" << OutlinerBenefitThreshold
                       << " B)\n");
 
+    // Remove all Linker Optimization Hints from the candidates since we did not
+    // check if the set of hints are the same for each of them.
+    // TODO: The intersection of the LOHs from all candidates should be legal in
+    // the outlined function.
+    SmallPtrSet<MachineInstr *, 2> MIs;
+    std::optional<size_t> MinRemovedLOHs;
+    for (Candidate &C : OF->Candidates) {
+      const TargetInstrInfo &TII = *C.getMF()->getSubtarget().getInstrInfo();
+      for (MachineInstr &MI : C)
+        MIs.insert(&MI);
+      size_t NumRemoved = TII.clearLOHs(MIs);
+      MIs.clear();
+      MinRemovedLOHs =
+          std::min(MinRemovedLOHs.value_or(NumRemoved), NumRemoved);
+    }
+    NumRemovedLOHs += MinRemovedLOHs.value_or(0);
+
     // It's beneficial. Create the function and outline its sequence's
     // occurrences.
     OF->MF = createOutlinedFunction(M, *OF, Mapper, OutlinedFunctionNum);

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

@@ -9641,14 +9641,20 @@ AArch64InstrInfo::getOutlinableRanges(MachineBasicBlock &MBB,
   return Ranges;
 }
 
+size_t
+AArch64InstrInfo::clearLOHs(const SmallPtrSetImpl<MachineInstr *> &MIs) const {
+  if (MIs.empty())
+    return 0;
+  auto *MI = *MIs.begin();
+  auto *FuncInfo = MI->getMF()->getInfo<AArch64FunctionInfo>();
+  return FuncInfo->clearLOHs(MIs);
+}
+
 outliner::InstrType
 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
@@ -9676,10 +9682,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/AArch64InstrInfo.h

@@ -493,6 +493,7 @@ class AArch64InstrInfo final : public AArch64GenInstrInfo {
   outliner::InstrType getOutliningTypeImpl(const MachineModuleInfo &MMI,
                                            MachineBasicBlock::iterator &MIT,
                                            unsigned Flags) const override;
+  size_t clearLOHs(const SmallPtrSetImpl<MachineInstr *> &MIs) const override;
   SmallVector<
       std::pair<MachineBasicBlock::iterator, MachineBasicBlock::iterator>>
   getOutlinableRanges(MachineBasicBlock &MBB, unsigned &Flags) const override;

llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h

@@ -504,6 +504,19 @@ class AArch64FunctionInfo final : public MachineFunctionInfo {
     LOHRelated.insert_range(Args);
   }
 
+  size_t clearLOHs(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/test/CodeGen/AArch64/machine-outliner-loh.ll

@@ -0,0 +1,63 @@
+; 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 i32 @a0(i32 %a) {
+  %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)
+  %addr2 = getelementptr inbounds i32, ptr @B, i32 4
+  store i32 0, 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:   mov w0, w20
+  ; CHECK-NEXT: [[L4:Lloh.+]]:
+  ; CHECK-NEXT:   str wzr, [x8, #16]
+  ret i32 %res
+  ; LOH-DAG:   .loh AdrpAdd [[L0]], [[L1]]
+  ; CHECK-DAG: .loh AdrpAddStr [[L2]], [[L3]], [[L4]]
+  ; CHECK:     .cfi_endproc
+}
+
+; CHECK-LABEL: _a1:
+define i32 @a1(i32 %a) {
+  %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