fix scheduling reasoning on physical cores directly

test/PIM/SpatialSchedulingTargetTest.cpp

@@ -3,15 +3,156 @@
 #include <string>
 #include <vector>
 
+#include "mlir/IR/BuiltinOps.h"
+
+#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.hpp"
 
 using namespace onnx_mlir::spatial;
 
+static void testPhysicalPeft() {
+  TransferCost transfer {.fixed = 50, .networkFlits = 4};
+  SchedulingTarget line;
+  line.processorCount = 3;
+  line.interProcessorLatencyNs = {
+    0,
+    1,
+    10,
+    1,
+    0,
+    1,
+    10,
+    1,
+    0,
+  };
+  line.averageInterProcessorLatencyNs = 4;
+  assert(getPeftTransferTime(transfer, 0, 1, line) == 54);
+  assert(getPeftTransferTime(transfer, 0, 2, line) == 90);
+  SchedulingTarget reversed = line;
+  reversed.interProcessorLatencyNs = {0, 10, 1, 10, 0, 1, 1, 1, 0};
+  assert(getPeftTransferTime(transfer, 0, 1, reversed) == 90);
+  assert(getPeftTransferTime(transfer, 0, 2, reversed) == 54);
+
+  mlir::MLIRContext context;
+  mlir::OwningOpRef<mlir::ModuleOp> owner =
+    mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
+  ComputeGraph chain;
+  chain.successors.resize(3);
+  chain.predecessors.resize(3);
+  for (uint32_t task = 0; task < 3; ++task) {
+    ComputeInstance instance {owner->getOperation(), task, 1};
+    ResidentWeight weight;
+    weight.opaqueLane = task;
+    chain.nodes.push_back({instance, 10, {weight}, task});
+    if (task != 0) {
+      chain.successors[task - 1].push_back({task, transfer});
+      chain.predecessors[task].push_back({task - 1, transfer});
+    }
+  }
+  line.residentWeightCapacity = reversed.residentWeightCapacity = 3;
+  for (const SchedulingTarget &target : {line, reversed}) {
+    auto scheduled = runPeftScheduler(chain, {target, &context});
+    size_t firstCore = scheduled.computeToCpuMap.lookup(chain.nodes[0].instance);
+    assert(firstCore == (target.interProcessorLatencyNs == line.interProcessorLatencyNs ? 1 : 2));
+    for (size_t task = 1; task < chain.nodes.size(); ++task) {
+      auto previous = chain.nodes[task - 1].instance;
+      auto current = chain.nodes[task].instance;
+      size_t source = scheduled.computeToCpuMap.lookup(previous);
+      size_t destination = scheduled.computeToCpuMap.lookup(current);
+      assert(source != destination);
+      assert(scheduled.computeToAestMap.lookup(current)
+             == scheduled.computeToAestMap.lookup(previous) + 10
+                  + getPeftTransferTime(transfer, source, destination, target));
+    }
+  }
+}
+
+static void testPipelineInputArrival() {
+  ComputeGraph graph;
+  graph.successors.resize(4);
+  graph.predecessors.resize(4);
+  MergeScheduleResult schedule;
+  schedule.processorCount = 4;
+  for (uint32_t task = 0; task < 4; ++task) {
+    ComputeInstance instance {nullptr, task, 1};
+    ResidentWeight weight;
+    weight.opaqueLane = task;
+    graph.nodes.push_back({instance, 1, {weight}, task});
+    schedule.dominanceOrderCompute.push_back(instance);
+    schedule.computeToCpuMap[instance] = task;
+    schedule.computeToCpuSlotMap[instance] = 0;
+    if (task != 0) {
+      TransferCost transfer {.fixed = 0, .networkFlits = 100};
+      graph.predecessors[task].push_back({task - 1, transfer});
+      graph.successors[task - 1].push_back({task, transfer});
+    }
+  }
+  SchedulingTarget target;
+  target.processorCount = 4;
+  target.residentWeightCapacity = 4;
+  target.interProcessorLatencyNs.assign(16, 1);
+  for (size_t core = 0; core < 4; ++core)
+    target.interProcessorLatencyNs[core * 4 + core] = 0;
+  std::string error;
+  assert(mlir::succeeded(applyPipelineScheduling(graph, schedule, 2, target, error)));
+  for (size_t first : {0, 2}) {
+    size_t source = schedule.computeToCpuMap.lookup(graph.nodes[first].instance);
+    size_t destination = schedule.computeToCpuMap.lookup(graph.nodes[first + 1].instance);
+    assert(schedule.processorStages[source] == first / 2);
+    // An idle remote core would finish later than waiting for the local core.
+    assert(source == destination);
+  }
+}
+
+static void testReadyCommunication() {
+  mlir::MLIRContext context;
+  mlir::OwningOpRef<mlir::ModuleOp> owner =
+    mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
+  auto type = mlir::RankedTensorType::get({1}, mlir::Float32Type::get(&context));
+  mlir::Block payloads;
+  auto payload = payloads.addArgument(type, owner->getLoc());
+  DeferredTransferPlan plan;
+  plan.stepCounts = {3, 1, 1};
+  plan.scheduled.resize(3);
+  for (auto [index, info] : llvm::enumerate(plan.scheduled)) {
+    info.op = owner->getOperation();
+    info.streamIds.push_back(index);
+  }
+  ProducedValue producer;
+  producer.scheduled = &plan.scheduled[0];
+  producer.payload = payload;
+  // Two consumers exercise removal of already-scheduled ready-queue entries.
+  for (unsigned target = 1; target <= 2; ++target) {
+    auto exchange = std::make_unique<DeferredExchangePlan>();
+    exchange->target = &plan.scheduled[target];
+    exchange->exchangeId = target;
+    RequirementFamily requirement;
+    requirement.producer = &producer;
+    requirement.publicationFragmentType = type;
+    exchange->requirements.push_back(requirement);
+    ExternalTransferFamily transfer;
+    transfer.requirement = &exchange->requirements.front();
+    transfer.sourceScheduled = producer.scheduled;
+    transfer.targetScheduled = exchange->target;
+    transfer.targetLanes = LaneSet::all(1);
+    transfer.targetStreams = onnx_mlir::StaticIntSequence::uniform(target, 1);
+    exchange->external.push_back(transfer);
+    plan.exchanges.push_back(std::move(exchange));
+  }
+  auto scheduled = scheduleDeferredCommunication({}, plan);
+  assert(mlir::succeeded(scheduled));
+  assert(scheduled->slices.size() == 2);
+  for (const auto &slice : scheduled->slices)
+    assert(slice.sourceInsertionStep == 1);
+}
+
 int main() {
+  testPipelineInputArrival();
+  testReadyCommunication();
   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);
@@ -56,43 +197,7 @@ int main() {
   assert(fast.getInterProcessorLatencyNs(0, 1) == 3);
   assert(slow.getInterProcessorLatencyNs(0, 1) == 10);
 
-  SchedulingTarget line;
-  line.processorCount = 3;
-  line.interProcessorLatencyNs = {
-    0,
-    1,
-    10,
-    1,
-    0,
-    1,
-    10,
-    1,
-    0,
-  };
-  std::vector<Cost> logicalTrafficFlits(9, 0);
-  logicalTrafficFlits[2] = 100;
-  assert(mapLogicalProcessorsToPhysicalCores(logicalTrafficFlits, line) == std::vector<size_t>({1, 0, 2}));
-
-  SchedulingTarget alreadyPlaced = line;
-  alreadyPlaced.interProcessorLatencyNs = {
-    0,
-    10,
-    1,
-    10,
-    0,
-    1,
-    1,
-    1,
-    0,
-  };
-  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]]);
+  testPhysicalPeft();
 
   ComputeGraph graph;
   graph.successors.resize(6);
@@ -108,16 +213,16 @@ int main() {
     graph.instanceToIndex[instance] = task;
   }
 
-  MergeScheduleResult logicalSchedule;
-  logicalSchedule.processorCount = 2;
-  logicalSchedule.dominanceOrderCompute.reserve(graph.nodes.size());
+  MergeScheduleResult physicalSchedule;
+  physicalSchedule.processorCount = 4;
+  physicalSchedule.dominanceOrderCompute.reserve(graph.nodes.size());
   for (size_t task = 0; task < graph.nodes.size(); ++task) {
     const ComputeInstance& instance = graph.nodes[task].instance;
-    logicalSchedule.dominanceOrderCompute.push_back(instance);
+    physicalSchedule.dominanceOrderCompute.push_back(instance);
     size_t cpu = task < 4 ? 0 : 1;
-    logicalSchedule.computeToCpuMap[instance] = cpu;
-    logicalSchedule.computeToCpuSlotMap[instance] = task < 4 ? task : task - 4;
-    logicalSchedule.computeToAestMap[instance] = task;
+    physicalSchedule.computeToCpuMap[instance] = cpu;
+    physicalSchedule.computeToCpuSlotMap[instance] = task < 4 ? task : task - 4;
+    physicalSchedule.computeToAestMap[instance] = task;
   }
 
   SchedulingTarget physical = fast;
@@ -129,24 +234,46 @@ int main() {
     3, 3, 0, 3,
     3, 3, 3, 0,
   };
+  SchedulingTarget clustered = physical;
+  clustered.interProcessorLatencyNs = {
+    0, 20, 1, 20,
+    20, 0, 20, 1,
+    1, 20, 0, 20,
+    20, 1, 20, 0,
+  };
+  PipelineCoreLayout clusteredLayout(std::vector<size_t> {2, 2}, clustered);
+  assert(clusteredLayout.getPhysicalCore(0) == 0);
+  assert(clusteredLayout.getPhysicalCore(1) == 2);
+  assert(clusteredLayout.getStageForCore(0) == 0);
+  assert(clusteredLayout.getStageForCore(2) == 0);
+  assert(clusteredLayout.getStageForCore(1) == 1);
+  assert(clusteredLayout.getStageForCore(3) == 1);
+  PipelineCoreLayout uniformLayout(std::vector<size_t> {2, 2}, physical);
+  assert(uniformLayout.getPhysicalCore(1) == 1);
   std::string pipelineError;
   ComputeGraph preparationGraph = graph;
   ResidentWeight extraWeight;
   extraWeight.opaqueLane = graph.nodes.size();
   preparationGraph.nodes[2].residentWeights.push_back(extraWeight);
   auto preparation = preparePipelineWorkload(
-    preparationGraph, logicalSchedule, 2, physical, pipelineError);
+    preparationGraph, 2, physical, pipelineError);
   assert(mlir::succeeded(preparation));
   assert(*preparation == PipelineWorkloadPreparation::Ready);
+  extraWeight.opaqueLane++;
+  preparationGraph.nodes[2].residentWeights.push_back(extraWeight);
+  assert(mlir::failed(preparePipelineWorkload(
+    preparationGraph, 2, physical, pipelineError)));
+  assert(pipelineError.find("physical core's crossbars") != std::string::npos);
+  pipelineError.clear();
   ComputeGraph emptyGraph;
   MergeScheduleResult emptySchedule;
-  emptySchedule.processorCount = 2;
+  emptySchedule.processorCount = 4;
   assert(mlir::succeeded(applyPipelineScheduling(
     emptyGraph, emptySchedule, 2, physical, pipelineError)));
   assert(emptySchedule.processorCount == physical.processorCount);
   assert(emptySchedule.processorStages == std::vector<size_t>({0, 0, 1, 1}));
 
-  MergeScheduleResult pipelineSchedule = logicalSchedule;
+  MergeScheduleResult pipelineSchedule = physicalSchedule;
   assert(mlir::succeeded(applyPipelineScheduling(
     graph, pipelineSchedule, 2, physical, pipelineError)));
   assert(pipelineSchedule.processorCount == 4);
@@ -175,7 +302,8 @@ int main() {
   fourStagePhysical.interProcessorLatencyNs.assign(64, 3);
   for (size_t core = 0; core < 8; ++core)
     fourStagePhysical.interProcessorLatencyNs[core * 8 + core] = 0;
-  MergeScheduleResult fourStageSchedule = logicalSchedule;
+  MergeScheduleResult fourStageSchedule = physicalSchedule;
+  fourStageSchedule.processorCount = 8;
   assert(mlir::succeeded(applyPipelineScheduling(
     graph, fourStageSchedule, 4, fourStagePhysical, pipelineError)));
   for (size_t task = 0; task < graph.nodes.size(); ++task)
@@ -199,7 +327,7 @@ int main() {
     {4, TransferCost {.fixed = 0, .networkFlits = 1}});
   const Cost communicationCosts[] = {6, 4, 6, 4, 1};
   MergeScheduleResult communicationSchedule;
-  communicationSchedule.processorCount = 2;
+  communicationSchedule.processorCount = 4;
   for (uint32_t task = 0; task < 5; ++task) {
     ComputeInstance instance {nullptr, task, 1};
     ResidentWeight weight;
@@ -222,6 +350,9 @@ int main() {
 
   SchedulingTarget slowPipeline = fastPipeline;
   slowPipeline.averageInterProcessorLatencyNs = 10;
+  slowPipeline.interProcessorLatencyNs.assign(16, 10);
+  for (size_t core = 0; core < 4; ++core)
+    slowPipeline.interProcessorLatencyNs[core * 4 + core] = 0;
   MergeScheduleResult slowCommunicationSchedule = communicationSchedule;
   assert(mlir::succeeded(applyPipelineScheduling(
     communicationGraph, slowCommunicationSchedule, 2, slowPipeline, pipelineError)));