Resnet is fast

This commit is contained in:
ilgeco
2026-07-20 11:34:59 +02:00
parent 5f42da36ae
commit 6bad9a8008
16 changed files with 883 additions and 138 deletions
@@ -191,6 +191,37 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
rewriter.replaceOp(planOp, computeOp.getResults());
continue;
}
if (auto planOp = dyn_cast<spatial::SpatMaxPool2DPlanOp>(&op)) {
auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
});
if (outputBlueprint == planOp.getResult().getUsers().end()) {
planOp.emitOpError("selected MaxPool plan requires a row-strip blueprint result");
signalPassFailure();
return;
}
FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
rewriter.setInsertionPoint(planOp);
FailureOr<Value> lowered = lowerSelectedMaxPool2DPlan(
planOp, succeeded(input) ? std::optional<Value> {input->storage} : std::nullopt, rewriter);
if (failed(lowered)) {
planOp.emitOpError("failed to lower selected row-strip Spatial MaxPool plan");
signalPassFailure();
return;
}
auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
if (failed(output)) {
signalPassFailure();
return;
}
rowStripValues[blueprint.getResult()] = *output;
eraseAfterLowering.insert(planOp);
eraseAfterLowering.insert(blueprint);
continue;
}
if (auto planOp = dyn_cast<spatial::SpatBiasAddPlanOp>(&op)) {
if (succeeded(getRowStripValue(rowStripValues, planOp.getInput()))) {
auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
@@ -386,6 +417,7 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
} else if (isa<spatial::SpatConv2DPlanOp,
spatial::SpatBiasAddPlanOp,
spatial::SpatReluPlanOp,
spatial::SpatMaxPool2DPlanOp,
spatial::SpatMaterializeLayoutOp>(op)
|| op->getDialect()->getNamespace() == "onnx") {
op->emitOpError("operation must not remain after LowerSpatialPlans");