better throughput in pipeline mode

test/PIM/SpatialSchedulingTargetTest.cpp

@@ -9,6 +9,36 @@
 using namespace onnx_mlir::spatial;
 
 int main() {
+  PipelineCoreLayout unevenLayout(138, 4);
+  assert(unevenLayout.isValid());
+  assert(unevenLayout.getLogicalProcessorCount() == 34);
+  assert(unevenLayout.getStageRange(0).begin == 0);
+  assert(unevenLayout.getStageRange(0).size == 35);
+  assert(unevenLayout.getStageRange(1).begin == 35);
+  assert(unevenLayout.getStageRange(1).size == 35);
+  assert(unevenLayout.getStageRange(2).begin == 70);
+  assert(unevenLayout.getStageRange(2).size == 34);
+  assert(unevenLayout.getStageRange(3).begin == 104);
+  assert(unevenLayout.getStageRange(3).size == 34);
+  assert(unevenLayout.getStageForCore(34) == 0);
+  assert(unevenLayout.getStageForCore(35) == 1);
+  assert(unevenLayout.getStageForCore(69) == 1);
+  assert(unevenLayout.getStageForCore(70) == 2);
+  assert(unevenLayout.getStageForCore(137) == 3);
+  assert(!unevenLayout.getStageForCore(138));
+
+  PipelineCoreLayout dynamicLayout(std::vector<size_t> {2, 4, 1, 3});
+  assert(dynamicLayout.isValid());
+  assert(dynamicLayout.getProcessorCount() == 10);
+  assert(dynamicLayout.getStageRange(0).begin == 0);
+  assert(dynamicLayout.getStageRange(1).begin == 2);
+  assert(dynamicLayout.getStageRange(2).begin == 6);
+  assert(dynamicLayout.getStageRange(3).begin == 7);
+  assert(dynamicLayout.getStageForCore(1) == 0);
+  assert(dynamicLayout.getStageForCore(2) == 1);
+  assert(dynamicLayout.getStageForCore(6) == 2);
+  assert(dynamicLayout.getStageForCore(9) == 3);
+
   TransferCost transfer {.fixed = 50, .networkFlits = 4};
 
   SchedulingTarget fast;
@@ -57,6 +87,13 @@ int main() {
   };
   assert(mapLogicalProcessorsToPhysicalCores(logicalTrafficFlits, alreadyPlaced) == std::vector<size_t>({0, 1, 2}));
 
+  std::vector<size_t> placementGroups {0, 1, 1};
+  std::vector<size_t> groupedPlacement = mapLogicalProcessorsToPhysicalCores(
+    logicalTrafficFlits, line, placementGroups);
+  for (size_t processor = 0; processor < groupedPlacement.size(); ++processor)
+    assert(placementGroups[processor]
+           == placementGroups[groupedPlacement[processor]]);
+
   ComputeGraph graph;
   graph.successors.resize(6);
   graph.predecessors.resize(6);
@@ -71,16 +108,16 @@ int main() {
     graph.instanceToIndex[instance] = task;
   }
 
-  MergeScheduleResult pipelineSchedule;
-  pipelineSchedule.processorCount = 2;
-  pipelineSchedule.dominanceOrderCompute.reserve(graph.nodes.size());
+  MergeScheduleResult logicalSchedule;
+  logicalSchedule.processorCount = 2;
+  logicalSchedule.dominanceOrderCompute.reserve(graph.nodes.size());
   for (size_t task = 0; task < graph.nodes.size(); ++task) {
     const ComputeInstance& instance = graph.nodes[task].instance;
-    pipelineSchedule.dominanceOrderCompute.push_back(instance);
+    logicalSchedule.dominanceOrderCompute.push_back(instance);
     size_t cpu = task < 4 ? 0 : 1;
-    pipelineSchedule.computeToCpuMap[instance] = cpu;
-    pipelineSchedule.computeToCpuSlotMap[instance] = task < 4 ? task : task - 4;
-    pipelineSchedule.computeToAestMap[instance] = task;
+    logicalSchedule.computeToCpuMap[instance] = cpu;
+    logicalSchedule.computeToCpuSlotMap[instance] = task < 4 ? task : task - 4;
+    logicalSchedule.computeToAestMap[instance] = task;
   }
 
   SchedulingTarget physical = fast;
@@ -93,20 +130,50 @@ int main() {
     3, 3, 3, 0,
   };
   std::string pipelineError;
+  MergeScheduleResult pipelineSchedule = logicalSchedule;
   assert(mlir::succeeded(applyPipelineScheduling(
     graph, pipelineSchedule, 2, physical, pipelineError)));
   assert(pipelineSchedule.processorCount == 4);
-  assert(pipelineSchedule.computeToCpuMap.lookup(graph.nodes[0].instance) == 0);
-  assert(pipelineSchedule.computeToCpuMap.lookup(graph.nodes[1].instance) == 0);
-  assert(pipelineSchedule.computeToCpuMap.lookup(graph.nodes[2].instance) == 2);
-  assert(pipelineSchedule.computeToCpuMap.lookup(graph.nodes[3].instance) == 2);
-  assert(pipelineSchedule.computeToCpuMap.lookup(graph.nodes[4].instance) == 1);
-  assert(pipelineSchedule.computeToCpuMap.lookup(graph.nodes[5].instance) == 3);
+  size_t predecessorCore = pipelineSchedule.computeToCpuMap.lookup(
+    graph.nodes[1].instance);
+  size_t successorCore = pipelineSchedule.computeToCpuMap.lookup(
+    graph.nodes[2].instance);
   assert(pipelineSchedule.computeToAestMap.lookup(graph.nodes[2].instance)
          >= pipelineSchedule.computeToAestMap.lookup(graph.nodes[1].instance)
-              + graph.nodes[1].cost + 4);
+              + graph.nodes[1].cost
+              + getPeftTransferTime(
+                  TransferCost {.fixed = 1, .networkFlits = 1},
+                  predecessorCore, successorCore, physical));
+  assert(pipelineSchedule.processorStages[predecessorCore]
+         <= pipelineSchedule.processorStages[successorCore]);
+  assert(pipelineSchedule.processorStages[successorCore]
+         <= pipelineSchedule.processorStages[predecessorCore] + 1);
   assert(pipelineSchedule.equivalentClass.empty());
 
+  graph.successors[0].push_back(
+    {5, TransferCost {.fixed = 1, .networkFlits = 1}});
+  graph.predecessors[5].push_back(
+    {0, TransferCost {.fixed = 1, .networkFlits = 1}});
+  SchedulingTarget fourStagePhysical = physical;
+  fourStagePhysical.processorCount = 8;
+  fourStagePhysical.interProcessorLatencyNs.assign(64, 3);
+  for (size_t core = 0; core < 8; ++core)
+    fourStagePhysical.interProcessorLatencyNs[core * 8 + core] = 0;
+  MergeScheduleResult fourStageSchedule = logicalSchedule;
+  assert(mlir::succeeded(applyPipelineScheduling(
+    graph, fourStageSchedule, 4, fourStagePhysical, pipelineError)));
+  for (size_t task = 0; task < graph.nodes.size(); ++task)
+    for (const auto &[predecessor, cost] : graph.predecessors[task]) {
+      (void)cost;
+      size_t sourceStage = fourStageSchedule.processorStages[
+        fourStageSchedule.computeToCpuMap.lookup(
+          graph.nodes[predecessor].instance)];
+      size_t targetStage = fourStageSchedule.processorStages[
+        fourStageSchedule.computeToCpuMap.lookup(graph.nodes[task].instance)];
+      assert(sourceStage <= targetStage);
+      assert(targetStage <= sourceStage + 1);
+    }
+
   ComputeGraph communicationGraph;
   communicationGraph.successors.resize(5);
   communicationGraph.predecessors.resize(5);
@@ -136,15 +203,19 @@ int main() {
   MergeScheduleResult fastCommunicationSchedule = communicationSchedule;
   assert(mlir::succeeded(applyPipelineScheduling(
     communicationGraph, fastCommunicationSchedule, 2, fastPipeline, pipelineError)));
-  assert(fastCommunicationSchedule.computeToCpuMap.lookup(
-           communicationGraph.nodes[2].instance) == 2);
 
   SchedulingTarget slowPipeline = fastPipeline;
   slowPipeline.averageInterProcessorLatencyNs = 10;
   MergeScheduleResult slowCommunicationSchedule = communicationSchedule;
   assert(mlir::succeeded(applyPipelineScheduling(
     communicationGraph, slowCommunicationSchedule, 2, slowPipeline, pipelineError)));
-  assert(slowCommunicationSchedule.computeToCpuMap.lookup(
-           communicationGraph.nodes[2].instance) < 2);
+  size_t sourceCore = slowCommunicationSchedule.computeToCpuMap.lookup(
+    communicationGraph.nodes[4].instance);
+  size_t targetCore = slowCommunicationSchedule.computeToCpuMap.lookup(
+    communicationGraph.nodes[3].instance);
+  assert(slowCommunicationSchedule.processorStages[sourceCore]
+         <= slowCommunicationSchedule.processorStages[targetCore]);
+  assert(slowCommunicationSchedule.processorStages[targetCore]
+         <= slowCommunicationSchedule.processorStages[sourceCore] + 1);
   return EXIT_SUCCESS;
 }