[llvm] [CGData] Outlined Hash Tree (PR #89792)

Kyungwoo Lee via llvm-commits llvm-commits at lists.llvm.org
Sat May 4 17:45:15 PDT 2024


https://github.com/kyulee-com updated https://github.com/llvm/llvm-project/pull/89792

>From dd8a1fd0125a0bad48bd83492cfe7f6c1fc4c1e0 Mon Sep 17 00:00:00 2001
From: Xuan Zhang <xuanzh at meta.com>
Date: Fri, 26 Apr 2024 12:44:02 -0700
Subject: [PATCH 1/5] efficient implementation of
 MachineOutliner::findCandidates()

---
 llvm/lib/CodeGen/MachineOutliner.cpp          | 22 ++++++++++---------
 .../AArch64/machine-outliner-overlap.mir      | 12 +++++-----
 2 files changed, 18 insertions(+), 16 deletions(-)

diff --git a/llvm/lib/CodeGen/MachineOutliner.cpp b/llvm/lib/CodeGen/MachineOutliner.cpp
index dc2f5ef15206e8..d553c0e6d24772 100644
--- a/llvm/lib/CodeGen/MachineOutliner.cpp
+++ b/llvm/lib/CodeGen/MachineOutliner.cpp
@@ -593,7 +593,11 @@ void MachineOutliner::findCandidates(
     unsigned NumDiscarded = 0;
     unsigned NumKept = 0;
 #endif
-    for (const unsigned &StartIdx : RS.StartIndices) {
+    // Sort the start indices so that we can efficiently check if candidates
+    // overlap with each other in MachineOutliner::findCandidates().
+    SmallVector<unsigned> SortedStartIndices(RS.StartIndices);
+    llvm::sort(SortedStartIndices);
+    for (const unsigned &StartIdx : SortedStartIndices) {
       // Trick: Discard some candidates that would be incompatible with the
       // ones we've already found for this sequence. This will save us some
       // work in candidate selection.
@@ -616,17 +620,15 @@ void MachineOutliner::findCandidates(
       // * End before the other starts
       // * Start after the other ends
       unsigned EndIdx = StartIdx + StringLen - 1;
-      auto FirstOverlap = find_if(
-          CandidatesForRepeatedSeq, [StartIdx, EndIdx](const Candidate &C) {
-            return EndIdx >= C.getStartIdx() && StartIdx <= C.getEndIdx();
-          });
-      if (FirstOverlap != CandidatesForRepeatedSeq.end()) {
+      if (CandidatesForRepeatedSeq.size() > 0 &&
+          StartIdx <= CandidatesForRepeatedSeq.back().getEndIdx()) {
 #ifndef NDEBUG
         ++NumDiscarded;
-        LLVM_DEBUG(dbgs() << "    .. DISCARD candidate @ [" << StartIdx
-                          << ", " << EndIdx << "]; overlaps with candidate @ ["
-                          << FirstOverlap->getStartIdx() << ", "
-                          << FirstOverlap->getEndIdx() << "]\n");
+        LLVM_DEBUG(dbgs() << "    .. DISCARD candidate @ [" << StartIdx << ", "
+                          << EndIdx << "]; overlaps with candidate @ ["
+                          << CandidatesForRepeatedSeq.back().getStartIdx()
+                          << ", " << CandidatesForRepeatedSeq.back().getEndIdx()
+                          << "]\n");
 #endif
         continue;
       }
diff --git a/llvm/test/CodeGen/AArch64/machine-outliner-overlap.mir b/llvm/test/CodeGen/AArch64/machine-outliner-overlap.mir
index 649bb33828c32c..c6bd4c1d04d871 100644
--- a/llvm/test/CodeGen/AArch64/machine-outliner-overlap.mir
+++ b/llvm/test/CodeGen/AArch64/machine-outliner-overlap.mir
@@ -8,27 +8,27 @@
 # CHECK-NEXT:    Candidates discarded: 0
 # CHECK-NEXT:    Candidates kept: 2
 # CHECK-DAG:  Sequence length: 8
-# CHECK-NEXT:    .. DISCARD candidate @ [5, 12]; overlaps with candidate @ [12, 19]
+# CHECK-NEXT:    .. DISCARD candidate @ [12, 19]; overlaps with candidate @ [5, 12]
 # CHECK-NEXT:    Candidates discarded: 1
 # CHECK-NEXT:    Candidates kept: 1
 # CHECK-DAG:   Sequence length: 9
-# CHECK-NEXT:    .. DISCARD candidate @ [4, 12]; overlaps with candidate @ [11, 19]
+# CHECK-NEXT:    .. DISCARD candidate @ [11, 19]; overlaps with candidate @ [4, 12]
 # CHECK-NEXT:    Candidates discarded: 1
 # CHECK-NEXT:    Candidates kept: 1
 # CHECK-DAG:   Sequence length: 10
-# CHECK-NEXT:    .. DISCARD candidate @ [3, 12]; overlaps with candidate @ [10, 19]
+# CHECK-NEXT:    .. DISCARD candidate @ [10, 19]; overlaps with candidate @ [3, 12]
 # CHECK-NEXT:    Candidates discarded: 1
 # CHECK-NEXT:    Candidates kept: 1
 # CHECK-DAG:   Sequence length: 11
-# CHECK-NEXT:    .. DISCARD candidate @ [2, 12]; overlaps with candidate @ [9, 19]
+# CHECK-NEXT:    .. DISCARD candidate @ [9, 19]; overlaps with candidate @ [2, 12]
 # CHECK-NEXT:    Candidates discarded: 1
 # CHECK-NEXT:    Candidates kept: 1
 # CHECK-DAG:   Sequence length: 12
-# CHECK-NEXT:    .. DISCARD candidate @ [1, 12]; overlaps with candidate @ [8, 19]
+# CHECK-NEXT:    .. DISCARD candidate @ [8, 19]; overlaps with candidate @ [1, 12]
 # CHECK-NEXT:    Candidates discarded: 1
 # CHECK-NEXT:    Candidates kept: 1
 # CHECK-DAG:   Sequence length: 13
-# CHECK-NEXT:    .. DISCARD candidate @ [0, 12]; overlaps with candidate @ [7, 19]
+# CHECK-NEXT:    .. DISCARD candidate @ [7, 19]; overlaps with candidate @ [0, 12]
 # CHECK-NEXT:    Candidates discarded: 1
 # CHECK-NEXT:    Candidates kept: 1
 

>From 5b8fb073f50490bd87d9be8eb36e5d7d0a9b9c96 Mon Sep 17 00:00:00 2001
From: Xuan Zhang <xuanzh at meta.com>
Date: Fri, 26 Apr 2024 13:04:38 -0700
Subject: [PATCH 2/5] outlining order by benefit-to-cost ratio

---
 llvm/lib/CodeGen/MachineOutliner.cpp          |  6 +-
 .../machine-outliner-sort-per-priority.ll     | 96 +++++++++++++++++++
 .../CodeGen/ARM/machine-outliner-calls.mir    | 80 ++++++----------
 .../CodeGen/ARM/machine-outliner-default.mir  | 33 +++----
 .../ARM/machine-outliner-stack-fixup-arm.mir  | 56 +++++------
 .../machine-outliner-stack-fixup-thumb.mir    | 74 ++++++--------
 6 files changed, 194 insertions(+), 151 deletions(-)
 create mode 100644 llvm/test/CodeGen/AArch64/machine-outliner-sort-per-priority.ll

diff --git a/llvm/lib/CodeGen/MachineOutliner.cpp b/llvm/lib/CodeGen/MachineOutliner.cpp
index d553c0e6d24772..bc21215a181f6c 100644
--- a/llvm/lib/CodeGen/MachineOutliner.cpp
+++ b/llvm/lib/CodeGen/MachineOutliner.cpp
@@ -830,10 +830,12 @@ bool MachineOutliner::outline(Module &M,
                     << "\n");
   bool OutlinedSomething = false;
 
-  // Sort by benefit. The most beneficial functions should be outlined first.
+  // Sort by priority where priority := getNotOutlinedCost / getOutliningCost.
+  // The function with highest priority should be outlined first.
   stable_sort(FunctionList,
               [](const OutlinedFunction &LHS, const OutlinedFunction &RHS) {
-                return LHS.getBenefit() > RHS.getBenefit();
+                return LHS.getNotOutlinedCost() * RHS.getOutliningCost() >
+                       RHS.getNotOutlinedCost() * LHS.getOutliningCost();
               });
 
   // Walk over each function, outlining them as we go along. Functions are
diff --git a/llvm/test/CodeGen/AArch64/machine-outliner-sort-per-priority.ll b/llvm/test/CodeGen/AArch64/machine-outliner-sort-per-priority.ll
new file mode 100644
index 00000000000000..00efc3c6e71c89
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/machine-outliner-sort-per-priority.ll
@@ -0,0 +1,96 @@
+; This tests the order in which functions are outlined in MachineOutliner
+; There are TWO key OutlinedFunction in FunctionList
+;
+; ===================== First One =====================
+;   ```
+;     mov     w0, #1
+;     mov     w1, #2
+;     mov     w2, #3
+;     mov     w3, #4
+;     mov     w4, #5
+;   ```
+; It has:
+;   - `SequenceSize=20` and `OccurrenceCount=6`
+;   - each Candidate has `CallOverhead=12` and `FrameOverhead=4`
+;   - `NotOutlinedCost=20*6=120` and `OutliningCost=12*6+20+4=96`
+;   - `Benefit=120-96=24` and `Priority=120/96=1.25`
+;
+; ===================== Second One =====================
+;   ```
+;     mov     w6, #6
+;     mov     w7, #7
+;     b
+;   ```
+; It has:
+;   - `SequenceSize=12` and `OccurrenceCount=4`
+;   - each Candidate has `CallOverhead=4` and `FrameOverhead=0`
+;   - `NotOutlinedCost=12*4=48` and `OutliningCost=4*4+12+0=28`
+;   - `Benefit=120-96=20` and `Priority=48/28=1.71`
+;
+; Note that the first one has higher benefit, but lower priority.
+; Hence, when outlining per priority, the second one will be outlined first.
+
+; RUN: llc %s -enable-machine-outliner=always -filetype=obj -o %t
+; RUN: llvm-objdump -d %t | FileCheck %s --check-prefix=CHECK-SORT-BY-PRIORITY
+
+; RUN: llc %s -enable-machine-outliner=always -outliner-benefit-threshold=22 -filetype=obj -o %t
+; RUN: llvm-objdump -d %t | FileCheck %s --check-prefix=CHECK-THRESHOLD
+
+
+target datalayout = "e-m:o-i64:64-i128:128-n32:64-S128"
+target triple = "arm64-apple-macosx14.0.0"
+
+declare i32 @_Z3fooiiii(i32 noundef, i32 noundef, i32 noundef, i32 noundef, i32 noundef, i32 noundef, i32 noundef, i32 noundef)
+
+define i32 @_Z2f1v() minsize {
+  %1 = tail call i32 @_Z3fooiiii(i32 noundef 1, i32 noundef 2, i32 noundef 3, i32 noundef 4, i32 noundef 5, i32 noundef 11, i32 noundef 6, i32 noundef 7)
+  ret i32 %1
+}
+
+define i32 @_Z2f2v() minsize {
+  %1 = tail call i32 @_Z3fooiiii(i32 noundef 1, i32 noundef 2, i32 noundef 3, i32 noundef 4, i32 noundef 5, i32 noundef 12, i32 noundef 6, i32 noundef 7)
+  ret i32 %1
+}
+
+define i32 @_Z2f3v() minsize {
+  %1 = tail call i32 @_Z3fooiiii(i32 noundef 1, i32 noundef 2, i32 noundef 3, i32 noundef 4, i32 noundef 5, i32 noundef 13, i32 noundef 6, i32 noundef 7)
+  ret i32 %1
+}
+
+define i32 @_Z2f4v() minsize {
+  %1 = tail call i32 @_Z3fooiiii(i32 noundef 1, i32 noundef 2, i32 noundef 3, i32 noundef 4, i32 noundef 5, i32 noundef 14, i32 noundef 6, i32 noundef 7)
+  ret i32 %1
+}
+
+define i32 @_Z2f5v() minsize {
+  %1 = tail call i32 @_Z3fooiiii(i32 noundef 1, i32 noundef 2, i32 noundef 3, i32 noundef 4, i32 noundef 5, i32 noundef 15, i32 noundef 8, i32 noundef 9)
+  ret i32 %1
+}
+
+define i32 @_Z2f6v() minsize {
+  %1 = tail call i32 @_Z3fooiiii(i32 noundef 1, i32 noundef 2, i32 noundef 3, i32 noundef 4, i32 noundef 5, i32 noundef 16, i32 noundef 9, i32 noundef 8)
+  ret i32 %1
+}
+
+; CHECK-SORT-BY-PRIORITY: <_OUTLINED_FUNCTION_0>:
+; CHECK-SORT-BY-PRIORITY-NEXT: mov     w6, #0x6
+; CHECK-SORT-BY-PRIORITY-NEXT: mov     w7, #0x7
+; CHECK-SORT-BY-PRIORITY-NEXT: b
+
+; CHECK-SORT-BY-PRIORITY: <_OUTLINED_FUNCTION_1>:
+; CHECK-SORT-BY-PRIORITY-NEXT: mov     w0, #0x1
+; CHECK-SORT-BY-PRIORITY-NEXT: mov     w1, #0x2
+; CHECK-SORT-BY-PRIORITY-NEXT: mov     w2, #0x3
+; CHECK-SORT-BY-PRIORITY-NEXT: mov     w3, #0x4
+; CHECK-SORT-BY-PRIORITY-NEXT: mov     w4, #0x5
+; CHECK-SORT-BY-PRIORITY-NEXT: ret
+
+; CHECK-THRESHOLD: <_OUTLINED_FUNCTION_0>:
+; CHECK-THRESHOLD-NEXT: mov     w0, #0x1
+; CHECK-THRESHOLD-NEXT: mov     w1, #0x2
+; CHECK-THRESHOLD-NEXT: mov     w2, #0x3
+; CHECK-THRESHOLD-NEXT: mov     w3, #0x4
+; CHECK-THRESHOLD-NEXT: mov     w4, #0x5
+; CHECK-THRESHOLD-NEXT: ret
+
+; CHECK-THRESHOLD-NOT: <_OUTLINED_FUNCTION_1>:
diff --git a/llvm/test/CodeGen/ARM/machine-outliner-calls.mir b/llvm/test/CodeGen/ARM/machine-outliner-calls.mir
index a92c9dd28be5ae..7634ecd6e863ae 100644
--- a/llvm/test/CodeGen/ARM/machine-outliner-calls.mir
+++ b/llvm/test/CodeGen/ARM/machine-outliner-calls.mir
@@ -26,15 +26,15 @@ body:             |
   ; CHECK:   frame-setup CFI_INSTRUCTION def_cfa_offset 8
   ; CHECK:   frame-setup CFI_INSTRUCTION offset $lr, -4
   ; CHECK:   frame-setup CFI_INSTRUCTION offset $r4, -8
-  ; CHECK:   BL @OUTLINED_FUNCTION_0
+  ; CHECK:   BL @OUTLINED_FUNCTION_2
   ; CHECK: bb.1:
-  ; CHECK:   BL @OUTLINED_FUNCTION_0
+  ; CHECK:   BL @OUTLINED_FUNCTION_2
   ; CHECK: bb.2:
-  ; CHECK:   BL @OUTLINED_FUNCTION_0
+  ; CHECK:   BL @OUTLINED_FUNCTION_2
   ; CHECK: bb.3:
-  ; CHECK:   BL @OUTLINED_FUNCTION_0
+  ; CHECK:   BL @OUTLINED_FUNCTION_2
   ; CHECK: bb.4:
-  ; CHECK:   BL @OUTLINED_FUNCTION_0
+  ; CHECK:   BL @OUTLINED_FUNCTION_2
   ; CHECK: bb.5:
   ; CHECK:   $sp = frame-destroy LDMIA_UPD $sp, 14 /* CC::al */, $noreg, def $r4, def $lr
   ; CHECK:   BX_RET 14 /* CC::al */, $noreg
@@ -139,13 +139,13 @@ body:             |
   ; CHECK:   frame-setup CFI_INSTRUCTION def_cfa_offset 8
   ; CHECK:   frame-setup CFI_INSTRUCTION offset $lr, -4
   ; CHECK:   frame-setup CFI_INSTRUCTION offset $r4, -8
-  ; CHECK:   BL @OUTLINED_FUNCTION_1
+  ; CHECK:   BL @OUTLINED_FUNCTION_0
   ; CHECK: bb.1:
-  ; CHECK:   BL @OUTLINED_FUNCTION_1
+  ; CHECK:   BL @OUTLINED_FUNCTION_0
   ; CHECK: bb.2:
-  ; CHECK:   BL @OUTLINED_FUNCTION_1
+  ; CHECK:   BL @OUTLINED_FUNCTION_0
   ; CHECK: bb.3:
-  ; CHECK:   BL @OUTLINED_FUNCTION_1
+  ; CHECK:   BL @OUTLINED_FUNCTION_0
   ; CHECK: bb.4:
   ; CHECK:   $sp = frame-destroy LDMIA_UPD $sp, 14 /* CC::al */, $noreg, def $r4, def $lr
   ; CHECK:   BX_RET 14 /* CC::al */, $noreg
@@ -245,19 +245,19 @@ body:             |
   ; CHECK:   frame-setup CFI_INSTRUCTION offset $lr, -4
   ; CHECK:   frame-setup CFI_INSTRUCTION offset $r4, -8
   ; CHECK:   BL @"\01mcount", csr_aapcs, implicit-def dead $lr, implicit $sp
-  ; CHECK:   BL @OUTLINED_FUNCTION_2
+  ; CHECK:   BL @OUTLINED_FUNCTION_1
   ; CHECK: bb.1:
   ; CHECK:   BL @"\01mcount", csr_aapcs, implicit-def dead $lr, implicit $sp
-  ; CHECK:   BL @OUTLINED_FUNCTION_2
+  ; CHECK:   BL @OUTLINED_FUNCTION_1
   ; CHECK: bb.2:
   ; CHECK:   BL @"\01mcount", csr_aapcs, implicit-def dead $lr, implicit $sp
-  ; CHECK:   BL @OUTLINED_FUNCTION_2
+  ; CHECK:   BL @OUTLINED_FUNCTION_1
   ; CHECK: bb.3:
   ; CHECK:   BL @"\01mcount", csr_aapcs, implicit-def dead $lr, implicit $sp
-  ; CHECK:   BL @OUTLINED_FUNCTION_2
+  ; CHECK:   BL @OUTLINED_FUNCTION_1
   ; CHECK: bb.4:
   ; CHECK:   BL @"\01mcount", csr_aapcs, implicit-def dead $lr, implicit $sp
-  ; CHECK:   BL @OUTLINED_FUNCTION_2
+  ; CHECK:   BL @OUTLINED_FUNCTION_1
   ; CHECK: bb.5:
   ; CHECK:   $sp = frame-destroy LDMIA_UPD $sp, 14 /* CC::al */, $noreg, def $r4, def $lr
   ; CHECK:   BX_RET 14 /* CC::al */, $noreg
@@ -307,38 +307,17 @@ body:             |
   bb.0:
     BX_RET 14, $noreg
 
-
   ; CHECK-LABEL: name: OUTLINED_FUNCTION_0
   ; CHECK: bb.0:
-  ; CHECK:   liveins: $r11, $r10, $r9, $r8, $r7, $r6, $r5, $d15, $d14, $d13, $d12, $d11, $d10, $d9, $d8, $lr
-  ; CHECK:   early-clobber $sp = frame-setup STR_PRE_IMM killed $lr, $sp, -8, 14 /* CC::al */, $noreg
-  ; CHECK:   frame-setup CFI_INSTRUCTION def_cfa_offset 8
-  ; CHECK:   frame-setup CFI_INSTRUCTION offset $lr, -8
-  ; CHECK:   BL @bar, implicit-def dead $lr, implicit $sp
-  ; CHECK:   $r0 = MOVi 1, 14 /* CC::al */, $noreg, $noreg
-  ; CHECK:   $r1 = MOVi 1, 14 /* CC::al */, $noreg, $noreg
-  ; CHECK:   $r2 = MOVi 1, 14 /* CC::al */, $noreg, $noreg
-  ; CHECK:   $r3 = MOVi 1, 14 /* CC::al */, $noreg, $noreg
-  ; CHECK:   $r4 = MOVi 1, 14 /* CC::al */, $noreg, $noreg
-  ; CHECK:   $lr, $sp = frame-destroy LDR_POST_IMM $sp, $noreg, 8, 14 /* CC::al */, $noreg
-  ; CHECK:   MOVPCLR 14 /* CC::al */, $noreg
-
-  ; CHECK-LABEL: name: OUTLINED_FUNCTION_1
-  ; CHECK: bb.0:
-  ; CHECK:   liveins: $r11, $r10, $r9, $r8, $r7, $r6, $r5, $d15, $d14, $d13, $d12, $d11, $d10, $d9, $d8, $lr
-  ; CHECK:   early-clobber $sp = frame-setup STR_PRE_IMM killed $lr, $sp, -8, 14 /* CC::al */, $noreg
-  ; CHECK:   frame-setup CFI_INSTRUCTION def_cfa_offset 8
-  ; CHECK:   frame-setup CFI_INSTRUCTION offset $lr, -8
-  ; CHECK:   BL @bar, implicit-def dead $lr, implicit $sp
+  ; CHECK:   liveins: $r11, $r10, $r9, $r8, $r7, $r6, $r5, $d15, $d14, $d13, $d12, $d11, $d10, $d9, $d8
   ; CHECK:   $r0 = MOVi 2, 14 /* CC::al */, $noreg, $noreg
   ; CHECK:   $r1 = MOVi 2, 14 /* CC::al */, $noreg, $noreg
   ; CHECK:   $r2 = MOVi 2, 14 /* CC::al */, $noreg, $noreg
   ; CHECK:   $r3 = MOVi 2, 14 /* CC::al */, $noreg, $noreg
   ; CHECK:   $r4 = MOVi 2, 14 /* CC::al */, $noreg, $noreg
-  ; CHECK:   $lr, $sp = frame-destroy LDR_POST_IMM $sp, $noreg, 8, 14 /* CC::al */, $noreg
   ; CHECK:   TAILJMPd @bar, implicit $sp
 
-  ; CHECK-LABEL: name: OUTLINED_FUNCTION_2
+  ; CHECK-LABEL: name: OUTLINED_FUNCTION_1
   ; CHECK: bb.0:
   ; CHECK:   liveins: $r11, $r10, $r9, $r8, $r7, $r6, $r5, $d15, $d14, $d13, $d12, $d11, $d10, $d9, $d8
   ; CHECK:   $r0 = MOVi 3, 14 /* CC::al */, $noreg, $noreg
@@ -348,31 +327,28 @@ body:             |
   ; CHECK:   $r4 = MOVi 3, 14 /* CC::al */, $noreg, $noreg
   ; CHECK:   MOVPCLR 14 /* CC::al */, $noreg
 
+  ; CHECK-LABEL: name: OUTLINED_FUNCTION_2
+  ; CHECK: bb.0:
+  ; CHECK:   liveins: $r11, $r10, $r9, $r8, $r7, $r6, $r5, $d15, $d14, $d13, $d12, $d11, $d10, $d9, $d8
+  ; CHECK:   $r0 = MOVi 1, 14 /* CC::al */, $noreg, $noreg
+  ; CHECK:   $r1 = MOVi 1, 14 /* CC::al */, $noreg, $noreg
+  ; CHECK:   $r2 = MOVi 1, 14 /* CC::al */, $noreg, $noreg
+  ; CHECK:   $r3 = MOVi 1, 14 /* CC::al */, $noreg, $noreg
+  ; CHECK:   $r4 = MOVi 1, 14 /* CC::al */, $noreg, $noreg
+  ; CHECK:   MOVPCLR 14 /* CC::al */, $noreg
+
   ; CHECK-LABEL: name: OUTLINED_FUNCTION_3
   ; CHECK: bb.0:
-  ; CHECK:   liveins: $r11, $r10, $r9, $r8, $r6, $r5, $r4, $d15, $d14, $d13, $d12, $d11, $d10, $d9, $d8, $lr
-  ; CHECK:   early-clobber $sp = frame-setup t2STR_PRE killed $lr, $sp, -8, 14 /* CC::al */, $noreg
-  ; CHECK:   frame-setup CFI_INSTRUCTION def_cfa_offset 8
-  ; CHECK:   frame-setup CFI_INSTRUCTION offset $lr, -8
-  ; CHECK:   tBL 14 /* CC::al */, $noreg, @bar, implicit-def dead $lr, implicit $sp
+  ; CHECK:   liveins: $r11, $r10, $r9, $r8, $r6, $r5, $r4, $d15, $d14, $d13, $d12, $d11, $d10, $d9, $d8
   ; CHECK:   $r0 = t2MOVi 2, 14 /* CC::al */, $noreg, $noreg
   ; CHECK:   $r1 = t2MOVi 2, 14 /* CC::al */, $noreg, $noreg
   ; CHECK:   $r2 = t2MOVi 2, 14 /* CC::al */, $noreg, $noreg
-  ; CHECK:   $lr, $sp = frame-destroy t2LDR_POST $sp, 8, 14 /* CC::al */, $noreg
   ; CHECK:   tTAILJMPdND @bar, 14 /* CC::al */, $noreg, implicit $sp
 
   ; CHECK-LABEL: name: OUTLINED_FUNCTION_4
   ; CHECK: bb.0:
-  ; CHECK:   liveins: $r11, $r10, $r9, $r8, $r6, $r5, $r4, $d15, $d14, $d13, $d12, $d11, $d10, $d9, $d8, $lr
-  ; CHECK:   early-clobber $sp = frame-setup t2STR_PRE killed $lr, $sp, -8, 14 /* CC::al */, $noreg
-  ; CHECK:   frame-setup CFI_INSTRUCTION def_cfa_offset 8
-  ; CHECK:   frame-setup CFI_INSTRUCTION offset $lr, -8
-  ; CHECK:   tBL 14 /* CC::al */, $noreg, @bar, implicit-def dead $lr, implicit $sp
+  ; CHECK:   liveins: $r11, $r10, $r9, $r8, $r6, $r5, $r4, $d15, $d14, $d13, $d12, $d11, $d10, $d9, $d8
   ; CHECK:   $r0 = t2MOVi 1, 14 /* CC::al */, $noreg, $noreg
   ; CHECK:   $r1 = t2MOVi 1, 14 /* CC::al */, $noreg, $noreg
   ; CHECK:   $r2 = t2MOVi 1, 14 /* CC::al */, $noreg, $noreg
-  ; CHECK:   $lr, $sp = frame-destroy t2LDR_POST $sp, 8, 14 /* CC::al */, $noreg
   ; CHECK:   tBX_RET 14 /* CC::al */, $noreg
-
-
-
diff --git a/llvm/test/CodeGen/ARM/machine-outliner-default.mir b/llvm/test/CodeGen/ARM/machine-outliner-default.mir
index 6d0218dbfe636d..de2b8f55969765 100644
--- a/llvm/test/CodeGen/ARM/machine-outliner-default.mir
+++ b/llvm/test/CodeGen/ARM/machine-outliner-default.mir
@@ -19,17 +19,17 @@ body:             |
   ; CHECK: bb.0:
   ; CHECK:   liveins: $lr
   ; CHECK:   early-clobber $sp = frame-setup STR_PRE_IMM killed $lr, $sp, -8, 14 /* CC::al */, $noreg
-  ; CHECK:   BL @OUTLINED_FUNCTION_0
+  ; CHECK:   BL @OUTLINED_FUNCTION_1
   ; CHECK:   $lr, $sp = frame-destroy LDR_POST_IMM $sp, $noreg, 8, 14 /* CC::al */, $noreg
   ; CHECK: bb.1:
   ; CHECK:   liveins: $lr, $r6, $r7, $r8, $r9, $r10, $r11
   ; CHECK:   early-clobber $sp = frame-setup STR_PRE_IMM killed $lr, $sp, -8, 14 /* CC::al */, $noreg
-  ; CHECK:   BL @OUTLINED_FUNCTION_0
+  ; CHECK:   BL @OUTLINED_FUNCTION_1
   ; CHECK:   $lr, $sp = frame-destroy LDR_POST_IMM $sp, $noreg, 8, 14 /* CC::al */, $noreg
   ; CHECK: bb.2:
   ; CHECK:   liveins: $lr, $r6, $r7, $r8, $r9, $r10, $r11
   ; CHECK:   early-clobber $sp = frame-setup STR_PRE_IMM killed $lr, $sp, -8, 14 /* CC::al */, $noreg
-  ; CHECK:   BL @OUTLINED_FUNCTION_0
+  ; CHECK:   BL @OUTLINED_FUNCTION_1
   ; CHECK:   $lr, $sp = frame-destroy LDR_POST_IMM $sp, $noreg, 8, 14 /* CC::al */, $noreg
   ; CHECK: bb.3:
   ; CHECK:   liveins: $lr, $r6, $r7, $r8, $r9, $r10, $r11
@@ -73,17 +73,17 @@ body:             |
   ; CHECK: bb.0:
   ; CHECK:   liveins: $lr
   ; CHECK:   early-clobber $sp = frame-setup t2STR_PRE killed $lr, $sp, -8, 14 /* CC::al */, $noreg
-  ; CHECK:   tBL 14 /* CC::al */, $noreg, @OUTLINED_FUNCTION_1
+  ; CHECK:   tBL 14 /* CC::al */, $noreg, @OUTLINED_FUNCTION_0
   ; CHECK:   $lr, $sp = frame-destroy t2LDR_POST $sp, 8, 14 /* CC::al */, $noreg
   ; CHECK: bb.1:
   ; CHECK:   liveins: $lr, $r4, $r5, $r6, $r7, $r8, $r9, $r10, $r11
   ; CHECK:   early-clobber $sp = frame-setup t2STR_PRE killed $lr, $sp, -8, 14 /* CC::al */, $noreg
-  ; CHECK:   tBL 14 /* CC::al */, $noreg, @OUTLINED_FUNCTION_1
+  ; CHECK:   tBL 14 /* CC::al */, $noreg, @OUTLINED_FUNCTION_0
   ; CHECK:   $lr, $sp = frame-destroy t2LDR_POST $sp, 8, 14 /* CC::al */, $noreg
   ; CHECK: bb.2:
   ; CHECK:   liveins: $lr, $r4, $r5, $r6, $r7, $r8, $r9, $r10, $r11
   ; CHECK:   early-clobber $sp = frame-setup t2STR_PRE killed $lr, $sp, -8, 14 /* CC::al */, $noreg
-  ; CHECK:   tBL 14 /* CC::al */, $noreg, @OUTLINED_FUNCTION_1
+  ; CHECK:   tBL 14 /* CC::al */, $noreg, @OUTLINED_FUNCTION_0
   ; CHECK:   $lr, $sp = frame-destroy t2LDR_POST $sp, 8, 14 /* CC::al */, $noreg
   ; CHECK: bb.3:
   ; CHECK:   liveins: $lr, $r4, $r5, $r6, $r7, $r8, $r9, $r10, $r11
@@ -114,6 +114,15 @@ body:             |
 
   ; CHECK-LABEL: name: OUTLINED_FUNCTION_0
   ; CHECK: bb.0:
+  ; CHECK:   liveins: $lr, $r4, $r5, $r6, $r7, $r8, $r9, $r10, $r11
+  ; CHECK:   $r0 = t2MOVi 1, 14 /* CC::al */, $noreg, $noreg
+  ; CHECK:   $r1 = t2MOVi 1, 14 /* CC::al */, $noreg, $noreg
+  ; CHECK:   $r2 = t2MOVi 1, 14 /* CC::al */, $noreg, $noreg
+  ; CHECK:   $r3 = t2MOVi 1, 14 /* CC::al */, $noreg, $noreg
+  ; CHECK:   tBX_RET 14 /* CC::al */, $noreg
+
+  ; CHECK-LABEL: name: OUTLINED_FUNCTION_1
+  ; CHECK: bb.0:
   ; CHECK:   liveins: $lr, $r6, $r7, $r8, $r9, $r10, $r11
   ; CHECK:   $r0 = MOVi 1, 14 /* CC::al */, $noreg, $noreg
   ; CHECK:   $r1 = MOVi 1, 14 /* CC::al */, $noreg, $noreg
@@ -122,15 +131,3 @@ body:             |
   ; CHECK:   $r4 = MOVi 1, 14 /* CC::al */, $noreg, $noreg
   ; CHECK:   $r5 = MOVi 1, 14 /* CC::al */, $noreg, $noreg
   ; CHECK:   MOVPCLR 14 /* CC::al */, $noreg
-
-  ; CHECK-LABEL: name: OUTLINED_FUNCTION_1
-  ; CHECK: bb.0:
-  ; CHECK:   liveins: $lr, $r4, $r5, $r6, $r7, $r8, $r9, $r10, $r11
-  ; CHECK:   $r0 = t2MOVi 1, 14 /* CC::al */, $noreg, $noreg
-  ; CHECK:   $r1 = t2MOVi 1, 14 /* CC::al */, $noreg, $noreg
-  ; CHECK:   $r2 = t2MOVi 1, 14 /* CC::al */, $noreg, $noreg
-  ; CHECK:   $r3 = t2MOVi 1, 14 /* CC::al */, $noreg, $noreg
-  ; CHECK:   tBX_RET 14 /* CC::al */, $noreg
-
-
-
diff --git a/llvm/test/CodeGen/ARM/machine-outliner-stack-fixup-arm.mir b/llvm/test/CodeGen/ARM/machine-outliner-stack-fixup-arm.mir
index ae5caa5b7c06db..e71edc8ceb3f60 100644
--- a/llvm/test/CodeGen/ARM/machine-outliner-stack-fixup-arm.mir
+++ b/llvm/test/CodeGen/ARM/machine-outliner-stack-fixup-arm.mir
@@ -18,6 +18,7 @@ body:             |
     liveins: $r0
     ; CHECK-LABEL: name:           CheckAddrMode_i12
     ; CHECK: $r1 = MOVr killed $r0, 14 /* CC::al */, $noreg, $noreg
+    ; CHECK-NEXT: BL @foo, implicit-def dead $lr, implicit $sp
     ; CHECK-NEXT: BL @OUTLINED_FUNCTION_[[I12:[0-9]+]]
     ; CHECK-NEXT: $r6 = LDRi12 $sp, 4088, 14 /* CC::al */, $noreg
     $r1 = MOVr killed $r0, 14, $noreg, $noreg
@@ -47,6 +48,7 @@ body:             |
     liveins: $r1
     ; CHECK-LABEL: name:           CheckAddrMode3
     ; CHECK: $r0 = MOVr killed $r1, 14 /* CC::al */, $noreg, $noreg
+    ; CHECK-NEXT: BL @foo, implicit-def dead $lr, implicit $sp
     ; CHECK-NEXT: BL @OUTLINED_FUNCTION_[[I3:[0-9]+]]
     ; CHECK-NEXT: $r6 = LDRSH $sp, $noreg, 248, 14 /* CC::al */, $noreg
     $r0 = MOVr killed $r1, 14, $noreg, $noreg
@@ -76,6 +78,7 @@ body:             |
     liveins: $r2
     ; CHECK-LABEL: name:           CheckAddrMode5
     ; CHECK: $r0 = MOVr killed $r2, 14 /* CC::al */, $noreg, $noreg
+    ; CHECK-NEXT: BL @foo, implicit-def dead $lr, implicit $sp
     ; CHECK-NEXT: BL @OUTLINED_FUNCTION_[[I5:[0-9]+]]
     ; CHECK-NEXT: $d5 = VLDRD $sp, 254, 14 /* CC::al */, $noreg
     $r0 = MOVr killed $r2, 14, $noreg, $noreg
@@ -110,6 +113,7 @@ body:             |
     liveins: $r3
     ; CHECK-LABEL: name:           CheckAddrMode5FP16
     ; CHECK: $r0 = MOVr killed $r3, 14 /* CC::al */, $noreg, $noreg
+    ; CHECK-NEXT: BL @foo, implicit-def dead $lr, implicit $sp
     ; CHECK-NEXT: BL @OUTLINED_FUNCTION_[[I5FP16:[0-9]+]]
     ; CHECK-NEXT: $s6 = VLDRH $sp, 252, 14, $noreg
     $r0 = MOVr killed $r3, 14, $noreg, $noreg
@@ -146,41 +150,29 @@ body:             |
     BX_RET 14, $noreg
 
     ;CHECK: name:           OUTLINED_FUNCTION_[[I5]]
-    ;CHECK: early-clobber $sp = frame-setup STR_PRE_IMM killed $lr, $sp, -8, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: frame-setup CFI_INSTRUCTION def_cfa_offset 8
-    ;CHECK-NEXT: frame-setup CFI_INSTRUCTION offset $lr, -8
-    ;CHECK-NEXT: BL @foo, implicit-def dead $lr, implicit $sp
-    ;CHECK-NEXT: $d0 = VLDRD $sp, 2, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: $d1 = VLDRD $sp, 10, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: $d4 = VLDRD $sp, 255, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: $lr, $sp = frame-destroy LDR_POST_IMM $sp, $noreg, 8, 14 /* CC::al */, $noreg
+    ;CHECK: liveins: $r10, $r9, $r8, $r7, $r6, $r5, $r4, $d15, $d14, $d13, $d12, $d11, $d10, $d9, $d8
+    ;CHECK: $d0 = VLDRD $sp, 0, 14 /* CC::al */, $noreg
+    ;CHECK-NEXT: $d1 = VLDRD $sp, 8, 14 /* CC::al */, $noreg
+    ;CHECK-NEXT: $d4 = VLDRD $sp, 253, 14 /* CC::al */, $noreg
+    ;CHECK-NEXT: BX_RET 14 /* CC::al */, $noreg
 
     ;CHECK: name:           OUTLINED_FUNCTION_[[I5FP16]]
-    ;CHECK: early-clobber $sp = frame-setup STR_PRE_IMM killed $lr, $sp, -8, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: frame-setup CFI_INSTRUCTION def_cfa_offset 8
-    ;CHECK-NEXT: frame-setup CFI_INSTRUCTION offset $lr, -8
-    ;CHECK-NEXT: BL @foo, implicit-def dead $lr, implicit $sp
-    ;CHECK-NEXT: $s1 = VLDRH $sp, 4, 14, $noreg
-    ;CHECK-NEXT: $s2 = VLDRH $sp, 12, 14, $noreg
-    ;CHECK-NEXT: $s5 = VLDRH $sp, 244, 14, $noreg
-    ;CHECK-NEXT: $lr, $sp = frame-destroy LDR_POST_IMM $sp, $noreg, 8, 14 /* CC::al */, $noreg
+    ;CHECK: liveins: $r10, $r9, $r8, $r7, $r6, $r5, $r4, $d15, $d14, $d13, $d12, $d11, $d10, $d9, $d8
+    ;CHECK: $s1 = VLDRH $sp, 0, 14, $noreg
+    ;CHECK-NEXT: $s2 = VLDRH $sp, 8, 14, $noreg
+    ;CHECK-NEXT: $s5 = VLDRH $sp, 240, 14, $noreg
+    ;CHECK-NEXT: BX_RET 14 /* CC::al */, $noreg
 
     ;CHECK: name:           OUTLINED_FUNCTION_[[I12]]
-    ;CHECK: early-clobber $sp = frame-setup STR_PRE_IMM killed $lr, $sp, -8, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: frame-setup CFI_INSTRUCTION def_cfa_offset 8
-    ;CHECK-NEXT: frame-setup CFI_INSTRUCTION offset $lr, -8
-    ;CHECK-NEXT: BL @foo, implicit-def dead $lr, implicit $sp
-    ;CHECK-NEXT: $r1 = LDRi12 $sp, 8, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: $r2 = LDRi12 $sp, 16, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: $r5 = LDRi12 $sp, 4094, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: $lr, $sp = frame-destroy LDR_POST_IMM $sp, $noreg, 8, 14 /* CC::al */, $noreg
+    ;CHECK: liveins: $r10, $r9, $r8, $r7, $d8, $r4, $d15, $d14, $d13, $d12, $d11, $d10, $d9
+    ;CHECK: $r1 = LDRi12 $sp, 0, 14 /* CC::al */, $noreg
+    ;CHECK-NEXT: $r2 = LDRi12 $sp, 8, 14 /* CC::al */, $noreg
+    ;CHECK-NEXT: $r5 = LDRi12 $sp, 4086, 14 /* CC::al */, $noreg
+    ;CHECK-NEXT: BX_RET 14 /* CC::al */, $noreg
 
     ;CHECK: name:           OUTLINED_FUNCTION_[[I3]]
-    ;CHECK: early-clobber $sp = frame-setup STR_PRE_IMM killed $lr, $sp, -8, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: frame-setup CFI_INSTRUCTION def_cfa_offset 8
-    ;CHECK-NEXT: frame-setup CFI_INSTRUCTION offset $lr, -8
-    ;CHECK-NEXT: BL @foo, implicit-def dead $lr, implicit $sp
-    ;CHECK-NEXT: $r1 = LDRSH $sp, $noreg, 8, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: $r2 = LDRSH $sp, $noreg, 16, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: $r5 = LDRSH $sp, $noreg, 255, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: $lr, $sp = frame-destroy LDR_POST_IMM $sp, $noreg, 8, 14 /* CC::al */, $noreg
+    ;CHECK: liveins: $r10, $r9, $r8, $r7, $d8, $r4, $d15, $d14, $d13, $d12, $d11, $d10, $d9
+    ;CHECK: $r1 = LDRSH $sp, $noreg, 0, 14 /* CC::al */, $noreg
+    ;CHECK-NEXT: $r2 = LDRSH $sp, $noreg, 8, 14 /* CC::al */, $noreg
+    ;CHECK-NEXT: $r5 = LDRSH $sp, $noreg, 247, 14 /* CC::al */, $noreg
+    ;CHECK-NEXT: BX_RET 14 /* CC::al */, $noreg
diff --git a/llvm/test/CodeGen/ARM/machine-outliner-stack-fixup-thumb.mir b/llvm/test/CodeGen/ARM/machine-outliner-stack-fixup-thumb.mir
index 56184448424489..f2cdaebdfa8baa 100644
--- a/llvm/test/CodeGen/ARM/machine-outliner-stack-fixup-thumb.mir
+++ b/llvm/test/CodeGen/ARM/machine-outliner-stack-fixup-thumb.mir
@@ -19,7 +19,7 @@ body:             |
   bb.0:
     liveins: $r1
     ;CHECK-LABEL: name:           CheckAddrModeT2_i12
-    ;CHECK: $r0 = tMOVr killed $r1, 14 /* CC::al */, $noreg
+    ;CHECK: tBL 14 /* CC::al */, $noreg, @OUTLINED_FUNCTION_[[SHARED:[0-9+]]]
     ;CHECK-NEXT: tBL 14 /* CC::al */, $noreg, @OUTLINED_FUNCTION_[[I12:[0-9]+]]
     ;CHECK-NEXT: $r0 = t2LDRi12 $sp, 4088, 14 /* CC::al */, $noreg
     $r0 = tMOVr killed $r1, 14, $noreg
@@ -49,7 +49,7 @@ body:             |
   bb.0:
     liveins: $r1
     ;CHECK-LABEL: name:           CheckAddrModeT2_i8
-    ;CHECK: $r0 = tMOVr $r1, 14 /* CC::al */, $noreg
+    ;CHECK: tBL 14 /* CC::al */, $noreg, @OUTLINED_FUNCTION_[[SHARED:[0-9+]]]
     ;CHECK-NEXT: tBL 14 /* CC::al */, $noreg, @OUTLINED_FUNCTION_[[I8:[0-9]+]]
     ;CHECK-NEXT: t2STRHT $r0, $sp, 248, 14 /* CC::al */, $noreg
     $r0 = tMOVr $r1, 14, $noreg
@@ -79,7 +79,7 @@ body:             |
   bb.0:
     liveins: $r1
     ;CHECK-LABEL: name:           CheckAddrModeT2_i8s4
-    ;CHECK: $r0 = tMOVr $r1, 14 /* CC::al */, $noreg
+    ;CHECK: tBL 14 /* CC::al */, $noreg, @OUTLINED_FUNCTION_[[SHARED:[0-9+]]]
     ;CHECK-NEXT: tBL 14 /* CC::al */, $noreg, @OUTLINED_FUNCTION_[[I8S4:[0-9]+]]
     ;CHECK-NEXT: t2STRDi8 $r0, $r1, $sp, 1020, 14 /* CC::al */, $noreg
     $r0 = tMOVr $r1, 14, $noreg
@@ -109,7 +109,7 @@ body:             |
   bb.0:
     liveins: $r1
     ;CHECK-LABEL: name:           CheckAddrModeT2_ldrex
-    ;CHECK: $r0 = tMOVr $r1, 14 /* CC::al */, $noreg
+    ;CHECK: tBL 14 /* CC::al */, $noreg, @OUTLINED_FUNCTION_[[SHARED:[0-9+]]]
     ;CHECK-NEXT: tBL 14 /* CC::al */, $noreg, @OUTLINED_FUNCTION_[[LDREX:[0-9]+]]
     ;CHECK-NEXT: $r1 = t2LDREX $sp, 254, 14 /* CC::al */, $noreg
     $r0 = tMOVr $r1, 14, $noreg
@@ -144,8 +144,10 @@ body:             |
   bb.0:
     liveins: $r0, $r1
     ;CHECK-LABEL: name:           CheckAddrModeT1_s
-    ;CHECK: $r0 = tMOVr $r1, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: tBL 14 /* CC::al */, $noreg, @OUTLINED_FUNCTION_[[T1_S:[0-9]+]]
+    ;CHECK: tBL 14 /* CC::al */, $noreg, @OUTLINED_FUNCTION_[[SHARED:[0-9+]]]
+    ;CHECK-NEXT: tSTRspi $r0, $sp, 0, 14 /* CC::al */, $noreg
+    ;CHECK-NEXT: tSTRspi $r0, $sp, 4, 14 /* CC::al */, $noreg
+    ;CHECK-NEXT: tSTRspi $r0, $sp, 253, 14 /* CC::al */, $noreg
     ;CHECK-NEXT: tSTRspi $r0, $sp, 254, 14 /* CC::al */, $noreg
     $r0 = tMOVr $r1, 14, $noreg
     tBL 14, $noreg, @foo, implicit-def dead $lr, implicit $sp
@@ -181,51 +183,29 @@ body:             |
     BX_RET 14, $noreg
 
     ;CHECK: name:           OUTLINED_FUNCTION_[[LDREX]]
-    ;CHECK: early-clobber $sp = frame-setup t2STR_PRE killed $lr, $sp, -8, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: frame-setup CFI_INSTRUCTION def_cfa_offset 8
-    ;CHECK-NEXT: frame-setup CFI_INSTRUCTION offset $lr, -8
-    ;CHECK-NEXT: tBL 14 /* CC::al */, $noreg, @foo, implicit-def dead $lr, implicit $sp
-    ;CHECK-NEXT: $r1 = t2LDREX $sp, 2, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: $r1 = t2LDREX $sp, 10, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: $r1 = t2LDREX $sp, 255, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: $lr, $sp = frame-destroy t2LDR_POST $sp, 8, 14 /* CC::al */, $noreg
+    ;CHECK: $r1 = t2LDREX $sp, 0, 14 /* CC::al */, $noreg
+    ;CHECK-NEXT: $r1 = t2LDREX $sp, 8, 14 /* CC::al */, $noreg
+    ;CHECK-NEXT: $r1 = t2LDREX $sp, 253, 14 /* CC::al */, $noreg
+    ;CHECK-NEXT: tBX_RET 14 /* CC::al */, $noreg
 
     ;CHECK: name:           OUTLINED_FUNCTION_[[I8]]
-    ;CHECK: early-clobber $sp = frame-setup t2STR_PRE killed $lr, $sp, -8, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: frame-setup CFI_INSTRUCTION def_cfa_offset 8
-    ;CHECK-NEXT: frame-setup CFI_INSTRUCTION offset $lr, -8
-    ;CHECK-NEXT: tBL 14 /* CC::al */, $noreg, @foo, implicit-def dead $lr, implicit $sp
-    ;CHECK-NEXT: t2STRHT $r0, $sp, 8, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: t2STRHT $r0, $sp, 12, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: t2STRHT $r0, $sp, 255, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: $lr, $sp = frame-destroy t2LDR_POST $sp, 8, 14 /* CC::al */, $noreg
+    ;CHECK: t2STRHT $r0, $sp, 0, 14 /* CC::al */, $noreg
+    ;CHECK-NEXT: t2STRHT $r0, $sp, 4, 14 /* CC::al */, $noreg
+    ;CHECK-NEXT: t2STRHT $r0, $sp, 247, 14 /* CC::al */, $noreg
+    ;CHECK-NEXT: tBX_RET 14 /* CC::al */, $noreg
 
     ;CHECK: name:           OUTLINED_FUNCTION_[[I8S4]]
-    ;CHECK: early-clobber $sp = frame-setup t2STR_PRE killed $lr, $sp, -8, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: frame-setup CFI_INSTRUCTION def_cfa_offset 8
-    ;CHECK-NEXT: frame-setup CFI_INSTRUCTION offset $lr, -8
-    ;CHECK-NEXT: tBL 14 /* CC::al */, $noreg, @foo, implicit-def dead $lr, implicit $sp
+    ;CHECK: t2STRDi8 $r0, $r1, $sp, 0, 14 /* CC::al */, $noreg
     ;CHECK-NEXT: t2STRDi8 $r0, $r1, $sp, 8, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: t2STRDi8 $r0, $r1, $sp, 16, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: t2STRDi8 $r0, $r1, $sp, 1020, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: $lr, $sp = frame-destroy t2LDR_POST $sp, 8, 14 /* CC::al */, $noreg
+    ;CHECK-NEXT: t2STRDi8 $r0, $r1, $sp, 1012, 14 /* CC::al */, $noreg
+    ;CHECK-NEXT: tBX_RET 14 /* CC::al */, $nore
 
     ;CHECK: name:           OUTLINED_FUNCTION_[[I12]]
-    ;CHECK: early-clobber $sp = frame-setup t2STR_PRE killed $lr, $sp, -8, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: frame-setup CFI_INSTRUCTION def_cfa_offset 8
-    ;CHECK-NEXT: frame-setup CFI_INSTRUCTION offset $lr, -8
-    ;CHECK-NEXT: tBL 14 /* CC::al */, $noreg, @foo, implicit-def dead $lr, implicit $sp
-    ;CHECK-NEXT: $r0 = t2LDRi12 $sp, 8, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: $r0 = t2LDRi12 $sp, 12, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: $r0 = t2LDRi12 $sp, 4094, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: $lr, $sp = frame-destroy t2LDR_POST $sp, 8, 14 /* CC::al */, $noreg
-
-    ;CHECK: name:           OUTLINED_FUNCTION_[[T1_S]]
-    ;CHECK: early-clobber $sp = frame-setup t2STR_PRE killed $lr, $sp, -8, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: frame-setup CFI_INSTRUCTION def_cfa_offset 8
-    ;CHECK-NEXT: frame-setup CFI_INSTRUCTION offset $lr, -8
-    ;CHECK-NEXT: tBL 14 /* CC::al */, $noreg, @foo, implicit-def dead $lr, implicit $sp
-    ;CHECK-NEXT: tSTRspi $r0, $sp, 2, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: tSTRspi $r0, $sp, 6, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: tSTRspi $r0, $sp, 255, 14 /* CC::al */, $noreg
-    ;CHECK-NEXT: $lr, $sp = frame-destroy t2LDR_POST $sp, 8, 14 /* CC::al */, $noreg
+    ;CHECK: $r0 = t2LDRi12 $sp, 0, 14 /* CC::al */, $noreg
+    ;CHECK-NEXT: $r0 = t2LDRi12 $sp, 4, 14 /* CC::al */, $noreg
+    ;CHECK-NEXT: $r0 = t2LDRi12 $sp, 4086, 14 /* CC::al */, $noreg
+    ;CHECK-NEXT: tBX_RET 14 /* CC::al */, $noreg
+
+    ;CHECK: name:           OUTLINED_FUNCTION_[[SHARED]]
+    ;CHECK: $r0 = tMOVr killed $r1, 14 /* CC::al */, $noreg
+    ;CHECK-NEXT: tTAILJMPdND @foo, 14 /* CC::al */, $noreg, implicit $sp

>From 75c515f4456237ac5c3eba3c3186b94c9503443c Mon Sep 17 00:00:00 2001
From: Xuan Zhang <xuanzh at meta.com>
Date: Fri, 12 Apr 2024 10:57:51 -0700
Subject: [PATCH 3/5] consider leaf descendants to include more candidates for
 outlining

---
 llvm/include/llvm/Support/SuffixTree.h        |  34 ++++-
 llvm/include/llvm/Support/SuffixTreeNode.h    |  25 +++-
 llvm/lib/CodeGen/MachineOutliner.cpp          |   8 +-
 llvm/lib/Support/SuffixTree.cpp               |  83 ++++++++++-
 llvm/lib/Support/SuffixTreeNode.cpp           |   5 +
 .../machine-outliner-cfi-tail-some.mir        |   2 +-
 .../machine-outliner-leaf-descendants.ll      | 124 ++++++++++++++++
 .../AArch64/machine-outliner-overlap.mir      |  24 +++-
 .../machine-outliner-retaddr-sign-sp-mod.mir  |   2 +-
 .../machine-outliner-retaddr-sign-thunk.ll    |   4 +-
 .../AArch64/machine-outliner-throw2.ll        |   4 +-
 .../CodeGen/AArch64/machine-outliner-thunk.ll |   2 +-
 .../test/CodeGen/AArch64/machine-outliner.mir |   2 +-
 .../RISCV/machineoutliner-pcrel-lo.mir        |   8 +-
 llvm/unittests/Support/SuffixTreeTest.cpp     | 134 ++++++++++++++++++
 15 files changed, 441 insertions(+), 20 deletions(-)
 create mode 100644 llvm/test/CodeGen/AArch64/machine-outliner-leaf-descendants.ll

diff --git a/llvm/include/llvm/Support/SuffixTree.h b/llvm/include/llvm/Support/SuffixTree.h
index 4940fbbf308d8b..37b73666404300 100644
--- a/llvm/include/llvm/Support/SuffixTree.h
+++ b/llvm/include/llvm/Support/SuffixTree.h
@@ -42,6 +42,9 @@ class SuffixTree {
   /// Each element is an integer representing an instruction in the module.
   ArrayRef<unsigned> Str;
 
+  /// Whether to consider leaf descendants or only leaf children.
+  bool OutlinerLeafDescendants;
+
   /// A repeated substring in the tree.
   struct RepeatedSubstring {
     /// The length of the string.
@@ -130,11 +133,27 @@ class SuffixTree {
   /// this step.
   unsigned extend(unsigned EndIdx, unsigned SuffixesToAdd);
 
+  /// This vector contains all leaf nodes of this suffix tree. These leaf nodes
+  /// are identified using post-order depth-first traversal, so that the order
+  /// of these leaf nodes in the vector matches the order of the leaves in the
+  /// tree from left to right if one were to draw the tree on paper.
+  std::vector<SuffixTreeLeafNode *> LeafNodes;
+
+  /// Perform a post-order depth-first traversal of the tree and perform two
+  /// tasks during the traversal. The first is to populate LeafNodes, adding
+  /// nodes in order of the traversal. The second is to keep track of the leaf
+  /// descendants of every internal node by assigning values to LeftLeafIndex
+  /// and RightLefIndex fields of SuffixTreeNode for all internal nodes.
+  void setLeafNodes();
+
 public:
   /// Construct a suffix tree from a sequence of unsigned integers.
   ///
   /// \param Str The string to construct the suffix tree for.
-  SuffixTree(const ArrayRef<unsigned> &Str);
+  /// \param OutlinerLeafDescendants Whether to consider leaf descendants or
+  /// only leaf children (used by Machine Outliner).
+  SuffixTree(const ArrayRef<unsigned> &Str,
+             bool OutlinerLeafDescendants = false);
 
   /// Iterator for finding all repeated substrings in the suffix tree.
   struct RepeatedSubstringIterator {
@@ -154,6 +173,12 @@ class SuffixTree {
     /// instruction lengths.
     const unsigned MinLength = 2;
 
+    /// Vector of leaf nodes of the suffix tree.
+    const std::vector<SuffixTreeLeafNode *> &LeafNodes;
+
+    /// Whether to consider leaf descendants or only leaf children.
+    bool OutlinerLeafDescendants = !LeafNodes.empty();
+
     /// Move the iterator to the next repeated substring.
     void advance();
 
@@ -179,7 +204,10 @@ class SuffixTree {
       return !(*this == Other);
     }
 
-    RepeatedSubstringIterator(SuffixTreeInternalNode *N) : N(N) {
+    RepeatedSubstringIterator(
+        SuffixTreeInternalNode *N,
+        const std::vector<SuffixTreeLeafNode *> &LeafNodes = {})
+        : N(N), LeafNodes(LeafNodes) {
       // Do we have a non-null node?
       if (!N)
         return;
@@ -191,7 +219,7 @@ class SuffixTree {
   };
 
   typedef RepeatedSubstringIterator iterator;
-  iterator begin() { return iterator(Root); }
+  iterator begin() { return iterator(Root, LeafNodes); }
   iterator end() { return iterator(nullptr); }
 };
 
diff --git a/llvm/include/llvm/Support/SuffixTreeNode.h b/llvm/include/llvm/Support/SuffixTreeNode.h
index 7d0d1cf0c58b95..84b590f2deb0cd 100644
--- a/llvm/include/llvm/Support/SuffixTreeNode.h
+++ b/llvm/include/llvm/Support/SuffixTreeNode.h
@@ -46,6 +46,17 @@ struct SuffixTreeNode {
   /// the root to this node.
   unsigned ConcatLen = 0;
 
+  /// These two indices give a range of indices for its leaf descendants.
+  /// Imagine drawing a tree on paper and assigning a unique index to each leaf
+  /// node in monotonically increasing order from left to right. This way of
+  /// numbering the leaf nodes allows us to associate a continuous range of
+  /// indices with each internal node. For example, if a node has leaf
+  /// descendants with indices i, i+1, ..., j, then its LeftLeafIdx is i and
+  /// its RightLeafIdx is j. These indices are for LeafNodes in the SuffixTree
+  /// class, which is constructed using post-order depth-first traversal.
+  unsigned LeftLeafIdx = EmptyIdx;
+  unsigned RightLeafIdx = EmptyIdx;
+
 public:
   // LLVM RTTI boilerplate.
   NodeKind getKind() const { return Kind; }
@@ -56,6 +67,18 @@ struct SuffixTreeNode {
   /// \returns the end index of this node.
   virtual unsigned getEndIdx() const = 0;
 
+  /// \return the index of this node's left most leaf node.
+  unsigned getLeftLeafIdx() const;
+
+  /// \return the index of this node's right most leaf node.
+  unsigned getRightLeafIdx() const;
+
+  /// Set the index of the left most leaf node of this node to \p Idx.
+  void setLeftLeafIdx(unsigned Idx);
+
+  /// Set the index of the right most leaf node of this node to \p Idx.
+  void setRightLeafIdx(unsigned Idx);
+
   /// Advance this node's StartIdx by \p Inc.
   void incrementStartIdx(unsigned Inc);
 
@@ -168,4 +191,4 @@ struct SuffixTreeLeafNode : SuffixTreeNode {
   virtual ~SuffixTreeLeafNode() = default;
 };
 } // namespace llvm
-#endif // LLVM_SUPPORT_SUFFIXTREE_NODE_H
\ No newline at end of file
+#endif // LLVM_SUPPORT_SUFFIXTREE_NODE_H
diff --git a/llvm/lib/CodeGen/MachineOutliner.cpp b/llvm/lib/CodeGen/MachineOutliner.cpp
index bc21215a181f6c..01a9966adb097a 100644
--- a/llvm/lib/CodeGen/MachineOutliner.cpp
+++ b/llvm/lib/CodeGen/MachineOutliner.cpp
@@ -121,6 +121,12 @@ static cl::opt<unsigned> OutlinerBenefitThreshold(
     cl::desc(
         "The minimum size in bytes before an outlining candidate is accepted"));
 
+static cl::opt<bool> OutlinerLeafDescendants(
+    "outliner-leaf-descendants", cl::init(true), cl::Hidden,
+    cl::desc("Consider all leaf descendants of internal nodes of the suffix "
+             "tree as candidates for outlining (if false, only leaf children "
+             "are considered)"));
+
 namespace {
 
 /// Maps \p MachineInstrs to unsigned integers and stores the mappings.
@@ -576,7 +582,7 @@ void MachineOutliner::emitOutlinedFunctionRemark(OutlinedFunction &OF) {
 void MachineOutliner::findCandidates(
     InstructionMapper &Mapper, std::vector<OutlinedFunction> &FunctionList) {
   FunctionList.clear();
-  SuffixTree ST(Mapper.UnsignedVec);
+  SuffixTree ST(Mapper.UnsignedVec, OutlinerLeafDescendants);
 
   // First, find all of the repeated substrings in the tree of minimum length
   // 2.
diff --git a/llvm/lib/Support/SuffixTree.cpp b/llvm/lib/Support/SuffixTree.cpp
index eaa653078e0900..4bd5cdfe8791dd 100644
--- a/llvm/lib/Support/SuffixTree.cpp
+++ b/llvm/lib/Support/SuffixTree.cpp
@@ -11,9 +11,11 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Support/SuffixTree.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/SuffixTreeNode.h"
+#include <stack>
 
 using namespace llvm;
 
@@ -26,7 +28,9 @@ static size_t numElementsInSubstring(const SuffixTreeNode *N) {
   return N->getEndIdx() - N->getStartIdx() + 1;
 }
 
-SuffixTree::SuffixTree(const ArrayRef<unsigned> &Str) : Str(Str) {
+SuffixTree::SuffixTree(const ArrayRef<unsigned> &Str,
+                       bool OutlinerLeafDescendants)
+    : Str(Str), OutlinerLeafDescendants(OutlinerLeafDescendants) {
   Root = insertRoot();
   Active.Node = Root;
 
@@ -46,6 +50,11 @@ SuffixTree::SuffixTree(const ArrayRef<unsigned> &Str) : Str(Str) {
   // Set the suffix indices of each leaf.
   assert(Root && "Root node can't be nullptr!");
   setSuffixIndices();
+
+  // Collect all leaf nodes of the suffix tree. And for each internal node,
+  // record the range of leaf nodes that are descendants of it.
+  if (OutlinerLeafDescendants)
+    setLeafNodes();
 }
 
 SuffixTreeNode *SuffixTree::insertLeaf(SuffixTreeInternalNode &Parent,
@@ -105,6 +114,68 @@ void SuffixTree::setSuffixIndices() {
   }
 }
 
+void SuffixTree::setLeafNodes() {
+  // A stack that keeps track of nodes to visit for post-order DFS traversal.
+  std::stack<SuffixTreeNode *> ToVisit;
+  ToVisit.push(Root);
+
+  // This keeps track of the index of the next leaf node to be added to
+  // the LeafNodes vector of the suffix tree.
+  unsigned LeafCounter = 0;
+
+  // This keeps track of nodes whose children have been added to the stack
+  // during the post-order depth-first traversal of the tree.
+  llvm::SmallPtrSet<SuffixTreeInternalNode *, 32> ChildrenAddedToStack;
+
+  // Traverse the tree in post-order.
+  while (!ToVisit.empty()) {
+    SuffixTreeNode *CurrNode = ToVisit.top();
+    ToVisit.pop();
+    if (auto *CurrInternalNode = dyn_cast<SuffixTreeInternalNode>(CurrNode)) {
+      // The current node is an internal node.
+      if (ChildrenAddedToStack.find(CurrInternalNode) !=
+          ChildrenAddedToStack.end()) {
+        // If the children of the current node has been added to the stack,
+        // then this is the second time we visit this node and at this point,
+        // all of its children have already been processed. Now, we can
+        // set its LeftLeafIdx and RightLeafIdx;
+        auto it = CurrInternalNode->Children.begin();
+        if (it != CurrInternalNode->Children.end()) {
+          // Get the first child to use its RightLeafIdx. The RightLeafIdx is
+          // used as the first child is the initial one added to the stack, so
+          // it's the last one to be processed. This implies that the leaf
+          // descendants of the first child are assigned the largest index
+          // numbers.
+          CurrNode->setRightLeafIdx(it->second->getRightLeafIdx());
+          // get the last child to use its LeftLeafIdx.
+          while (std::next(it) != CurrInternalNode->Children.end())
+            it = std::next(it);
+          CurrNode->setLeftLeafIdx(it->second->getLeftLeafIdx());
+          assert(CurrNode->getLeftLeafIdx() <= CurrNode->getRightLeafIdx() &&
+                 "LeftLeafIdx should not be larger than RightLeafIdx");
+        }
+      } else {
+        // This is the first time we visit this node. This means that its
+        // children have not been added to the stack yet. Hence, we will add
+        // the current node back to the stack and add its children to the
+        // stack for processing.
+        ToVisit.push(CurrNode);
+        for (auto &ChildPair : CurrInternalNode->Children)
+          ToVisit.push(ChildPair.second);
+        ChildrenAddedToStack.insert(CurrInternalNode);
+      }
+    } else {
+      // The current node is a leaf node.
+      // We can simplyset its LeftLeafIdx and RightLeafIdx.
+      CurrNode->setLeftLeafIdx(LeafCounter);
+      CurrNode->setRightLeafIdx(LeafCounter);
+      LeafCounter++;
+      auto *CurrLeafNode = cast<SuffixTreeLeafNode>(CurrNode);
+      LeafNodes.push_back(CurrLeafNode);
+    }
+  }
+}
+
 unsigned SuffixTree::extend(unsigned EndIdx, unsigned SuffixesToAdd) {
   SuffixTreeInternalNode *NeedsLink = nullptr;
 
@@ -230,6 +301,7 @@ void SuffixTree::RepeatedSubstringIterator::advance() {
 
   // Each leaf node represents a repeat of a string.
   SmallVector<unsigned> RepeatedSubstringStarts;
+  SmallVector<SuffixTreeLeafNode *> LeafDescendants;
 
   // Continue visiting nodes until we find one which repeats more than once.
   while (!InternalNodesToVisit.empty()) {
@@ -252,7 +324,7 @@ void SuffixTree::RepeatedSubstringIterator::advance() {
         continue;
       }
 
-      if (Length < MinLength)
+      if (Length < MinLength || OutlinerLeafDescendants)
         continue;
 
       // Have an occurrence of a potentially repeated string. Save it.
@@ -260,6 +332,13 @@ void SuffixTree::RepeatedSubstringIterator::advance() {
       RepeatedSubstringStarts.push_back(Leaf->getSuffixIdx());
     }
 
+    if (OutlinerLeafDescendants && Length >= MinLength) {
+      LeafDescendants.assign(LeafNodes.begin() + Curr->getLeftLeafIdx(),
+                             LeafNodes.begin() + Curr->getRightLeafIdx() + 1);
+      for (SuffixTreeLeafNode *Leaf : LeafDescendants)
+        RepeatedSubstringStarts.push_back(Leaf->getSuffixIdx());
+    }
+
     // The root never represents a repeated substring. If we're looking at
     // that, then skip it.
     if (Curr->isRoot())
diff --git a/llvm/lib/Support/SuffixTreeNode.cpp b/llvm/lib/Support/SuffixTreeNode.cpp
index 113b990fd352fc..9f1f94a39895e8 100644
--- a/llvm/lib/Support/SuffixTreeNode.cpp
+++ b/llvm/lib/Support/SuffixTreeNode.cpp
@@ -38,3 +38,8 @@ unsigned SuffixTreeLeafNode::getEndIdx() const {
 
 unsigned SuffixTreeLeafNode::getSuffixIdx() const { return SuffixIdx; }
 void SuffixTreeLeafNode::setSuffixIdx(unsigned Idx) { SuffixIdx = Idx; }
+
+unsigned SuffixTreeNode::getLeftLeafIdx() const { return LeftLeafIdx; }
+unsigned SuffixTreeNode::getRightLeafIdx() const { return RightLeafIdx; }
+void SuffixTreeNode::setLeftLeafIdx(unsigned Idx) { LeftLeafIdx = Idx; }
+void SuffixTreeNode::setRightLeafIdx(unsigned Idx) { RightLeafIdx = Idx; }
diff --git a/llvm/test/CodeGen/AArch64/machine-outliner-cfi-tail-some.mir b/llvm/test/CodeGen/AArch64/machine-outliner-cfi-tail-some.mir
index 67d411962ce4f7..3afa1d5559a585 100644
--- a/llvm/test/CodeGen/AArch64/machine-outliner-cfi-tail-some.mir
+++ b/llvm/test/CodeGen/AArch64/machine-outliner-cfi-tail-some.mir
@@ -1,5 +1,5 @@
 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
-# RUN: llc -mtriple=aarch64-apple-unknown -run-pass=machine-outliner -verify-machineinstrs %s -o - | FileCheck %s
+# RUN: llc -mtriple=aarch64-apple-unknown -run-pass=machine-outliner -verify-machineinstrs -outliner-leaf-descendants=false %s -o - | FileCheck %s
 
 # Outlining CFI instructions is unsafe if we cannot outline all of the CFI
 # instructions from a function.  This shows that we choose not to outline the
diff --git a/llvm/test/CodeGen/AArch64/machine-outliner-leaf-descendants.ll b/llvm/test/CodeGen/AArch64/machine-outliner-leaf-descendants.ll
new file mode 100644
index 00000000000000..bdaf653a79566d
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/machine-outliner-leaf-descendants.ll
@@ -0,0 +1,124 @@
+; This test is mainly for the -outliner-leaf-descendants flag for MachineOutliner.
+;
+; ===================== -outliner-leaf-descendants=false =====================
+; MachineOutliner finds THREE key `OutlinedFunction` and outlines them. They are:
+;   ```
+;     mov     w0, #1
+;     mov     w1, #2
+;     mov     w2, #3
+;     mov     w3, #4
+;     mov     w4, #5
+;     mov     w5, #6 or #7 or #8
+;     b
+;   ```
+; Each has:
+;   - `SequenceSize=28` and `OccurrenceCount=2`
+;   - each Candidate has `CallOverhead=4` and `FrameOverhead=0`
+;   - `NotOutlinedCost=28*2=56` and `OutliningCost=4*2+28+0=36`
+;   - `Benefit=56-36=20` and `Priority=56/36=1.56`
+;
+; ===================== -outliner-leaf-descendants=false =====================
+; MachineOutliner finds a FOURTH key `OutlinedFunction`, which is:
+;   ```
+;   mov     w0, #1
+;   mov     w1, #2
+;   mov     w2, #3
+;   mov     w3, #4
+;   mov     w4, #5
+;   ```
+; This corresponds to an internal node that has ZERO leaf children, but SIX leaf descendants.
+; It has:
+;   - `SequenceSize=20` and `OccurrenceCount=6`
+;   - each Candidate has `CallOverhead=12` and `FrameOverhead=4`
+;   - `NotOutlinedCost=20*6=120` and `OutliningCost=12*6+20+4=96`
+;   - `Benefit=120-96=24` and `Priority=120/96=1.25`
+;
+; The FOURTH `OutlinedFunction` has lower _priority_ compared to the first THREE `OutlinedFunction`
+; Hence, if we additionally include the `-sort-per-priority` flag,  the first THREE `OutlinedFunction` are outlined.
+
+; RUN: llc %s -enable-machine-outliner=always -outliner-leaf-descendants=false -filetype=obj -o %t
+; RUN: llvm-objdump -d %t | FileCheck %s --check-prefix=CHECK-BASELINE
+
+; RUN: llc %s -enable-machine-outliner=always -outliner-leaf-descendants=false -outliner-benefit-threshold=22 -filetype=obj -o %t
+; RUN: llvm-objdump -d %t | FileCheck %s --check-prefix=CHECK-NO-CANDIDATE
+
+; RUN: llc %s -enable-machine-outliner=always -outliner-leaf-descendants=true -filetype=obj -o %t
+; RUN: llvm-objdump -d %t | FileCheck %s --check-prefix=CHECK-BASELINE
+
+; RUN: llc %s -enable-machine-outliner=always -outliner-leaf-descendants=true -outliner-benefit-threshold=22 -filetype=obj -o %t
+; RUN: llvm-objdump -d %t | FileCheck %s --check-prefix=CHECK-LEAF-DESCENDANTS
+
+
+target datalayout = "e-m:o-i64:64-i128:128-n32:64-S128"
+target triple = "arm64-apple-macosx14.0.0"
+
+declare i32 @_Z3fooiiii(i32 noundef, i32 noundef, i32 noundef, i32 noundef, i32 noundef, i32 noundef)
+
+define i32 @_Z2f1v() minsize {
+  %1 = tail call i32 @_Z3fooiiii(i32 noundef 1, i32 noundef 2, i32 noundef 3, i32 noundef 4, i32 noundef 5, i32 noundef 6)
+  ret i32 %1
+}
+
+define i32 @_Z2f2v() minsize {
+  %1 = tail call i32 @_Z3fooiiii(i32 noundef 1, i32 noundef 2, i32 noundef 3, i32 noundef 4, i32 noundef 5, i32 noundef 6)
+  ret i32 %1
+}
+
+define i32 @_Z2f3v() minsize {
+  %1 = tail call i32 @_Z3fooiiii(i32 noundef 1, i32 noundef 2, i32 noundef 3, i32 noundef 4, i32 noundef 5, i32 noundef 7)
+  ret i32 %1
+}
+
+define i32 @_Z2f4v() minsize {
+  %1 = tail call i32 @_Z3fooiiii(i32 noundef 1, i32 noundef 2, i32 noundef 3, i32 noundef 4, i32 noundef 5, i32 noundef 7)
+  ret i32 %1
+}
+
+define i32 @_Z2f5v() minsize {
+  %1 = tail call i32 @_Z3fooiiii(i32 noundef 1, i32 noundef 2, i32 noundef 3, i32 noundef 4, i32 noundef 5, i32 noundef 8)
+  ret i32 %1
+}
+
+define i32 @_Z2f6v() minsize {
+  %1 = tail call i32 @_Z3fooiiii(i32 noundef 1, i32 noundef 2, i32 noundef 3, i32 noundef 4, i32 noundef 5, i32 noundef 8)
+  ret i32 %1
+}
+
+; CHECK-BASELINE: <_OUTLINED_FUNCTION_0>:
+; CHECK-BASELINE-NEXT: mov     w0, #0x1
+; CHECK-BASELINE-NEXT: mov     w1, #0x2
+; CHECK-BASELINE-NEXT: mov     w2, #0x3
+; CHECK-BASELINE-NEXT: mov     w3, #0x4
+; CHECK-BASELINE-NEXT: mov     w4, #0x5
+; CHECK-BASELINE-NEXT: mov     w5, #0x6
+; CHECK-BASELINE-NEXT: b
+
+; CHECK-BASELINE: <_OUTLINED_FUNCTION_1>:
+; CHECK-BASELINE-NEXT: mov     w0, #0x1
+; CHECK-BASELINE-NEXT: mov     w1, #0x2
+; CHECK-BASELINE-NEXT: mov     w2, #0x3
+; CHECK-BASELINE-NEXT: mov     w3, #0x4
+; CHECK-BASELINE-NEXT: mov     w4, #0x5
+; CHECK-BASELINE-NEXT: mov     w5, #0x8
+; CHECK-BASELINE-NEXT: b
+
+; CHECK-BASELINE: <_OUTLINED_FUNCTION_2>:
+; CHECK-BASELINE-NEXT: mov     w0, #0x1
+; CHECK-BASELINE-NEXT: mov     w1, #0x2
+; CHECK-BASELINE-NEXT: mov     w2, #0x3
+; CHECK-BASELINE-NEXT: mov     w3, #0x4
+; CHECK-BASELINE-NEXT: mov     w4, #0x5
+; CHECK-BASELINE-NEXT: mov     w5, #0x7
+; CHECK-BASELINE-NEXT: b
+
+; CHECK-LEAF-DESCENDANTS: <_OUTLINED_FUNCTION_0>:
+; CHECK-LEAF-DESCENDANTS-NEXT: mov     w0, #0x1
+; CHECK-LEAF-DESCENDANTS-NEXT: mov     w1, #0x2
+; CHECK-LEAF-DESCENDANTS-NEXT: mov     w2, #0x3
+; CHECK-LEAF-DESCENDANTS-NEXT: mov     w3, #0x4
+; CHECK-LEAF-DESCENDANTS-NEXT: mov     w4, #0x5
+; CHECK-LEAF-DESCENDANTS-NEXT: ret
+
+; CHECK-LEAF-DESCENDANTS-NOT: <_OUTLINED_FUNCTION_1>:
+
+; CHECK-NO-CANDIDATE-NOT: <_OUTLINED_FUNCTION_0>:
diff --git a/llvm/test/CodeGen/AArch64/machine-outliner-overlap.mir b/llvm/test/CodeGen/AArch64/machine-outliner-overlap.mir
index c6bd4c1d04d871..8b0e03dbee8d8d 100644
--- a/llvm/test/CodeGen/AArch64/machine-outliner-overlap.mir
+++ b/llvm/test/CodeGen/AArch64/machine-outliner-overlap.mir
@@ -1,5 +1,6 @@
 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --include-generated-funcs
-# RUN: llc %s -mtriple aarch64 -debug-only=machine-outliner -run-pass=machine-outliner -o - 2>&1 | FileCheck %s
+# RUN: llc %s -mtriple aarch64 -outliner-leaf-descendants=false -debug-only=machine-outliner -run-pass=machine-outliner -o - 2>&1 | FileCheck %s
+# RUN: llc %s -mtriple aarch64 -debug-only=machine-outliner -run-pass=machine-outliner -o - 2>&1 | FileCheck %s --check-prefix=CHECK-LEAF
 # REQUIRES: asserts
 
 # CHECK: *** Discarding overlapping candidates ***
@@ -54,6 +55,27 @@ body:             |
     ; CHECK-NEXT: BL @OUTLINED_FUNCTION_0, implicit-def $lr, implicit $sp, implicit-def $lr, implicit-def $x8, implicit-def $x9, implicit $sp, implicit $x0, implicit $x9
     ; CHECK-NEXT: RET undef $x9
 
+    ; CHECK-LEAF-LABEL: name: overlap
+    ; CHECK-LEAF: liveins: $x0, $x9
+    ; CHECK-LEAF-NEXT: {{  $}}
+    ; CHECK-LEAF-NEXT: BL @OUTLINED_FUNCTION_0
+    ; CHECK-LEAF-NEXT: BL @OUTLINED_FUNCTION_0
+    ; CHECK-LEAF-NEXT: $x8 = ADDXri $x0, 3, 0
+    ; CHECK-LEAF-NEXT: BL @OUTLINED_FUNCTION_0
+    ; CHECK-LEAF-NEXT: BL @OUTLINED_FUNCTION_0
+    ; CHECK-LEAF-NEXT: $x8 = ADDXri $x0, 3, 0
+    ; CHECK-LEAF-NEXT: BL @OUTLINED_FUNCTION_0
+    ; CHECK-LEAF-NEXT: BL @OUTLINED_FUNCTION_0
+    ; CHECK-LEAF-NEXT: RET undef $x9
+
+    ; CHECK-LEAF-LABEL: name: OUTLINED_FUNCTION_0
+    ; CHECK-LEAF: liveins: $x0, $x9, $lr
+    ; CHECK-LEAF-NEXT: {{  $}}
+    ; CHECK-LEAF-NEXT: $x9 = ADDXri $x9, 16, 0
+    ; CHECK-LEAF-NEXT: $x9 = ADDXri $x9, 16, 0
+    ; CHECK-LEAF-NEXT: $x9 = ADDXri $x9, 16, 0
+    ; CHECK-LEAF-NEXT: RET $lr
+
     ; fixme: outline!
     $x9 = ADDXri $x9, 16, 0
     $x9 = ADDXri $x9, 16, 0
diff --git a/llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-sp-mod.mir b/llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-sp-mod.mir
index c1c2720dec6ad6..22e5edef2a9395 100644
--- a/llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-sp-mod.mir
+++ b/llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-sp-mod.mir
@@ -1,4 +1,4 @@
-# RUN: llc -verify-machineinstrs -run-pass=machine-outliner -run-pass=aarch64-ptrauth %s -o - | FileCheck %s
+# RUN: llc -verify-machineinstrs -run-pass=machine-outliner -run-pass=aarch64-ptrauth -outliner-leaf-descendants=false %s -o - | FileCheck %s
 
 --- |
   target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
diff --git a/llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-thunk.ll b/llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-thunk.ll
index 9250718fc0d585..618973b9368d1d 100644
--- a/llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-thunk.ll
+++ b/llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-thunk.ll
@@ -1,6 +1,6 @@
-; RUN: llc -mtriple aarch64-arm-linux-gnu --enable-machine-outliner \
+; RUN: llc -mtriple aarch64-arm-linux-gnu --enable-machine-outliner -outliner-leaf-descendants=false \
 ; RUN: -verify-machineinstrs %s -o - | FileCheck --check-prefixes CHECK,V8A %s
-; RUN-V83A: llc -mtriple aarch64 -enable-machine-outliner \
+; RUN-V83A: llc -mtriple aarch64 -enable-machine-outliner -outliner-leaf-descendants=false \
 ; RUN-V83A: -verify-machineinstrs -mattr=+v8.3a %s -o - > %t
 ; RUN-V83A: FileCheck --check-prefixes CHECK,V83A < %t %s
 
diff --git a/llvm/test/CodeGen/AArch64/machine-outliner-throw2.ll b/llvm/test/CodeGen/AArch64/machine-outliner-throw2.ll
index aa6e31d6ff21d7..538e1165e39c1d 100644
--- a/llvm/test/CodeGen/AArch64/machine-outliner-throw2.ll
+++ b/llvm/test/CodeGen/AArch64/machine-outliner-throw2.ll
@@ -1,5 +1,5 @@
-; RUN: llc -verify-machineinstrs -enable-machine-outliner -mtriple=aarch64 -frame-pointer=non-leaf < %s | FileCheck %s --check-prefix=NOOMIT
-; RUN: llc -verify-machineinstrs -enable-machine-outliner -mtriple=aarch64 -frame-pointer=none < %s | FileCheck %s --check-prefix=OMITFP
+; RUN: llc -verify-machineinstrs -enable-machine-outliner -outliner-leaf-descendants=false -mtriple=aarch64 -frame-pointer=non-leaf < %s | FileCheck %s --check-prefix=NOOMIT
+; RUN: llc -verify-machineinstrs -enable-machine-outliner -outliner-leaf-descendants=false -mtriple=aarch64 -frame-pointer=none < %s | FileCheck %s --check-prefix=OMITFP
 
 define void @_Z1giii(i32 %x, i32 %y, i32 %z) minsize {
 ; NOOMIT-LABEL: _Z1giii:
diff --git a/llvm/test/CodeGen/AArch64/machine-outliner-thunk.ll b/llvm/test/CodeGen/AArch64/machine-outliner-thunk.ll
index 8740aac0549eec..7e34adf16d25d3 100644
--- a/llvm/test/CodeGen/AArch64/machine-outliner-thunk.ll
+++ b/llvm/test/CodeGen/AArch64/machine-outliner-thunk.ll
@@ -1,5 +1,5 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
-; RUN: llc < %s -enable-machine-outliner -verify-machineinstrs | FileCheck %s
+; RUN: llc < %s -enable-machine-outliner -outliner-leaf-descendants=false -verify-machineinstrs | FileCheck %s
 
 target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
 target triple = "aarch64-pc-linux-gnu"
diff --git a/llvm/test/CodeGen/AArch64/machine-outliner.mir b/llvm/test/CodeGen/AArch64/machine-outliner.mir
index 83eda744d24a9e..66779addaff0a8 100644
--- a/llvm/test/CodeGen/AArch64/machine-outliner.mir
+++ b/llvm/test/CodeGen/AArch64/machine-outliner.mir
@@ -1,4 +1,4 @@
-# RUN: llc -mtriple=aarch64--- -run-pass=prologepilog -run-pass=machine-outliner -verify-machineinstrs -frame-pointer=non-leaf %s -o - | FileCheck %s
+# RUN: llc -mtriple=aarch64--- -run-pass=prologepilog -run-pass=machine-outliner -verify-machineinstrs -frame-pointer=non-leaf -outliner-leaf-descendants=false %s -o - | FileCheck %s
 --- |
 
   @x = common global i32 0, align 4
diff --git a/llvm/test/CodeGen/RISCV/machineoutliner-pcrel-lo.mir b/llvm/test/CodeGen/RISCV/machineoutliner-pcrel-lo.mir
index 34f7a93b6a168c..8a83543b0280fd 100644
--- a/llvm/test/CodeGen/RISCV/machineoutliner-pcrel-lo.mir
+++ b/llvm/test/CodeGen/RISCV/machineoutliner-pcrel-lo.mir
@@ -1,11 +1,11 @@
 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
-# RUN: llc -mtriple=riscv32 -x mir -run-pass=machine-outliner -simplify-mir -verify-machineinstrs < %s \
+# RUN: llc -mtriple=riscv32 -x mir -run-pass=machine-outliner -simplify-mir -verify-machineinstrs -outliner-leaf-descendants=false < %s \
 # RUN: | FileCheck %s
-# RUN: llc -mtriple=riscv64 -x mir -run-pass=machine-outliner -simplify-mir -verify-machineinstrs < %s \
+# RUN: llc -mtriple=riscv64 -x mir -run-pass=machine-outliner -simplify-mir -verify-machineinstrs -outliner-leaf-descendants=false < %s \
 # RUN: | FileCheck %s
-# RUN: llc -mtriple=riscv32 -x mir -run-pass=machine-outliner -simplify-mir --function-sections -verify-machineinstrs < %s \
+# RUN: llc -mtriple=riscv32 -x mir -run-pass=machine-outliner -simplify-mir --function-sections -verify-machineinstrs -outliner-leaf-descendants=false < %s \
 # RUN: | FileCheck -check-prefix=CHECK-FS %s
-# RUN: llc -mtriple=riscv64 -x mir -run-pass=machine-outliner -simplify-mir --function-sections -verify-machineinstrs < %s \
+# RUN: llc -mtriple=riscv64 -x mir -run-pass=machine-outliner -simplify-mir --function-sections -verify-machineinstrs -outliner-leaf-descendants=false < %s \
 # RUN: | FileCheck -check-prefix=CHECK-FS %s
 
 --- |
diff --git a/llvm/unittests/Support/SuffixTreeTest.cpp b/llvm/unittests/Support/SuffixTreeTest.cpp
index f5d8112ccf5cec..59f6dc9bf963fa 100644
--- a/llvm/unittests/Support/SuffixTreeTest.cpp
+++ b/llvm/unittests/Support/SuffixTreeTest.cpp
@@ -140,4 +140,138 @@ TEST(SuffixTreeTest, TestExclusion) {
   }
 }
 
+// Tests that the SuffixTree is able to find the following substrings:
+// {1, 1} at indices 0, 1, 2, 3, and 4;
+// {1, 1, 1} at indices 0, 1, 2, and 3;
+// {1, 1, 1, 1}  at indices 0, 1, and 2; and
+// {1, 1, 1, 1, 1} at indices 0 and 1.
+//
+// This is a FIX to the Test TestSingleCharacterRepeat
+TEST(SuffixTreeTest, TestSingleCharacterRepeatWithLeafDescendants) {
+  std::vector<unsigned> RepeatedRepetitionData = {1, 1, 1, 1, 1, 1, 2};
+  std::vector<unsigned>::iterator RRDIt, RRDIt2;
+  SuffixTree ST(RepeatedRepetitionData, true);
+  std::vector<SuffixTree::RepeatedSubstring> SubStrings;
+  for (auto It = ST.begin(); It != ST.end(); It++)
+    SubStrings.push_back(*It);
+  EXPECT_EQ(SubStrings.size(), 4u);
+  for (SuffixTree::RepeatedSubstring &RS : SubStrings) {
+    EXPECT_EQ(RS.StartIndices.size(),
+              RepeatedRepetitionData.size() - RS.Length);
+    for (unsigned StartIdx : SubStrings[0].StartIndices) {
+      RRDIt = RRDIt2 = RepeatedRepetitionData.begin();
+      std::advance(RRDIt, StartIdx);
+      std::advance(RRDIt2, StartIdx + SubStrings[0].Length);
+      ASSERT_TRUE(
+          all_of(make_range<std::vector<unsigned>::iterator>(RRDIt, RRDIt2),
+                 [](unsigned Elt) { return Elt == 1; }));
+    }
+  }
+}
+
+// Tests that the SuffixTree is able to find three substrings
+// {1, 2, 3} at indices 6 and 10;
+// {2, 3} at indices 7 and 11; and
+// {1, 2} at indicies 0 and 3.
+//
+// FIXME: {1, 2} has indices 6 and 10 missing as it is a substring of {1, 2, 3}
+// See Test TestSubstringRepeatsWithLeafDescendants for the FIX
+TEST(SuffixTreeTest, TestSubstringRepeats) {
+  std::vector<unsigned> RepeatedRepetitionData = {1, 2, 100, 1, 2, 101, 1,
+                                                  2, 3, 103, 1, 2, 3,   104};
+  SuffixTree ST(RepeatedRepetitionData);
+  std::vector<SuffixTree::RepeatedSubstring> SubStrings;
+  for (auto It = ST.begin(); It != ST.end(); It++)
+    SubStrings.push_back(*It);
+  EXPECT_EQ(SubStrings.size(), 3u);
+  unsigned Len;
+  for (SuffixTree::RepeatedSubstring &RS : SubStrings) {
+    Len = RS.Length;
+    bool IsExpectedLen = (Len == 3u || Len == 2u);
+    ASSERT_TRUE(IsExpectedLen);
+    bool IsExpectedIndex;
+
+    if (Len == 3u) { // {1, 2, 3}
+      EXPECT_EQ(RS.StartIndices.size(), 2u);
+      for (unsigned StartIdx : RS.StartIndices) {
+        IsExpectedIndex = (StartIdx == 6u || StartIdx == 10u);
+        EXPECT_TRUE(IsExpectedIndex);
+        EXPECT_EQ(RepeatedRepetitionData[StartIdx], 1u);
+        EXPECT_EQ(RepeatedRepetitionData[StartIdx + 1], 2u);
+        EXPECT_EQ(RepeatedRepetitionData[StartIdx + 2], 3u);
+      }
+    } else {
+      if (RepeatedRepetitionData[RS.StartIndices[0]] == 1u) { // {1, 2}
+        EXPECT_EQ(RS.StartIndices.size(), 2u);
+        for (unsigned StartIdx : RS.StartIndices) {
+          IsExpectedIndex = (StartIdx == 0u || StartIdx == 3u);
+          EXPECT_TRUE(IsExpectedIndex);
+          EXPECT_EQ(RepeatedRepetitionData[StartIdx + 1], 2u);
+        }
+      } else { // {2, 3}
+        EXPECT_EQ(RS.StartIndices.size(), 2u);
+        for (unsigned StartIdx : RS.StartIndices) {
+          IsExpectedIndex = (StartIdx == 7u || StartIdx == 11u);
+          EXPECT_TRUE(IsExpectedIndex);
+          EXPECT_EQ(RepeatedRepetitionData[StartIdx], 2u);
+          EXPECT_EQ(RepeatedRepetitionData[StartIdx + 1], 3u);
+        }
+      }
+    }
+  }
+}
+
+// Tests that the SuffixTree is able to find three substrings
+// {1, 2, 3} at indices 6 and 10;
+// {2, 3} at indices 7 and 11; and
+// {1, 2} at indicies 0, 3, 6, and 10.
+//
+// This is a FIX to the Test TestSubstringRepeats
+
+TEST(SuffixTreeTest, TestSubstringRepeatsWithLeafDescendants) {
+  std::vector<unsigned> RepeatedRepetitionData = {1, 2, 100, 1, 2, 101, 1,
+                                                  2, 3, 103, 1, 2, 3,   104};
+  SuffixTree ST(RepeatedRepetitionData, true);
+  std::vector<SuffixTree::RepeatedSubstring> SubStrings;
+  for (auto It = ST.begin(); It != ST.end(); It++)
+    SubStrings.push_back(*It);
+  EXPECT_EQ(SubStrings.size(), 3u);
+  unsigned Len;
+  for (SuffixTree::RepeatedSubstring &RS : SubStrings) {
+    Len = RS.Length;
+    bool IsExpectedLen = (Len == 3u || Len == 2u);
+    ASSERT_TRUE(IsExpectedLen);
+    bool IsExpectedIndex;
+
+    if (Len == 3u) { // {1, 2, 3}
+      EXPECT_EQ(RS.StartIndices.size(), 2u);
+      for (unsigned StartIdx : RS.StartIndices) {
+        IsExpectedIndex = (StartIdx == 6u || StartIdx == 10u);
+        EXPECT_TRUE(IsExpectedIndex);
+        EXPECT_EQ(RepeatedRepetitionData[StartIdx], 1u);
+        EXPECT_EQ(RepeatedRepetitionData[StartIdx + 1], 2u);
+        EXPECT_EQ(RepeatedRepetitionData[StartIdx + 2], 3u);
+      }
+    } else {
+      if (RepeatedRepetitionData[RS.StartIndices[0]] == 1u) { // {1, 2}
+        EXPECT_EQ(RS.StartIndices.size(), 4u);
+        for (unsigned StartIdx : RS.StartIndices) {
+          IsExpectedIndex = (StartIdx == 0u || StartIdx == 3u ||
+                             StartIdx == 6u || StartIdx == 10u);
+          EXPECT_TRUE(IsExpectedIndex);
+          EXPECT_EQ(RepeatedRepetitionData[StartIdx + 1], 2u);
+        }
+      } else { // {2, 3}
+        EXPECT_EQ(RS.StartIndices.size(), 2u);
+        for (unsigned StartIdx : RS.StartIndices) {
+          IsExpectedIndex = (StartIdx == 7u || StartIdx == 11u);
+          EXPECT_TRUE(IsExpectedIndex);
+          EXPECT_EQ(RepeatedRepetitionData[StartIdx], 2u);
+          EXPECT_EQ(RepeatedRepetitionData[StartIdx + 1], 3u);
+        }
+      }
+    }
+  }
+}
+
 } // namespace

>From 2c2601941e1b532c121b53d522541f5316a38db9 Mon Sep 17 00:00:00 2001
From: Kyungwoo Lee <kyulee at meta.com>
Date: Mon, 22 Apr 2024 15:29:25 -0700
Subject: [PATCH 4/5] [CGData] Outlined Hash Tree

This defines the OutlinedHashTree class.
It contains sequences of stable hash values of instructions that have been outlined.
This OutlinedHashTree can be used to track the outlined instruction sequences across modules.
A trie structure is used in its implementation, allowing for a compact sharing of common prefixes.
---
 .../llvm/CodeGenData/OutlinedHashTree.h       | 107 +++++++++++
 .../llvm/CodeGenData/OutlinedHashTreeRecord.h |  67 +++++++
 llvm/lib/CMakeLists.txt                       |   1 +
 llvm/lib/CodeGenData/CMakeLists.txt           |  14 ++
 llvm/lib/CodeGenData/OutlinedHashTree.cpp     | 131 ++++++++++++++
 .../CodeGenData/OutlinedHashTreeRecord.cpp    | 168 ++++++++++++++++++
 llvm/unittests/CMakeLists.txt                 |   1 +
 llvm/unittests/CodeGenData/CMakeLists.txt     |  14 ++
 .../OutlinedHashTreeRecordTest.cpp            | 118 ++++++++++++
 .../CodeGenData/OutlinedHashTreeTest.cpp      |  81 +++++++++
 10 files changed, 702 insertions(+)
 create mode 100644 llvm/include/llvm/CodeGenData/OutlinedHashTree.h
 create mode 100644 llvm/include/llvm/CodeGenData/OutlinedHashTreeRecord.h
 create mode 100644 llvm/lib/CodeGenData/CMakeLists.txt
 create mode 100644 llvm/lib/CodeGenData/OutlinedHashTree.cpp
 create mode 100644 llvm/lib/CodeGenData/OutlinedHashTreeRecord.cpp
 create mode 100644 llvm/unittests/CodeGenData/CMakeLists.txt
 create mode 100644 llvm/unittests/CodeGenData/OutlinedHashTreeRecordTest.cpp
 create mode 100644 llvm/unittests/CodeGenData/OutlinedHashTreeTest.cpp

diff --git a/llvm/include/llvm/CodeGenData/OutlinedHashTree.h b/llvm/include/llvm/CodeGenData/OutlinedHashTree.h
new file mode 100644
index 00000000000000..875e1a78bb4010
--- /dev/null
+++ b/llvm/include/llvm/CodeGenData/OutlinedHashTree.h
@@ -0,0 +1,107 @@
+//===- OutlinedHashTree.h --------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===---------------------------------------------------------------------===//
+//
+// This defines the OutlinedHashTree class. It contains sequences of stable
+// hash values of instructions that have been outlined. This OutlinedHashTree
+// can be used to track the outlined instruction sequences across modules.
+//
+//===---------------------------------------------------------------------===//
+
+#ifndef LLVM_CODEGENDATA_OUTLINEDHASHTREE_H
+#define LLVM_CODEGENDATA_OUTLINEDHASHTREE_H
+
+#include "llvm/ADT/StableHashing.h"
+#include "llvm/ObjectYAML/YAML.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include <unordered_map>
+#include <vector>
+
+namespace llvm {
+
+/// A HashNode is an entry in an OutlinedHashTree, holding a hash value
+/// and a collection of Successors (other HashNodes). If a HashNode has
+/// a positive terminal value (Terminals > 0), it signifies the end of
+/// a hash sequence with that occurrence count.
+struct HashNode {
+  /// The hash value of the node.
+  stable_hash Hash;
+  /// The number of terminals in the sequence ending at this node.
+  unsigned Terminals;
+  /// The successors of this node.
+  std::unordered_map<stable_hash, std::unique_ptr<HashNode>> Successors;
+};
+
+/// HashNodeStable is the serialized, stable, and compact representation
+/// of a HashNode.
+struct HashNodeStable {
+  llvm::yaml::Hex64 Hash;
+  unsigned Terminals;
+  std::vector<unsigned> SuccessorIds;
+};
+
+class OutlinedHashTree {
+
+  using EdgeCallbackFn =
+      std::function<void(const HashNode *, const HashNode *)>;
+  using NodeCallbackFn = std::function<void(const HashNode *)>;
+
+  using HashSequence = std::vector<stable_hash>;
+  using HashSequencePair = std::pair<std::vector<stable_hash>, unsigned>;
+
+public:
+  /// Walks every edge and node in the OutlinedHashTree and calls CallbackEdge
+  /// for the edges and CallbackNode for the nodes with the stable_hash for
+  /// the source and the stable_hash of the sink for an edge. These generic
+  /// callbacks can be used to traverse a OutlinedHashTree for the purpose of
+  /// print debugging or serializing it.
+  void walkGraph(NodeCallbackFn CallbackNode,
+                 EdgeCallbackFn CallbackEdge = nullptr,
+                 bool SortedWalk = false) const;
+
+  /// Release all hash nodes except the root hash node.
+  void clear() {
+    assert(getRoot()->Hash == 0 && getRoot()->Terminals == 0);
+    getRoot()->Successors.clear();
+  }
+
+  /// \returns true if the hash tree has only the root node.
+  bool empty() { return size() == 1; }
+
+  /// \returns the size of a OutlinedHashTree by traversing it. If
+  /// \p GetTerminalCountOnly is true, it only counts the terminal nodes
+  /// (meaning it returns the the number of hash sequences in the
+  /// OutlinedHashTree).
+  size_t size(bool GetTerminalCountOnly = false) const;
+
+  /// \returns the depth of a OutlinedHashTree by traversing it.
+  size_t depth() const;
+
+  /// \returns the root hash node of a OutlinedHashTree.
+  const HashNode *getRoot() const { return Root.get(); }
+  HashNode *getRoot() { return Root.get(); }
+
+  /// Inserts a \p Sequence into the this tree. The last node in the sequence
+  /// will increase Terminals.
+  void insert(const HashSequencePair &SequencePair);
+
+  /// Merge a \p OtherTree into this Tree.
+  void merge(const OutlinedHashTree *OtherTree);
+
+  /// \returns the matching count if \p Sequence exists in the OutlinedHashTree.
+  unsigned find(const HashSequence &Sequence) const;
+
+  OutlinedHashTree() { Root = std::make_unique<HashNode>(); }
+
+private:
+  std::unique_ptr<HashNode> Root;
+};
+
+} // namespace llvm
+
+#endif
diff --git a/llvm/include/llvm/CodeGenData/OutlinedHashTreeRecord.h b/llvm/include/llvm/CodeGenData/OutlinedHashTreeRecord.h
new file mode 100644
index 00000000000000..ccd2ad26dd0871
--- /dev/null
+++ b/llvm/include/llvm/CodeGenData/OutlinedHashTreeRecord.h
@@ -0,0 +1,67 @@
+//===- OutlinedHashTreeRecord.h --------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===---------------------------------------------------------------------===//
+//
+// This defines the OutlinedHashTreeRecord class. This class holds the outlined
+// hash tree for both serialization and deserialization processes. It utilizes
+// two data formats for serialization: raw binary data and YAML.
+// These two formats can be used interchangeably.
+//
+//===---------------------------------------------------------------------===//
+
+#ifndef LLVM_CODEGENDATA_OUTLINEDHASHTREERECORD_H
+#define LLVM_CODEGENDATA_OUTLINEDHASHTREERECORD_H
+
+#include "llvm/CodeGenData/OutlinedHashTree.h"
+
+namespace llvm {
+
+using IdHashNodeStableMapTy = std::map<unsigned, HashNodeStable>;
+using IdHashNodeMapTy = std::map<unsigned, HashNode *>;
+using HashNodeIdMapTy = std::unordered_map<const HashNode *, unsigned>;
+
+struct OutlinedHashTreeRecord {
+  std::unique_ptr<OutlinedHashTree> HashTree;
+
+  OutlinedHashTreeRecord() { HashTree = std::make_unique<OutlinedHashTree>(); }
+  OutlinedHashTreeRecord(std::unique_ptr<OutlinedHashTree> HashTree)
+      : HashTree(std::move(HashTree)){};
+
+  /// Serialize the outlined hash tree to a raw_ostream.
+  void serialize(raw_ostream &OS) const;
+  /// Deserialize the outlined hash tree from a raw_ostream.
+  void deserialize(const unsigned char *&Ptr);
+  /// Serialize the outlined hash tree to a YAML stream.
+  void serializeYAML(yaml::Output &YOS) const;
+  /// Deserialize the outlined hash tree from a YAML stream.
+  void deserializeYAML(yaml::Input &YIS);
+
+  /// Merge the other outlined hash tree into this one.
+  void merge(const OutlinedHashTreeRecord &Other) {
+    HashTree->merge(Other.HashTree.get());
+  }
+
+  /// \returns true if the outlined hash tree is empty.
+  bool empty() const { return HashTree->empty(); }
+
+  /// Print the outlined hash tree in a YAML format.
+  void print(raw_ostream &OS = llvm::errs()) const {
+    yaml::Output YOS(OS);
+    serializeYAML(YOS);
+  }
+
+private:
+  /// Convert the outlined hash tree to stable data.
+  void convertToStableData(IdHashNodeStableMapTy &IdNodeStableMap) const;
+
+  /// Convert the stable data back to the outlined hash tree.
+  void convertFromStableData(const IdHashNodeStableMapTy &IdNodeStableMap);
+};
+
+} // end namespace llvm
+
+#endif // LLVM_CODEGENDATA_OUTLINEDHASHTREERECORD_H
diff --git a/llvm/lib/CMakeLists.txt b/llvm/lib/CMakeLists.txt
index 74e2d03c07953d..2ac0b0dc026e16 100644
--- a/llvm/lib/CMakeLists.txt
+++ b/llvm/lib/CMakeLists.txt
@@ -10,6 +10,7 @@ add_subdirectory(InterfaceStub)
 add_subdirectory(IRPrinter)
 add_subdirectory(IRReader)
 add_subdirectory(CodeGen)
+add_subdirectory(CodeGenData)
 add_subdirectory(CodeGenTypes)
 add_subdirectory(BinaryFormat)
 add_subdirectory(Bitcode)
diff --git a/llvm/lib/CodeGenData/CMakeLists.txt b/llvm/lib/CodeGenData/CMakeLists.txt
new file mode 100644
index 00000000000000..3ba90f96cc86d4
--- /dev/null
+++ b/llvm/lib/CodeGenData/CMakeLists.txt
@@ -0,0 +1,14 @@
+add_llvm_component_library(LLVMCodeGenData
+  OutlinedHashTree.cpp
+  OutlinedHashTreeRecord.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${LLVM_MAIN_INCLUDE_DIR}/llvm/CodeGenData
+
+  DEPENDS
+  intrinsics_gen
+
+  LINK_COMPONENTS
+  Core
+  Support
+  )
diff --git a/llvm/lib/CodeGenData/OutlinedHashTree.cpp b/llvm/lib/CodeGenData/OutlinedHashTree.cpp
new file mode 100644
index 00000000000000..032993ded60ead
--- /dev/null
+++ b/llvm/lib/CodeGenData/OutlinedHashTree.cpp
@@ -0,0 +1,131 @@
+//===-- OutlinedHashTree.cpp ----------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// An OutlinedHashTree is a Trie that contains sequences of stable hash values
+// of instructions that have been outlined. This OutlinedHashTree can be used
+// to understand the outlined instruction sequences collected across modules.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGenData/OutlinedHashTree.h"
+
+#include <stack>
+#include <tuple>
+
+#define DEBUG_TYPE "outlined-hash-tree"
+
+using namespace llvm;
+
+void OutlinedHashTree::walkGraph(NodeCallbackFn CallbackNode,
+                                 EdgeCallbackFn CallbackEdge,
+                                 bool SortedWalk) const {
+  std::stack<const HashNode *> Stack;
+  Stack.push(getRoot());
+
+  while (!Stack.empty()) {
+    const auto *Current = Stack.top();
+    Stack.pop();
+    if (CallbackNode)
+      CallbackNode(Current);
+
+    auto HandleNext = [&](const HashNode *Next) {
+      if (CallbackEdge)
+        CallbackEdge(Current, Next);
+      Stack.push(Next);
+    };
+    if (SortedWalk) {
+      std::map<stable_hash, const HashNode *> SortedSuccessors;
+      for (const auto &P : Current->Successors)
+        SortedSuccessors[P.first] = P.second.get();
+      for (const auto &P : SortedSuccessors)
+        HandleNext(P.second);
+    } else {
+      for (const auto &P : Current->Successors)
+        HandleNext(P.second.get());
+    }
+  }
+}
+
+size_t OutlinedHashTree::size(bool GetTerminalCountOnly) const {
+  size_t Size = 0;
+  walkGraph([&Size, GetTerminalCountOnly](const HashNode *N) {
+    Size += (N && (!GetTerminalCountOnly || N->Terminals));
+  });
+  return Size;
+}
+
+size_t OutlinedHashTree::depth() const {
+  size_t Size = 0;
+  std::unordered_map<const HashNode *, size_t> DepthMap;
+  walkGraph([&Size, &DepthMap](
+                const HashNode *N) { Size = std::max(Size, DepthMap[N]); },
+            [&DepthMap](const HashNode *Src, const HashNode *Dst) {
+              size_t Depth = DepthMap[Src];
+              DepthMap[Dst] = Depth + 1;
+            });
+  return Size;
+}
+
+void OutlinedHashTree::insert(const HashSequencePair &SequencePair) {
+  const auto &Sequence = SequencePair.first;
+  unsigned Count = SequencePair.second;
+  HashNode *Current = getRoot();
+
+  for (stable_hash StableHash : Sequence) {
+    auto I = Current->Successors.find(StableHash);
+    if (I == Current->Successors.end()) {
+      std::unique_ptr<HashNode> Next = std::make_unique<HashNode>();
+      HashNode *NextPtr = Next.get();
+      NextPtr->Hash = StableHash;
+      Current->Successors.emplace(StableHash, std::move(Next));
+      Current = NextPtr;
+    } else
+      Current = I->second.get();
+  }
+  Current->Terminals += Count;
+}
+
+void OutlinedHashTree::merge(const OutlinedHashTree *Tree) {
+  HashNode *Dst = getRoot();
+  const HashNode *Src = Tree->getRoot();
+  std::stack<std::pair<HashNode *, const HashNode *>> Stack;
+  Stack.push({Dst, Src});
+
+  while (!Stack.empty()) {
+    auto [DstNode, SrcNode] = Stack.top();
+    Stack.pop();
+    if (!SrcNode)
+      continue;
+    DstNode->Terminals += SrcNode->Terminals;
+
+    for (auto &[Hash, NextSrcNode] : SrcNode->Successors) {
+      HashNode *NextDstNode;
+      auto I = DstNode->Successors.find(Hash);
+      if (I == DstNode->Successors.end()) {
+        auto NextDst = std::make_unique<HashNode>();
+        NextDstNode = NextDst.get();
+        NextDstNode->Hash = Hash;
+        DstNode->Successors.emplace(Hash, std::move(NextDst));
+      } else
+        NextDstNode = I->second.get();
+
+      Stack.push({NextDstNode, NextSrcNode.get()});
+    }
+  }
+}
+
+unsigned OutlinedHashTree::find(const HashSequence &Sequence) const {
+  const HashNode *Current = getRoot();
+  for (stable_hash StableHash : Sequence) {
+    const auto I = Current->Successors.find(StableHash);
+    if (I == Current->Successors.end())
+      return 0;
+    Current = I->second.get();
+  }
+  return Current->Terminals;
+}
diff --git a/llvm/lib/CodeGenData/OutlinedHashTreeRecord.cpp b/llvm/lib/CodeGenData/OutlinedHashTreeRecord.cpp
new file mode 100644
index 00000000000000..0d5dd864c89c55
--- /dev/null
+++ b/llvm/lib/CodeGenData/OutlinedHashTreeRecord.cpp
@@ -0,0 +1,168 @@
+//===-- OutlinedHashTreeRecord.cpp ----------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This defines the OutlinedHashTreeRecord class. This class holds the outlined
+// hash tree for both serialization and deserialization processes. It utilizes
+// two data formats for serialization: raw binary data and YAML.
+// These two formats can be used interchangeably.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGenData/OutlinedHashTreeRecord.h"
+#include "llvm/CodeGenData/OutlinedHashTree.h"
+#include "llvm/ObjectYAML/YAML.h"
+#include "llvm/Support/Endian.h"
+#include "llvm/Support/EndianStream.h"
+
+#define DEBUG_TYPE "outlined-hash-tree"
+
+using namespace llvm;
+using namespace llvm::support;
+
+namespace llvm {
+namespace yaml {
+
+template <> struct MappingTraits<HashNodeStable> {
+  static void mapping(IO &io, HashNodeStable &res) {
+    io.mapRequired("Hash", res.Hash);
+    io.mapRequired("Terminals", res.Terminals);
+    io.mapRequired("SuccessorIds", res.SuccessorIds);
+  }
+};
+
+template <> struct CustomMappingTraits<IdHashNodeStableMapTy> {
+  static void inputOne(IO &io, StringRef Key, IdHashNodeStableMapTy &V) {
+    HashNodeStable NodeStable;
+    io.mapRequired(Key.str().c_str(), NodeStable);
+    unsigned Id;
+    if (Key.getAsInteger(0, Id)) {
+      io.setError("Id not an integer");
+      return;
+    }
+    V.insert({Id, NodeStable});
+  }
+
+  static void output(IO &io, IdHashNodeStableMapTy &V) {
+    for (auto Iter = V.begin(); Iter != V.end(); ++Iter)
+      io.mapRequired(utostr(Iter->first).c_str(), Iter->second);
+  }
+};
+
+} // namespace yaml
+} // namespace llvm
+
+void OutlinedHashTreeRecord::serialize(raw_ostream &OS) const {
+  IdHashNodeStableMapTy IdNodeStableMap;
+  convertToStableData(IdNodeStableMap);
+  support::endian::Writer Writer(OS, endianness::little);
+  Writer.write<uint32_t>(IdNodeStableMap.size());
+
+  for (const auto &[Id, NodeStable] : IdNodeStableMap) {
+    Writer.write<uint32_t>(Id);
+    Writer.write<uint64_t>(NodeStable.Hash);
+    Writer.write<uint32_t>(NodeStable.Terminals);
+    Writer.write<uint32_t>(NodeStable.SuccessorIds.size());
+    for (auto SuccessorId : NodeStable.SuccessorIds)
+      Writer.write<uint32_t>(SuccessorId);
+  }
+}
+
+void OutlinedHashTreeRecord::deserialize(const unsigned char *&Ptr) {
+  IdHashNodeStableMapTy IdNodeStableMap;
+  auto NumIdNodeStableMap =
+      endian::readNext<uint32_t, endianness::little, unaligned>(Ptr);
+
+  for (unsigned I = 0; I < NumIdNodeStableMap; ++I) {
+    auto Id = endian::readNext<uint32_t, endianness::little, unaligned>(Ptr);
+    HashNodeStable NodeStable;
+    NodeStable.Hash =
+        endian::readNext<uint64_t, endianness::little, unaligned>(Ptr);
+    NodeStable.Terminals =
+        endian::readNext<uint32_t, endianness::little, unaligned>(Ptr);
+    auto NumSuccessorIds =
+        endian::readNext<uint32_t, endianness::little, unaligned>(Ptr);
+    for (unsigned J = 0; J < NumSuccessorIds; ++J)
+      NodeStable.SuccessorIds.push_back(
+          endian::readNext<uint32_t, endianness::little, unaligned>(Ptr));
+
+    IdNodeStableMap[Id] = std::move(NodeStable);
+  }
+
+  convertFromStableData(IdNodeStableMap);
+}
+
+void OutlinedHashTreeRecord::serializeYAML(yaml::Output &YOS) const {
+  IdHashNodeStableMapTy IdNodeStableMap;
+  convertToStableData(IdNodeStableMap);
+
+  YOS << IdNodeStableMap;
+}
+
+void OutlinedHashTreeRecord::deserializeYAML(yaml::Input &YIS) {
+  IdHashNodeStableMapTy IdNodeStableMap;
+
+  YIS >> IdNodeStableMap;
+  YIS.nextDocument();
+
+  convertFromStableData(IdNodeStableMap);
+}
+
+void OutlinedHashTreeRecord::convertToStableData(
+    IdHashNodeStableMapTy &IdNodeStableMap) const {
+  // Build NodeIdMap
+  HashNodeIdMapTy NodeIdMap;
+  HashTree->walkGraph(
+      [&NodeIdMap](const HashNode *Current) {
+        size_t Index = NodeIdMap.size();
+        NodeIdMap[Current] = Index;
+        assert(Index = NodeIdMap.size() + 1 &&
+                       "Expected size of NodeMap to increment by 1");
+      },
+      /*EdgeCallbackFn=*/nullptr, /*SortedWork=*/true);
+
+  // Convert NodeIdMap to NodeStableMap
+  for (auto &P : NodeIdMap) {
+    auto *Node = P.first;
+    auto Id = P.second;
+    HashNodeStable NodeStable;
+    NodeStable.Hash = Node->Hash;
+    NodeStable.Terminals = Node->Terminals;
+    for (auto &P : Node->Successors)
+      NodeStable.SuccessorIds.push_back(NodeIdMap[P.second.get()]);
+    IdNodeStableMap[Id] = NodeStable;
+  }
+
+  // Sort the Successors so that they come out in the same order as in the map.
+  for (auto &P : IdNodeStableMap)
+    std::sort(P.second.SuccessorIds.begin(), P.second.SuccessorIds.end());
+}
+
+void OutlinedHashTreeRecord::convertFromStableData(
+    const IdHashNodeStableMapTy &IdNodeStableMap) {
+  IdHashNodeMapTy IdNodeMap;
+  // Initialize the root node at 0.
+  IdNodeMap[0] = HashTree->getRoot();
+  assert(IdNodeMap[0]->Successors.empty());
+
+  for (auto &P : IdNodeStableMap) {
+    auto Id = P.first;
+    const HashNodeStable &NodeStable = P.second;
+    assert(IdNodeMap.count(Id));
+    HashNode *Curr = IdNodeMap[Id];
+    Curr->Hash = NodeStable.Hash;
+    Curr->Terminals = NodeStable.Terminals;
+    auto &Successors = Curr->Successors;
+    assert(Successors.empty());
+    for (auto SuccessorId : NodeStable.SuccessorIds) {
+      auto Sucessor = std::make_unique<HashNode>();
+      IdNodeMap[SuccessorId] = Sucessor.get();
+      auto Hash = IdNodeStableMap.at(SuccessorId).Hash;
+      Successors[Hash] = std::move(Sucessor);
+    }
+  }
+}
diff --git a/llvm/unittests/CMakeLists.txt b/llvm/unittests/CMakeLists.txt
index 46f30ff398e10d..cb4b8513e6d02e 100644
--- a/llvm/unittests/CMakeLists.txt
+++ b/llvm/unittests/CMakeLists.txt
@@ -21,6 +21,7 @@ add_subdirectory(BinaryFormat)
 add_subdirectory(Bitcode)
 add_subdirectory(Bitstream)
 add_subdirectory(CodeGen)
+add_subdirectory(CodeGenData)
 add_subdirectory(DebugInfo)
 add_subdirectory(Debuginfod)
 add_subdirectory(Demangle)
diff --git a/llvm/unittests/CodeGenData/CMakeLists.txt b/llvm/unittests/CodeGenData/CMakeLists.txt
new file mode 100644
index 00000000000000..3d821b87e29d8c
--- /dev/null
+++ b/llvm/unittests/CodeGenData/CMakeLists.txt
@@ -0,0 +1,14 @@
+set(LLVM_LINK_COMPONENTS
+  ${LLVM_TARGETS_TO_BUILD}
+  CodeGen
+  CodeGenData
+  Core
+  Support
+  )
+
+add_llvm_unittest(CodeGenDataTests
+  OutlinedHashTreeRecordTest.cpp
+  OutlinedHashTreeTest.cpp
+  )
+
+target_link_libraries(CodeGenDataTests PRIVATE LLVMTestingSupport)
diff --git a/llvm/unittests/CodeGenData/OutlinedHashTreeRecordTest.cpp b/llvm/unittests/CodeGenData/OutlinedHashTreeRecordTest.cpp
new file mode 100644
index 00000000000000..aa7ad4a33754ff
--- /dev/null
+++ b/llvm/unittests/CodeGenData/OutlinedHashTreeRecordTest.cpp
@@ -0,0 +1,118 @@
+//===- OutlinedHashTreeRecordTest.cpp -------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGenData/OutlinedHashTreeRecord.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+namespace {
+
+TEST(OutlinedHashTreeRecordTest, Empty) {
+  OutlinedHashTreeRecord HashTreeRecord;
+  ASSERT_TRUE(HashTreeRecord.empty());
+}
+
+TEST(OutlinedHashTreeRecordTest, Print) {
+  OutlinedHashTreeRecord HashTreeRecord;
+  HashTreeRecord.HashTree->insert({{1, 2}, 3});
+
+  const char *ExpectedTreeStr = R"(---
+0:
+  Hash:            0x0
+  Terminals:       0
+  SuccessorIds:    [ 1 ]
+1:
+  Hash:            0x1
+  Terminals:       0
+  SuccessorIds:    [ 2 ]
+2:
+  Hash:            0x2
+  Terminals:       3
+  SuccessorIds:    [  ]
+...
+)";
+  std::string TreeDump;
+  raw_string_ostream OS(TreeDump);
+  HashTreeRecord.print(OS);
+  EXPECT_EQ(ExpectedTreeStr, TreeDump);
+}
+
+TEST(OutlinedHashTreeRecordTest, Stable) {
+  OutlinedHashTreeRecord HashTreeRecord1;
+  HashTreeRecord1.HashTree->insert({{1, 2}, 4});
+  HashTreeRecord1.HashTree->insert({{1, 3}, 5});
+
+  OutlinedHashTreeRecord HashTreeRecord2;
+  HashTreeRecord2.HashTree->insert({{1, 3}, 5});
+  HashTreeRecord2.HashTree->insert({{1, 2}, 4});
+
+  // Output is stable regardless of insertion order.
+  std::string TreeDump1;
+  raw_string_ostream OS1(TreeDump1);
+  HashTreeRecord1.print(OS1);
+  std::string TreeDump2;
+  raw_string_ostream OS2(TreeDump2);
+  HashTreeRecord2.print(OS2);
+
+  EXPECT_EQ(TreeDump1, TreeDump2);
+}
+
+TEST(OutlinedHashTreeRecordTest, Serialize) {
+  OutlinedHashTreeRecord HashTreeRecord1;
+  HashTreeRecord1.HashTree->insert({{1, 2}, 4});
+  HashTreeRecord1.HashTree->insert({{1, 3}, 5});
+
+  // Serialize and deserialize the tree.
+  SmallVector<char> Out;
+  raw_svector_ostream OS(Out);
+  HashTreeRecord1.serialize(OS);
+
+  OutlinedHashTreeRecord HashTreeRecord2;
+  const uint8_t *Data = reinterpret_cast<const uint8_t *>(Out.data());
+  HashTreeRecord2.deserialize(Data);
+
+  // Two trees should be identical.
+  std::string TreeDump1;
+  raw_string_ostream OS1(TreeDump1);
+  HashTreeRecord1.print(OS1);
+  std::string TreeDump2;
+  raw_string_ostream OS2(TreeDump2);
+  HashTreeRecord2.print(OS2);
+
+  EXPECT_EQ(TreeDump1, TreeDump2);
+}
+
+TEST(OutlinedHashTreeRecordTest, SerializeYAML) {
+  OutlinedHashTreeRecord HashTreeRecord1;
+  HashTreeRecord1.HashTree->insert({{1, 2}, 4});
+  HashTreeRecord1.HashTree->insert({{1, 3}, 5});
+
+  // Serialize and deserialize the tree in a YAML format.
+  std::string Out;
+  raw_string_ostream OS(Out);
+  yaml::Output YOS(OS);
+  HashTreeRecord1.serializeYAML(YOS);
+
+  OutlinedHashTreeRecord HashTreeRecord2;
+  yaml::Input YIS(StringRef(Out.data(), Out.size()));
+  HashTreeRecord2.deserializeYAML(YIS);
+
+  // Two trees should be identical.
+  std::string TreeDump1;
+  raw_string_ostream OS1(TreeDump1);
+  HashTreeRecord1.print(OS1);
+  std::string TreeDump2;
+  raw_string_ostream OS2(TreeDump2);
+  HashTreeRecord2.print(OS2);
+
+  EXPECT_EQ(TreeDump1, TreeDump2);
+}
+
+} // end namespace
diff --git a/llvm/unittests/CodeGenData/OutlinedHashTreeTest.cpp b/llvm/unittests/CodeGenData/OutlinedHashTreeTest.cpp
new file mode 100644
index 00000000000000..d11618cf8e4fae
--- /dev/null
+++ b/llvm/unittests/CodeGenData/OutlinedHashTreeTest.cpp
@@ -0,0 +1,81 @@
+//===- OutlinedHashTreeTest.cpp -------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGenData/OutlinedHashTree.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+namespace {
+
+TEST(OutlinedHashTreeTest, Empty) {
+  OutlinedHashTree HashTree;
+  ASSERT_TRUE(HashTree.empty());
+  // The header node is always present.
+  ASSERT_TRUE(HashTree.size() == 1);
+  ASSERT_TRUE(HashTree.depth() == 0);
+}
+
+TEST(OutlinedHashTreeTest, Insert) {
+  OutlinedHashTree HashTree;
+  HashTree.insert({{1, 2, 3}, 1});
+  // The node count is 4 (including the root node).
+  ASSERT_TRUE(HashTree.size() == 4);
+  // The terminal count is 1.
+  ASSERT_TRUE(HashTree.size(/*GetTerminalCountOnly=*/true) == 1);
+  // The depth is 3.
+  ASSERT_TRUE(HashTree.depth() == 3);
+
+  HashTree.clear();
+  ASSERT_TRUE(HashTree.empty());
+
+  HashTree.insert({{1, 2, 3}, 1});
+  HashTree.insert({{1, 2, 4}, 2});
+  // The nodes of 1 and 2 are shared with the same prefix.
+  // The nodes are root, 1, 2, 3 and 4, whose counts are 5.
+  ASSERT_TRUE(HashTree.size() == 5);
+}
+
+TEST(OutlinedHashTreeTest, Find) {
+  OutlinedHashTree HashTree;
+  HashTree.insert({{1, 2, 3}, 1});
+  HashTree.insert({{1, 2, 3}, 2});
+
+  // The node count does not change as the same sequences are added.
+  ASSERT_TRUE(HashTree.size() == 4);
+  // The terminal counts are accumulated from two same sequences.
+  ASSERT_TRUE(HashTree.find({1, 2, 3}) == 3);
+  ASSERT_TRUE(HashTree.find({1, 2}) == 0);
+}
+
+TEST(OutlinedHashTreeTest, Merge) {
+  // Build HashTree1 inserting 2 sequences.
+  OutlinedHashTree HashTree1;
+
+  HashTree1.insert({{1, 2}, 20});
+  HashTree1.insert({{1, 4}, 30});
+
+  // Build HashTree2 and HashTree3 for each
+  OutlinedHashTree HashTree2;
+  HashTree2.insert({{1, 2}, 20});
+  OutlinedHashTree HashTree3;
+  HashTree3.insert({{1, 4}, 30});
+
+  // Merge HashTree3 into HashTree2.
+  HashTree2.merge(&HashTree3);
+
+  // Compare HashTree1 and HashTree2.
+  EXPECT_EQ(HashTree1.size(), HashTree2.size());
+  EXPECT_EQ(HashTree1.depth(), HashTree2.depth());
+  EXPECT_EQ(HashTree1.find({1, 2}), HashTree2.find({1, 2}));
+  EXPECT_EQ(HashTree1.find({1, 4}), HashTree2.find({1, 4}));
+  EXPECT_EQ(HashTree1.find({1, 3}), HashTree2.find({1, 3}));
+}
+
+} // end namespace

>From 338a5d32849af27c3835b97fe098b59beba88e98 Mon Sep 17 00:00:00 2001
From: Kyungwoo Lee <kyulee at meta.com>
Date: Sat, 4 May 2024 17:24:33 -0700
Subject: [PATCH 5/5] Address comments from Ellis

---
 .../llvm/CodeGenData/OutlinedHashTree.h       | 29 +++++++------------
 .../llvm/CodeGenData/OutlinedHashTreeRecord.h | 13 +++++++--
 llvm/lib/CodeGenData/OutlinedHashTree.cpp     | 29 ++++++++++---------
 .../CodeGenData/OutlinedHashTreeRecord.cpp    |  4 +--
 .../CodeGenData/OutlinedHashTreeTest.cpp      |  5 ++--
 5 files changed, 41 insertions(+), 39 deletions(-)

diff --git a/llvm/include/llvm/CodeGenData/OutlinedHashTree.h b/llvm/include/llvm/CodeGenData/OutlinedHashTree.h
index 875e1a78bb4010..c40038cd8c5174 100644
--- a/llvm/include/llvm/CodeGenData/OutlinedHashTree.h
+++ b/llvm/include/llvm/CodeGenData/OutlinedHashTree.h
@@ -30,29 +30,22 @@ namespace llvm {
 /// a hash sequence with that occurrence count.
 struct HashNode {
   /// The hash value of the node.
-  stable_hash Hash;
+  stable_hash Hash = 0;
   /// The number of terminals in the sequence ending at this node.
-  unsigned Terminals;
+  std::optional<unsigned> Terminals;
   /// The successors of this node.
+  /// We don't use DenseMap as a stable_hash value can be tombstone.
   std::unordered_map<stable_hash, std::unique_ptr<HashNode>> Successors;
 };
 
-/// HashNodeStable is the serialized, stable, and compact representation
-/// of a HashNode.
-struct HashNodeStable {
-  llvm::yaml::Hex64 Hash;
-  unsigned Terminals;
-  std::vector<unsigned> SuccessorIds;
-};
-
 class OutlinedHashTree {
 
   using EdgeCallbackFn =
       std::function<void(const HashNode *, const HashNode *)>;
   using NodeCallbackFn = std::function<void(const HashNode *)>;
 
-  using HashSequence = std::vector<stable_hash>;
-  using HashSequencePair = std::pair<std::vector<stable_hash>, unsigned>;
+  using HashSequence = SmallVector<stable_hash>;
+  using HashSequencePair = std::pair<HashSequence, unsigned>;
 
 public:
   /// Walks every edge and node in the OutlinedHashTree and calls CallbackEdge
@@ -66,7 +59,7 @@ class OutlinedHashTree {
 
   /// Release all hash nodes except the root hash node.
   void clear() {
-    assert(getRoot()->Hash == 0 && getRoot()->Terminals == 0);
+    assert(getRoot()->Hash == 0 && !getRoot()->Terminals);
     getRoot()->Successors.clear();
   }
 
@@ -83,8 +76,8 @@ class OutlinedHashTree {
   size_t depth() const;
 
   /// \returns the root hash node of a OutlinedHashTree.
-  const HashNode *getRoot() const { return Root.get(); }
-  HashNode *getRoot() { return Root.get(); }
+  const HashNode *getRoot() const { return &Root; }
+  HashNode *getRoot() { return &Root; }
 
   /// Inserts a \p Sequence into the this tree. The last node in the sequence
   /// will increase Terminals.
@@ -94,12 +87,10 @@ class OutlinedHashTree {
   void merge(const OutlinedHashTree *OtherTree);
 
   /// \returns the matching count if \p Sequence exists in the OutlinedHashTree.
-  unsigned find(const HashSequence &Sequence) const;
-
-  OutlinedHashTree() { Root = std::make_unique<HashNode>(); }
+  std::optional<unsigned> find(const HashSequence &Sequence) const;
 
 private:
-  std::unique_ptr<HashNode> Root;
+  HashNode Root;
 };
 
 } // namespace llvm
diff --git a/llvm/include/llvm/CodeGenData/OutlinedHashTreeRecord.h b/llvm/include/llvm/CodeGenData/OutlinedHashTreeRecord.h
index ccd2ad26dd0871..2960e319604489 100644
--- a/llvm/include/llvm/CodeGenData/OutlinedHashTreeRecord.h
+++ b/llvm/include/llvm/CodeGenData/OutlinedHashTreeRecord.h
@@ -16,13 +16,22 @@
 #ifndef LLVM_CODEGENDATA_OUTLINEDHASHTREERECORD_H
 #define LLVM_CODEGENDATA_OUTLINEDHASHTREERECORD_H
 
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/CodeGenData/OutlinedHashTree.h"
 
 namespace llvm {
 
+/// HashNodeStable is the serialized, stable, and compact representation
+/// of a HashNode.
+struct HashNodeStable {
+  llvm::yaml::Hex64 Hash;
+  unsigned Terminals;
+  std::vector<unsigned> SuccessorIds;
+};
+
 using IdHashNodeStableMapTy = std::map<unsigned, HashNodeStable>;
-using IdHashNodeMapTy = std::map<unsigned, HashNode *>;
-using HashNodeIdMapTy = std::unordered_map<const HashNode *, unsigned>;
+using IdHashNodeMapTy = DenseMap<unsigned, HashNode *>;
+using HashNodeIdMapTy = DenseMap<const HashNode *, unsigned>;
 
 struct OutlinedHashTreeRecord {
   std::unique_ptr<OutlinedHashTree> HashTree;
diff --git a/llvm/lib/CodeGenData/OutlinedHashTree.cpp b/llvm/lib/CodeGenData/OutlinedHashTree.cpp
index 032993ded60ead..3111f7e910dc30 100644
--- a/llvm/lib/CodeGenData/OutlinedHashTree.cpp
+++ b/llvm/lib/CodeGenData/OutlinedHashTree.cpp
@@ -24,19 +24,18 @@ using namespace llvm;
 void OutlinedHashTree::walkGraph(NodeCallbackFn CallbackNode,
                                  EdgeCallbackFn CallbackEdge,
                                  bool SortedWalk) const {
-  std::stack<const HashNode *> Stack;
-  Stack.push(getRoot());
+  SmallVector<const HashNode *> Stack;
+  Stack.emplace_back(getRoot());
 
   while (!Stack.empty()) {
-    const auto *Current = Stack.top();
-    Stack.pop();
+    const auto *Current = Stack.pop_back_val();
     if (CallbackNode)
       CallbackNode(Current);
 
     auto HandleNext = [&](const HashNode *Next) {
       if (CallbackEdge)
         CallbackEdge(Current, Next);
-      Stack.push(Next);
+      Stack.emplace_back(Next);
     };
     if (SortedWalk) {
       std::map<stable_hash, const HashNode *> SortedSuccessors;
@@ -87,22 +86,23 @@ void OutlinedHashTree::insert(const HashSequencePair &SequencePair) {
     } else
       Current = I->second.get();
   }
-  Current->Terminals += Count;
+  if (Count)
+    Current->Terminals = (Current->Terminals ? *Current->Terminals : 0) + Count;
 }
 
 void OutlinedHashTree::merge(const OutlinedHashTree *Tree) {
   HashNode *Dst = getRoot();
   const HashNode *Src = Tree->getRoot();
-  std::stack<std::pair<HashNode *, const HashNode *>> Stack;
-  Stack.push({Dst, Src});
+  SmallVector<std::pair<HashNode *, const HashNode *>> Stack;
+  Stack.emplace_back(Dst, Src);
 
   while (!Stack.empty()) {
-    auto [DstNode, SrcNode] = Stack.top();
-    Stack.pop();
+    auto [DstNode, SrcNode] = Stack.pop_back_val();
     if (!SrcNode)
       continue;
-    DstNode->Terminals += SrcNode->Terminals;
-
+    if (SrcNode->Terminals)
+      DstNode->Terminals =
+          (DstNode->Terminals ? *DstNode->Terminals : 0) + *SrcNode->Terminals;
     for (auto &[Hash, NextSrcNode] : SrcNode->Successors) {
       HashNode *NextDstNode;
       auto I = DstNode->Successors.find(Hash);
@@ -114,12 +114,13 @@ void OutlinedHashTree::merge(const OutlinedHashTree *Tree) {
       } else
         NextDstNode = I->second.get();
 
-      Stack.push({NextDstNode, NextSrcNode.get()});
+      Stack.emplace_back(NextDstNode, NextSrcNode.get());
     }
   }
 }
 
-unsigned OutlinedHashTree::find(const HashSequence &Sequence) const {
+std::optional<unsigned>
+OutlinedHashTree::find(const HashSequence &Sequence) const {
   const HashNode *Current = getRoot();
   for (stable_hash StableHash : Sequence) {
     const auto I = Current->Successors.find(StableHash);
diff --git a/llvm/lib/CodeGenData/OutlinedHashTreeRecord.cpp b/llvm/lib/CodeGenData/OutlinedHashTreeRecord.cpp
index 0d5dd864c89c55..b57c035c8a841f 100644
--- a/llvm/lib/CodeGenData/OutlinedHashTreeRecord.cpp
+++ b/llvm/lib/CodeGenData/OutlinedHashTreeRecord.cpp
@@ -131,7 +131,7 @@ void OutlinedHashTreeRecord::convertToStableData(
     auto Id = P.second;
     HashNodeStable NodeStable;
     NodeStable.Hash = Node->Hash;
-    NodeStable.Terminals = Node->Terminals;
+    NodeStable.Terminals = Node->Terminals ? *Node->Terminals : 0;
     for (auto &P : Node->Successors)
       NodeStable.SuccessorIds.push_back(NodeIdMap[P.second.get()]);
     IdNodeStableMap[Id] = NodeStable;
@@ -139,7 +139,7 @@ void OutlinedHashTreeRecord::convertToStableData(
 
   // Sort the Successors so that they come out in the same order as in the map.
   for (auto &P : IdNodeStableMap)
-    std::sort(P.second.SuccessorIds.begin(), P.second.SuccessorIds.end());
+    llvm::sort(P.second.SuccessorIds);
 }
 
 void OutlinedHashTreeRecord::convertFromStableData(
diff --git a/llvm/unittests/CodeGenData/OutlinedHashTreeTest.cpp b/llvm/unittests/CodeGenData/OutlinedHashTreeTest.cpp
index d11618cf8e4fae..5fdfa60673b7fc 100644
--- a/llvm/unittests/CodeGenData/OutlinedHashTreeTest.cpp
+++ b/llvm/unittests/CodeGenData/OutlinedHashTreeTest.cpp
@@ -50,8 +50,9 @@ TEST(OutlinedHashTreeTest, Find) {
   // The node count does not change as the same sequences are added.
   ASSERT_TRUE(HashTree.size() == 4);
   // The terminal counts are accumulated from two same sequences.
-  ASSERT_TRUE(HashTree.find({1, 2, 3}) == 3);
-  ASSERT_TRUE(HashTree.find({1, 2}) == 0);
+  ASSERT_TRUE(HashTree.find({1, 2, 3}));
+  ASSERT_TRUE(HashTree.find({1, 2, 3}).value() == 3);
+  ASSERT_FALSE(HashTree.find({1, 2}));
 }
 
 TEST(OutlinedHashTreeTest, Merge) {



More information about the llvm-commits mailing list