1 files changed, 9 insertions, 5 deletions
diff --git a/llvm/lib/CodeGen/ModuloSchedule.cpp b/llvm/lib/CodeGen/ModuloSchedule.cpp
index 3fa2285..65dc11e 100644
--- a/llvm/lib/CodeGen/ModuloSchedule.cpp
+++ b/llvm/lib/CodeGen/ModuloSchedule.cpp
@@ -1440,11 +1440,15 @@ Register KernelRewriter::remapUse(Register Reg, MachineInstr &MI) {
     // immediately prior to pruning.
     auto RC = MRI.getRegClass(Reg);
     Register R = MRI.createVirtualRegister(RC);
-    BuildMI(*BB, MI, DebugLoc(), TII->get(TargetOpcode::PHI), R)
-        .addReg(IllegalPhiDefault.getValue())
-        .addMBB(PreheaderBB) // Block choice is arbitrary and has no effect.
-        .addReg(LoopReg)
-        .addMBB(BB); // Block choice is arbitrary and has no effect.
+    MachineInstr *IllegalPhi =
+        BuildMI(*BB, MI, DebugLoc(), TII->get(TargetOpcode::PHI), R)
+            .addReg(IllegalPhiDefault.getValue())
+            .addMBB(PreheaderBB) // Block choice is arbitrary and has no effect.
+            .addReg(LoopReg)
+            .addMBB(BB); // Block choice is arbitrary and has no effect.
+    // Illegal phi should belong to the producer stage so that it can be
+    // filtered correctly during peeling.
+    S.setStage(IllegalPhi, LoopProducerStage);
     return R;
   }