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

even faster on pimcomp models
Validate Operations / validate-operations (push) Has been cancelled
Details

Browse Source
This commit is contained in:
NiccoloN
2026-07-31 21:15:28 +02:00
parent 9ca1a0ed9f
commit f4a3b012cc
49 changed files with 1923 additions and 583 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp
+51
View File
@@ -58,6 +58,11 @@ lowerRowStripRelu(const RowStripPhysicalValue& input, spatial::SpatReluPlanOp pl
return applyRowStripRelu(input, rewriter, planOp.getLoc());
}
static FailureOr<Value>
lowerRowStripSilu(const RowStripPhysicalValue& input, spatial::SpatSiluPlanOp planOp, PatternRewriter& rewriter) {
return applyRowStripSilu(input, rewriter, planOp.getLoc());
}
static FailureOr<Value> lowerRowStripBiasAdd(const RowStripPhysicalValue& input,
spatial::SpatBiasAddPlanOp planOp,
PatternRewriter& rewriter) {
@@ -349,6 +354,51 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
rewriter.replaceOp(planOp, computeOp.getResults());
continue;
}
if (auto planOp = dyn_cast<spatial::SpatSiluPlanOp>(&op)) {
if (succeeded(getRowStripValue(rowStripValues, planOp.getInput()))) {
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("row-strip SiLU plan requires a row-strip blueprint result");
signalPassFailure();
return;
}
FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
rewriter.setInsertionPoint(planOp);
FailureOr<Value> lowered = lowerRowStripSilu(*input, planOp, rewriter);
if (failed(lowered)) {
planOp.emitOpError("failed to lower selected row-strip Spatial SiLU 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;
}
rewriter.setInsertionPoint(planOp);
auto computeOp = createSpatCompute<1>(
rewriter, planOp.getLoc(), planOp.getOutput().getType(), {}, planOp.getInput(), [&](Value x) {
Value sigmoid = spatial::SpatSigmoidOp::create(
rewriter, planOp.getLoc(), planOp.getOutput().getType(), x).getResult();
Value silu = spatial::SpatVMulOp::create(
rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, sigmoid).getResult();
spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), silu);
});
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);
@@ -650,6 +700,7 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
spatial::SpatBiasAddPlanOp,
spatial::SpatAddPlanOp,
spatial::SpatReluPlanOp,
spatial::SpatSiluPlanOp,
spatial::SpatMaxPool2DPlanOp,
spatial::SpatGlobalAveragePoolPlanOp,
spatial::SpatMaterializeLayoutOp>(op)