Unexpected invariant now it's clear (batched in the first tensor rank)
Validate Operations / validate-operations (push) Has been cancelled
Validate Operations / validate-operations (push) Has been cancelled
This commit is contained in:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user