nnicolosi/Raptor
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

better throughput scheduling with steady state cost model, now limited to spatial

Browse Source
This commit is contained in:
NiccoloN
2026-10-26 12:07:01 +01:00
parent fd493ae6e8
commit 4b53f6cf10
62 changed files with 4293 additions and 871 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/PIM/PimPipelineSelectionTest.cpp
+284 -40
View File
@@ -1,11 +1,259 @@
#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 "TestCheck.hpp"
#include <limits>
#include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
#include "src/Accelerators/PIM/Passes/PIMPasses.h"
#include "src/Accelerators/PIM/Compiler/PimPipelineSelection.hpp"
#include "src/Accelerators/PIM/Compiler/PimCodeGen.hpp"
#include "src/Accelerators/PIM/Compiler/PimCompilerUtils.hpp"
#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
#include "src/Accelerators/PIM/Common/Support/FileSystemUtils.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/Passes/Transforms/Bufferization/OpBufferizationInterfaces.hpp"
#include "mlir/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/IR/Verifier.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/MemoryBuffer.h"
using namespace onnx_mlir;
using namespace onnx_mlir::spatial;
static WholeScheduleCost scheduleCost(double interval) {
WholeScheduleCost cost;
cost.primary = cost.alternative = interval;
cost.selectionMargin = 0.021;
cost.loweringSupport = StaticLoweringSupport::Supported;
for (auto &recurrence : cost.recurrence)
recurrence = {ScheduleRecurrenceEvidence::Proved, 1, interval, {}};
return cost;
}
static void testEmptyGraphPreparation() {
using namespace mlir;
MLIRContext context;
context.loadDialect<func::FuncDialect, SpatialDialect>();
OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
OpBuilder builder(module->getBodyRegion());
auto function = func::FuncOp::create(builder, module->getLoc(), "main_graph",
builder.getFunctionType({}, {}));
builder.setInsertionPointToStart(function.addEntryBlock());
func::ReturnOp::create(builder, module->getLoc());
SchedulingTarget target;
target.processorCount = target.residentWeightCapacity = 1;
target.interProcessorLatencyNs = {0};
for (bool collectCandidates : {false, true}) {
auto prepared = prepareSpatialSchedules(*module, target, 1, collectCandidates);
TEST_CHECK(succeeded(prepared) && prepared->candidates.size() == 1);
TEST_CHECK(prepared->candidates.front().dominanceOrderCompute.empty());
}
}
static void testSelectedLoweringFailure() {
using namespace mlir;
MLIRContext context(MLIRContext::Threading::DISABLED);
context.loadDialect<SpatialDialect, pim::PimDialect, func::FuncDialect,
arith::ArithDialect, tensor::TensorDialect, memref::MemRefDialect,
scf::SCFDialect, linalg::LinalgDialect, affine::AffineDialect,
bufferization::BufferizationDialect>();
OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
OpBuilder builder(module->getBodyRegion());
auto loc = module->getLoc();
auto type = RankedTensorType::get({1, 536870912}, builder.getF32Type());
auto function = func::FuncOp::create(builder, loc, "main_graph", builder.getFunctionType({type}, {type}));
builder.setInsertionPointToStart(function.addEntryBlock());
auto compute = SpatGraphCompute::create(builder, loc, TypeRange{type}, ValueRange{},
ValueRange{function.getArgument(0)});
compute.getBody().push_back(new Block);
Value input = compute.getBody().front().addArgument(type, loc);
builder.setInsertionPointToStart(&compute.getBody().front());
Value result = SpatReluOp::create(builder, loc, type, input);
SpatYieldOp::create(builder, loc, ValueRange{result});
builder.setInsertionPointAfter(compute);
func::ReturnOp::create(builder, loc, compute.getResults());
TEST_CHECK(succeeded(verify(*module)));
SchedulingTarget target;
target.processorCount = 2;
target.residentWeightCapacity = 1;
target.synchronizationRegisterCount = 64;
target.interProcessorLatencyNs = {0, 2, 2, 0};
target.averageInterProcessorLatencyNs = 2;
target.estimateRealizedScheduleCost = [](ModuleOp, llvm::json::Object &) {
TEST_CHECK(false && "single-candidate compilation must skip costs");
return WholeScheduleCost{};
};
SpatialTargetResources resources{{128, 128}, 1, 2, 16};
std::string diagnostics;
ScopedDiagnosticHandler handler(&context, [&](Diagnostic &diagnostic) {
diagnostics += diagnostic.str(); return success();
});
ScopedCompilerOutputSuppression silence;
// Match the compiler's preloaded backend registry before nested pipelines run.
PassManager backend(&context);
addPimLoweringPasses(backend, resources, EmitPimCodegen, false);
DialectRegistry registry;
backend.getDependentDialects(registry);
pim::registerOpBufferizationInterfaces(registry);
arith::registerBufferizableOpInterfaceExternalModels(registry);
bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(registry);
context.appendDialectRegistry(registry);
context.loadAllAvailableDialects();
for (bool collectCandidates : {false, true}) {
OwningOpRef<ModuleOp> checkpoint = cast<ModuleOp>((*module)->clone());
auto prepared = prepareSpatialSchedules(*checkpoint, target, 2, collectCandidates);
TEST_CHECK(succeeded(prepared) && !prepared->candidates.empty());
}
PassManager manager(&context);
manager.addPass(createPimPipelineSelectionPass(target, resources, 2, SpatialDataflowExportStage::None));
TEST_CHECK(failed(manager.run(*module)));
TEST_CHECK(diagnostics.find("backend selection failed") != std::string::npos);
TEST_CHECK(diagnostics.find("outside representable range") != std::string::npos);
}
static void testResourceRetries() {
using namespace mlir;
auto previousCores = coresCount.getValue();
auto previousSize = crossbarSize.getValue();
auto previousCrossbars = crossbarCountInCore.getValue();
coresCount = 2; crossbarSize = 128; crossbarCountInCore = 1;
MLIRContext context;
context.loadDialect<pim::PimDialect, memref::MemRefDialect, func::FuncDialect,
arith::ArithDialect>();
OwningOpRef<ModuleOp> checkpoint = ModuleOp::create(UnknownLoc::get(&context));
OpBuilder builder(checkpoint->getBodyRegion());
auto loc = checkpoint->getLoc();
auto function = func::FuncOp::create(builder, loc, "main_graph", builder.getFunctionType({}, {}));
builder.setInsertionPointToStart(function.addEntryBlock());
auto core = pim::PimCoreOp::create(builder, loc, ValueRange{}, uint32_t{0});
core.getBody().push_back(new Block);
builder.setInsertionPointToStart(&core.getBody().front());
auto allocation = memref::AllocOp::create(builder, loc,
MemRefType::get({int64_t{1} << 31}, builder.getI8Type()));
auto smallType = MemRefType::get({1}, builder.getI8Type());
auto destination = memref::AllocOp::create(builder, loc, smallType);
Value zero = arith::ConstantIndexOp::create(builder, loc, 0);
pim::PimMemCopyOp::create(builder, loc, smallType, zero, zero, destination,
allocation, builder.getI32IntegerAttr(1));
pim::PimHaltOp::create(builder, loc);
builder.setInsertionPointAfter(core);
func::ReturnOp::create(builder, loc);
TEST_CHECK(succeeded(verify(*checkpoint)));
std::string before;
llvm::raw_string_ostream stream(before); checkpoint->print(stream, OpPrintingFlags().printGenericOpForm());
auto original = scheduleCost(100);
auto faster = scheduleCost(90);
std::vector<WholeScheduleCost> costs{original, faster};
std::string reason;
SpatialTargetResources resources{{128, 128}, 1, 2, 16};
PassManager backend(&context);
addPimLoweringPasses(backend, resources, EmitPimCodegen, false);
DialectRegistry registry;
backend.getDependentDialects(registry);
pim::registerOpBufferizationInterfaces(registry);
arith::registerBufferizableOpInterfaceExternalModels(registry);
bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(registry);
context.appendDialectRegistry(registry);
context.loadAllAvailableDialects();
auto emit = [&](ModuleOp module) {
llvm::SmallString<128> temporary;
TEST_CHECK(!llvm::sys::fs::createUniqueDirectory("raptor-retry-program", temporary));
std::string output = temporary.str().str();
ScopedCompilerOutputSuppression silence;
TEST_CHECK(compileToPimCode(module, output) == CompilerSuccess);
auto binary = llvm::MemoryBuffer::getFile(output + "/core_0.pim");
TEST_CHECK(binary && (*binary)->getBufferSize() > 12);
std::string program = (*binary)->getBuffer().str();
TEST_CHECK(!llvm::sys::fs::remove_directories(temporary));
return program;
};
std::string emitted;
for (bool allFail : {false, true}) {
std::vector<size_t> attempted;
OwningOpRef<ModuleOp> winner;
std::vector<llvm::json::Object> summaries;
auto selected = lowerSelectedPimCandidate(costs, [&](size_t index) {
attempted.push_back(index);
IRMapping mapping;
OwningOpRef<ModuleOp> candidate = cast<ModuleOp>((*checkpoint)->clone(mapping));
if (!index && !allFail)
cast<memref::AllocOp>(mapping.lookup(allocation.getResult()).getDefiningOp())
.getResult().setType(MemRefType::get({16}, builder.getI8Type()));
llvm::json::Object summary;
auto result = lowerPimPipelineCandidate(*candidate, resources, summary);
if (result == PimCandidateResult::Success) winner = std::move(candidate);
summaries.push_back(std::move(summary));
return result;
}, reason);
TEST_CHECK(attempted == std::vector<size_t>({1, 0}));
TEST_CHECK(allFail ? !selected : selected == 0);
TEST_CHECK(summaries.front().getString("status") == "resource-limit");
TEST_CHECK(summaries.front().getString("reason")->contains("signed int32 address range"));
if (selected) emitted = emit(*winner);
if (allFail) TEST_CHECK(reason.find("all candidate schedules failed") != std::string::npos);
std::string after;
llvm::raw_string_ostream out(after); checkpoint->print(out, OpPrintingFlags().printGenericOpForm());
TEST_CHECK(before == after);
}
TEST_CHECK(!emitted.empty());
IRMapping mapping;
OwningOpRef<ModuleOp> alternative = cast<ModuleOp>((*checkpoint)->clone(mapping));
cast<memref::AllocOp>(mapping.lookup(allocation.getResult()).getDefiningOp())
.getResult().setType(MemRefType::get({16}, builder.getI8Type()));
llvm::json::Object summary;
TEST_CHECK(lowerPimPipelineCandidate(*alternative, resources, summary) == PimCandidateResult::Success);
TEST_CHECK(emit(*alternative) == emitted);
size_t attempts = 0;
auto empty = lowerSelectedPimCandidate({}, [&](size_t) {
++attempts;
return PimCandidateResult::Success;
}, reason);
TEST_CHECK(!empty && attempts == 0 && reason.find("at least one candidate") != std::string::npos);
auto selected = lowerSelectedPimCandidate(costs, [&](size_t) {
++attempts;
OwningOpRef<ModuleOp> invalid = cast<ModuleOp>((*alternative)->clone());
invalid->walk([&](pim::PimMemCopyOp copy) {
copy->setAttr("size", builder.getI64IntegerAttr(1));
});
llvm::json::Object failure;
auto result = lowerPimPipelineCandidate(*invalid, resources, failure);
TEST_CHECK(result == PimCandidateResult::Failure);
TEST_CHECK(failure.getString("status") == "systemic-or-unclassified-error");
return result;
}, reason);
TEST_CHECK(!selected && attempts == 1);
TEST_CHECK(reason.find("recovery is not justified") != std::string::npos);
for (unsigned unknown = 0; unknown < 4; ++unknown) {
auto invalid = faster;
if (unknown == 0) invalid.primary.reset();
if (unknown == 1) invalid.recurrence[0].status = ScheduleRecurrenceEvidence::Unproved;
if (unknown == 2) invalid.loweringSupport = StaticLoweringSupport::Unknown;
if (unknown == 3) invalid.loweringSupport = StaticLoweringSupport::Unsupported;
attempts = 0;
std::vector<WholeScheduleCost> incomplete{original, invalid};
selected = lowerSelectedPimCandidate(incomplete, [&](size_t) {
++attempts;
return PimCandidateResult::Success;
}, reason);
TEST_CHECK(!selected && attempts == 0);
}
coresCount = previousCores; crossbarSize = previousSize; crossbarCountInCore = previousCrossbars;
}
int main() {
testEmptyGraphPreparation();
testSelectedLoweringFailure();
testResourceRetries();
mlir::MLIRContext context;
context.allowUnregisteredDialects();
mlir::OwningOpRef<mlir::ModuleOp> checkpoint = mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
@@ -25,56 +273,52 @@ int main() {
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);
TEST_CHECK(copy.op != op && copy.laneStart == 3 && copy.laneCount == 2);
TEST_CHECK(remapped.computeToCpuMap.lookup(copy) == 7 && remapped.computeToCpuSlotMap.lookup(copy) == 4);
TEST_CHECK(remapped.computeToAestMap.lookup(copy) == 42 && remapped.isLastComputeOfCpu.contains(copy));
TEST_CHECK(remapped.cpuToLastComputeMap.lookup(7) == copy);
copy.op->erase();
assert(schedule.dominanceOrderCompute.front().op == op && op->getBlock() == checkpoint->getBody());
TEST_CHECK(schedule.dominanceOrderCompute.front().op == op && op->getBlock() == checkpoint->getBody());
}
auto score = [](double value, std::string id) { return WholeScheduleCost{value, value, {}, id, 0.021}; };
auto score = [](double value, std::string) {
return scheduleCost(value);
};
std::string reason;
auto original = score(100, "0"), faster = score(90, "1");
for (auto support : {StaticLoweringSupport::Unknown,
StaticLoweringSupport::Supported, StaticLoweringSupport::Unsupported}) {
auto cost = faster;
cost.loweringSupport = support;
TEST_CHECK(cost.available());
TEST_CHECK(cost.eligible() == (support == StaticLoweringSupport::Supported));
TEST_CHECK(cost.loweringSupported().has_value() == (support != StaticLoweringSupport::Unknown));
}
auto incompleteProof = faster;
incompleteProof.recurrence[1].periodGenerations = 0;
TEST_CHECK(incompleteProof.available() && !incompleteProof.eligible());
incompleteProof.recurrence[1] = {};
TEST_CHECK(incompleteProof.convergenceProved() == false && !incompleteProof.eligible());
incompleteProof.recurrence[0] = {};
TEST_CHECK(!incompleteProof.convergenceProved().has_value());
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);
TEST_CHECK(select({original, faster}) == 1);
TEST_CHECK(select({original, score(97, "2")}) == 0);
TEST_CHECK(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);
TEST_CHECK(select({original, sensitive}) == 0);
TEST_CHECK(select({original, faster, score(91, "2")}) == 1);
TEST_CHECK(select({original, faster, score(90, "1")}) == 1);
auto unavailable = WholeScheduleCost{};
assert(select({unavailable, faster}) == 0);
assert(select({original, unavailable}) == 0);
TEST_CHECK(select({unavailable, faster}) == 0);
TEST_CHECK(select({original, unavailable}) == 0);
auto unknownUncertainty = faster;
unknownUncertainty.relativeUncertainty.reset();
assert(select({original, unknownUncertainty}) == 0);
unknownUncertainty.selectionMargin.reset();
TEST_CHECK(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);
TEST_CHECK(select({original, score(value, "2")}) == 0);
for (int i = 0; i < 8; ++i) TEST_CHECK(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");
}
}