3 files changed, 128 insertions, 0 deletions

diff --git a/llvm/lib/CodeGen/CodeGen.cpp b/llvm/lib/CodeGen/CodeGen.cpp
index c6376290dab..51be4e3af01 100644
--- a/llvm/lib/CodeGen/CodeGen.cpp
+++ b/llvm/lib/CodeGen/CodeGen.cpp
@@ -68,6 +68,7 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeMachineOptimizationRemarkEmitterPassPass(Registry);
   initializeMachineOutlinerPass(Registry);
   initializeMachinePipelinerPass(Registry);
+  initializeModuloScheduleTestPass(Registry);
   initializeMachinePostDominatorTreePass(Registry);
   initializeMachineRegionInfoPassPass(Registry);
   initializeMachineSchedulerPass(Registry);
diff --git a/llvm/lib/CodeGen/MachinePipeliner.cpp b/llvm/lib/CodeGen/MachinePipeliner.cpp
index b138342c784..ff3d2feab5e 100644
--- a/llvm/lib/CodeGen/MachinePipeliner.cpp
+++ b/llvm/lib/CodeGen/MachinePipeliner.cpp
@@ -154,6 +154,12 @@ static cl::opt<bool> SwpShowResMask("pipeliner-show-mask", cl::Hidden,
 static cl::opt<bool> SwpDebugResource("pipeliner-dbg-res", cl::Hidden,
                                       cl::init(false));
 
+static cl::opt<bool> EmitTestAnnotations(
+    "pipeliner-annotate-for-testing", cl::Hidden, cl::init(false),
+    cl::desc("Instead of emitting the pipelined code, annotate instructions "
+             "with the generated schedule for feeding into the "
+             "-modulo-schedule-test pass"));
+
 namespace llvm {
 
 // A command line option to enable the CopyToPhi DAG mutation.
@@ -536,6 +542,13 @@ void SwingSchedulerDAG::schedule() {
 
   ModuloSchedule MS(MF, &Loop, std::move(OrderedInsts), std::move(Cycles),
                     std::move(Stages));
+  if (EmitTestAnnotations) {
+    assert(NewInstrChanges.empty() &&
+           "Cannot serialize a schedule with InstrChanges!");
+    ModuloScheduleTestAnnotater MSTI(MF, MS);
+    MSTI.annotate();
+    return;
+  }
   ModuloScheduleExpander MSE(MF, MS, LIS, std::move(NewInstrChanges));
   MSE.expand();
   ++NumPipelined;
diff --git a/llvm/lib/CodeGen/ModuloSchedule.cpp b/llvm/lib/CodeGen/ModuloSchedule.cpp
index 703d7926a69..80b022e232f 100644
--- a/llvm/lib/CodeGen/ModuloSchedule.cpp
+++ b/llvm/lib/CodeGen/ModuloSchedule.cpp
@@ -7,14 +7,21 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/CodeGen/ModuloSchedule.h"
+#include "llvm/Support/raw_ostream.h"
+#include "llvm/ADT/StringExtras.h"
 #include "llvm/CodeGen/LiveIntervals.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
+#include "llvm/MC/MCContext.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "pipeliner"
 using namespace llvm;
 
+//===----------------------------------------------------------------------===//
+// ModuloScheduleExpander implementation
+//===----------------------------------------------------------------------===//
+
 /// Return the register values for  the operands of a Phi instruction.
 /// This function assume the instruction is a Phi.
 static void getPhiRegs(MachineInstr &Phi, MachineBasicBlock *Loop,
@@ -1188,3 +1195,110 @@ bool ModuloScheduleExpander::isLoopCarried(MachineInstr &Phi) {
   int LoopStage = Schedule.getStage(Use);
   return (LoopCycle > DefCycle) || (LoopStage <= DefStage);
 }
+
+//===----------------------------------------------------------------------===//
+// ModuloScheduleTestPass implementation
+//===----------------------------------------------------------------------===//
+// This pass constructs a ModuloSchedule from its module and runs
+// ModuloScheduleExpander.
+//
+// The module is expected to contain a single-block analyzable loop.
+// The total order of instructions is taken from the loop as-is.
+// Instructions are expected to be annotated with a PostInstrSymbol.
+// This PostInstrSymbol must have the following format:
+//  "Stage=%d Cycle=%d".
+//===----------------------------------------------------------------------===//
+
+class ModuloScheduleTest : public MachineFunctionPass {
+public:
+  static char ID;
+
+  ModuloScheduleTest() : MachineFunctionPass(ID) {
+    initializeModuloScheduleTestPass(*PassRegistry::getPassRegistry());
+  }
+
+  bool runOnMachineFunction(MachineFunction &MF) override;
+  void runOnLoop(MachineFunction &MF, MachineLoop &L);
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<MachineLoopInfo>();
+    AU.addRequired<LiveIntervals>();
+    MachineFunctionPass::getAnalysisUsage(AU);
+  }
+};
+
+char ModuloScheduleTest::ID = 0;
+
+INITIALIZE_PASS_BEGIN(ModuloScheduleTest, "modulo-schedule-test",
+                      "Modulo Schedule test pass", false, false)
+INITIALIZE_PASS_DEPENDENCY(MachineLoopInfo)
+INITIALIZE_PASS_DEPENDENCY(LiveIntervals)
+INITIALIZE_PASS_END(ModuloScheduleTest, "modulo-schedule-test",
+                    "Modulo Schedule test pass", false, false)
+
+bool ModuloScheduleTest::runOnMachineFunction(MachineFunction &MF) {
+  MachineLoopInfo &MLI = getAnalysis<MachineLoopInfo>();
+  for (auto *L : MLI) {
+    if (L->getTopBlock() != L->getBottomBlock())
+      continue;
+    runOnLoop(MF, *L);
+    return false;
+  }
+  return false;
+}
+
+static void parseSymbolString(StringRef S, int &Cycle, int &Stage) {
+  std::pair<StringRef, StringRef> StageAndCycle = getToken(S, "_");
+  std::pair<StringRef, StringRef> StageTokenAndValue =
+      getToken(StageAndCycle.first, "-");
+  std::pair<StringRef, StringRef> CycleTokenAndValue =
+      getToken(StageAndCycle.second, "-");
+  if (StageTokenAndValue.first != "Stage" ||
+      CycleTokenAndValue.first != "_Cycle") {
+    llvm_unreachable(
+        "Bad post-instr symbol syntax: see comment in ModuloScheduleTest");
+    return;
+  }
+
+  StageTokenAndValue.second.drop_front().getAsInteger(10, Stage);
+  CycleTokenAndValue.second.drop_front().getAsInteger(10, Cycle);
+
+  dbgs() << "  Stage=" << Stage << ", Cycle=" << Cycle << "\n";
+}
+
+void ModuloScheduleTest::runOnLoop(MachineFunction &MF, MachineLoop &L) {
+  LiveIntervals &LIS = getAnalysis<LiveIntervals>();
+  MachineBasicBlock *BB = L.getTopBlock();
+  dbgs() << "--- ModuloScheduleTest running on BB#" << BB->getNumber() << "\n";
+
+  DenseMap<MachineInstr *, int> Cycle, Stage;
+  std::vector<MachineInstr *> Instrs;
+  for (MachineInstr &MI : *BB) {
+    if (MI.isTerminator())
+      continue;
+    Instrs.push_back(&MI);
+    if (MCSymbol *Sym = MI.getPostInstrSymbol()) {
+      dbgs() << "Parsing post-instr symbol for " << MI;
+      parseSymbolString(Sym->getName(), Cycle[&MI], Stage[&MI]);
+    }
+  }
+
+  ModuloSchedule MS(MF, &L, std::move(Instrs), std::move(Cycle), std::move(Stage));
+  ModuloScheduleExpander MSE(
+      MF, MS, LIS, /*InstrChanges=*/ModuloScheduleExpander::InstrChangesTy());
+  MSE.expand();
+}
+
+//===----------------------------------------------------------------------===//
+// ModuloScheduleTestAnnotater implementation
+//===----------------------------------------------------------------------===//
+
+void ModuloScheduleTestAnnotater::annotate() {
+  for (MachineInstr *MI : S.getInstructions()) {
+    SmallVector<char, 16> SV;
+    raw_svector_ostream OS(SV);
+    OS << "Stage-" << S.getStage(MI) << "_Cycle-" << S.getCycle(MI);
+    MCSymbol *Sym = MF.getContext().getOrCreateSymbol(OS.str());
+    MI->setPostInstrSymbol(MF, Sym);
+  }
+}