[llvm][CodeGen] Resolve issues when updating live intervals in window scheduler

Corrupted live interval information can cause window scheduling to crash
in some cases. By adding the missing MBB's live interval information in
the ModuloScheduleExpander, the information can be correctly analyzed
in the window scheduler.

Author: huaatian

llvm/include/llvm/CodeGen/WindowScheduler.h

@@ -150,8 +150,6 @@ class WindowScheduler {
   /// Using the scheduling infrastructure to expand the results of window
   /// scheduling. It is usually necessary to add prologue and epilogue MBBs.
   virtual void expand();
-  /// Reanalyze the live intervals in some scenarios.
-  virtual void reanalyzeLiveIntervals();
   /// Update the live intervals for all registers used within MBB.
   virtual void updateLiveIntervals();
   /// Estimate a II value at which all MIs will be scheduled successfully.

llvm/lib/CodeGen/ModuloSchedule.cpp

@@ -130,6 +130,7 @@ void ModuloScheduleExpander::generatePipelinedLoop() {
   // Generate the prolog instructions that set up the pipeline.
   generateProlog(MaxStageCount, KernelBB, VRMap, PrologBBs);
   MF.insert(BB->getIterator(), KernelBB);
+  LIS.insertMBBInMaps(KernelBB);
 
   // Rearrange the instructions to generate the new, pipelined loop,
   // and update register names as needed.
@@ -210,6 +211,7 @@ void ModuloScheduleExpander::generateProlog(unsigned LastStage,
     NewBB->transferSuccessors(PredBB);
     PredBB->addSuccessor(NewBB);
     PredBB = NewBB;
+    LIS.insertMBBInMaps(NewBB);
 
     // Generate instructions for each appropriate stage. Process instructions
     // in original program order.
@@ -283,6 +285,7 @@ void ModuloScheduleExpander::generateEpilog(
 
     PredBB->replaceSuccessor(LoopExitBB, NewBB);
     NewBB->addSuccessor(LoopExitBB);
+    LIS.insertMBBInMaps(NewBB);
 
     if (EpilogStart == LoopExitBB)
       EpilogStart = NewBB;

llvm/lib/CodeGen/WindowScheduler.cpp

@@ -242,7 +242,6 @@ bool WindowScheduler::initialize() {
     LLVM_DEBUG(dbgs() << "There are too few MIs in the window region!\n");
     return false;
   }
-  reanalyzeLiveIntervals();
   return true;
 }
 
@@ -644,20 +643,6 @@ void WindowScheduler::expand() {
   MSE.cleanup();
 }
 
-void WindowScheduler::reanalyzeLiveIntervals() {
-  auto SlotIndex = Context->LIS->getSlotIndexes();
-  // Check if the SlotIndex infomation is missing.
-  for (auto &MBB : *MF)
-    for (auto &MI : MBB)
-      if (!SlotIndex->hasIndex(MI)) {
-        // The slot index and live intervals of MF have been corrupted and need
-        // to be reanalyzed.
-        SlotIndex->reanalyze(*MF);
-        Context->LIS->reanalyze(*MF);
-        return;
-      }
-}
-
 void WindowScheduler::updateLiveIntervals() {
   SmallVector<Register, 128> UsedRegs;
   for (MachineInstr &MI : *MBB)

llvm/test/CodeGen/Hexagon/swp-ws-live-intervals.mir

@@ -8,8 +8,21 @@
 # The bug was reported at https://github.com/llvm/llvm-project/pull/99454. 
 # It is caused by the corruption of live intervals in certain scenarios.
 #
+# We check the newly generated MBBs after successful scheduling here.
 # CHECK: Best window offset is {{[0-9]+}} and Best II is {{[0-9]+}}.
+# CHECK: prolog:
+# CHECK: bb.5:
+# CHECK: New block
+# CHECK: bb.6:
+# CHECK: epilog:
+# CHECK: bb.7:
 # CHECK: Best window offset is {{[0-9]+}} and Best II is {{[0-9]+}}.
+# CHECK: prolog:
+# CHECK: bb.8:
+# CHECK: New block
+# CHECK: bb.9:
+# CHECK: epilog:
+# CHECK: bb.10:
 
 --- |
   target triple = "hexagon"