1 files changed, 44 insertions, 30 deletions

diff --git a/llvm/lib/CodeGen/ModuloSchedule.cpp b/llvm/lib/CodeGen/ModuloSchedule.cpp
index 163e52d9199d1..d85b1b7988cec 100644
--- a/llvm/lib/CodeGen/ModuloSchedule.cpp
+++ b/llvm/lib/CodeGen/ModuloSchedule.cpp
@@ -8,6 +8,7 @@
 
 #include "llvm/CodeGen/ModuloSchedule.h"
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/Analysis/MemoryLocation.h"
 #include "llvm/CodeGen/LiveIntervals.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineLoopUtils.h"
@@ -420,7 +421,7 @@ void ModuloScheduleExpander::generateExistingPhis(
     unsigned NewReg = 0;
     unsigned AccessStage = (LoopValStage != -1) ? LoopValStage : StageScheduled;
     // In the epilog, we may need to look back one stage to get the correct
-    // Phi name because the epilog and prolog blocks execute the same stage.
+    // Phi name, because the epilog and prolog blocks execute the same stage.
     // The correct name is from the previous block only when the Phi has
     // been completely scheduled prior to the epilog, and Phi value is not
     // needed in multiple stages.
@@ -913,7 +914,12 @@ bool ModuloScheduleExpander::computeDelta(MachineInstr &MI, unsigned &Delta) {
   const TargetRegisterInfo *TRI = MF.getSubtarget().getRegisterInfo();
   const MachineOperand *BaseOp;
   int64_t Offset;
-  if (!TII->getMemOperandWithOffset(MI, BaseOp, Offset, TRI))
+  bool OffsetIsScalable;
+  if (!TII->getMemOperandWithOffset(MI, BaseOp, Offset, OffsetIsScalable, TRI))
+    return false;
+
+  // FIXME: This algorithm assumes instructions have fixed-size offsets.
+  if (OffsetIsScalable)
     return false;
 
   if (!BaseOp->isReg())
@@ -1435,11 +1441,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;
   }
 
@@ -1620,18 +1630,21 @@ void PeelingModuloScheduleExpander::moveStageBetweenBlocks(
     MachineInstr *MI = &*I++;
     if (MI->isPHI()) {
       // This is an illegal PHI. If we move any instructions using an illegal
-      // PHI, we need to create a legal Phi
-      Register PhiR = MI->getOperand(0).getReg();
-      auto RC = MRI.getRegClass(PhiR);
-      Register NR = MRI.createVirtualRegister(RC);
-      MachineInstr *NI = BuildMI(*DestBB, DestBB->getFirstNonPHI(), DebugLoc(),
-                                 TII->get(TargetOpcode::PHI), NR)
-                             .addReg(PhiR)
-                             .addMBB(SourceBB);
-      BlockMIs[{DestBB, CanonicalMIs[MI]}] = NI;
-      CanonicalMIs[NI] = CanonicalMIs[MI];
-      Remaps[PhiR] = NR;
-      continue;
+      // PHI, we need to create a legal Phi.
+      if (getStage(MI) != Stage) {
+        // The legal Phi is not necessary if the illegal phi's stage
+        // is being moved.
+        Register PhiR = MI->getOperand(0).getReg();
+        auto RC = MRI.getRegClass(PhiR);
+        Register NR = MRI.createVirtualRegister(RC);
+        MachineInstr *NI = BuildMI(*DestBB, DestBB->getFirstNonPHI(),
+                                   DebugLoc(), TII->get(TargetOpcode::PHI), NR)
+                               .addReg(PhiR)
+                               .addMBB(SourceBB);
+        BlockMIs[{DestBB, CanonicalMIs[MI]}] = NI;
+        CanonicalMIs[NI] = CanonicalMIs[MI];
+        Remaps[PhiR] = NR;
+      }
     }
     if (getStage(MI) != Stage)
       continue;
@@ -1649,8 +1662,8 @@ void PeelingModuloScheduleExpander::moveStageBetweenBlocks(
     // we don't need the phi anymore.
     if (getStage(Def) == Stage) {
       Register PhiReg = MI.getOperand(0).getReg();
-      MRI.replaceRegWith(MI.getOperand(0).getReg(),
-                         Def->getOperand(0).getReg());
+      assert(Def->findRegisterDefOperandIdx(MI.getOperand(1).getReg()) != -1);
+      MRI.replaceRegWith(MI.getOperand(0).getReg(), MI.getOperand(1).getReg());
       MI.getOperand(0).setReg(PhiReg);
       PhiToDelete.push_back(&MI);
     }
@@ -1698,16 +1711,17 @@ PeelingModuloScheduleExpander::getPhiCanonicalReg(MachineInstr *CanonicalPhi,
                                                   MachineInstr *Phi) {
   unsigned distance = PhiNodeLoopIteration[Phi];
   MachineInstr *CanonicalUse = CanonicalPhi;
+  Register CanonicalUseReg = CanonicalUse->getOperand(0).getReg();
   for (unsigned I = 0; I < distance; ++I) {
     assert(CanonicalUse->isPHI());
     assert(CanonicalUse->getNumOperands() == 5);
     unsigned LoopRegIdx = 3, InitRegIdx = 1;
     if (CanonicalUse->getOperand(2).getMBB() == CanonicalUse->getParent())
       std::swap(LoopRegIdx, InitRegIdx);
-    CanonicalUse =
-        MRI.getVRegDef(CanonicalUse->getOperand(LoopRegIdx).getReg());
+    CanonicalUseReg = CanonicalUse->getOperand(LoopRegIdx).getReg();
+    CanonicalUse = MRI.getVRegDef(CanonicalUseReg);
   }
-  return CanonicalUse->getOperand(0).getReg();
+  return CanonicalUseReg;
 }
 
 void PeelingModuloScheduleExpander::peelPrologAndEpilogs() {
@@ -1933,7 +1947,7 @@ void PeelingModuloScheduleExpander::fixupBranches() {
     SmallVector<MachineOperand, 4> Cond;
     TII->removeBranch(*Prolog);
     Optional<bool> StaticallyGreater =
-        Info->createTripCountGreaterCondition(TC, *Prolog, Cond);
+        LoopInfo->createTripCountGreaterCondition(TC, *Prolog, Cond);
     if (!StaticallyGreater.hasValue()) {
       LLVM_DEBUG(dbgs() << "Dynamic: TC > " << TC << "\n");
       // Dynamically branch based on Cond.
@@ -1961,10 +1975,10 @@ void PeelingModuloScheduleExpander::fixupBranches() {
   }
 
   if (!KernelDisposed) {
-    Info->adjustTripCount(-(Schedule.getNumStages() - 1));
-    Info->setPreheader(Prologs.back());
+    LoopInfo->adjustTripCount(-(Schedule.getNumStages() - 1));
+    LoopInfo->setPreheader(Prologs.back());
   } else {
-    Info->disposed();
+    LoopInfo->disposed();
   }
 }
 
@@ -1977,8 +1991,8 @@ void PeelingModuloScheduleExpander::expand() {
   BB = Schedule.getLoop()->getTopBlock();
   Preheader = Schedule.getLoop()->getLoopPreheader();
   LLVM_DEBUG(Schedule.dump());
-  Info = TII->analyzeLoopForPipelining(BB);
-  assert(Info);
+  LoopInfo = TII->analyzeLoopForPipelining(BB);
+  assert(LoopInfo);
 
   rewriteKernel();
   peelPrologAndEpilogs();