[MachinePipeliner] Factor experimental codegen into a class

[jmolloy]

Summary

This new class, PipelineEmitter, is abstract and allows custom code generators to plug into MachinePipeliner. A single implementation is provided, SinglEpilogPipelineEmitter, that is equivalent to the current experimental code generator.

The name of this class serves to distinguish it from other code generators that will appear in the future. Quoting from the header:

/ PipelinerEmitter implementation that produces prolog blocks linked to a
/ single shared set of epilog blocks. This reduces the number of epilog blocks
/ required, but assumes that it is correct to emit stages in non-FIFO order:
/
/ Prolog0 -- executes A[0]
/ | \
/ Prolog1 \ | executes A[1], B[0]
/ | \ |
/ Prolog2 \ || executes A[2], B[1], C[0]
/ | |||
/ Kernel ||| executes X[3], Y[2], Z[1], W[0]
/ | / ||
/ Epilog0 || executes (A or Y)[3]
/ | / |
/ Epilog1 / executes (A or Z)[2, 3]
/ | /
/ Epilog2 executes (A or W)[1, 2, 3]
/
/ Note that assuming the kernel is taken the epilog order is:
/ Y[3], Z[2], Z[3], W[1], W[2], W[3]
/ Whereas a FIFO order for the epilog would be:
/ Y[3], Z[2], W[1], Z[3], W[2], W[3]
/
/ For many targets this ordering does not matter. However, if a hardware
/ functional unit is being used that depends on FIFO order (for example a
/ streaming hardware unit), wrong code may be generated.

[jmolloy]

Abandoning as this patch has been superseded.

include/llvm/CodeGen/MachinePipeliner.h

public:
   }
 
   void addMutation(std::unique_ptr<ScheduleDAGMutation> Mutation) {
     Mutations.push_back(std::move(Mutation));
   }
 
   static bool classof(const ScheduleDAGInstrs *DAG) { return true; }
 
+  void removeDeadInstructions(MachineBasicBlock *KernelBB,
+                              MBBVectorTy &EpilogBBs);
+  void splitLifetimes(MachineBasicBlock *KernelBB,
+                      ArrayRef<MachineBasicBlock *> EpilogBBs,
+                      SMSchedule &Schedule);
+
 private:
   void addLoopCarriedDependences(AliasAnalysis *AA);
   void updatePhiDependences();
   void changeDependences();
   unsigned calculateResMII();
   unsigned calculateRecMII(NodeSetType &RecNodeSets);
   void findCircuits(NodeSetType &NodeSets);
   void fuseRecs(NodeSetType &NodeSets);
void generateExistingPhis(MachineBasicBlock *NewBB, MachineBasicBlock *BB1,
                             SMSchedule &Schedule, ValueMapTy *VRMap,
                             InstrMapTy &InstrMap, unsigned LastStageNum,
                             unsigned CurStageNum, bool IsLast);
   void generatePhis(MachineBasicBlock *NewBB, MachineBasicBlock *BB1,
                     MachineBasicBlock *BB2, MachineBasicBlock *KernelBB,
                     SMSchedule &Schedule, ValueMapTy *VRMap,
                     InstrMapTy &InstrMap, unsigned LastStageNum,
                     unsigned CurStageNum, bool IsLast);
-  void removeDeadInstructions(MachineBasicBlock *KernelBB,
-                              MBBVectorTy &EpilogBBs);
-  void splitLifetimes(MachineBasicBlock *KernelBB,
-                      ArrayRef<MachineBasicBlock*> EpilogBBs,
-                      SMSchedule &Schedule);
   void addBranches(MachineBasicBlock &PreheaderBB, MBBVectorTy &PrologBBs,
                    MachineBasicBlock *KernelBB, MBBVectorTy &EpilogBBs,
                    SMSchedule &Schedule, ValueMapTy *VRMap);
   bool computeDelta(MachineInstr &MI, unsigned &Delta);
   void updateMemOperands(MachineInstr &NewMI, MachineInstr &OldMI,
                          unsigned Num);
   MachineInstr *cloneInstr(MachineInstr *OldMI, unsigned CurStageNum,
                            unsigned InstStageNum);

 /// map from scheduled cycle to instructions.  During scheduling, the
 /// data structure explicitly represents all stages/iterations.   When
 /// the algorithm finshes, the schedule is collapsed into a single stage,
 /// which represents instructions from different loop iterations.
 ///
 /// The SMS algorithm allows negative values for cycles, so the first cycle
 /// in the schedule is the smallest cycle value.
 class SMSchedule {
-private:
+protected:
   /// Map from execution cycle to instructions.
   DenseMap<int, std::deque<SUnit *>> ScheduledInstrs;
 
   /// Map from instruction to execution cycle.
   std::map<SUnit *, int> InstrToCycle;
 
   /// Map for each register and the max difference between its uses and def.
   /// The first element in the pair is the max difference in stages. The
void orderDependence(SwingSchedulerDAG *SSD, SUnit *SU,
                        std::deque<SUnit *> &Insts);
   bool isLoopCarried(SwingSchedulerDAG *SSD, MachineInstr &Phi);
   bool isLoopCarriedDefOfUse(SwingSchedulerDAG *SSD, MachineInstr *Def,
                              MachineOperand &MO);
   void print(raw_ostream &os) const;
   void dump() const;
 };
 
+/// Once SwingSchedulerDAG has generated a correct schedule, it calls an
+/// implementation of this class to generate the code for that schedule.
+///
+/// The PipelineEmitter class itself is abstract, but defines a suite of helper
+/// objects that can be used to generate the schedule.
+class PipelineEmitter {
+protected:
+  /// A simple class that takes identical copies of a single basic block and
+  /// allows
+  /// querying of equivalent instructions at a later point when the blocks are
+  /// no longer identical.
+  class EquivalenceRegistry {
+  public:
+    EquivalenceRegistry() {}
+    void insert(MachineBasicBlock *MBB) {
+      unsigned Idx = 0;
+      for (auto &MI : *MBB) {
+        Insts[{MBB, Idx}] = &MI;
+        Indexes[&MI] = Idx++;
+      }
+    }
+
+    MachineInstr *lookup(MachineInstr *MI, MachineBasicBlock *MBB) {
+      unsigned Idx = Indexes[MI];
+      return Insts[{MBB, Idx}];
+    }
+
+    MachineOperand *lookup(MachineOperand *MO, MachineBasicBlock *MBB) {
+      unsigned Idx = MO->getParent()->findRegisterDefOperandIdx(MO->getReg());
+      return &lookup(MO->getParent(), MBB)->getOperand(Idx);
+    }
+
+    bool contains(MachineInstr *MI) { return Indexes.count(MI); }
+
+  private:
+    DenseMap<std::pair<MachineBasicBlock *, unsigned>, MachineInstr *> Insts;
+    DenseMap<MachineInstr *, unsigned> Indexes;
+  };
+
+  /// A CGBlock wraps a MachineBasicBlock with a copy of the scheduled kernel.
+  /// One prototypical CGBlock is created from the original unpipelined loop and
+  /// the SMSchedule. This can then be duplicated to form the prolog and epilog
+  /// code.
+  class CGBlock {
+  public:
+    /// Create a ScheduledBlock from the original loop in BB. BB is rewritten
+    /// in-place to adhere to Schedule. Where schedule finalization has updated
+    /// instructions to use pre-inc values (NewMIs) the rewritten block uses
+    /// these updates.
+    CGBlock(MachineBasicBlock *BB, SMSchedule &Schedule, PipelineEmitter &PE,
+            LiveIntervals &LIS, EquivalenceRegistry &ER);
+
+    /// Peel this loop. If PeelFront is true, peel the first iteration otherwise
+    /// peel the last iteration.
+    CGBlock peel(bool PeelFront);
+
+    /// Mark all instructions that are not in the set LiveStages as Pruned.
+    /// Pruning does not happen in this function.
+    void setLiveStages(ArrayRef<unsigned> LiveStages);
+    /// Rewrite all uses of pruned-out instructions.
+    void rewritePrunedInstUses(MachineBasicBlock *LoopExit);
+    /// Actually prune all instructions marked as Pruned by setLiveStages.
+    void prune();
+    void addPredecessor(CGBlock &PredCGB);
+    void setSinglePredecessor(CGBlock &PredCGB);
+
+    MachineBasicBlock *getBB() { return BB; }
+
+  private:
+    struct InstrInfo {
+      InstrInfo() = default;
+      InstrInfo(int Stage, int Cycle, unsigned Index, bool Prune)
+          : Stage(Stage), Cycle(Cycle), Index(Index), Prune(Prune) {}
+      /// The stage this instruction is scheduled in.
+      int Stage;
+      /// The cycle this instruction is scheduled in.
+      int Cycle;
+      /// The linear order of this instruction. Cycle only provides a partial
+      /// order - this provides a total order.
+      unsigned Index;
+      /// True if this instruction should be pruned out of this block.
+      bool Prune = false;
+    };
+
+    /// Private constructor for use by peel().
+    CGBlock(MachineBasicBlock *BB, const CGBlock &CGB,
+            DenseMap<MachineInstr *, InstrInfo> InstrInfos);
+
+    /// Reg is used by MI. Return the new register MI should use to adhere to
+    /// the schedule. Insert phis as necessary.
+    Register remapUse(Register Reg, MachineInstr &MI);
+    /// Insert a phi that carries LoopReg from the loop body and InitReg
+    /// otherwise. 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);
+
+    MachineBasicBlock *BB;
+    SMSchedule &Schedule;
+    MachineRegisterInfo &MRI;
+    const TargetInstrInfo *TII;
+    LiveIntervals &LIS;
+    EquivalenceRegistry &ER;
+
+    MachineBasicBlock *PreheaderBB;
+    DenseMap<MachineInstr *, InstrInfo> InstrInfos;
+    // Map from register class to canonical undef register for that class.
+    DenseMap<const TargetRegisterClass *, Register> Undefs;
+    // Map from <LoopReg, InitReg> to phi register for all created phis. Note
+    // that this map is only used when InitReg is non-undef.
+    DenseMap<std::pair<unsigned, unsigned>, Register> Phis;
+    // Map from LoopReg to phi register where the InitReg is undef.
+    DenseMap<unsigned, Register> UndefPhis;
+  };
+
+  SMSchedule &Schedule;
+  MachineBasicBlock *BB;
+  MachineLoop &L;
+  LiveIntervals &LIS;
+  SwingSchedulerDAG *DAG;
+  const TargetInstrInfo *TII;
+  MachineRegisterInfo &MRI;
+
+public:
+  PipelineEmitter(SMSchedule &Schedule, MachineBasicBlock *BB,
+                  LiveIntervals &LIS, MachineLoop &L,
+                  SwingSchedulerDAG *DAG = nullptr)
+      : Schedule(Schedule), BB(BB), L(L), LIS(LIS), DAG(DAG),
+        TII(BB->getParent()->getSubtarget().getInstrInfo()),
+        MRI(BB->getParent()->getRegInfo()) {}
+  virtual ~PipelineEmitter();
+
+  // Rewrite BB to adhere to Schedule and emit any prolog or epilog code.
+  virtual void emit() = 0;
+};
+
+/// PipelinerEmitter implementation that produces prolog blocks linked to a
+/// single shared set of epilog blocks. This reduces the number of epilog blocks
+/// required, but assumes that it is correct to emit stages in non-FIFO order:
+///
+/// Prolog0 --    executes A[0]
+///    |      \
+/// Prolog1 \  |  executes A[1], B[0]
+///    |     \ |
+/// Prolog2 \ ||  executes A[2], B[1], C[0]
+///    |     |||
+/// Kernel   |||  executes X[3], Y[2], Z[1], W[0]
+///    |    / ||
+/// Epilog0   ||  executes (A or Y)[3]
+///    |     / |
+/// Epilog1   /   executes (A or Z)[2, 3]
+///    |     /
+/// Epilog2       executes (A or W)[1, 2, 3]
+///
+/// Note that assuming the kernel is taken the epilog order is:
+///   Y[3], Z[2], Z[3], W[1], W[2], W[3]
+/// Whereas a FIFO order for the epilog would be:
+///   Y[3], Z[2], W[1], Z[3], W[2], W[3]
+///
+/// For many targets this ordering does not matter. However, if a hardware
+/// functional unit is being used that depends on FIFO order (for example a
+/// streaming hardware unit), wrong code may be generated.
+class SingleEpilogPipelineEmitter : public PipelineEmitter {
+public:
+  using PipelineEmitter::PipelineEmitter;
+
+  void emit() override;
+};
+
 } // end namespace llvm
 
 #endif // LLVM_LIB_CODEGEN_MACHINEPIPELINER_H

lib/CodeGen/MachinePipeliner.cpp

SwpEnableCopyToPhi("pipeliner-enable-copytophi", cl::ReallyHidden,
                        cl::init(true), cl::ZeroOrMore,
                        cl::desc("Enable CopyToPhi DAG Mutation"));
 } // end namespace llvm
 
 /// Enables the experimental code generator.
 static cl::opt<bool> EnableExperimentalCodeGen(
     "swp-experimental-cg",
     cl::desc("Enable experimental CG code in MachinePipeliner."), cl::Hidden,
-    cl::init(false));
+    cl::init(true));
 
 unsigned SwingSchedulerDAG::Circuits::MaxPaths = 5;
 char MachinePipeliner::ID = 0;
 #ifndef NDEBUG
 int MachinePipeliner::NumTries = 0;
 #endif
 char &llvm::MachinePipelinerID = MachinePipeliner::ID;
 
void SwingSchedulerDAG::schedule() {
   if (SwpMaxStages > -1 && (int)numStages > SwpMaxStages) {
     LLVM_DEBUG(dbgs() << "numStages:" << numStages << ">" << SwpMaxStages
                       << " : too many stages, abort\n");
     NumFailLargeMaxStage++;
     return;
   }
 
   bool CanUseExperimentalCodeGen = InstrChanges.empty();
-  if (EnableExperimentalCodeGen && CanUseExperimentalCodeGen)
-    experimentalGeneratePipeline(Schedule);
-  else
+  if (EnableExperimentalCodeGen && CanUseExperimentalCodeGen) {
+    SingleEpilogPipelineEmitter Emitter(Schedule, BB, LIS, Loop, this);
+    Emitter.emit();
+    // NewMIs have been added to BB, so don't delete them.
+    NewMIs.clear();
+  } else {
     generatePipelinedLoop(Schedule);
+  }
   ++NumPipelined;
 }
 
 /// Clean up after the software pipeliner runs.
 void SwingSchedulerDAG::finishBlock() {
   for (MachineInstr *I : NewMIs)
     MF.DeleteMachineInstr(I);
   NewMIs.clear();
for (auto I = MBB->begin(); I != MBB->getFirstNonPHI();) {
         LIS.RemoveMachineInstrFromMaps(MI);
         MI.eraseFromParent();
         Changed = true;
       }
     }
   }
 }
 
-// A simple class that takes identical copies of a single basic block and allows
-// querying of equivalent instructions at a later point when the blocks are no
-// longer identical.
-class EquivalenceRegistry {
-public:
-  EquivalenceRegistry() {}
-  void insert(MachineBasicBlock *MBB) {
-    unsigned Idx = 0;
-    for (auto &MI : *MBB) {
-      Insts[{MBB, Idx}] = &MI;
-      Indexes[&MI] = Idx++;
-    }
-  }
-
-  MachineInstr *lookup(MachineInstr *MI, MachineBasicBlock *MBB) {
-    unsigned Idx = Indexes[MI];
-    return Insts[{MBB, Idx}];
-  }
-
-  MachineOperand *lookup(MachineOperand *MO, MachineBasicBlock *MBB) {
-    unsigned Idx = MO->getParent()->findRegisterDefOperandIdx(MO->getReg());
-    return &lookup(MO->getParent(), MBB)->getOperand(Idx);
-  }
-
-  bool contains(MachineInstr *MI) { return Indexes.count(MI); }
-
-private:
-  DenseMap<std::pair<MachineBasicBlock *, unsigned>, MachineInstr *> Insts;
-  DenseMap<MachineInstr *, unsigned> Indexes;
-};
-
-// A CGBlock wraps a MachineBasicBlock with a copy of the scheduled kernel.
-// One prototypical CGBlock is created from the original unpipelined loop and
-// the SMSchedule. This can then be duplicated to form the prolog and epilog
-// code.
-class CGBlock {
-public:
-  // Create a ScheduledBlock from the original loop in BB. BB is rewritten
-  // in-place to adhere to Schedule. Where schedule finalization has updated
-  // instructions to use pre-inc values (NewMIs) the rewritten block uses these
-  // updates.
-  CGBlock(MachineBasicBlock *BB, SMSchedule &Schedule, SwingSchedulerDAG &DAG,
-          LiveIntervals &LIS, EquivalenceRegistry &ER);
-
-  // Peel this loop in the direction given by LPD.
-  CGBlock peel(LoopPeelDirection LPD);
-
-  // Mark all instructions that are not in the set LiveStages as Pruned. Pruning
-  // does not happen in this function.
-  void setLiveStages(ArrayRef<unsigned> LiveStages);
-  // Rewrite all uses of pruned-out instructions.
-  void rewritePrunedInstUses(MachineBasicBlock *LoopExit);
-  // Actually prune all instructions marked as Pruned by setLiveStages.
-  void prune();
-  void addPredecessor(CGBlock &PredCGB);
-  void setSinglePredecessor(CGBlock &PredCGB);
-  void insertBranchToSuccessors(ArrayRef<MachineOperand> Cond);
-
-  MachineBasicBlock *getBB() { return BB; }
-
-private:
-  struct InstrInfo {
-    InstrInfo() = default;
-    InstrInfo(int Stage, int Cycle, unsigned Index, bool Prune) :
-        Stage(Stage), Cycle(Cycle), Index(Index), Prune(Prune) {}
-    // The stage this instruction is scheduled in.
-    int Stage;
-    // The cycle this instruction is scheduled in.
-    int Cycle;
-    // The linear order of this instruction. Cycle only provides a partial
-    // order - this provides a total order.
-    unsigned Index;
-    // True if this instruction should be pruned out of this block.
-    bool Prune = false;
-  };
-
-  // Private constructor for use by peel().
-  CGBlock(MachineBasicBlock *BB, const CGBlock &CGB,
-          DenseMap<MachineInstr *, InstrInfo> InstrInfos);
-
-  // Reg is used by MI. Return the new register MI should use to adhere to the
-  // schedule. Insert phis as necessary.
-  Register remapUse(Register Reg, MachineInstr &MI);
-  // Insert a phi that carries LoopReg from the loop body and InitReg otherwise.
-  // 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);
-
-  MachineBasicBlock *BB;
-  SMSchedule &Schedule;
-  SwingSchedulerDAG &DAG;
-  MachineRegisterInfo &MRI;
-  const TargetInstrInfo *TII;
-  LiveIntervals &LIS;
-  EquivalenceRegistry &ER;
-
-  MachineBasicBlock *PreheaderBB;
-  DenseMap<MachineInstr *, InstrInfo> InstrInfos;
-  // Map from register class to canonical undef register for that class.
-  DenseMap<const TargetRegisterClass *, Register> Undefs;
-  // Map from <LoopReg, InitReg> to phi register for all created phis. Note that
-  // this map is only used when InitReg is non-undef.
-  DenseMap<std::pair<unsigned, unsigned>, Register> Phis;
-  // Map from LoopReg to phi register where the InitReg is undef.
-  DenseMap<Register, Register> UndefPhis;
-};
-
+/// Encapsulates the state required for interacting with TII->reduceLoopCount.
+/// FIXME: The reduceLoopCount API is overfit for hardware loop instructions.
+/// Architectures that use "normal" conditional branches with induction
+/// variables don't get a way to indicate instructions that are necessary for
+/// calculating the loop count (which must be in stage0!) and don't get enough
+/// information to reconstruct conditional branches in the kernel and prologs.
 struct ReduceLoopCountState {
-  ReduceLoopCountState(MachineBasicBlock *Preheader,
-                       MachinePipeliner::LoopInfo &LI, unsigned NumPeeled)
-      : Preheader(Preheader), LI(LI), NumPeeled(NumPeeled) {}
+  ReduceLoopCountState(MachineBasicBlock *Preheader, MachineInstr *LoopCompare,
+                       MachineInstr *LoopInductionVar, unsigned NumPeeled)
+      : Preheader(Preheader), LoopCompare(LoopCompare),
+        LoopInductionVar(LoopInductionVar), NumPeeled(NumPeeled) {}
   SmallVector<MachineOperand, 4> Cond;
 
   MachineBasicBlock *Preheader;
-  MachinePipeliner::LoopInfo &LI;
+  MachineInstr *LoopCompare;
+  MachineInstr *LoopInductionVar;
   unsigned PrevLC = UINT_MAX;
   unsigned FirstLC = UINT_MAX;
   SmallVector<MachineInstr *, 4> PrevInsts;
   unsigned NumPeeled;
   unsigned I = 0;
 };
 
 Optional<unsigned> DecrementLoopCount(MachineBasicBlock *BB,
                                       ReduceLoopCountState &S,
                                       const TargetInstrInfo *TII) {
   if (S.PrevLC == 0)
     return S.FirstLC;
   S.Cond.clear();
-  unsigned LC = TII->reduceLoopCount(*BB, *S.Preheader, S.LI.LoopInductionVar,
-                                     *S.LI.LoopCompare, S.Cond, S.PrevInsts,
+  unsigned LC = TII->reduceLoopCount(*BB, *S.Preheader, S.LoopInductionVar,
+                                     *S.LoopCompare, S.Cond, S.PrevInsts,
                                      S.NumPeeled - 1 - S.I++, S.NumPeeled - 1);
   if (S.FirstLC == UINT_MAX)
     S.FirstLC = LC;
   S.PrevLC = LC;
   return Register::isVirtualRegister(LC) ? Optional<unsigned>() : S.FirstLC;
 }
 } // namespace
 
-void SwingSchedulerDAG::experimentalGeneratePipeline(SMSchedule &Schedule) {
+PipelineEmitter::~PipelineEmitter() {}
+
+void SingleEpilogPipelineEmitter::emit() {
   EquivalenceRegistry ER;
   MachineBasicBlock *PreheaderBB = *BB->pred_begin();
   if (PreheaderBB == BB)
     PreheaderBB = *std::next(BB->pred_begin());
-  MachineBasicBlock *LoopExitBB = Loop.getExitBlock();
+  MachineBasicBlock *LoopExitBB = *BB->succ_begin();
+  if (LoopExitBB == BB)
+    LoopExitBB = *std::next(BB->succ_begin());
 
   std::list<CGBlock> CGBlocks;
   SmallVector<CGBlock *, 4> Prologs, Epilogs;
   std::list<CGBlock>::iterator Kernel;
   unsigned NumStages = Schedule.getMaxStageCount() + 1;
   unsigned NumPrologs = NumStages - 1;
 
+  MachineInstr *LoopInductionVar = nullptr;
+  MachineInstr *LoopCompare = nullptr;
+  assert(!TII->analyzeLoop(L, LoopInductionVar, LoopCompare));
+
   // Create the first CGBlock, which is a rescheduled loop BB.
   CGBlocks.emplace_back(BB, Schedule, *this, LIS, ER);
   Kernel = CGBlocks.begin();
   LLVM_DEBUG({
     dbgs() << "Initial rewritten kernel: " << *Kernel->getBB();
   });
 
   // Create the prologs, peeling in reverse.
   SmallVector<unsigned, 4> LiveStages;
   for (unsigned I = 0; I < NumPrologs; ++I) {
     LiveStages.push_back(I);
-    auto It = CGBlocks.insert(Kernel, Kernel->peel(LPD_Front));
+    auto It = CGBlocks.insert(Kernel, Kernel->peel(/*PeelFront=*/true));
     It->setLiveStages(LiveStages);
     Prologs.push_back(&*It);
   }
 
   // The epilogs end up peeling off the last iterations. So Epilog 0 performs
   // stage NumStages-1, Epilog 1 performs stage NumStages-2 THEN NumStages-1.
   for (unsigned I = 0; I < NumPrologs; ++I) {
     unsigned StartStage = NumStages - 1 - I;
-    CGBlocks.insert(std::next(Kernel), Kernel->peel(LPD_Back));
+    CGBlocks.insert(std::next(Kernel), Kernel->peel(/*PeelFront=*/false));
     for (unsigned Stage = StartStage + 1; Stage < NumStages; ++Stage) {
-      CGBlocks.insert(std::next(Kernel), Kernel->peel(LPD_Back));
+      CGBlocks.insert(std::next(Kernel), Kernel->peel(/*PeelFront=*/false));
     }
   }
   auto EBI = std::next(Kernel);
   for (unsigned I = 0; I < NumPrologs; ++I) {
     unsigned StartStage = NumStages - 1 - I;
     Epilogs.push_back(&*EBI);
     EBI++->setLiveStages({StartStage});
     for (unsigned Stage = StartStage + 1; Stage < NumStages; ++Stage) {
       EBI++->setLiveStages({Stage});
     }
   }
 
   // Stitch the epilogs to the prologs, working outwards from the kernel.
   auto EI = Epilogs.begin();
   auto PI = Prologs.rbegin();
   unsigned N = 0;
-  ReduceLoopCountState RLCS(PreheaderBB, Pass.LI, NumPrologs);
+  ReduceLoopCountState RLCS(PreheaderBB, LoopCompare, LoopInductionVar,
+                            NumPrologs);
   for (; EI != Epilogs.end(); ++EI, ++PI, ++N) {
     CGBlock &PB = **PI;
     CGBlock &EB = **EI;
     TII->removeBranch(*PB.getBB());
     Optional<unsigned> TripCount = DecrementLoopCount(PB.getBB(), RLCS, TII);
     if (!TripCount.hasValue()) {
       // Unknown trip count; plant predecessors.
       PB.getBB()->addSuccessor(EB.getBB());
LLVM_DEBUG({
     dbgs() << "Kernel:\n";
     It++->getBB()->dump();
     dbgs() << "Epilogs:\n";
     while (It != CGBlocks.end())
       It++->getBB()->dump();
   });
 
   // Do a final optimization pass on the kernel.
+  if (DAG) {
     SmallVector<MachineBasicBlock *, 4> EpilogBBs;
     for (auto I = std::next(Kernel); I != CGBlocks.end(); ++I)
       EpilogBBs.push_back(I->getBB());
-  splitLifetimes(Kernel->getBB(), EpilogBBs, Schedule);
-  removeDeadInstructions(Kernel->getBB(), EpilogBBs);
+    DAG->splitLifetimes(Kernel->getBB(), EpilogBBs, Schedule);
+    DAG->removeDeadInstructions(Kernel->getBB(), EpilogBBs);
+  }
 
-  for (CGBlock &CGB : CGBlocks) {
-    if (CGB.getBB()->pred_empty() || CGB.getBB()->succ_empty())
-      CGB.getBB()->eraseFromParent();
+  SmallPtrSet<MachineBasicBlock *, 4> CreatedBlocks;
+  for (CGBlock &CGB : CGBlocks)
+    CreatedBlocks.insert(CGB.getBB());
+  bool Changed = true;
+  while (Changed) {
+    Changed = false;
+    for (MachineBasicBlock *BB : CreatedBlocks) {
+      if (BB->pred_empty()) {
+        while (!BB->succ_empty())
+          BB->removeSuccessor(BB->succ_begin());
+        BB->eraseFromParent();
+        CreatedBlocks.erase(BB);
+        Changed = true;
+      }
+    }
   }
   if (Kernel->getBB()->succ_size() == 1)
     Kernel->getBB()->eraseFromParent();
-
-  // NewMIs have been added to BB, so don't delete them.
-  NewMIs.clear();
 }
 
-CGBlock::CGBlock(MachineBasicBlock *BB, const CGBlock &CGB,
-                 DenseMap<MachineInstr *, InstrInfo> NewInstrInfos)
-    : BB(BB), Schedule(CGB.Schedule), DAG(CGB.DAG), MRI(DAG.MRI), TII(DAG.TII),
-      LIS(CGB.LIS), ER(CGB.ER), InstrInfos(std::move(NewInstrInfos)) {
+PipelineEmitter::CGBlock::CGBlock(
+    MachineBasicBlock *BB, const CGBlock &CGB,
+    : BB(BB), Schedule(CGB.Schedule), MRI(CGB.MRI), TII(CGB.TII), LIS(CGB.LIS),
+      ER(CGB.ER), InstrInfos(std::move(NewInstrInfos)) {
   ER.insert(BB);
   for (auto &II : InstrInfos)
     II.second.Prune = false;
 }
 
-CGBlock::CGBlock(MachineBasicBlock *BB, SMSchedule &Schedule,
-                 SwingSchedulerDAG &DAG, LiveIntervals &LIS,
+PipelineEmitter::CGBlock::CGBlock(MachineBasicBlock *BB, SMSchedule &Schedule,
+                                  PipelineEmitter &PE, LiveIntervals &LIS,
                                   EquivalenceRegistry &ER)
-    : BB(BB), Schedule(Schedule), DAG(DAG), MRI(DAG.MRI), TII(DAG.TII),
-      LIS(LIS), ER(ER) {
+    : BB(BB), Schedule(Schedule), MRI(PE.MRI), TII(PE.TII), LIS(LIS), ER(ER) {
   // Discover this block's preheader.
   PreheaderBB = *BB->pred_begin();
   if (PreheaderBB == BB)
     PreheaderBB = *std::next(BB->pred_begin());
 
   // Rearrange the loop to be in schedule order. Note that the schedule may
   // contain instructions that are not owned by the loop block (InstrChanges and
   // friends), so we gracefully handle unowned instructions and delete any
   // instructions that weren't in the schedule.
   auto InsertPt = BB->getFirstTerminator();
   unsigned Index = 0;
   MachineInstr *FirstMI = nullptr;
-
   for (int Cycle = Schedule.getFirstCycle(); Cycle <= Schedule.getFinalCycle();
        ++Cycle)
     for (SUnit *SU : Schedule.getInstructions(Cycle)) {
       auto *MI = SU->getInstr();
-      InstrInfos[MI] = {Schedule.stageScheduled(SU), Cycle, Index++, false};
       if (MI->isPHI())
         continue;
+      InstrInfos[MI] = {Schedule.stageScheduled(SU), Cycle, Index++, false};
       if (MI->getParent())
         MI->removeFromParent();
       BB->insert(InsertPt, MI);
       if (!FirstMI)
         FirstMI = MI;
     }
 
   // At this point all of the scheduled instructions are between FirstMI
   // and the end of the block. Kill from the first non-phi to FirstMI.
   for (auto I = BB->getFirstNonPHI(); I != FirstMI->getIterator();) {
     LIS.RemoveMachineInstrFromMaps(*I);
+    InstrInfos.erase(&*I);
     (I++)->eraseFromParent();
   }
 
   // Now remap every instruction in the loop.
   for (MachineInstr &MI : *BB) {
     if (MI.isPHI())
       continue;
     for (MachineOperand &MO : MI.uses()) {
for (MachineOperand &Def : MI->defs()) {
         }
       }
     }
   }
 
   ER.insert(BB);
 }
 
-Register CGBlock::remapUse(Register Reg, MachineInstr &MI) {
+Register PipelineEmitter::CGBlock::remapUse(Register Reg, MachineInstr &MI) {
   MachineInstr *Producer = MRI.getUniqueVRegDef(Reg);
   if (!Producer)
     return Reg;
 
   int ConsumerStage = InstrInfos[&MI].Stage;
   if (!Producer->isPHI()) {
     // Non-phi producers are simple to remap. Insert as many phis as the
     // difference between the consumer and producer stages.
Register PipelineEmitter::CGBlock::remapUse(Register Reg, MachineInstr &MI) {
   }
 
   // Now we know the number of stages to jump back, insert the phi chain.
   while (DefaultI != Defaults.rend())
     LoopReg = phi(LoopReg, *DefaultI++, MRI.getRegClass(Reg));
   return LoopReg;
 }
 
-Register CGBlock::phi(Register LoopReg, Optional<Register> InitReg,
+Register PipelineEmitter::CGBlock::phi(Register LoopReg,
+                                       Optional<Register> InitReg,
                                        const TargetRegisterClass *RC) {
   // If the init register is not undef, try and find an existing phi.
   if (InitReg.hasValue()) {
     auto I = Phis.find({LoopReg, InitReg.getValue()});
     if (I != Phis.end())
       return I->second;
   } else {
     for (auto &KV : Phis) {
       if (KV.first.first == LoopReg)
BuildMI(*BB, BB->getFirstNonPHI(), DebugLoc(), TII->get(TargetOpcode::PHI), R)
       .addMBB(BB);
   if (!InitReg.hasValue())
     UndefPhis[LoopReg] = R;
   else
     Phis[{LoopReg, *InitReg}] = R;
   return R;
 }
 
-Register CGBlock::undef(const TargetRegisterClass *RC) {
+Register PipelineEmitter::CGBlock::undef(const TargetRegisterClass *RC) {
   Register &R = Undefs[RC];
   if (R == 0) {
     // Create an IMPLICIT_DEF that defines this register if we need it.
     // All uses of this should be removed by the time we have finished unrolling
     // prologs and epilogs.
     R = MRI.createVirtualRegister(RC);
     auto *InsertBB = &PreheaderBB->getParent()->front();
     BuildMI(*InsertBB, InsertBB->getFirstTerminator(), DebugLoc(),
             TII->get(TargetOpcode::IMPLICIT_DEF), R);
   }
   return R;
 }
 
-CGBlock CGBlock::peel(LoopPeelDirection LPD) {
+PipelineEmitter::CGBlock PipelineEmitter::CGBlock::peel(bool PeelFront) {
+  auto LPD = PeelFront ? LPD_Front : LPD_Back;
   MachineBasicBlock *NewBB = PeelSingleBlockLoop(LPD, BB, MRI, TII);
   auto OI = BB->begin();
   auto NI = NewBB->begin();
   DenseMap<MachineInstr *, InstrInfo> NewInstrInfos;
   for (; OI != BB->end(); ++OI, ++NI) {
     if (InstrInfos.count(&*OI))
       NewInstrInfos.insert({&*NI, InstrInfos[&*OI]});
   }
   return CGBlock(NewBB, *this, std::move(NewInstrInfos));
 }
 
-void CGBlock::setLiveStages(ArrayRef<unsigned> LiveStages) {
+void PipelineEmitter::CGBlock::setLiveStages(ArrayRef<unsigned> LiveStages) {
   for (auto &KV : InstrInfos) {
     if (KV.second.Stage != -1 &&
         find(LiveStages, KV.second.Stage) == LiveStages.end())
       KV.second.Prune = true;
   }
 }
 
 // Return a PHI that uses R in BB.
 static MachineInstr *findPhiProviding(Register R, MachineBasicBlock *BB) {
   for (auto I = BB->begin(); I != BB->end() && I->isPHI(); ++I) {
     if (I->findRegisterUseOperand(R) != nullptr)
       return &*I;
   }
   llvm_unreachable("findPhiProviding failed to find an existing phi!");
 }
 
-void CGBlock::rewritePrunedInstUses(MachineBasicBlock *LoopExit) {
+void PipelineEmitter::CGBlock::rewritePrunedInstUses(
+    MachineBasicBlock *LoopExit) {
   // Rewrite all uses of predicated-out instructions.
   for (MachineInstr &MI : *BB) {
     if (MI.isPHI() || InstrInfos.count(&MI) == 0)
       continue;
     if (!InstrInfos[&MI].Prune)
       // Predicate is live.
       continue;
     // This instruction is predicated out. Any uses outside this block must be
if (!PredMI)
       continue;
     MachineInstr *Phi = findPhiProviding(PredMI->getOperand(0).getReg(), BB);
     // Update the loop-carried value of the illegal PHI. Don't erase it yet
     // until the next pass when all values have been rewritten.
     MI->getOperand(3).setReg(Phi->getOperand(0).getReg());
   }
 }
 
-void CGBlock::prune() {
+void PipelineEmitter::CGBlock::prune() {
   // As a final pass, fix up any temporary illegal phis we created. These phis
   // were created in the middle of the block to switch between either a
   // loop-carried value computed in this block if available, or a default value
   // if the loop-carried value is pruned out.
   for (auto I = BB->getFirstNonPHI(); I != BB->end();) {
     MachineInstr *MI = &*I++;
     if (!MI->isPHI())
       continue;
     Register PhiR = MI->getOperand(0).getReg();
     // The second PHI input is the desired value.
     Register R = MI->getOperand(3).getReg();
     MachineInstr *RMI = MRI.getUniqueVRegDef(R);
     assert(RMI->isPHI() || InstrInfos.count(RMI) &&
                                "We shouldn't have created a phi in this case!");
     if (!RMI->isPHI() && InstrInfos[RMI].Prune)
       // Desired register will be pruned out, so use the init register instead.
       R = MI->getOperand(1).getReg();
     MRI.replaceRegWith(PhiR, R);
     MRI.setRegClass(R, MRI.getRegClass(PhiR));
     MI->eraseFromParent();
   }
 
   for (auto &KV : InstrInfos) {
-    if (!KV.first->isPHI() && KV.second.Prune) {
+    if (KV.second.Prune) {
       LIS.RemoveMachineInstrFromMaps(*KV.first);
       KV.first->eraseFromParent();
     }
   }
-  EliminateDeadPhis(BB, MRI, LIS);
 }
 
-void CGBlock::addPredecessor(CGBlock &PredCGB) {
+void PipelineEmitter::CGBlock::addPredecessor(CGBlock &PredCGB) {
   // PredBB is our current unique predecessor.
   MachineBasicBlock *PredBB = *BB->pred_begin();
   // OtherPredBB is the BB that we are adding as a predecessor.
   MachineBasicBlock *OtherPredBB = PredCGB.getBB();
 
   // For every PHI, we have only one value which must come from PredBB due to
   // how the loop blocks were peeled. Find the equivalent value in OtherPredBB.
   // This is easy because all peeled blocks are identical, so we can just use
if (PredDefMI->getParent() == PredBB) {
       unsigned MOOffset = PredDefMI->findRegisterDefOperandIdx(PredReg);
       Reg = OtherDefMI->getOperand(MOOffset).getReg();
     }
     MI.addOperand(MachineOperand::CreateReg(Reg, /*isDef=*/false));
     MI.addOperand(MachineOperand::CreateMBB(OtherPredBB));
   }
 }
 
-void CGBlock::setSinglePredecessor(CGBlock &PredCGB) {
+void PipelineEmitter::CGBlock::setSinglePredecessor(CGBlock &PredCGB) {
   // PredBB is our current unique predecessor.
   MachineBasicBlock *PredBB = *BB->pred_begin();
   // OtherPredBB is the BB that we are adding as a predecessor.
   MachineBasicBlock *OtherPredBB = PredCGB.getBB();
   PredBB->removeSuccessor(BB);
 
   // For every PHI, we have only one value which must come from PredBB due to
   // how the loop blocks were peeled. Find the equivalent value in OtherPredBB.