[llvm] efficient implementation of MachineOutliner::findCandidates() (PR #90260)

[Xuan Zhang via llvm-commits]

https://github.com/xuanzh-meta created https://github.com/llvm/llvm-project/pull/90260

None

>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] 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