Test Spatial Scheduling

2026-08-06 14:32:57 +02:00
parent a39fdba366
commit aeedf2f566
1 changed files with 92 additions and 0 deletions
@@ -1,8 +1,10 @@
 #include <cassert>
 #include <cstdlib>
+#include <string>
 #include <vector>

 #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;

@@ -54,5 +56,95 @@ int main() {
    0,
  };
  assert(mapLogicalProcessorsToPhysicalCores(logicalTrafficFlits, alreadyPlaced) == std::vector<size_t>({0, 1, 2}));
+
+  ComputeGraph graph;
+  graph.successors.resize(6);
+  graph.predecessors.resize(6);
+  graph.successors[1].push_back({2, TransferCost {.fixed = 1, .networkFlits = 1}});
+  graph.predecessors[2].push_back({1, TransferCost {.fixed = 1, .networkFlits = 1}});
+  const Cost costs[] = {6, 4, 6, 4, 8, 8};
+  for (uint32_t task = 0; task < 6; ++task) {
+    ComputeInstance instance {nullptr, task, 1};
+    ResidentWeight weight;
+    weight.opaqueLane = task;
+    graph.nodes.push_back({instance, costs[task], {weight}, task});
+    graph.instanceToIndex[instance] = task;
+  }
+
+  MergeScheduleResult pipelineSchedule;
+  pipelineSchedule.processorCount = 2;
+  pipelineSchedule.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);
+    size_t cpu = task < 4 ? 0 : 1;
+    pipelineSchedule.computeToCpuMap[instance] = cpu;
+    pipelineSchedule.computeToCpuSlotMap[instance] = task < 4 ? task : task - 4;
+    pipelineSchedule.computeToAestMap[instance] = task;
+  }
+
+  SchedulingTarget physical = fast;
+  physical.processorCount = 4;
+  physical.residentWeightCapacity = 2;
+  physical.interProcessorLatencyNs = {
+    0, 3, 3, 3,
+    3, 0, 3, 3,
+    3, 3, 0, 3,
+    3, 3, 3, 0,
+  };
+  std::string pipelineError;
+  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);
+  assert(pipelineSchedule.computeToAestMap.lookup(graph.nodes[2].instance)
+         >= pipelineSchedule.computeToAestMap.lookup(graph.nodes[1].instance)
+              + graph.nodes[1].cost + 4);
+  assert(pipelineSchedule.equivalentClass.empty());
+
+  ComputeGraph communicationGraph;
+  communicationGraph.successors.resize(5);
+  communicationGraph.predecessors.resize(5);
+  communicationGraph.successors[4].push_back(
+    {3, TransferCost {.fixed = 0, .networkFlits = 1}});
+  communicationGraph.predecessors[3].push_back(
+    {4, TransferCost {.fixed = 0, .networkFlits = 1}});
+  const Cost communicationCosts[] = {6, 4, 6, 4, 1};
+  MergeScheduleResult communicationSchedule;
+  communicationSchedule.processorCount = 2;
+  for (uint32_t task = 0; task < 5; ++task) {
+    ComputeInstance instance {nullptr, task, 1};
+    ResidentWeight weight;
+    weight.opaqueLane = task;
+    communicationGraph.nodes.push_back(
+      {instance, communicationCosts[task], {weight}, task});
+    communicationGraph.instanceToIndex[instance] = task;
+    communicationSchedule.dominanceOrderCompute.push_back(instance);
+    size_t cpu = task < 4 ? 0 : 1;
+    communicationSchedule.computeToCpuMap[instance] = cpu;
+    communicationSchedule.computeToCpuSlotMap[instance] = task < 4 ? task : 0;
+    communicationSchedule.computeToAestMap[instance] = task;
+  }
+
+  SchedulingTarget fastPipeline = physical;
+  fastPipeline.residentWeightCapacity = 4;
+  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);
  return EXIT_SUCCESS;
 }