diff --git a/test/PIM/SpatialSchedulingTargetTest.cpp b/test/PIM/SpatialSchedulingTargetTest.cpp index e50c296..dce2d35 100644 --- a/test/PIM/SpatialSchedulingTargetTest.cpp +++ b/test/PIM/SpatialSchedulingTargetTest.cpp @@ -1,8 +1,10 @@ #include #include +#include #include #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({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; }