Resnet is fast
This commit is contained in:
@@ -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");
|
||||
|
||||
Reference in New Issue
Block a user