Unexpected invariant now it's clear (batched in the first tensor rank)
Validate Operations / validate-operations (push) Has been cancelled

This commit is contained in:
ilgeco
2026-07-13 12:05:59 +02:00
parent fed6d343e5
commit 61e3ea9996
29 changed files with 2791 additions and 707 deletions
@@ -1448,12 +1448,10 @@ static FailureOr<Value> reconstructDepthwiseGemmRows(Value pieces,
Value laneIndex = createOrFoldAffineApply(
rewriter, tileLoc, (d0 * tiling.totalPatches) + d1, ValueRange {channelTileIndex, patchIndex}, anchorOp);
auto rowTileType = RankedTensorType::get({1, tiling.tileOutputChannels}, piecesType.getElementType());
SmallVector<OpFoldResult> pieceOffsets {laneIndex, rewriter.getIndexAttr(0)};
SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(1),
rewriter.getIndexAttr(tiling.tileOutputChannels)};
Value rowTile = tensor::ExtractSliceOp::create(
rewriter, tileLoc, rowTileType, piecesArg, pieceOffsets, pieceSizes, getUnitStrides(rewriter, 2));
Value rowNext = insertOutputTile(rowTile, rowAcc, channelTileIndex, tiling, rewriter, tileLoc);
FailureOr<Value> rowTile = extractGraphBatchPhysicalFragment(rewriter, tileLoc, piecesArg, laneIndex, rowTileType);
if (failed(rowTile))
return failure();
Value rowNext = insertOutputTile(*rowTile, rowAcc, channelTileIndex, tiling, rewriter, tileLoc);
tileYielded.push_back(rowNext);
return success();
});
@@ -1556,8 +1554,9 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter&
auto gemmOutType =
RankedTensorType::get({tiling->totalPatches, state.outType.getDimSize(1)}, state.outType.getElementType());
auto piecesType = RankedTensorType::get({tiling->totalPatches * tiling->numChannelTiles, tiling->tileOutputChannels},
state.outType.getElementType());
auto rowTileType = RankedTensorType::get({1, tiling->tileOutputChannels}, state.outType.getElementType());
auto piecesType = spatial::getGraphBatchPhysicalResultType(
tiling->totalPatches * tiling->numChannelTiles, rowTileType);
auto paddedInputType = cast<RankedTensorType>(paddedInput.getType());
auto inputTileType =
RankedTensorType::get({1, tiling->channelsPerTile, state.wType.getDimSize(2), state.wType.getDimSize(3)},
@@ -1626,7 +1625,6 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter&
*tiling,
rewriter,
loc);
auto rowTileType = RankedTensorType::get({1, tiling->tileOutputChannels}, state.outType.getElementType());
Value rowTile = spatial::SpatVMMOp::create(rewriter, loc, rowTileType, weightTile, inputTile).getResult();
if (args.inputs.size() > 1) {
Value biasArg = pickInputByRank(/*rank=*/2);
@@ -1637,11 +1635,7 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter&
Value biasTile = tiling->numChannelTiles == 1 ? biasArg : createBiasTile(biasArg, channelTileIndex, *tiling, rewriter, loc);
rowTile = spatial::SpatVAddOp::create(rewriter, loc, rowTileType, rowTile, biasTile).getResult();
}
SmallVector<OpFoldResult> outputOffsets {args.lane, rewriter.getIndexAttr(0)};
SmallVector<OpFoldResult> outputSizes {rewriter.getIndexAttr(1),
rewriter.getIndexAttr(tiling->tileOutputChannels)};
createParallelInsertSliceIntoBatchOutput(
rewriter, loc, rowTile, args.outputs.front(), outputOffsets, outputSizes, getUnitStrides(rewriter, 2));
publishGraphBatchPhysicalFragment(rewriter, loc, rowTile, args.outputs.front(), args.lane);
return success();
});
if (failed(batchOp))
@@ -1650,14 +1644,10 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter&
auto nhwcType = RankedTensorType::get(
{state.xType.getDimSize(0), state.outType.getDimSize(2), state.outType.getDimSize(3), state.outType.getDimSize(1)},
state.outType.getElementType());
Value collectedRows = batchOp->getResult(0);
if (tiling->numChannelTiles != 1) {
auto reconstructedRows =
reconstructDepthwiseGemmRows(batchOp->getResult(0), piecesType, gemmOutType, *tiling, rewriter, loc);
if (failed(reconstructedRows))
return failure();
collectedRows = *reconstructedRows;
}
auto reconstructedRows = reconstructDepthwiseGemmRows(batchOp->getResult(0), piecesType, gemmOutType, *tiling, rewriter, loc);
if (failed(reconstructedRows))
return failure();
Value collectedRows = *reconstructedRows;
return createCollectedConvOutput(ValueRange {collectedRows},
state.outType,