NiccoloN
81 lines
4.0 KiB
C++
81 lines
4.0 KiB
C++
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ScheduleEvaluation.hpp"
|
|
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.hpp"
|
|
#include "llvm/Support/JSON.h"
|
|
#include <cassert>
|
|
#include <limits>
|
|
using namespace onnx_mlir;
|
|
using namespace onnx_mlir::spatial;
|
|
int main() {
|
|
mlir::MLIRContext context;
|
|
context.allowUnregisteredDialects();
|
|
mlir::OwningOpRef<mlir::ModuleOp> checkpoint = mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
|
|
mlir::OperationState state(checkpoint->getLoc(), "test.compute");
|
|
auto *op = mlir::Operation::create(state);
|
|
checkpoint->getBody()->push_back(op);
|
|
spatial::ComputeInstance instance{op, 3, 2};
|
|
spatial::MergeScheduleResult schedule;
|
|
schedule.dominanceOrderCompute.push_back(instance);
|
|
schedule.computeToCpuMap[instance] = 7;
|
|
schedule.computeToCpuSlotMap[instance] = 4;
|
|
schedule.computeToAestMap[instance] = 42;
|
|
schedule.isLastComputeOfCpu.insert(instance);
|
|
schedule.cpuToLastComputeMap[7] = instance;
|
|
for (int i = 0; i < 3; ++i) {
|
|
mlir::IRMapping mapping;
|
|
mlir::OwningOpRef<mlir::ModuleOp> trial = mlir::cast<mlir::ModuleOp>((*checkpoint)->clone(mapping));
|
|
auto remapped = spatial::remapSpatialSchedule(schedule, mapping);
|
|
auto copy = remapped.dominanceOrderCompute.front();
|
|
assert(copy.op != op && copy.laneStart == 3 && copy.laneCount == 2);
|
|
assert(remapped.computeToCpuMap.lookup(copy) == 7 && remapped.computeToCpuSlotMap.lookup(copy) == 4);
|
|
assert(remapped.computeToAestMap.lookup(copy) == 42 && remapped.isLastComputeOfCpu.contains(copy));
|
|
assert(remapped.cpuToLastComputeMap.lookup(7) == copy);
|
|
copy.op->erase();
|
|
assert(schedule.dominanceOrderCompute.front().op == op && op->getBlock() == checkpoint->getBody());
|
|
}
|
|
|
|
auto score = [](double value, std::string id) { return WholeScheduleCost{value, value, {}, id, 0.021}; };
|
|
std::string reason;
|
|
auto original = score(100, "0"), faster = score(90, "1");
|
|
auto select = [&](std::initializer_list<WholeScheduleCost> values) {
|
|
return selectPipelineCandidate(llvm::ArrayRef<WholeScheduleCost>(values), reason);
|
|
};
|
|
assert(select({original, faster}) == 1);
|
|
assert(select({original, score(97, "2")}) == 0);
|
|
assert(select({original, score(105, "2")}) == 0);
|
|
auto sensitive = faster; sensitive.alternative = 101;
|
|
assert(select({original, sensitive}) == 0);
|
|
assert(select({original, faster, score(91, "2")}) == 1);
|
|
assert(select({original, faster, score(90, "1")}) == 1);
|
|
auto unavailable = WholeScheduleCost{};
|
|
assert(select({unavailable, faster}) == 0);
|
|
assert(select({original, unavailable}) == 0);
|
|
auto unknownUncertainty = faster;
|
|
unknownUncertainty.relativeUncertainty.reset();
|
|
assert(select({original, unknownUncertainty}) == 0);
|
|
for (double value : {0., -1., std::numeric_limits<double>::infinity(), std::numeric_limits<double>::quiet_NaN()})
|
|
assert(select({original, score(value, "2")}) == 0);
|
|
for (int i = 0; i < 8; ++i) assert(select({original, faster}) == 1);
|
|
|
|
// The generic query accepts target-lowered IR without a target dialect API.
|
|
for (double interval : {90., 97.}) {
|
|
SchedulingTarget target;
|
|
assert(!target.prepareScheduleCost);
|
|
target.prepareScheduleCost = [interval](mlir::ModuleOp module,
|
|
llvm::json::Object &) -> mlir::FailureOr<ScheduleCostQuery> {
|
|
assert(module.getBody()->getOperations().size() == 1);
|
|
ScheduleCostQuery query;
|
|
query.identity = "static-program";
|
|
query.evaluate = [interval](llvm::json::Object &) {
|
|
return WholeScheduleCost{interval, interval, {}, "static-program", 0.021};
|
|
};
|
|
return query;
|
|
};
|
|
llvm::json::Object summary;
|
|
auto query = target.prepareScheduleCost(*checkpoint, summary);
|
|
assert(mlir::succeeded(query) && query->identity == "static-program");
|
|
auto cost = query->evaluate(summary);
|
|
assert(cost.eligible() && select({original, cost}) == (interval == 90. ? 1 : 0));
|
|
assert(checkpoint->getBody()->front().getName().getStringRef() == "test.compute");
|
|
}
|
|
}
|