Files
Raptor/test/PIM/SpatialSchedulingTargetTest.cpp
NiccoloN add20e56eb
Validate Operations / validate-operations (push) Waiting to run
minor fix
2026-08-19 16:33:04 +02:00

230 lines
9.3 KiB
C++

#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;
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;
fast.processorCount = 2;
fast.interProcessorLatencyNs = {0, 3, 3, 0};
fast.averageInterProcessorLatencyNs = 3;
SchedulingTarget slow = fast;
slow.interProcessorLatencyNs = {0, 10, 10, 0};
slow.averageInterProcessorLatencyNs = 10;
assert(getPeftTransferTime(transfer, 0, 0, fast) == 0);
assert(getPeftTransferTime(transfer, 0, 1, fast) == 62);
assert(getPeftTransferTime(transfer, 0, 1, slow) == 90);
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]]);
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 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;
logicalSchedule.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;
}
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;
ComputeGraph emptyGraph;
MergeScheduleResult emptySchedule;
emptySchedule.processorCount = 2;
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;
assert(mlir::succeeded(applyPipelineScheduling(
graph, pipelineSchedule, 2, physical, pipelineError)));
assert(pipelineSchedule.processorCount == 4);
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
+ 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);
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)));
SchedulingTarget slowPipeline = fastPipeline;
slowPipeline.averageInterProcessorLatencyNs = 10;
MergeScheduleResult slowCommunicationSchedule = communicationSchedule;
assert(mlir::succeeded(applyPipelineScheduling(
communicationGraph, slowCommunicationSchedule, 2, slowPipeline, pipelineError)));
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;
}