NiccoloN
369 lines
16 KiB
C++
369 lines
16 KiB
C++
#include <cassert>
|
|
#include <cstdlib>
|
|
#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.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);
|
|
|
|
testPhysicalPeft();
|
|
|
|
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 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;
|
|
physicalSchedule.dominanceOrderCompute.push_back(instance);
|
|
size_t cpu = task < 4 ? 0 : 1;
|
|
physicalSchedule.computeToCpuMap[instance] = cpu;
|
|
physicalSchedule.computeToCpuSlotMap[instance] = task < 4 ? task : task - 4;
|
|
physicalSchedule.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,
|
|
};
|
|
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, 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 = 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 = physicalSchedule;
|
|
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 = physicalSchedule;
|
|
fourStageSchedule.processorCount = 8;
|
|
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 = 4;
|
|
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;
|
|
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)));
|
|
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;
|
|
}
|