Index: llvm/lib/Target/AMDGPU/SIOptimizeVGPRLiveRange.cpp
===================================================================
--- llvm/lib/Target/AMDGPU/SIOptimizeVGPRLiveRange.cpp
+++ llvm/lib/Target/AMDGPU/SIOptimizeVGPRLiveRange.cpp
@@ -384,12 +384,24 @@
   assert(ThenEntry && "No successor in Then region?");
 
   LiveVariables::VarInfo &OldVarInfo = LV->getVarInfo(Reg);
-  df_iterator_default_set<MachineBasicBlock *, 16> Visited;
-
-  for (MachineBasicBlock *MBB : depth_first_ext(ThenEntry, Visited)) {
-    if (MBB == Flow)
-      continue;
+  SetVector<MachineBasicBlock *> Blocks;
+  SmallVector<MachineBasicBlock *> WorkList;
+
+  WorkList.push_back(If);
+
+  // Collect all successors until we see the flow block, where we should
+  // reconverge.
+  while (!WorkList.empty()) {
+    auto *MBB = WorkList.pop_back_val();
+    for (auto *Succ : MBB->successors()) {
+      if (Succ != Flow && !Blocks.contains(Succ)) {
+        WorkList.push_back(Succ);
+        Blocks.insert(Succ);
+      }
+    }
+  }
 
+  for (MachineBasicBlock *MBB : Blocks) {
     // Clear Live bit, as we will recalculate afterwards
     LLVM_DEBUG(dbgs() << "Clear AliveBlock " << printMBBReference(*MBB)
                       << '\n');
@@ -401,17 +413,12 @@
        ++I) {
     auto *UseMI = I->getParent();
     if (UseMI->isPHI() && I->readsReg()) {
-      if (Visited.contains(UseMI->getParent()))
+      if (Blocks.contains(UseMI->getParent()))
         PHIIncoming.insert(UseMI->getOperand(I.getOperandNo() + 1).getMBB());
     }
   }
 
-  Visited.clear();
-
-  for (MachineBasicBlock *MBB : depth_first_ext(ThenEntry, Visited)) {
-    if (MBB == Flow)
-      break;
-
+  for (MachineBasicBlock *MBB : Blocks) {
     SmallVector<MachineInstr *> Uses;
     // PHI instructions has been processed before.
     findNonPHIUsesInBlock(Reg, MBB, Uses);
@@ -438,7 +445,7 @@
 
   // Set the isKilled flag if we get new Kills in the THEN region.
   for (auto *MI : OldVarInfo.Kills) {
-    if (Visited.contains(MI->getParent()))
+    if (Blocks.contains(MI->getParent()))
       MI->addRegisterKilled(Reg, TRI);
   }
 }