diff --git a/llvm/include/llvm/CodeGen/SlotIndexes.h b/llvm/include/llvm/CodeGen/SlotIndexes.h
index 1517d33e65a38..828165c8b887b 100644
--- a/llvm/include/llvm/CodeGen/SlotIndexes.h
+++ b/llvm/include/llvm/CodeGen/SlotIndexes.h
@@ -472,18 +472,25 @@ class raw_ostream;
     /// begin and basic block)
     using MBBIndexIterator = SmallVectorImpl<IdxMBBPair>::const_iterator;
 
-    /// Move iterator to the next IdxMBBPair where the SlotIndex is greater or
-    /// equal to \p To.
-    MBBIndexIterator advanceMBBIndex(MBBIndexIterator I, SlotIndex To) const {
-      return std::partition_point(
-          I, idx2MBBMap.end(),
-          [=](const IdxMBBPair &IM) { return IM.first < To; });
+    /// Get an iterator pointing to the first IdxMBBPair with SlotIndex greater
+    /// than or equal to \p Idx. If \p Start is provided, only search the range
+    /// from \p Start to the end of the function.
+    MBBIndexIterator getMBBLowerBound(MBBIndexIterator Start,
+                                      SlotIndex Idx) const {
+      return std::lower_bound(
+          Start, MBBIndexEnd(), Idx,
+          [](const IdxMBBPair &IM, SlotIndex Idx) { return IM.first < Idx; });
+    }
+    MBBIndexIterator getMBBLowerBound(SlotIndex Idx) const {
+      return getMBBLowerBound(MBBIndexBegin(), Idx);
     }
 
-    /// Get an iterator pointing to the IdxMBBPair with the biggest SlotIndex
-    /// that is greater or equal to \p Idx.
-    MBBIndexIterator findMBBIndex(SlotIndex Idx) const {
-      return advanceMBBIndex(idx2MBBMap.begin(), Idx);
+    /// Get an iterator pointing to the first IdxMBBPair with SlotIndex greater
+    /// than \p Idx.
+    MBBIndexIterator getMBBUpperBound(SlotIndex Idx) const {
+      return std::upper_bound(
+          MBBIndexBegin(), MBBIndexEnd(), Idx,
+          [](SlotIndex Idx, const IdxMBBPair &IM) { return Idx < IM.first; });
     }
 
     /// Returns an iterator for the begin of the idx2MBBMap.
@@ -501,16 +508,11 @@ class raw_ostream;
       if (MachineInstr *MI = getInstructionFromIndex(index))
         return MI->getParent();
 
-      MBBIndexIterator I = findMBBIndex(index);
-      // Take the pair containing the index
-      MBBIndexIterator J =
-        ((I != MBBIndexEnd() && I->first > index) ||
-         (I == MBBIndexEnd() && !idx2MBBMap.empty())) ? std::prev(I) : I;
-
-      assert(J != MBBIndexEnd() && J->first <= index &&
-             index < getMBBEndIdx(J->second) &&
+      MBBIndexIterator I = std::prev(getMBBUpperBound(index));
+      assert(I != MBBIndexEnd() && I->first <= index &&
+             index < getMBBEndIdx(I->second) &&
              "index does not correspond to an MBB");
-      return J->second;
+      return I->second;
     }
 
     /// Insert the given machine instruction into the mapping. Returns the
diff --git a/llvm/lib/CodeGen/VirtRegMap.cpp b/llvm/lib/CodeGen/VirtRegMap.cpp
index 5bdd86aebcd03..48f4ee29fbe95 100644
--- a/llvm/lib/CodeGen/VirtRegMap.cpp
+++ b/llvm/lib/CodeGen/VirtRegMap.cpp
@@ -312,7 +312,7 @@ void VirtRegRewriter::addLiveInsForSubRanges(const LiveInterval &LI,
 
   // Check all mbb start positions between First and Last while
   // simultaneously advancing an iterator for each subrange.
-  for (SlotIndexes::MBBIndexIterator MBBI = Indexes->findMBBIndex(First);
+  for (SlotIndexes::MBBIndexIterator MBBI = Indexes->getMBBLowerBound(First);
        MBBI != Indexes->MBBIndexEnd() && MBBI->first <= Last; ++MBBI) {
     SlotIndex MBBBegin = MBBI->first;
     // Advance all subrange iterators so that their end position is just
@@ -363,7 +363,7 @@ void VirtRegRewriter::addMBBLiveIns() {
       // sorted by slot indexes.
       SlotIndexes::MBBIndexIterator I = Indexes->MBBIndexBegin();
       for (const auto &Seg : LI) {
-        I = Indexes->advanceMBBIndex(I, Seg.start);
+        I = Indexes->getMBBLowerBound(I, Seg.start);
         for (; I != Indexes->MBBIndexEnd() && I->first < Seg.end; ++I) {
           MachineBasicBlock *MBB = I->second;
           MBB->addLiveIn(PhysReg);