[ModuloSchedule] Fix a bug in experimental expander during prologue/epilogue stitching.

[ThomasRaoux]

Summary

Fix two problems that popped up after my last patch. One is that the
stitching of prologue/epilogue can be wrong when reading a value from a
previous stage. Also changed how we duplicate phi instructions to avoid
generating extra phi that we delete later.

[jmolloy]

llvm/lib/CodeGen/ModuloSchedule.cpp

1670–1671

What's the rationale for this change? Is this just to avoid creating a phi that we destroy later?

If so, I'd prefer to keep the logic simple over optimal. We can always clean up dead nodes easily.

1791

The nesting is getting too deep here, let's move this into a helper function.

[ThomasRaoux]

llvm/lib/CodeGen/ModuloSchedule.cpp

1670–1671

Yes the goal is to avoid inserting useless Phi. The problem is that when we have a large number of stages we can run into combinatorial explosion since we move the phi instructions block by block. I moved the logic in a lambda to make it easier to read.

[ThomasRaoux]

@jmolloy, could you take one more look when you get a chance?

[jmolloy]

LGTM, sorry for the delay.

[ThomasRaoux]

In D70213#1756873, @jmolloy wrote:

LGTM, sorry for the delay.

Thanks James!

llvm/include/llvm/CodeGen/ModuloSchedule.h

class PeelingModuloScheduleExpander {
   /// All prolog and epilog blocks.
   SmallVector<MachineBasicBlock *, 4> Prologs, Epilogs;
   /// For every block, the stages that are produced.
   DenseMap<MachineBasicBlock *, BitVector> LiveStages;
   /// For every block, the stages that are available. A stage can be available
   /// but not produced (in the epilog) or produced but not available (in the
   /// prolog).
   DenseMap<MachineBasicBlock *, BitVector> AvailableStages;
+  /// When peeling the epilogue keep track of the distance between the phi
+  /// nodes and the kernel.
+  DenseMap<MachineInstr *, unsigned> PhiNodeLoopIteration;
 
   /// CanonicalMIs and BlockMIs form a bidirectional map between any of the
   /// loop kernel clones.
   DenseMap<MachineInstr *, MachineInstr *> CanonicalMIs;
   DenseMap<std::pair<MachineBasicBlock *, MachineInstr *>, MachineInstr *>
       BlockMIs;
 
   /// State passed from peelKernel to peelPrologAndEpilogs().
private:
   /// the loop exiting block as any other kernel clone.
   MachineBasicBlock *CreateLCSSAExitingBlock();
   /// Helper to get the stage of an instruction in the schedule.
   unsigned getStage(MachineInstr *MI) {
     if (CanonicalMIs.count(MI))
       MI = CanonicalMIs[MI];
     return Schedule.getStage(MI);
   }
+  /// Helper function to find the right canonical register for a phi instruction
+  /// coming from a peeled out prologue.
+  Register getPhiCanonicalReg(MachineInstr* CanonicalPhi, MachineInstr* Phi);
   /// Target loop info before kernel peeling.
   std::unique_ptr<TargetInstrInfo::PipelinerLoopInfo> Info;
 };
 
 /// Expander that simply annotates each scheduled instruction with a post-instr
 /// symbol that can be consumed by the ModuloScheduleTest pass.
 ///
 /// The post-instr symbol is a way of annotating an instruction that can be

llvm/lib/CodeGen/ModuloSchedule.cpp

 // This is a reimplementation of ModuloScheduleExpander that works by creating
 // a fully correct steady-state kernel and peeling off the prolog and epilogs.
 //===----------------------------------------------------------------------===//
 
 namespace {
 // Remove any dead phis in MBB. Dead phis either have only one block as input
 // (in which case they are the identity) or have no uses.
 void EliminateDeadPhis(MachineBasicBlock *MBB, MachineRegisterInfo &MRI,
-                       LiveIntervals *LIS) {
+                       LiveIntervals *LIS, bool KeepSingleSrcPhi = false) {
   bool Changed = true;
   while (Changed) {
     Changed = false;
     for (auto I = MBB->begin(); I != MBB->getFirstNonPHI();) {
       MachineInstr &MI = *I++;
       assert(MI.isPHI());
       if (MRI.use_empty(MI.getOperand(0).getReg())) {
         if (LIS)
           LIS->RemoveMachineInstrFromMaps(MI);
         MI.eraseFromParent();
         Changed = true;
-      } else if (MI.getNumExplicitOperands() == 3) {
+      } else if (!KeepSingleSrcPhi && MI.getNumExplicitOperands() == 3) {
         MRI.constrainRegClass(MI.getOperand(1).getReg(),
                               MRI.getRegClass(MI.getOperand(0).getReg()));
         MRI.replaceRegWith(MI.getOperand(0).getReg(),
                            MI.getOperand(1).getReg());
         if (LIS)
           LIS->RemoveMachineInstrFromMaps(MI);
         MI.eraseFromParent();
         Changed = true;
if (getStage(Def) == Stage) {
                          Def->getOperand(0).getReg());
       MI.getOperand(0).setReg(PhiReg);
       PhiToDelete.push_back(&MI);
     }
   }
   for (auto *P : PhiToDelete)
     P->eraseFromParent();
   InsertPt = DestBB->getFirstNonPHI();
-  for (MachineInstr &MI : SourceBB->phis()) {
-    MachineInstr *NewMI = MF.CloneMachineInstr(&MI);
+  // Helper to clone Phi instructions into the destination block. We clone Phi
+  // greedily to avoid combinatorial explosion of Phi instructions.
+  auto clonePhi = [&](MachineInstr *Phi) {
+    MachineInstr *NewMI = MF.CloneMachineInstr(Phi);
     DestBB->insert(InsertPt, NewMI);
-    Register OrigR = MI.getOperand(0).getReg();
+    Register OrigR = Phi->getOperand(0).getReg();
     Register R = MRI.createVirtualRegister(MRI.getRegClass(OrigR));
     NewMI->getOperand(0).setReg(R);
     NewMI->getOperand(1).setReg(OrigR);
     NewMI->getOperand(2).setMBB(*DestBB->pred_begin());
     Remaps[OrigR] = R;
-    CanonicalMIs[NewMI] = CanonicalMIs[&MI];
-    BlockMIs[{DestBB, CanonicalMIs[&MI]}] = NewMI;
-  }
-  for (auto I = DestBB->getFirstNonPHI(); I != DestBB->end(); ++I)
-    for (MachineOperand &MO : I->uses())
-      if (MO.isReg() && Remaps.count(MO.getReg()))
+    CanonicalMIs[NewMI] = CanonicalMIs[Phi];
+    BlockMIs[{DestBB, CanonicalMIs[Phi]}] = NewMI;
+    PhiNodeLoopIteration[NewMI] = PhiNodeLoopIteration[Phi];
+    return R;
+  };
+  for (auto I = DestBB->getFirstNonPHI(); I != DestBB->end(); ++I) {
+    for (MachineOperand &MO : I->uses()) {
+      if (!MO.isReg())
+        continue;
+      if (Remaps.count(MO.getReg()))
         MO.setReg(Remaps[MO.getReg()]);
+      else {
+        // If we are using a phi from the source block we need to add a new phi
+        // pointing to the old one.
+        MachineInstr *Use = MRI.getUniqueVRegDef(MO.getReg());
+        if (Use && Use->isPHI() && Use->getParent() == SourceBB) {
+          Register R = clonePhi(Use);
+          MO.setReg(R);
+        }
+      }
+    }
+  }
+}
+
+Register
+PeelingModuloScheduleExpander::getPhiCanonicalReg(MachineInstr *CanonicalPhi,
+                                                  MachineInstr *Phi) {
+  unsigned distance = PhiNodeLoopIteration[Phi];
+  MachineInstr *CanonicalUse = CanonicalPhi;
+  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());
+  }
+  return CanonicalUse->getOperand(0).getReg();
 }
 
 void PeelingModuloScheduleExpander::peelPrologAndEpilogs() {
   BitVector LS(Schedule.getNumStages(), true);
   BitVector AS(Schedule.getNumStages(), true);
   LiveStages[BB] = LS;
   AvailableStages[BB] = AS;
 
   // Peel out the prologs.
   LS.reset();
   for (int I = 0; I < Schedule.getNumStages() - 1; ++I) {
     LS[I] = 1;
     Prologs.push_back(peelKernel(LPD_Front));
     LiveStages[Prologs.back()] = LS;
     AvailableStages[Prologs.back()] = LS;
   }
 
   // Create a block that will end up as the new loop exiting block (dominated by
   // all prologs and epilogs). It will only contain PHIs, in the same order as
   // BB's PHIs. This gives us a poor-man's LCSSA with the inductive property
   // that the exiting block is a (sub) clone of BB. This in turn gives us the
   // property that any value deffed in BB but used outside of BB is used by a
   // PHI in the exiting block.
   MachineBasicBlock *ExitingBB = CreateLCSSAExitingBlock();
-
+  EliminateDeadPhis(ExitingBB, MRI, LIS, /*KeepSingleSrcPhi=*/true);
   // Push out the epilogs, again in reverse order.
   // We can't assume anything about the minumum loop trip count at this point,
   // so emit a fairly complex epilog.
 
   // We first peel number of stages minus one epilogue. Then we remove dead
   // stages and reorder instructions based on their stage. If we have 3 stages
   // we generate first:
   // E0[3, 2, 1]
   // E1[3', 2']
   // E2[3'']
   // And then we move instructions based on their stages to have:
   // E0[3]
   // E1[2, 3']
   // E2[1, 2', 3'']
   // The transformation is legal because we only move instructions past
   // instructions of a previous loop iteration.
   for (int I = 1; I <= Schedule.getNumStages() - 1; ++I) {
     Epilogs.push_back(peelKernel(LPD_Back));
-    filterInstructions(Epilogs.back(), Schedule.getNumStages() - I);
+    MachineBasicBlock *B = Epilogs.back();
+    filterInstructions(B, Schedule.getNumStages() - I);
+    // Keep track at which iteration each phi belongs to. We need it to know
+    // what version of the variable to use during prologue/epilogue stitching.
+    EliminateDeadPhis(B, MRI, LIS, /*KeepSingleSrcPhi=*/true);
+    for (auto Phi = B->begin(), IE = B->getFirstNonPHI(); Phi != IE; ++Phi)
+      PhiNodeLoopIteration[&*Phi] = Schedule.getNumStages() - I;
   }
   for (size_t I = 0; I < Epilogs.size(); I++) {
     LS.reset();
     for (size_t J = I; J < Epilogs.size(); J++) {
       int Iteration = J;
       unsigned Stage = Schedule.getNumStages() - 1 + I - J;
       // Move stage one block at a time so that Phi nodes are updated correctly.
       for (size_t K = Iteration; K > I; K--)
void PeelingModuloScheduleExpander::peelPrologAndEpilogs() {
   auto EI = Epilogs.begin();
   assert(Prologs.size() == Epilogs.size());
   for (; PI != Prologs.end(); ++PI, ++EI) {
     MachineBasicBlock *Pred = *(*EI)->pred_begin();
     (*PI)->addSuccessor(*EI);
     for (MachineInstr &MI : (*EI)->phis()) {
       Register Reg = MI.getOperand(1).getReg();
       MachineInstr *Use = MRI.getUniqueVRegDef(Reg);
-      if (Use && Use->getParent() == Pred)
+      if (Use && Use->getParent() == Pred) {
+        MachineInstr *CanonicalUse = CanonicalMIs[Use];
+        if (CanonicalUse->isPHI()) {
+          // If the use comes from a phi we need to skip as many phi as the
+          // distance between the epilogue and the kernel. Trace through the phi
+          // chain to find the right value.
+          Reg = getPhiCanonicalReg(CanonicalUse, Use);
+        }
         Reg = getEquivalentRegisterIn(Reg, *PI);
+      }
       MI.addOperand(MachineOperand::CreateReg(Reg, /*isDef=*/false));
       MI.addOperand(MachineOperand::CreateMBB(*PI));
     }
   }
 
   // Create a list of all blocks in order.
   SmallVector<MachineBasicBlock *, 8> Blocks;
   llvm::copy(PeeledFront, std::back_inserter(Blocks));

llvm/test/CodeGen/Hexagon/swp-epilog-phi12.ll

 ; RUN: llc -march=hexagon -hexagon-initial-cfg-cleanup=0 -pipeliner-experimental-cg=true < %s | FileCheck %s
 
 ; Test epilogue generation when reading loop-carried dependency from a previous
 ; stage. The first epilogue should read value from iteration N-1 of the kernel.
 
 ; CHECK: loop0
 ; CHECK: r{{[0-9]+}} = add([[REG0:r([0-9]+)]],#8)
-; CHECK: [[REG0:r([0-9]+)]] = [[REG1:r([0-9]+)]]
+; CHECK: [[REG0]] = [[REG1:r([0-9]+)]]
 ; CHECK: endloop0
 ; CHECK: = add([[REG1]],#8)
 
 ; Function Attrs: nounwind
 define i32* @f0(i16* nocapture readonly %a0, i32 %a1) #0 {
 b0:
   %v0 = alloca [129 x i32], align 8
   br i1 undef, label %b1, label %b3

llvm/test/CodeGen/Hexagon/swp-stages4.ll

 ; RUN: llc -march=hexagon -mcpu=hexagonv5 -enable-pipeliner -pipeliner-max-stages=2 -disable-block-placement=0 -hexagon-bit=0 < %s -pipeliner-experimental-cg=true | FileCheck %s
 
 ; Test that we rename registers correctly for multiple stages when there is a
 ; Phi and depends upon another Phi.
 
 ; CHECK: = and
 ; CHECK: = and
+; CHECK: r[[REGA:[0-9]+]] = memub(r{{[0-9]+}}+#1)
 ; CHECK: r[[REG0:[0-9]+]] = and(r[[REG1:[0-9]+]],#255)
 ; CHECK-NOT: r[[REG0]] = and(r[[REG1]],#255)
 ; CHECK: loop0(.LBB0_[[LOOP:.]],
 ; CHECK: .LBB0_[[LOOP]]:
 ; CHECK: = add(r{{[0-9]+}},r[[REG0]])
 ; CHECK: = and
 ; CHECK: r[[REG0]] = and
 ; CHECK: endloop