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

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

Browse Source
This commit is contained in:
NiccoloN
2026-07-30 14:07:01 +02:00
parent 1b4f070bef
commit c12f69133d
4 changed files with 140 additions and 111 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp
+25 -16
View File
@@ -2511,6 +2511,11 @@ static bool rowStripOutputTileFitsOneCore(const ConvGeometry& geometry) {
<= static_cast<int64_t>(crossbarCountInCore.getValue());
}
static bool rowStripOutputChannelTileFitsOneCore(const ConvGeometry& geometry) {
return ceilIntegerDivide(geometry.k, geometry.xbarSize)
<= static_cast<int64_t>(crossbarCountInCore.getValue());
}
static bool canConsumePixelMajorRowStripFragments(const ConvLoweringState& state, StringRef& failureReason) {
if (state.batchSize != 1) {
failureReason = "batch_not_one";
@@ -2528,24 +2533,16 @@ static bool canConsumePixelMajorRowStripFragments(const ConvLoweringState& state
failureReason = "non_float_input";
return false;
}
if (state.strideHeight != 1 || state.strideWidth != 1) {
failureReason = "stride_not_one";
return false;
}
if (state.dilationHeight != 1 || state.dilationWidth != 1) {
failureReason = "dilation_not_one";
return false;
}
if (state.outHeight != state.xHeight || state.outWidth != state.xWidth) {
failureReason = "not_same_spatial_shape";
return false;
}
if (!getHostConstDenseElementsAttr(state.w)) {
failureReason = "non_constant_weight";
return false;
}
if (!rowStripOutputTileFitsOneCore(buildConvGeometry(state))) {
failureReason = "output_row_does_not_fit_one_core";
if (!rowStripOutputChannelTileFitsOneCore(buildConvGeometry(state))) {
failureReason = "output_channel_tile_does_not_fit_one_core";
return false;
}
if (state.hasBias && !isSupportedBiasAddValue(state.b, state.outType)) {
@@ -3306,7 +3303,7 @@ static FailureOr<Value> createOutputChannelTiledRowStripConvOutput(const ConvLow
static FailureOr<Value>
createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRewriter& rewriter, Location loc) {
ConvGeometry geometry = buildConvGeometry(state);
if (state.group != 1 || state.batchSize != 1 || !rowStripOutputTileFitsOneCore(geometry))
if (state.group != 1 || state.batchSize != 1 || !rowStripOutputChannelTileFitsOneCore(geometry))
return failure();
auto weightDenseAttr = getHostConstDenseElementsAttr(state.w);
@@ -3320,11 +3317,16 @@ static FailureOr<Value>
const int64_t numKSlices = ceilIntegerDivide(patchSize, xbarDim);
const int64_t paddedK = numKSlices * xbarDim;
const int64_t paddedOutputChannels =
ceilIntegerDivide(state.numChannelsOut, xbarDim) * xbarDim;
Value paddedWeights = standard::createPaddedPixelMajorWeightConstant(
weightDenseAttr, state, paddedK, paddedOutputChannels, rewriter);
if (!rowStripOutputTileFitsOneCore(geometry)) {
Value tiledWeights =
standard::createPaddedOutputChannelTiledWeightConstant(weightDenseAttr, state, paddedK, xbarDim, rewriter);
return createOutputChannelTiledRowStripConvOutput(
state, state.x, tiledWeights, paddedK, numKSlices, xbarDim, rewriter, loc);
}
const int64_t paddedOutputChannels = ceilIntegerDivide(state.numChannelsOut, xbarDim) * xbarDim;
Value paddedWeights =
standard::createPaddedPixelMajorWeightConstant(weightDenseAttr, state, paddedK, paddedOutputChannels, rewriter);
FailureOr<Value> bias = failure();
if (state.hasBias)
bias = createBiasRowConstant(state, rewriter);
@@ -3355,6 +3357,13 @@ static FailureOr<Value> createConvOutputFromPixelMajorRowStripFragments(Value ro
auto weightDenseAttr = getHostConstDenseElementsAttr(state.w);
if (!weightDenseAttr)
return failure();
if (!rowStripOutputTileFitsOneCore(geometry)) {
Value tiledWeights =
standard::createPaddedOutputChannelTiledWeightConstant(weightDenseAttr, state, paddedK, xbarDim, rewriter);
return createOutputChannelTiledRowStripConvOutput(
state, rowStripStorage, tiledWeights, paddedK, numKSlices, xbarDim, rewriter, loc);
}
const int64_t paddedOutputChannels =
ceilIntegerDivide(state.numChannelsOut, xbarDim) * xbarDim;
Value paddedWeights = standard::createPaddedPixelMajorWeightConstant(
@@ -4298,7 +4307,7 @@ LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp) {
analysis.barrierKind = DistributedConvBarrierKind::UnsupportedConsumer;
analysis.barrierDetail = "selected row-strip layout";
ConvGeometry geometry = buildConvGeometry(*state);
if (!rowStripOutputTileFitsOneCore(geometry))
if (!rowStripOutputChannelTileFitsOneCore(geometry))
return failure();
ConvLoweringDecision decision = chooseConvLoweringStrategy(geometry, *requestedStrategy, analysis);
if (decision.strategy == PimConvLoweringDepthwise && !depthwise::canUseStructuredRewrite(*state)