[ModuloSchedule] Pass loop block explicitly to kernel rewriter.

[hgreving]

Summary

This change is NFC upstream. We pass in the loop's block to the kernel
rewriter explicitly, instead of assuming it's the loop's top block. This
change is made for downstream targets where this assumption doesn't hold.

llvm/lib/CodeGen/ModuloSchedule.cpp

class KernelRewriter {
   // If InitReg is not given it is chosen arbitrarily. It will either be undef
   // or will be chosen so as to share another phi.
   Register phi(Register LoopReg, Optional<Register> InitReg = {},
                const TargetRegisterClass *RC = nullptr);
   // Create an undef register of the given register class.
   Register undef(const TargetRegisterClass *RC);
 
 public:
-  KernelRewriter(MachineLoop &L, ModuloSchedule &S,
+  KernelRewriter(MachineLoop &L, ModuloSchedule &S, MachineBasicBlock *LoopBB,
                  LiveIntervals *LIS = nullptr);
   void rewrite();
 };
 } // namespace
 
 KernelRewriter::KernelRewriter(MachineLoop &L, ModuloSchedule &S,
-                               LiveIntervals *LIS)
-    : S(S), BB(L.getTopBlock()), PreheaderBB(L.getLoopPreheader()),
+                               MachineBasicBlock *LoopBB, LiveIntervals *LIS)
+    : S(S), BB(LoopBB), PreheaderBB(L.getLoopPreheader()),
       ExitBB(L.getExitBlock()), MRI(BB->getParent()->getRegInfo()),
       TII(BB->getParent()->getSubtarget().getInstrInfo()), LIS(LIS) {
   PreheaderBB = *BB->pred_begin();
   if (PreheaderBB == BB)
     PreheaderBB = *std::next(BB->pred_begin());
 }
 
 void KernelRewriter::rewrite() {
if (!KernelDisposed) {
     LoopInfo->adjustTripCount(-(Schedule.getNumStages() - 1));
     LoopInfo->setPreheader(Prologs.back());
   } else {
     LoopInfo->disposed();
   }
 }
 
 void PeelingModuloScheduleExpander::rewriteKernel() {
-  KernelRewriter KR(*Schedule.getLoop(), Schedule);
+  KernelRewriter KR(*Schedule.getLoop(), Schedule, BB);
   KR.rewrite();
 }
 
 void PeelingModuloScheduleExpander::expand() {
   BB = Schedule.getLoop()->getTopBlock();
   Preheader = Schedule.getLoop()->getLoopPreheader();
   LLVM_DEBUG(Schedule.dump());
   LoopInfo = TII->analyzeLoopForPipelining(BB);
if (!ExpandedKernel) {
     // The expander optimized away the kernel. We can't do any useful checking.
     MSE.cleanup();
     return;
   }
   // Before running the KernelRewriter, re-add BB into the CFG.
   Preheader->addSuccessor(BB);
 
   // Now run the new expansion algorithm.
-  KernelRewriter KR(*Schedule.getLoop(), Schedule);
+  KernelRewriter KR(*Schedule.getLoop(), Schedule, BB);
   KR.rewrite();
   peelPrologAndEpilogs();
 
   // Collect all illegal phis that the new algorithm created. We'll give these
   // to KernelOperandInfo.
   SmallPtrSet<MachineInstr *, 4> IllegalPhis;
   for (auto NI = BB->getFirstNonPHI(); NI != BB->end(); ++NI) {
     if (NI->isPHI())