From 36b5d3f1af9ad132917c79c6a28ec6689385608e Mon Sep 17 00:00:00 2001
From: Hendrik Greving <hgreving@google.com>
Date: Wed, 23 Jun 2021 12:32:04 -0700
Subject: [PATCH] [ModuloSchedule] Pass loop block explicitly to kernel
 rewriter.

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.

Differential Revision: https://reviews.llvm.org/D104811
---
 lib/CodeGen/ModuloSchedule.cpp | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/lib/CodeGen/ModuloSchedule.cpp b/lib/CodeGen/ModuloSchedule.cpp
index 095da09ea82..b5517c40a28 100644
--- a/lib/CodeGen/ModuloSchedule.cpp
+++ b/lib/CodeGen/ModuloSchedule.cpp
@@ -1275,15 +1275,15 @@ class KernelRewriter {
   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();
@@ -1981,7 +1981,7 @@ void PeelingModuloScheduleExpander::fixupBranches() {
 }
 
 void PeelingModuloScheduleExpander::rewriteKernel() {
-  KernelRewriter KR(*Schedule.getLoop(), Schedule);
+  KernelRewriter KR(*Schedule.getLoop(), Schedule, BB);
   KR.rewrite();
 }
 
@@ -2024,7 +2024,7 @@ void PeelingModuloScheduleExpander::validateAgainstModuloScheduleExpander() {
   Preheader->addSuccessor(BB);
 
   // Now run the new expansion algorithm.
-  KernelRewriter KR(*Schedule.getLoop(), Schedule);
+  KernelRewriter KR(*Schedule.getLoop(), Schedule, BB);
   KR.rewrite();
   peelPrologAndEpilogs();