relu_conv_relu Faster on Arch-A
This commit is contained in:
ilgeco
@@ -58,8 +58,18 @@ static mlir::Value resolveForYieldedAliasToInit(mlir::scf::ForOp forOp,
|
||||
mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnowledge* knowledge) {
|
||||
value = resolveAlias(value, knowledge);
|
||||
|
||||
if (mlir::isa<mlir::BlockArgument>(value))
|
||||
if (auto blockArgument = mlir::dyn_cast<mlir::BlockArgument>(value)) {
|
||||
auto forOp = mlir::dyn_cast_or_null<mlir::scf::ForOp>(blockArgument.getOwner()->getParentOp());
|
||||
if (forOp && blockArgument.getArgNumber() > 0) {
|
||||
const unsigned iterArgIndex = blockArgument.getArgNumber() - 1;
|
||||
auto yieldOp = mlir::dyn_cast<mlir::scf::YieldOp>(forOp.getBody()->getTerminator());
|
||||
if (iterArgIndex < forOp.getInitArgs().size() && yieldOp
|
||||
&& iterArgIndex < yieldOp.getNumOperands()
|
||||
&& resolveAlias(yieldOp.getOperand(iterArgIndex), knowledge) == blockArgument)
|
||||
return resolveLoopCarriedAliasImpl(forOp.getInitArgs()[iterArgIndex], knowledge);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
mlir::Operation* definingOp = value.getDefiningOp();
|
||||
if (!definingOp)
|
||||
|
||||
@@ -18,10 +18,10 @@ FailureOr<RowStripPhysicalValue> describeRowStripPhysicalValue(Value storage, Ra
|
||||
if (!storageType || !storageType.hasStaticShape() || !logicalType || !logicalType.hasStaticShape()
|
||||
|| storageType.getRank() != 5 || logicalType.getRank() != 4 || logicalType.getDimSize(0) != 1
|
||||
|| storageType.getElementType() != logicalType.getElementType()
|
||||
|| storageType.getDimSize(1) != 1 || storageType.getDimSize(2) <= 0
|
||||
|| storageType.getDimSize(3) != 1 || storageType.getDimSize(4) != logicalType.getDimSize(3))
|
||||
|| storageType.getDimSize(1) != 1 || storageType.getDimSize(2) != 1
|
||||
|| storageType.getDimSize(3) != logicalType.getDimSize(3) || storageType.getDimSize(4) <= 0)
|
||||
return failure();
|
||||
const int64_t tilesPerRow = ceilIntegerDivide(logicalType.getDimSize(1), storageType.getDimSize(2));
|
||||
const int64_t tilesPerRow = ceilIntegerDivide(logicalType.getDimSize(1), storageType.getDimSize(4));
|
||||
if (storageType.getDimSize(0) != logicalType.getDimSize(2) * tilesPerRow)
|
||||
return failure();
|
||||
return RowStripPhysicalValue {storage, logicalType,
|
||||
@@ -30,7 +30,8 @@ FailureOr<RowStripPhysicalValue> describeRowStripPhysicalValue(Value storage, Ra
|
||||
}
|
||||
|
||||
RankedTensorType getRowStripFragmentType(RankedTensorType logicalType) {
|
||||
return RankedTensorType::get({logicalType.getDimSize(0), logicalType.getDimSize(1), 1, logicalType.getDimSize(3)},
|
||||
return RankedTensorType::get({logicalType.getDimSize(0), 1, logicalType.getDimSize(3),
|
||||
logicalType.getDimSize(1)},
|
||||
logicalType.getElementType(),
|
||||
logicalType.getEncoding());
|
||||
}
|
||||
@@ -78,16 +79,18 @@ void insertRowStripFragment(Value fragment,
|
||||
FailureOr<Value> createPerChannelConstantFragment(DenseElementsAttr denseAttr,
|
||||
RankedTensorType fragmentType,
|
||||
PatternRewriter& rewriter) {
|
||||
FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, fragmentType);
|
||||
auto logicalType = RankedTensorType::get(
|
||||
{1, fragmentType.getDimSize(3), 1, 1}, fragmentType.getElementType());
|
||||
FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, logicalType);
|
||||
if (failed(channelValues))
|
||||
return failure();
|
||||
|
||||
SmallVector<Attribute> values;
|
||||
values.reserve(fragmentType.getNumElements());
|
||||
for (int64_t n = 0; n < fragmentType.getDimSize(0); ++n)
|
||||
for (int64_t channel = 0; channel < fragmentType.getDimSize(1); ++channel)
|
||||
for (int64_t h = 0; h < fragmentType.getDimSize(2); ++h)
|
||||
for (int64_t w = 0; w < fragmentType.getDimSize(3); ++w)
|
||||
for (int64_t h = 0; h < fragmentType.getDimSize(1); ++h)
|
||||
for (int64_t w = 0; w < fragmentType.getDimSize(2); ++w)
|
||||
for (int64_t channel = 0; channel < fragmentType.getDimSize(3); ++channel)
|
||||
values.push_back((*channelValues)[channel]);
|
||||
|
||||
auto attr = DenseElementsAttr::get(fragmentType, values);
|
||||
@@ -117,7 +120,6 @@ FailureOr<Value> createRowStripStorageFromRows(Value rows,
|
||||
return failure();
|
||||
|
||||
auto rowSliceType = RankedTensorType::get({width, channels}, logicalType.getElementType(), rowsType.getEncoding());
|
||||
auto channelWidthType = RankedTensorType::get({channels, width}, logicalType.getElementType(), rowsType.getEncoding());
|
||||
auto fragmentType = getRowStripFragmentType(logicalType);
|
||||
auto storageType = getRowStripStorageType(logicalType);
|
||||
auto batchOp = createSpatComputeBatch(
|
||||
@@ -128,10 +130,8 @@ FailureOr<Value> createRowStripStorageFromRows(Value rows,
|
||||
SmallVector<OpFoldResult> rowSizes {rewriter.getIndexAttr(width), rewriter.getIndexAttr(channels)};
|
||||
Value rowSlice = tensor::ExtractSliceOp::create(
|
||||
rewriter, loc, rowSliceType, args.inputs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 2));
|
||||
Value channelWidth = ONNXTransposeOp::create(
|
||||
rewriter, loc, channelWidthType, rowSlice, rewriter.getI64ArrayAttr({1, 0})).getResult();
|
||||
Value fragment = tensor::ExpandShapeOp::create(
|
||||
rewriter, loc, fragmentType, channelWidth, SmallVector<ReassociationIndices> {{0, 1}, {2, 3}});
|
||||
rewriter, loc, fragmentType, rowSlice, SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
|
||||
insertRowStripFragment(fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
|
||||
return success();
|
||||
});
|
||||
@@ -143,8 +143,28 @@ FailureOr<Value> createRowStripStorageFromRows(Value rows,
|
||||
FailureOr<Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& value,
|
||||
PatternRewriter& rewriter,
|
||||
Location loc) {
|
||||
const int64_t laneCount = cast<RankedTensorType>(value.storage.getType()).getDimSize(0);
|
||||
const int64_t tileChannels = value.fragmentType.getDimSize(3);
|
||||
auto nchwFragmentType = RankedTensorType::get(
|
||||
{1, tileChannels, 1, value.logicalType.getDimSize(3)}, value.fragmentType.getElementType(),
|
||||
value.fragmentType.getEncoding());
|
||||
auto nchwStorageType = spatial::getGraphBatchPhysicalResultType(laneCount, nchwFragmentType);
|
||||
auto transposed = createSpatComputeBatch(
|
||||
rewriter, loc, TypeRange {nchwStorageType}, laneCount, {}, ValueRange {value.storage},
|
||||
[&](detail::SpatComputeBatchBodyArgs args) {
|
||||
FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(
|
||||
rewriter, loc, args.inputs.front(), args.lane, value.fragmentType);
|
||||
if (failed(fragment))
|
||||
return failure();
|
||||
Value nchw = ONNXTransposeOp::create(
|
||||
rewriter, loc, nchwFragmentType, *fragment, rewriter.getI64ArrayAttr({0, 3, 1, 2}));
|
||||
publishGraphBatchPhysicalFragment(rewriter, loc, nchw, args.outputs.front(), args.lane);
|
||||
return success();
|
||||
});
|
||||
if (failed(transposed))
|
||||
return failure();
|
||||
|
||||
SmallVector<FragmentAssemblyEntry> entries;
|
||||
const int64_t tileChannels = value.fragmentType.getDimSize(1);
|
||||
for (int64_t row = 0; row < value.logicalType.getDimSize(2); ++row)
|
||||
for (int64_t tile = 0; tile < value.tilesPerRow; ++tile) {
|
||||
const int64_t channelOffset = tile * tileChannels;
|
||||
@@ -152,7 +172,7 @@ FailureOr<Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& va
|
||||
{1, std::min(tileChannels, value.logicalType.getDimSize(1) - channelOffset), 1,
|
||||
value.logicalType.getDimSize(3)}});
|
||||
}
|
||||
return createFragmentAssemblyBlueprint(value.storage, value.logicalType, entries, "nchw_row_strip",
|
||||
return createFragmentAssemblyBlueprint(transposed->getResult(0), value.logicalType, entries, "nhwc_row_strip",
|
||||
kRowStripIndexMap, rewriter, loc);
|
||||
}
|
||||
|
||||
@@ -193,11 +213,12 @@ FailureOr<Value> applyRowStripBiasAdd(const RowStripPhysicalValue& value,
|
||||
auto biasStorageType = spatial::getGraphBatchPhysicalResultType(value.tilesPerRow, value.fragmentType);
|
||||
SmallVector<Attribute> biasValues(
|
||||
biasStorageType.getNumElements(), cast<Attribute>(rewriter.getZeroAttr(value.fragmentType.getElementType())));
|
||||
const int64_t tileChannels = value.fragmentType.getDimSize(1);
|
||||
const int64_t width = value.fragmentType.getDimSize(3);
|
||||
const int64_t tileChannels = value.fragmentType.getDimSize(3);
|
||||
const int64_t width = value.fragmentType.getDimSize(2);
|
||||
for (int64_t channel = 0; channel < value.logicalType.getDimSize(1); ++channel)
|
||||
for (int64_t w = 0; w < width; ++w)
|
||||
biasValues[((channel / tileChannels) * tileChannels + channel % tileChannels) * width + w] =
|
||||
biasValues[(channel / tileChannels) * width * tileChannels + w * tileChannels
|
||||
+ channel % tileChannels] =
|
||||
(*channelValues)[channel];
|
||||
Value biasStorage = getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(),
|
||||
DenseElementsAttr::get(biasStorageType, biasValues), biasStorageType);
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
|
||||
namespace onnx_mlir {
|
||||
|
||||
inline constexpr llvm::StringLiteral kRowStripIndexMap = "nchw_row_strip_fragments";
|
||||
inline constexpr llvm::StringLiteral kRowStripIndexMap = "nhwc_row_strip_fragments";
|
||||
|
||||
struct RowStripPhysicalValue {
|
||||
mlir::Value storage;
|
||||
|
||||
@@ -30,7 +30,7 @@ namespace onnx_mlir {
|
||||
namespace {
|
||||
|
||||
static constexpr StringLiteral kDenseLayout = "dense_nchw";
|
||||
static constexpr StringLiteral kRowStripLayout = "nchw_row_strip";
|
||||
static constexpr StringLiteral kRowStripLayout = "nhwc_row_strip";
|
||||
|
||||
static FailureOr<RowStripPhysicalValue> getRowStripValue(llvm::DenseMap<Value, RowStripPhysicalValue>& rowStripValues,
|
||||
Value value) {
|
||||
@@ -251,20 +251,8 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
|
||||
FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
|
||||
rewriter.setInsertionPoint(planOp);
|
||||
std::optional<Value> physicalInput;
|
||||
if (succeeded(input)) {
|
||||
if (input->tilesPerRow == 1) {
|
||||
physicalInput = input->storage;
|
||||
}
|
||||
else {
|
||||
FailureOr<Value> denseInput = materializeRowStripToDense(*input, planOp.getLoc(), rewriter);
|
||||
if (failed(denseInput)) {
|
||||
planOp.emitOpError("failed to materialize tiled row-strip input for MaxPool");
|
||||
signalPassFailure();
|
||||
return;
|
||||
}
|
||||
planOp.getInputMutable().assign(*denseInput);
|
||||
}
|
||||
}
|
||||
if (succeeded(input))
|
||||
physicalInput = input->storage;
|
||||
FailureOr<Value> lowered = lowerSelectedMaxPool2DPlan(
|
||||
planOp, physicalInput, rewriter);
|
||||
if (failed(lowered)) {
|
||||
|
||||
@@ -1834,6 +1834,29 @@ static Value createPaddedInputKTiledWeightConstant(DenseElementsAttr sourceAttr,
|
||||
return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), paddedAttr, paddedType);
|
||||
}
|
||||
|
||||
static Value createPaddedPixelMajorWeightConstant(DenseElementsAttr sourceAttr,
|
||||
const ConvLoweringState& state,
|
||||
int64_t paddedK,
|
||||
int64_t paddedC,
|
||||
PatternRewriter& rewriter) {
|
||||
auto paddedType = RankedTensorType::get({paddedK, paddedC}, state.wType.getElementType());
|
||||
SmallVector<Attribute> sourceValues(sourceAttr.getValues<Attribute>());
|
||||
SmallVector<Attribute> paddedValues(
|
||||
paddedType.getNumElements(), cast<Attribute>(rewriter.getZeroAttr(paddedType.getElementType())));
|
||||
for (int64_t outChannel = 0; outChannel < state.numChannelsOut; ++outChannel)
|
||||
for (int64_t kernelH = 0; kernelH < state.wHeight; ++kernelH)
|
||||
for (int64_t kernelW = 0; kernelW < state.wWidth; ++kernelW)
|
||||
for (int64_t inChannel = 0; inChannel < state.numChannelsIn; ++inChannel) {
|
||||
const int64_t sourceFlatIndex =
|
||||
(((outChannel * state.numChannelsIn) + inChannel) * state.wHeight + kernelH) * state.wWidth + kernelW;
|
||||
const int64_t patchIndex =
|
||||
((kernelH * state.wWidth) + kernelW) * state.numChannelsIn + inChannel;
|
||||
paddedValues[patchIndex * paddedC + outChannel] = sourceValues[sourceFlatIndex];
|
||||
}
|
||||
return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(),
|
||||
DenseElementsAttr::get(paddedType, paddedValues), paddedType);
|
||||
}
|
||||
|
||||
static Value createPaddedOutputChannelTiledWeightConstant(DenseElementsAttr sourceAttr,
|
||||
const ConvLoweringState& state,
|
||||
int64_t paddedK,
|
||||
@@ -1853,7 +1876,8 @@ static Value createPaddedOutputChannelTiledWeightConstant(DenseElementsAttr sour
|
||||
for (int64_t kernelW = 0; kernelW < state.wWidth; ++kernelW) {
|
||||
const int64_t sourceFlatIndex =
|
||||
(((outChannel * state.numChannelsIn) + inChannel) * state.wHeight + kernelH) * state.wWidth + kernelW;
|
||||
const int64_t patchIndex = ((inChannel * state.wHeight) + kernelH) * state.wWidth + kernelW;
|
||||
const int64_t patchIndex =
|
||||
((kernelH * state.wWidth) + kernelW) * state.numChannelsIn + inChannel;
|
||||
const int64_t destinationFlatIndex =
|
||||
((outputTile * paddedK) + patchIndex) * xbarDim + tileChannel;
|
||||
paddedValues[destinationFlatIndex] = sourceValues[sourceFlatIndex];
|
||||
@@ -2467,7 +2491,7 @@ static Value createZeroGemmBias(RankedTensorType gemmResultType, PatternRewriter
|
||||
return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), zeroAttr, gemmResultType);
|
||||
}
|
||||
|
||||
static bool canConsumeNchwRowStripFragments(const ConvLoweringState& state, StringRef& failureReason) {
|
||||
static bool canConsumePixelMajorRowStripFragments(const ConvLoweringState& state, StringRef& failureReason) {
|
||||
if (state.batchSize != 1) {
|
||||
failureReason = "batch_not_one";
|
||||
return false;
|
||||
@@ -2539,9 +2563,8 @@ static Value createZeroTensorConstant(RankedTensorType type, PatternRewriter& re
|
||||
return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), zeroAttr, type);
|
||||
}
|
||||
|
||||
static FailureOr<Value> createPaddedBiasRowConstant(const ConvLoweringState& state,
|
||||
int64_t paddedChannels,
|
||||
PatternRewriter& rewriter) {
|
||||
static FailureOr<Value> createBiasRowConstant(const ConvLoweringState& state,
|
||||
PatternRewriter& rewriter) {
|
||||
DenseElementsAttr denseAttr;
|
||||
if (!isSupportedBiasAddValue(state.b, state.outType, &denseAttr))
|
||||
return failure();
|
||||
@@ -2549,12 +2572,11 @@ static FailureOr<Value> createPaddedBiasRowConstant(const ConvLoweringState& sta
|
||||
if (failed(channelValues))
|
||||
return failure();
|
||||
|
||||
auto biasType = RankedTensorType::get({1, paddedChannels}, state.outType.getElementType());
|
||||
SmallVector<Attribute> values(biasType.getNumElements(), cast<Attribute>(rewriter.getZeroAttr(biasType.getElementType())));
|
||||
for (int64_t channel = 0; channel < state.numChannelsOut; ++channel)
|
||||
values[channel] = (*channelValues)[channel];
|
||||
auto biasAttr = DenseElementsAttr::get(biasType, values);
|
||||
return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), biasAttr, biasType);
|
||||
auto biasType = RankedTensorType::get({1, state.numChannelsOut}, state.outType.getElementType());
|
||||
return getOrCreateConstant(rewriter,
|
||||
rewriter.getInsertionBlock()->getParentOp(),
|
||||
DenseElementsAttr::get(biasType, *channelValues),
|
||||
biasType);
|
||||
}
|
||||
|
||||
static FailureOr<Value> createPaddedBiasTileConstant(const ConvLoweringState& state,
|
||||
@@ -2581,13 +2603,13 @@ static Value createHorizontallyPaddedRowStripFragment(Value fragment,
|
||||
PatternRewriter& rewriter,
|
||||
Location loc) {
|
||||
auto paddedType = RankedTensorType::get(
|
||||
{1, state.numChannelsIn, 1, state.xWidth + state.padWidthBegin + state.padWidthEnd},
|
||||
{1, 1, state.xWidth + state.padWidthBegin + state.padWidthEnd, state.numChannelsIn},
|
||||
state.xType.getElementType(),
|
||||
state.xType.getEncoding());
|
||||
return createZeroPaddedTensor(fragment,
|
||||
paddedType,
|
||||
{0, 0, 0, state.padWidthBegin},
|
||||
{0, 0, 0, state.padWidthEnd},
|
||||
{0, 0, state.padWidthBegin, 0},
|
||||
{0, 0, state.padWidthEnd, 0},
|
||||
rewriter,
|
||||
loc);
|
||||
}
|
||||
@@ -2643,6 +2665,8 @@ static Value extractDenseConvWindowRow(Value denseInput,
|
||||
Location loc) {
|
||||
Value tableIndex = createRowStripWindowTableIndex(outputHeight, kernelRow, state, rewriter, loc);
|
||||
Value sourceRow = tensor::ExtractOp::create(rewriter, loc, sourceRowTable, ValueRange {tableIndex}).getResult();
|
||||
auto nchwType = RankedTensorType::get(
|
||||
{1, state.numChannelsIn, 1, state.xWidth}, state.xType.getElementType(), state.xType.getEncoding());
|
||||
auto fragmentType = getRowStripFragmentType(state.xType);
|
||||
SmallVector<OpFoldResult> offsets {
|
||||
rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), sourceRow, rewriter.getIndexAttr(0)};
|
||||
@@ -2650,8 +2674,10 @@ static Value extractDenseConvWindowRow(Value denseInput,
|
||||
rewriter.getIndexAttr(state.numChannelsIn),
|
||||
rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(state.xWidth)};
|
||||
return tensor::ExtractSliceOp::create(
|
||||
rewriter, loc, fragmentType, denseInput, offsets, sizes, getUnitStrides(rewriter, 4));
|
||||
Value nchw = tensor::ExtractSliceOp::create(
|
||||
rewriter, loc, nchwType, denseInput, offsets, sizes, getUnitStrides(rewriter, 4));
|
||||
return ONNXTransposeOp::create(
|
||||
rewriter, loc, fragmentType, nchw, rewriter.getI64ArrayAttr({0, 2, 3, 1}));
|
||||
}
|
||||
|
||||
static FailureOr<Value> createRowStripWindowMaskTable(const ConvLoweringState& state, PatternRewriter& rewriter) {
|
||||
@@ -2661,7 +2687,7 @@ static FailureOr<Value> createRowStripWindowMaskTable(const ConvLoweringState& s
|
||||
return failure();
|
||||
|
||||
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
|
||||
auto tableType = RankedTensorType::get({state.outHeight * state.wHeight, state.numChannelsIn, 1, state.xWidth},
|
||||
auto tableType = RankedTensorType::get({state.outHeight * state.wHeight, 1, state.xWidth, state.numChannelsIn},
|
||||
elementType,
|
||||
state.xType.getEncoding());
|
||||
Attribute zero = rewriter.getZeroAttr(elementType);
|
||||
@@ -2673,8 +2699,8 @@ static FailureOr<Value> createRowStripWindowMaskTable(const ConvLoweringState& s
|
||||
int64_t sourceRow =
|
||||
outputRow * state.strideHeight + kernelRow * state.dilationHeight - state.padHeightBegin;
|
||||
Attribute value = (sourceRow < 0 || sourceRow >= state.xHeight) ? zero : one;
|
||||
for (int64_t channel = 0; channel < state.numChannelsIn; ++channel)
|
||||
for (int64_t width = 0; width < state.xWidth; ++width)
|
||||
for (int64_t width = 0; width < state.xWidth; ++width)
|
||||
for (int64_t channel = 0; channel < state.numChannelsIn; ++channel)
|
||||
values.push_back(value);
|
||||
}
|
||||
}
|
||||
@@ -2693,9 +2719,9 @@ static Value extractProjectedRowStripWindowMask(Value maskTable,
|
||||
SmallVector<OpFoldResult> offsets {
|
||||
tableIndex, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
|
||||
SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(state.numChannelsIn),
|
||||
rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(state.xWidth)};
|
||||
rewriter.getIndexAttr(state.xWidth),
|
||||
rewriter.getIndexAttr(state.numChannelsIn)};
|
||||
return tensor::ExtractSliceOp::create(rewriter,
|
||||
loc,
|
||||
fragmentType,
|
||||
@@ -2716,7 +2742,7 @@ static FailureOr<Value> createConvInputWindow(Value input,
|
||||
if (!denseInput && inputType != getRowStripStorageType(state.xType))
|
||||
return failure();
|
||||
auto paddedWindowType = RankedTensorType::get(
|
||||
{1, state.numChannelsIn, state.wHeight, state.xWidth + state.padWidthBegin + state.padWidthEnd},
|
||||
{1, state.wHeight, state.xWidth + state.padWidthBegin + state.padWidthEnd, state.numChannelsIn},
|
||||
state.xType.getElementType(),
|
||||
state.xType.getEncoding());
|
||||
Value sourceRowTable = createRowStripWindowSourceRowTable(state, rewriter);
|
||||
@@ -2744,120 +2770,130 @@ static FailureOr<Value> createConvInputWindow(Value input,
|
||||
paddedRow,
|
||||
window,
|
||||
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0),
|
||||
rewriter.getIndexAttr(0),
|
||||
rewriter.getIndexAttr(kernelRowIndex),
|
||||
rewriter.getIndexAttr(0),
|
||||
rewriter.getIndexAttr(0)},
|
||||
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(state.numChannelsIn),
|
||||
rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(
|
||||
state.xWidth + state.padWidthBegin
|
||||
+ state.padWidthEnd)},
|
||||
+ state.padWidthEnd),
|
||||
rewriter.getIndexAttr(state.numChannelsIn)},
|
||||
getUnitStrides(rewriter, 4));
|
||||
}
|
||||
return window;
|
||||
}
|
||||
|
||||
static FailureOr<Value> createNchwRowStripConvPatchRow(Value paddedWindow,
|
||||
const ConvLoweringState& state,
|
||||
Value outputWidth,
|
||||
PatternRewriter& rewriter,
|
||||
Location loc) {
|
||||
static FailureOr<Value> createPixelMajorConvPatchRow(Value paddedWindow,
|
||||
const ConvLoweringState& state,
|
||||
Value outputWidth,
|
||||
PatternRewriter& rewriter,
|
||||
Location loc) {
|
||||
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
|
||||
const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
|
||||
auto patchType = RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.wWidth},
|
||||
auto patchType = RankedTensorType::get({1, state.wHeight, state.wWidth, state.numChannelsIn},
|
||||
state.xType.getElementType(),
|
||||
state.xType.getEncoding());
|
||||
auto rowType = RankedTensorType::get({1, patchSize}, state.xType.getElementType(), state.xType.getEncoding());
|
||||
Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
|
||||
Value inputWidthOffset = affineMulConst(rewriter, loc, outputWidth, state.strideWidth, anchorOp);
|
||||
Value patch = createConvInputPatch(paddedWindow,
|
||||
patchType,
|
||||
c0,
|
||||
c0,
|
||||
c0,
|
||||
inputWidthOffset,
|
||||
state.dilationHeight,
|
||||
state.dilationWidth,
|
||||
rewriter,
|
||||
loc);
|
||||
SmallVector<OpFoldResult> offsets {
|
||||
rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), inputWidthOffset, rewriter.getIndexAttr(0)};
|
||||
SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(state.wHeight),
|
||||
rewriter.getIndexAttr(state.wWidth),
|
||||
rewriter.getIndexAttr(state.numChannelsIn)};
|
||||
SmallVector<OpFoldResult> strides {rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(state.dilationWidth),
|
||||
rewriter.getIndexAttr(1)};
|
||||
Value patch = tensor::ExtractSliceOp::create(
|
||||
rewriter, loc, patchType, paddedWindow, offsets, sizes, strides);
|
||||
return tensor::CollapseShapeOp::create(
|
||||
rewriter, loc, rowType, patch, SmallVector<ReassociationIndices> {{0}, {1, 2, 3}})
|
||||
.getResult();
|
||||
}
|
||||
|
||||
static FailureOr<Value> createPaddedConvOutputTile(Value paddedPatchRow,
|
||||
Value tileWeights,
|
||||
int64_t numKSlices,
|
||||
int64_t xbarDim,
|
||||
PatternRewriter& rewriter,
|
||||
Location loc) {
|
||||
auto elementType = cast<RankedTensorType>(paddedPatchRow.getType()).getElementType();
|
||||
static FailureOr<Value> createConvOutputTile(Value patchRow,
|
||||
Value partialInputScratch,
|
||||
Value tileWeights,
|
||||
int64_t patchSize,
|
||||
int64_t numKSlices,
|
||||
int64_t xbarDim,
|
||||
PatternRewriter& rewriter,
|
||||
Location loc) {
|
||||
auto elementType = cast<RankedTensorType>(patchRow.getType()).getElementType();
|
||||
auto paddedRowType = RankedTensorType::get({1, xbarDim}, elementType);
|
||||
auto weightElementType = cast<RankedTensorType>(tileWeights.getType()).getElementType();
|
||||
auto paddedWeightTileType = RankedTensorType::get({xbarDim, xbarDim}, weightElementType);
|
||||
|
||||
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
|
||||
Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
|
||||
Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
|
||||
Value cNumKSlices = getOrCreateIndexConstant(rewriter, anchorOp, numKSlices);
|
||||
Value cXbar = getOrCreateIndexConstant(rewriter, anchorOp, xbarDim);
|
||||
auto createPiece = [&](Value kSlice, Location pieceLoc) -> Value {
|
||||
Value kOffset = arith::MulIOp::create(rewriter, pieceLoc, kSlice, cXbar);
|
||||
SmallVector<OpFoldResult> aOffsets {rewriter.getIndexAttr(0), kOffset};
|
||||
SmallVector<OpFoldResult> aSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)};
|
||||
Value aTile = extractStaticSliceOrIdentity(
|
||||
rewriter, pieceLoc, paddedPatchRow, paddedRowType, aOffsets, aSizes, getUnitStrides(rewriter, 2));
|
||||
SmallVector<OpFoldResult> bOffsets {kOffset, rewriter.getIndexAttr(0)};
|
||||
Value tileResult;
|
||||
for (int64_t kSlice = 0; kSlice < numKSlices; ++kSlice) {
|
||||
const int64_t kOffset = kSlice * xbarDim;
|
||||
const int64_t sliceSize = std::min(xbarDim, patchSize - kOffset);
|
||||
Value inputTile;
|
||||
if (sliceSize == xbarDim) {
|
||||
inputTile = extractStaticSliceOrIdentity(
|
||||
rewriter,
|
||||
loc,
|
||||
patchRow,
|
||||
paddedRowType,
|
||||
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(kOffset)},
|
||||
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)},
|
||||
getUnitStrides(rewriter, 2));
|
||||
}
|
||||
else {
|
||||
if (!partialInputScratch)
|
||||
return failure();
|
||||
auto partialType = RankedTensorType::get({1, sliceSize}, elementType);
|
||||
Value partial = extractStaticSliceOrIdentity(
|
||||
rewriter,
|
||||
loc,
|
||||
patchRow,
|
||||
partialType,
|
||||
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(kOffset)},
|
||||
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(sliceSize)},
|
||||
getUnitStrides(rewriter, 2));
|
||||
inputTile = tensor::InsertSliceOp::create(
|
||||
rewriter,
|
||||
loc,
|
||||
partial,
|
||||
partialInputScratch,
|
||||
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
|
||||
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(sliceSize)},
|
||||
getUnitStrides(rewriter, 2));
|
||||
}
|
||||
SmallVector<OpFoldResult> bOffsets {
|
||||
rewriter.getIndexAttr(kOffset), rewriter.getIndexAttr(0)};
|
||||
SmallVector<OpFoldResult> bSizes {rewriter.getIndexAttr(xbarDim), rewriter.getIndexAttr(xbarDim)};
|
||||
Value bTile = extractStaticSliceOrIdentity(
|
||||
rewriter, pieceLoc, tileWeights, paddedWeightTileType, bOffsets, bSizes, getUnitStrides(rewriter, 2));
|
||||
return spatial::SpatVMMOp::create(rewriter, pieceLoc, paddedRowType, bTile, aTile).getResult();
|
||||
};
|
||||
|
||||
Value tileResult = createPiece(c0, loc);
|
||||
if (numKSlices == 1)
|
||||
return tileResult;
|
||||
|
||||
auto kLoop = buildNormalizedScfFor(
|
||||
rewriter,
|
||||
loc,
|
||||
c1,
|
||||
cNumKSlices,
|
||||
c1,
|
||||
ValueRange {tileResult},
|
||||
[&](OpBuilder&, Location reduceLoc, Value kSlice, ValueRange reduceIterArgs, SmallVectorImpl<Value>& reduceYielded) {
|
||||
Value piece = createPiece(kSlice, reduceLoc);
|
||||
reduceYielded.push_back(
|
||||
spatial::SpatVAddOp::create(rewriter, reduceLoc, paddedRowType, reduceIterArgs.front(), piece).getResult());
|
||||
return success();
|
||||
});
|
||||
if (failed(kLoop))
|
||||
return failure();
|
||||
return kLoop->results.front();
|
||||
rewriter, loc, tileWeights, paddedWeightTileType, bOffsets, bSizes, getUnitStrides(rewriter, 2));
|
||||
Value piece = spatial::SpatVMMOp::create(
|
||||
rewriter, loc, paddedRowType, bTile, inputTile).getResult();
|
||||
tileResult = tileResult
|
||||
? spatial::SpatVAddOp::create(
|
||||
rewriter, loc, paddedRowType, tileResult, piece).getResult()
|
||||
: piece;
|
||||
}
|
||||
return tileResult;
|
||||
}
|
||||
|
||||
static FailureOr<Value> createPaddedConvOutputRow(Value patchRow,
|
||||
const ConvLoweringState& state,
|
||||
Value paddedWeights,
|
||||
Value paddedBias,
|
||||
int64_t paddedK,
|
||||
int64_t numKSlices,
|
||||
int64_t xbarDim,
|
||||
PatternRewriter& rewriter,
|
||||
Location loc) {
|
||||
const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
|
||||
auto elementType = state.outType.getElementType();
|
||||
auto rowType = RankedTensorType::get({1, state.numChannelsOut}, elementType);
|
||||
auto paddedRowType = RankedTensorType::get({1, xbarDim}, elementType);
|
||||
auto paddedPatchRowType = RankedTensorType::get({1, paddedK}, elementType);
|
||||
auto tileWeightsType = RankedTensorType::get({paddedK, xbarDim}, state.wType.getElementType());
|
||||
const int64_t outputTileCount = ceilIntegerDivide(state.numChannelsOut, xbarDim);
|
||||
|
||||
Value paddedPatchRow = patchRow;
|
||||
if (patchSize != paddedK)
|
||||
paddedPatchRow = createZeroPaddedTensor(
|
||||
paddedPatchRow, paddedPatchRowType, {0, 0}, {0, paddedK - patchSize}, rewriter, loc);
|
||||
static FailureOr<Value> createConvOutputRow(Value patchRow,
|
||||
Value partialInputScratch,
|
||||
int64_t patchSize,
|
||||
int64_t paddedK,
|
||||
int64_t outputChannels,
|
||||
Value paddedWeights,
|
||||
Value bias,
|
||||
int64_t numKSlices,
|
||||
int64_t xbarDim,
|
||||
PatternRewriter& rewriter,
|
||||
Location loc) {
|
||||
auto elementType = cast<RankedTensorType>(patchRow.getType()).getElementType();
|
||||
auto rowType = RankedTensorType::get({1, outputChannels}, elementType);
|
||||
auto tileWeightsType =
|
||||
RankedTensorType::get({paddedK, xbarDim},
|
||||
cast<RankedTensorType>(paddedWeights.getType()).getElementType());
|
||||
const int64_t outputTileCount = ceilIntegerDivide(outputChannels, xbarDim);
|
||||
|
||||
auto getTileWeights = [&](int64_t outputTile) {
|
||||
if (outputTileCount == 1)
|
||||
@@ -2871,29 +2907,31 @@ static FailureOr<Value> createPaddedConvOutputRow(Value patchRow,
|
||||
};
|
||||
|
||||
if (outputTileCount == 1) {
|
||||
FailureOr<Value> rowResult = createPaddedConvOutputTile(
|
||||
paddedPatchRow, getTileWeights(0), numKSlices, xbarDim, rewriter, loc);
|
||||
FailureOr<Value> rowResult = createConvOutputTile(
|
||||
patchRow, partialInputScratch, getTileWeights(0), patchSize, numKSlices, xbarDim, rewriter, loc);
|
||||
if (failed(rowResult))
|
||||
return failure();
|
||||
if (paddedBias)
|
||||
rowResult = spatial::SpatVAddOp::create(rewriter, loc, paddedRowType, *rowResult, paddedBias).getResult();
|
||||
if (state.numChannelsOut == xbarDim)
|
||||
return *rowResult;
|
||||
|
||||
SmallVector<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
|
||||
SmallVector<OpFoldResult> outputSizes {
|
||||
rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut)};
|
||||
return tensor::ExtractSliceOp::create(
|
||||
rewriter, loc, rowType, *rowResult, outputOffsets, outputSizes, getUnitStrides(rewriter, 2))
|
||||
.getResult();
|
||||
Value validRow = *rowResult;
|
||||
if (outputChannels != xbarDim)
|
||||
validRow = tensor::ExtractSliceOp::create(
|
||||
rewriter,
|
||||
loc,
|
||||
rowType,
|
||||
validRow,
|
||||
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
|
||||
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(outputChannels)},
|
||||
getUnitStrides(rewriter, 2));
|
||||
if (bias)
|
||||
validRow = spatial::SpatVAddOp::create(rewriter, loc, rowType, validRow, bias).getResult();
|
||||
return validRow;
|
||||
}
|
||||
|
||||
const int64_t paddedOutputChannels = outputTileCount * xbarDim;
|
||||
auto paddedOutputType = RankedTensorType::get({1, paddedOutputChannels}, elementType);
|
||||
Value paddedOutput = tensor::EmptyOp::create(rewriter, loc, paddedOutputType.getShape(), elementType);
|
||||
for (int64_t outputTile = 0; outputTile < outputTileCount; ++outputTile) {
|
||||
FailureOr<Value> tileResult = createPaddedConvOutputTile(
|
||||
paddedPatchRow, getTileWeights(outputTile), numKSlices, xbarDim, rewriter, loc);
|
||||
FailureOr<Value> tileResult = createConvOutputTile(
|
||||
patchRow, partialInputScratch, getTileWeights(outputTile), patchSize, numKSlices, xbarDim, rewriter, loc);
|
||||
if (failed(tileResult))
|
||||
return failure();
|
||||
SmallVector<OpFoldResult> tileOffsets {
|
||||
@@ -2904,7 +2942,7 @@ static FailureOr<Value> createPaddedConvOutputRow(Value patchRow,
|
||||
}
|
||||
|
||||
SmallVector<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
|
||||
SmallVector<OpFoldResult> outputSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut)};
|
||||
SmallVector<OpFoldResult> outputSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(outputChannels)};
|
||||
return tensor::ExtractSliceOp::create(
|
||||
rewriter, loc, rowType, paddedOutput, outputOffsets, outputSizes, getUnitStrides(rewriter, 2))
|
||||
.getResult();
|
||||
@@ -2927,8 +2965,8 @@ static FailureOr<Value> createOutputChannelTiledRowStripConvOutput(const ConvLow
|
||||
auto elementType = state.outType.getElementType();
|
||||
auto paddedPatchRowType = RankedTensorType::get({1, paddedK}, elementType);
|
||||
auto paddedRowType = RankedTensorType::get({1, xbarDim}, elementType);
|
||||
auto tilePixelType = RankedTensorType::get({1, xbarDim, 1, 1}, elementType);
|
||||
auto tileFragmentType = RankedTensorType::get({1, xbarDim, 1, state.outWidth}, elementType);
|
||||
auto tilePixelType = RankedTensorType::get({1, 1, 1, xbarDim}, elementType);
|
||||
auto tileFragmentType = RankedTensorType::get({1, 1, state.outWidth, xbarDim}, elementType);
|
||||
auto tileWeightsType = RankedTensorType::get({paddedK, xbarDim}, state.wType.getElementType());
|
||||
const int64_t laneCount = state.outHeight * outputTileCount;
|
||||
auto tileStorageType = spatial::getGraphBatchPhysicalResultType(laneCount, tileFragmentType);
|
||||
@@ -2963,26 +3001,35 @@ static FailureOr<Value> createOutputChannelTiledRowStripConvOutput(const ConvLow
|
||||
if (failed(inputWindow))
|
||||
return failure();
|
||||
Value fragmentInit = tensor::EmptyOp::create(rewriter, loc, tileFragmentType.getShape(), elementType);
|
||||
SmallVector<Value> widthLoopInit {fragmentInit};
|
||||
if (patchSize != paddedK)
|
||||
widthLoopInit.push_back(createZeroTensorConstant(paddedPatchRowType, rewriter));
|
||||
auto widthLoop = buildNormalizedScfFor(
|
||||
rewriter,
|
||||
loc,
|
||||
c0,
|
||||
cOutWidth,
|
||||
c1,
|
||||
ValueRange {fragmentInit},
|
||||
widthLoopInit,
|
||||
[&](OpBuilder&,
|
||||
Location widthLoc,
|
||||
Value widthIndex,
|
||||
ValueRange widthIterArgs,
|
||||
SmallVectorImpl<Value>& widthYielded) {
|
||||
FailureOr<Value> patchRow =
|
||||
createNchwRowStripConvPatchRow(*inputWindow, state, widthIndex, rewriter, widthLoc);
|
||||
createPixelMajorConvPatchRow(*inputWindow, state, widthIndex, rewriter, widthLoc);
|
||||
if (failed(patchRow))
|
||||
return failure();
|
||||
Value paddedPatchRow = *patchRow;
|
||||
if (patchSize != paddedK)
|
||||
paddedPatchRow = createZeroPaddedTensor(
|
||||
paddedPatchRow, paddedPatchRowType, {0, 0}, {0, paddedK - patchSize}, rewriter, widthLoc);
|
||||
paddedPatchRow = tensor::InsertSliceOp::create(
|
||||
rewriter,
|
||||
widthLoc,
|
||||
paddedPatchRow,
|
||||
widthIterArgs[1],
|
||||
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
|
||||
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(patchSize)},
|
||||
getUnitStrides(rewriter, 2));
|
||||
FailureOr<Value> paddedOutputRow = createPaddedConvOutputTile(
|
||||
paddedPatchRow, tileWeights, numKSlices, xbarDim, rewriter, widthLoc);
|
||||
if (failed(paddedOutputRow))
|
||||
@@ -2991,13 +3038,13 @@ static FailureOr<Value> createOutputChannelTiledRowStripConvOutput(const ConvLow
|
||||
paddedOutputRow = spatial::SpatVAddOp::create(
|
||||
rewriter, widthLoc, paddedRowType, *paddedOutputRow, *biasTile).getResult();
|
||||
Value outputPixel = tensor::ExpandShapeOp::create(
|
||||
rewriter, widthLoc, tilePixelType, *paddedOutputRow, SmallVector<ReassociationIndices> {{0}, {1, 2, 3}});
|
||||
rewriter, widthLoc, tilePixelType, *paddedOutputRow, SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
|
||||
SmallVector<OpFoldResult> rowOffsets {
|
||||
rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex};
|
||||
rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex, rewriter.getIndexAttr(0)};
|
||||
SmallVector<OpFoldResult> rowSizes {rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(xbarDim),
|
||||
rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(1)};
|
||||
rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(xbarDim)};
|
||||
Value nextFragment = tensor::InsertSliceOp::create(rewriter,
|
||||
widthLoc,
|
||||
outputPixel,
|
||||
@@ -3006,6 +3053,8 @@ static FailureOr<Value> createOutputChannelTiledRowStripConvOutput(const ConvLow
|
||||
rowSizes,
|
||||
getUnitStrides(rewriter, 4));
|
||||
widthYielded.push_back(nextFragment);
|
||||
if (patchSize != paddedK)
|
||||
widthYielded.push_back(paddedPatchRow);
|
||||
return success();
|
||||
});
|
||||
if (failed(widthLoop))
|
||||
@@ -3036,12 +3085,13 @@ createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRe
|
||||
const int64_t numKSlices = ceilIntegerDivide(patchSize, xbarDim);
|
||||
const int64_t paddedK = numKSlices * xbarDim;
|
||||
auto elementType = state.outType.getElementType();
|
||||
auto paddedPatchRowType = RankedTensorType::get({1, paddedK}, elementType);
|
||||
auto fragmentType = getRowStripFragmentType(state.outType);
|
||||
auto outputPixelType = RankedTensorType::get({1, state.numChannelsOut, 1, 1}, elementType);
|
||||
auto outputPixelType = RankedTensorType::get({1, 1, 1, state.numChannelsOut}, elementType);
|
||||
auto outputStorageType = getRowStripStorageType(state.outType);
|
||||
|
||||
Value paddedWeights = state.numChannelsOut <= xbarDim
|
||||
? standard::createPaddedInputKTiledWeightConstant(
|
||||
? standard::createPaddedPixelMajorWeightConstant(
|
||||
weightDenseAttr, state, paddedK, xbarDim, rewriter)
|
||||
: standard::createPaddedOutputChannelTiledWeightConstant(
|
||||
weightDenseAttr, state, paddedK, xbarDim, rewriter);
|
||||
@@ -3049,10 +3099,10 @@ createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRe
|
||||
return createOutputChannelTiledRowStripConvOutput(
|
||||
state, paddedWeights, paddedK, numKSlices, xbarDim, rewriter, loc);
|
||||
|
||||
FailureOr<Value> paddedBias = failure();
|
||||
FailureOr<Value> bias = failure();
|
||||
if (state.hasBias)
|
||||
paddedBias = createPaddedBiasRowConstant(state, xbarDim, rewriter);
|
||||
if (state.hasBias && failed(paddedBias))
|
||||
bias = createBiasRowConstant(state, rewriter);
|
||||
if (state.hasBias && failed(bias))
|
||||
return failure();
|
||||
|
||||
auto batchOp = createSpatComputeBatch(
|
||||
@@ -3061,7 +3111,7 @@ createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRe
|
||||
TypeRange {outputStorageType},
|
||||
state.outHeight,
|
||||
ValueRange {paddedWeights},
|
||||
state.hasBias ? ValueRange {state.x, *paddedBias} : ValueRange {state.x},
|
||||
state.hasBias ? ValueRange {state.x, *bias} : ValueRange {state.x},
|
||||
[&](detail::SpatComputeBatchBodyArgs args) {
|
||||
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
|
||||
Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
|
||||
@@ -3072,23 +3122,35 @@ createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRe
|
||||
if (failed(inputWindow))
|
||||
return failure();
|
||||
Value fragmentInit = tensor::EmptyOp::create(rewriter, loc, fragmentType.getShape(), elementType);
|
||||
SmallVector<Value> widthLoopInit {fragmentInit};
|
||||
if (patchSize != paddedK)
|
||||
widthLoopInit.push_back(createZeroTensorConstant(paddedPatchRowType, rewriter));
|
||||
auto widthLoop = buildNormalizedScfFor(
|
||||
rewriter,
|
||||
loc,
|
||||
c0,
|
||||
cOutWidth,
|
||||
c1,
|
||||
ValueRange {fragmentInit},
|
||||
widthLoopInit,
|
||||
[&](OpBuilder&, Location widthLoc, Value widthIndex, ValueRange widthIterArgs, SmallVectorImpl<Value>& widthYielded) {
|
||||
FailureOr<Value> patchRow =
|
||||
createNchwRowStripConvPatchRow(*inputWindow, state, widthIndex, rewriter, widthLoc);
|
||||
createPixelMajorConvPatchRow(*inputWindow, state, widthIndex, rewriter, widthLoc);
|
||||
if (failed(patchRow))
|
||||
return failure();
|
||||
FailureOr<Value> outputRow = createPaddedConvOutputRow(*patchRow,
|
||||
state,
|
||||
Value paddedPatchRow = *patchRow;
|
||||
if (patchSize != paddedK)
|
||||
paddedPatchRow = tensor::InsertSliceOp::create(
|
||||
rewriter,
|
||||
widthLoc,
|
||||
paddedPatchRow,
|
||||
widthIterArgs[1],
|
||||
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
|
||||
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(patchSize)},
|
||||
getUnitStrides(rewriter, 2));
|
||||
FailureOr<Value> outputRow = createPaddedConvOutputRow(paddedPatchRow,
|
||||
state.numChannelsOut,
|
||||
args.weights.front(),
|
||||
state.hasBias ? args.inputs[1] : Value(),
|
||||
paddedK,
|
||||
numKSlices,
|
||||
xbarDim,
|
||||
rewriter,
|
||||
@@ -3100,15 +3162,17 @@ createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRe
|
||||
widthLoc,
|
||||
outputPixelType,
|
||||
*outputRow,
|
||||
SmallVector<ReassociationIndices> {{0}, {1, 2, 3}});
|
||||
SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
|
||||
SmallVector<OpFoldResult> rowOffsets {
|
||||
rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex};
|
||||
rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex, rewriter.getIndexAttr(0)};
|
||||
SmallVector<OpFoldResult> rowSizes {
|
||||
rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut), rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(1)};
|
||||
rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(state.numChannelsOut)};
|
||||
Value nextFragment = tensor::InsertSliceOp::create(
|
||||
rewriter, widthLoc, outputFragment, widthIterArgs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 4));
|
||||
widthYielded.push_back(nextFragment);
|
||||
if (patchSize != paddedK)
|
||||
widthYielded.push_back(paddedPatchRow);
|
||||
return success();
|
||||
});
|
||||
if (failed(widthLoop))
|
||||
@@ -3122,16 +3186,16 @@ createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRe
|
||||
return batchOp->getResult(0);
|
||||
}
|
||||
|
||||
static FailureOr<Value> createConvOutputFromNchwRowStripFragments(Value rowStripStorage,
|
||||
const ConvLoweringState& state,
|
||||
PatternRewriter& rewriter,
|
||||
Location loc) {
|
||||
static FailureOr<Value> createConvOutputFromPixelMajorRowStripFragments(Value rowStripStorage,
|
||||
const ConvLoweringState& state,
|
||||
PatternRewriter& rewriter,
|
||||
Location loc) {
|
||||
auto inputType = dyn_cast<RankedTensorType>(rowStripStorage.getType());
|
||||
if (!inputType || inputType != getRowStripStorageType(state.xType))
|
||||
return failure();
|
||||
|
||||
StringRef failureReason;
|
||||
if (!canConsumeNchwRowStripFragments(state, failureReason))
|
||||
if (!canConsumePixelMajorRowStripFragments(state, failureReason))
|
||||
return failure();
|
||||
|
||||
ConvGeometry geometry = buildConvGeometry(state);
|
||||
@@ -3140,16 +3204,18 @@ static FailureOr<Value> createConvOutputFromNchwRowStripFragments(Value rowStrip
|
||||
const int64_t numKSlices = ceilIntegerDivide(patchSize, xbarDim);
|
||||
const int64_t paddedK = numKSlices * xbarDim;
|
||||
auto elementType = state.outType.getElementType();
|
||||
auto outputPixelType = RankedTensorType::get({1, state.numChannelsOut, 1, 1}, elementType);
|
||||
auto paddedPatchRowType = RankedTensorType::get({1, paddedK}, elementType);
|
||||
auto outputPixelType = RankedTensorType::get({1, 1, 1, state.numChannelsOut}, elementType);
|
||||
auto outputStorageType = getRowStripStorageType(state.outType);
|
||||
auto weightDenseAttr = getHostConstDenseElementsAttr(state.w);
|
||||
if (!weightDenseAttr)
|
||||
return failure();
|
||||
Value paddedWeights = standard::createPaddedInputKTiledWeightConstant(weightDenseAttr, state, paddedK, xbarDim, rewriter);
|
||||
FailureOr<Value> paddedBias = failure();
|
||||
Value paddedWeights =
|
||||
standard::createPaddedPixelMajorWeightConstant(weightDenseAttr, state, paddedK, xbarDim, rewriter);
|
||||
FailureOr<Value> bias = failure();
|
||||
if (state.hasBias)
|
||||
paddedBias = createPaddedBiasRowConstant(state, xbarDim, rewriter);
|
||||
if (state.hasBias && failed(paddedBias))
|
||||
bias = createBiasRowConstant(state, rewriter);
|
||||
if (state.hasBias && failed(bias))
|
||||
return failure();
|
||||
|
||||
auto batchOp = createSpatComputeBatch(
|
||||
@@ -3158,7 +3224,7 @@ static FailureOr<Value> createConvOutputFromNchwRowStripFragments(Value rowStrip
|
||||
TypeRange {outputStorageType},
|
||||
state.outHeight,
|
||||
ValueRange {paddedWeights},
|
||||
state.hasBias ? ValueRange {rowStripStorage, *paddedBias} : ValueRange {rowStripStorage},
|
||||
state.hasBias ? ValueRange {rowStripStorage, *bias} : ValueRange {rowStripStorage},
|
||||
[&](detail::SpatComputeBatchBodyArgs args) {
|
||||
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
|
||||
Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
|
||||
@@ -3169,24 +3235,36 @@ static FailureOr<Value> createConvOutputFromNchwRowStripFragments(Value rowStrip
|
||||
if (failed(inputWindow))
|
||||
return failure();
|
||||
Value fragmentInit = tensor::EmptyOp::create(rewriter, loc, fragmentType.getShape(), elementType);
|
||||
SmallVector<Value> widthLoopInit {fragmentInit};
|
||||
if (patchSize != paddedK)
|
||||
widthLoopInit.push_back(createZeroTensorConstant(paddedPatchRowType, rewriter));
|
||||
auto widthLoop = buildNormalizedScfFor(
|
||||
rewriter,
|
||||
loc,
|
||||
c0,
|
||||
cOutWidth,
|
||||
c1,
|
||||
ValueRange {fragmentInit},
|
||||
widthLoopInit,
|
||||
[&](OpBuilder&, Location widthLoc, Value widthIndex, ValueRange widthIterArgs, SmallVectorImpl<Value>& widthYielded) {
|
||||
FailureOr<Value> patchRow =
|
||||
createNchwRowStripConvPatchRow(*inputWindow, state, widthIndex, rewriter, widthLoc);
|
||||
createPixelMajorConvPatchRow(*inputWindow, state, widthIndex, rewriter, widthLoc);
|
||||
if (failed(patchRow))
|
||||
return failure();
|
||||
|
||||
FailureOr<Value> outputRow = createPaddedConvOutputRow(*patchRow,
|
||||
state,
|
||||
Value paddedPatchRow = *patchRow;
|
||||
if (patchSize != paddedK)
|
||||
paddedPatchRow = tensor::InsertSliceOp::create(
|
||||
rewriter,
|
||||
widthLoc,
|
||||
paddedPatchRow,
|
||||
widthIterArgs[1],
|
||||
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
|
||||
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(patchSize)},
|
||||
getUnitStrides(rewriter, 2));
|
||||
FailureOr<Value> outputRow = createPaddedConvOutputRow(paddedPatchRow,
|
||||
state.numChannelsOut,
|
||||
args.weights.front(),
|
||||
state.hasBias ? args.inputs[1] : Value(),
|
||||
paddedK,
|
||||
numKSlices,
|
||||
xbarDim,
|
||||
rewriter,
|
||||
@@ -3198,15 +3276,17 @@ static FailureOr<Value> createConvOutputFromNchwRowStripFragments(Value rowStrip
|
||||
widthLoc,
|
||||
outputPixelType,
|
||||
*outputRow,
|
||||
SmallVector<ReassociationIndices> {{0}, {1, 2, 3}});
|
||||
SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
|
||||
SmallVector<OpFoldResult> rowOffsets {
|
||||
rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex};
|
||||
rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex, rewriter.getIndexAttr(0)};
|
||||
SmallVector<OpFoldResult> rowSizes {
|
||||
rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut), rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(1)};
|
||||
rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(state.numChannelsOut)};
|
||||
Value nextFragment = tensor::InsertSliceOp::create(
|
||||
rewriter, widthLoc, outputFragment, widthIterArgs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 4));
|
||||
widthYielded.push_back(nextFragment);
|
||||
if (patchSize != paddedK)
|
||||
widthYielded.push_back(paddedPatchRow);
|
||||
return success();
|
||||
});
|
||||
if (failed(widthLoop))
|
||||
@@ -3228,7 +3308,7 @@ static FailureOr<Value> createPointwiseOutputFromRowStripFragments(Value rowStri
|
||||
if (failed(input)) return failure();
|
||||
ConvGeometry geometry = buildConvGeometry(state);
|
||||
const int64_t xbarDim = geometry.xbarSize;
|
||||
const int64_t inputFragmentChannels = input->fragmentType.getDimSize(1);
|
||||
const int64_t inputFragmentChannels = input->fragmentType.getDimSize(3);
|
||||
if (inputFragmentChannels % xbarDim != 0 || state.numChannelsIn % xbarDim != 0)
|
||||
return failure();
|
||||
auto weightDenseAttr = getHostConstDenseElementsAttr(state.w);
|
||||
@@ -3241,7 +3321,7 @@ static FailureOr<Value> createPointwiseOutputFromRowStripFragments(Value rowStri
|
||||
auto inputRowType = RankedTensorType::get({1, inputFragmentChannels}, elementType);
|
||||
auto weightTileType = RankedTensorType::get({state.numChannelsIn, xbarDim}, state.wType.getElementType());
|
||||
auto weightSliceType = RankedTensorType::get({xbarDim, xbarDim}, state.wType.getElementType());
|
||||
auto outputFragmentType = RankedTensorType::get({1, xbarDim, 1, 1}, elementType);
|
||||
auto outputFragmentType = RankedTensorType::get({1, 1, 1, xbarDim}, elementType);
|
||||
auto outputStorageType = spatial::getGraphBatchPhysicalResultType(outputTileCount, outputFragmentType);
|
||||
Value paddedWeights = standard::createPaddedOutputChannelTiledWeightConstant(
|
||||
weightDenseAttr, state, state.numChannelsIn, xbarDim, rewriter);
|
||||
@@ -3272,7 +3352,7 @@ static FailureOr<Value> createPointwiseOutputFromRowStripFragments(Value rowStri
|
||||
rewriter, pieceLoc, args.inputs.front(), sourceSlot, input->fragmentType);
|
||||
if (failed(fragment)) return failure();
|
||||
Value inputRow = tensor::CollapseShapeOp::create(rewriter, pieceLoc, inputRowType, *fragment,
|
||||
SmallVector<ReassociationIndices> {{0}, {1, 2, 3}});
|
||||
SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
|
||||
Value inputSlice = tensor::ExtractSliceOp::create(rewriter, pieceLoc, paddedRowType, inputRow,
|
||||
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), sourceOffset},
|
||||
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)},
|
||||
@@ -3305,7 +3385,7 @@ static FailureOr<Value> createPointwiseOutputFromRowStripFragments(Value rowStri
|
||||
result = spatial::SpatVAddOp::create(rewriter, loc, paddedRowType, *result, *bias).getResult();
|
||||
}
|
||||
Value fragment = tensor::ExpandShapeOp::create(rewriter, loc, outputFragmentType, *result,
|
||||
SmallVector<ReassociationIndices> {{0}, {1, 2, 3}});
|
||||
SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
|
||||
publishGraphBatchPhysicalFragment(rewriter, loc, fragment, args.outputs.front(), args.lane);
|
||||
return success();
|
||||
});
|
||||
@@ -3320,7 +3400,7 @@ static FailureOr<Value> createConvOutputFromRowStripInput(const ConvLoweringStat
|
||||
Location loc) {
|
||||
if (state.xHeight == 1 && state.xWidth == 1 && state.wHeight == 1 && state.wWidth == 1)
|
||||
return createPointwiseOutputFromRowStripFragments(rowStripInput, state, rewriter, loc);
|
||||
return createConvOutputFromNchwRowStripFragments(rowStripInput, state, rewriter, loc);
|
||||
return createConvOutputFromPixelMajorRowStripFragments(rowStripInput, state, rewriter, loc);
|
||||
}
|
||||
|
||||
static Value createFragmentConstant(const DistributedTensorStep& step,
|
||||
@@ -3350,9 +3430,9 @@ static Value createFragmentReciprocalConstant(const DistributedTensorStep& step,
|
||||
channelValues.push_back(value);
|
||||
values.reserve(fragmentType.getNumElements());
|
||||
for (int64_t n = 0; n < fragmentType.getDimSize(0); ++n)
|
||||
for (int64_t channel = 0; channel < fragmentType.getDimSize(1); ++channel)
|
||||
for (int64_t h = 0; h < fragmentType.getDimSize(2); ++h)
|
||||
for (int64_t w = 0; w < fragmentType.getDimSize(3); ++w) {
|
||||
for (int64_t h = 0; h < fragmentType.getDimSize(1); ++h)
|
||||
for (int64_t w = 0; w < fragmentType.getDimSize(2); ++w)
|
||||
for (int64_t channel = 0; channel < fragmentType.getDimSize(3); ++channel) {
|
||||
APFloat reciprocal = channelValues[channel];
|
||||
APFloat one(reciprocal.getSemantics(), 1);
|
||||
[[maybe_unused]] APFloat::opStatus status = one.divide(reciprocal, APFloat::rmNearestTiesToEven);
|
||||
@@ -4183,7 +4263,7 @@ LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp) {
|
||||
return failure();
|
||||
|
||||
StringRef failureReason;
|
||||
return canConsumeNchwRowStripFragments(*state, failureReason) ? success() : failure();
|
||||
return canConsumePixelMajorRowStripFragments(*state, failureReason) ? success() : failure();
|
||||
}
|
||||
|
||||
FailureOr<Value>
|
||||
|
||||
@@ -481,9 +481,12 @@ FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
|
||||
const int64_t kernelWidth = planOp.getKernelShape()[1];
|
||||
Value input = rowStripInput.value_or(planOp.getInput());
|
||||
auto actualInputType = dyn_cast<RankedTensorType>(input.getType());
|
||||
const bool physicalInput = actualInputType == getRowStripStorageType(inputType);
|
||||
FailureOr<RowStripPhysicalValue> physicalValue = describeRowStripPhysicalValue(input, inputType);
|
||||
const bool physicalInput = succeeded(physicalValue);
|
||||
if (!physicalInput && actualInputType != inputType)
|
||||
return failure();
|
||||
const int64_t tilesPerRow = physicalInput ? physicalValue->tilesPerRow : 1;
|
||||
const int64_t tileChannels = physicalInput ? physicalValue->fragmentType.getDimSize(3) : channels;
|
||||
|
||||
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
|
||||
Value rowTable = createClampedPoolIndexTable(rewriter,
|
||||
@@ -502,51 +505,78 @@ FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
|
||||
planOp.getDilations()[1],
|
||||
planOp.getPads()[1],
|
||||
inputWidth);
|
||||
auto inputFragmentType = getRowStripFragmentType(inputType);
|
||||
auto outputFragmentType = getRowStripFragmentType(outputType);
|
||||
auto outputStorageType = getRowStripStorageType(outputType);
|
||||
auto tileType = RankedTensorType::get({1, channels, 1, 1}, outputType.getElementType());
|
||||
auto inputFragmentType =
|
||||
physicalInput ? physicalValue->fragmentType : getRowStripFragmentType(inputType);
|
||||
auto nchwInputFragmentType = RankedTensorType::get(
|
||||
{1, channels, 1, inputWidth}, inputType.getElementType(), inputType.getEncoding());
|
||||
auto outputFragmentType = RankedTensorType::get(
|
||||
{1, 1, outputWidth, tileChannels}, outputType.getElementType(), outputType.getEncoding());
|
||||
auto outputStorageType =
|
||||
spatial::getGraphBatchPhysicalResultType(outputHeight * tilesPerRow, outputFragmentType);
|
||||
auto tileType = RankedTensorType::get({1, 1, 1, tileChannels}, outputType.getElementType());
|
||||
auto batch = createSpatComputeBatch(
|
||||
rewriter,
|
||||
loc,
|
||||
TypeRange {outputStorageType},
|
||||
outputHeight,
|
||||
outputHeight * tilesPerRow,
|
||||
{},
|
||||
ValueRange {input},
|
||||
[&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
|
||||
SmallVector<Value> inputRows;
|
||||
inputRows.reserve(kernelHeight);
|
||||
Value outputRow = tilesPerRow == 1
|
||||
? args.lane
|
||||
: affineFloorDivConst(rewriter, loc, args.lane, tilesPerRow, anchorOp);
|
||||
Value channelTile = tilesPerRow == 1
|
||||
? getOrCreateIndexConstant(rewriter, anchorOp, 0)
|
||||
: affineModConst(rewriter, loc, args.lane, tilesPerRow, anchorOp);
|
||||
for (int64_t kernelRow = 0; kernelRow < kernelHeight; ++kernelRow) {
|
||||
Value sourceRow =
|
||||
extractPoolIndex(rewriter, loc, anchorOp, rowTable, args.lane, kernelRow, kernelHeight);
|
||||
extractPoolIndex(rewriter, loc, anchorOp, rowTable, outputRow, kernelRow, kernelHeight);
|
||||
if (physicalInput) {
|
||||
inputRows.push_back(
|
||||
extractRowStripFragment(args.inputs.front(), inputType, sourceRow, rewriter, loc));
|
||||
Value sourceSlot = sourceRow;
|
||||
if (tilesPerRow != 1) {
|
||||
sourceSlot = arith::AddIOp::create(
|
||||
rewriter,
|
||||
loc,
|
||||
arith::MulIOp::create(rewriter,
|
||||
loc,
|
||||
sourceRow,
|
||||
getOrCreateIndexConstant(rewriter, anchorOp, tilesPerRow)),
|
||||
channelTile);
|
||||
}
|
||||
FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(
|
||||
rewriter, loc, args.inputs.front(), sourceSlot, inputFragmentType);
|
||||
if (failed(fragment))
|
||||
return failure();
|
||||
inputRows.push_back(*fragment);
|
||||
}
|
||||
else {
|
||||
SmallVector<OpFoldResult> offsets {
|
||||
rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), sourceRow, rewriter.getIndexAttr(0)};
|
||||
inputRows.push_back(tensor::ExtractSliceOp::create(rewriter,
|
||||
loc,
|
||||
inputFragmentType,
|
||||
args.inputs.front(),
|
||||
offsets,
|
||||
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(channels),
|
||||
rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(inputWidth)},
|
||||
getUnitStrides(rewriter, 4)));
|
||||
Value nchw = tensor::ExtractSliceOp::create(rewriter,
|
||||
loc,
|
||||
nchwInputFragmentType,
|
||||
args.inputs.front(),
|
||||
offsets,
|
||||
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(channels),
|
||||
rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(inputWidth)},
|
||||
getUnitStrides(rewriter, 4));
|
||||
inputRows.push_back(ONNXTransposeOp::create(
|
||||
rewriter, loc, inputFragmentType, nchw, rewriter.getI64ArrayAttr({0, 2, 3, 1})));
|
||||
}
|
||||
}
|
||||
|
||||
auto windowType = RankedTensorType::get(
|
||||
{1, channels, kernelHeight, inputWidth}, inputType.getElementType(), inputType.getEncoding());
|
||||
{1, kernelHeight, inputWidth, tileChannels}, inputType.getElementType(), inputType.getEncoding());
|
||||
Value window = tensor::EmptyOp::create(
|
||||
rewriter, loc, windowType.getShape(), windowType.getElementType());
|
||||
for (int64_t kernelRow = 0; kernelRow < kernelHeight; ++kernelRow) {
|
||||
SmallVector<OpFoldResult> offsets {rewriter.getIndexAttr(0),
|
||||
rewriter.getIndexAttr(0),
|
||||
rewriter.getIndexAttr(kernelRow),
|
||||
rewriter.getIndexAttr(0),
|
||||
rewriter.getIndexAttr(0)};
|
||||
window = tensor::InsertSliceOp::create(rewriter,
|
||||
loc,
|
||||
@@ -554,9 +584,9 @@ FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
|
||||
window,
|
||||
offsets,
|
||||
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(channels),
|
||||
rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(inputWidth)},
|
||||
rewriter.getIndexAttr(inputWidth),
|
||||
rewriter.getIndexAttr(tileChannels)},
|
||||
getUnitStrides(rewriter, 4));
|
||||
}
|
||||
|
||||
@@ -585,43 +615,43 @@ FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
|
||||
kernelColumn,
|
||||
kernelWidth);
|
||||
SmallVector<OpFoldResult> offsets {
|
||||
rewriter.getIndexAttr(0),
|
||||
rewriter.getIndexAttr(0),
|
||||
rewriter.getIndexAttr(kernelRow),
|
||||
sourceColumn};
|
||||
sourceColumn,
|
||||
rewriter.getIndexAttr(0)};
|
||||
Value point = tensor::ExtractSliceOp::create(rewriter,
|
||||
nestedLoc,
|
||||
tileType,
|
||||
window,
|
||||
offsets,
|
||||
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(channels),
|
||||
rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(1)},
|
||||
rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(tileChannels)},
|
||||
getUnitStrides(rewriter, 4));
|
||||
reduced = reduced ? spatial::SpatVMaxOp::create(rewriter, nestedLoc, tileType, reduced, point).getResult()
|
||||
: materializeTileTensor(rewriter, nestedLoc, point);
|
||||
}
|
||||
}
|
||||
SmallVector<OpFoldResult> outputOffsets {
|
||||
rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), outputColumn};
|
||||
rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), outputColumn, rewriter.getIndexAttr(0)};
|
||||
Value updated = tensor::InsertSliceOp::create(rewriter,
|
||||
nestedLoc,
|
||||
reduced,
|
||||
iterArgs.front(),
|
||||
outputOffsets,
|
||||
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(channels),
|
||||
rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(1)},
|
||||
rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(tileChannels)},
|
||||
getUnitStrides(rewriter, 4));
|
||||
yielded.push_back(updated);
|
||||
return success();
|
||||
});
|
||||
if (failed(outputLoop))
|
||||
return failure();
|
||||
insertRowStripFragment(
|
||||
outputLoop->results.front(), args.outputs.front(), outputType, args.lane, rewriter, loc);
|
||||
publishGraphBatchPhysicalFragment(
|
||||
rewriter, loc, outputLoop->results.front(), args.outputs.front(), args.lane);
|
||||
return success();
|
||||
});
|
||||
if (failed(batch))
|
||||
|
||||
@@ -19,11 +19,11 @@ namespace {
|
||||
|
||||
static constexpr StringLiteral kLogicalLayout = "nchw";
|
||||
static constexpr StringLiteral kDenseLayout = "dense_nchw";
|
||||
static constexpr StringLiteral kRowStripLayout = "nchw_row_strip";
|
||||
static constexpr StringLiteral kRowStripLayout = "nhwc_row_strip";
|
||||
|
||||
enum class SelectedLayout {
|
||||
DenseNchw,
|
||||
NchwRowStrip,
|
||||
PixelMajorRowStrip,
|
||||
};
|
||||
|
||||
static SelectedLayout getSelectedLayout(llvm::DenseMap<Value, SelectedLayout>& layouts, Value value) {
|
||||
@@ -33,13 +33,13 @@ static SelectedLayout getSelectedLayout(llvm::DenseMap<Value, SelectedLayout>& l
|
||||
|
||||
static bool usesSelectedRowStrip(Operation* user, llvm::DenseMap<Value, SelectedLayout>& layouts) {
|
||||
if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(user))
|
||||
return getSelectedLayout(layouts, reluPlan.getResult()) == SelectedLayout::NchwRowStrip;
|
||||
return getSelectedLayout(layouts, reluPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
|
||||
if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(user))
|
||||
return getSelectedLayout(layouts, biasAddPlan.getResult()) == SelectedLayout::NchwRowStrip;
|
||||
return getSelectedLayout(layouts, biasAddPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
|
||||
if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(user))
|
||||
return getSelectedLayout(layouts, convPlan.getResult()) == SelectedLayout::NchwRowStrip;
|
||||
return getSelectedLayout(layouts, convPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
|
||||
if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(user))
|
||||
return getSelectedLayout(layouts, maxPoolPlan.getResult()) == SelectedLayout::NchwRowStrip;
|
||||
return getSelectedLayout(layouts, maxPoolPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
|
||||
return false;
|
||||
}
|
||||
|
||||
@@ -77,7 +77,7 @@ static bool hasRowStripConsumer(Value value) {
|
||||
static bool canSelectConvRowStrip(spatial::SpatConv2DPlanOp convPlan,
|
||||
llvm::DenseMap<Value, SelectedLayout>& layouts) {
|
||||
SelectedLayout inputLayout = getSelectedLayout(layouts, convPlan.getInput());
|
||||
if (inputLayout == SelectedLayout::NchwRowStrip)
|
||||
if (inputLayout == SelectedLayout::PixelMajorRowStrip)
|
||||
return succeeded(canConsumeAndProduceRowStrip(convPlan));
|
||||
return succeeded(canLowerConvPlanToRowStrip(convPlan));
|
||||
}
|
||||
@@ -88,23 +88,23 @@ static SelectedLayout chooseConvLayout(spatial::SpatConv2DPlanOp convPlan,
|
||||
return SelectedLayout::DenseNchw;
|
||||
if (!allUsersCanHandleRowStrip(convPlan.getResult(), layouts))
|
||||
return SelectedLayout::DenseNchw;
|
||||
return SelectedLayout::NchwRowStrip;
|
||||
return SelectedLayout::PixelMajorRowStrip;
|
||||
}
|
||||
|
||||
static SelectedLayout chooseReluLayout(spatial::SpatReluPlanOp reluPlan,
|
||||
llvm::DenseMap<Value, SelectedLayout>& layouts) {
|
||||
if (getSelectedLayout(layouts, reluPlan.getInput()) != SelectedLayout::NchwRowStrip)
|
||||
if (getSelectedLayout(layouts, reluPlan.getInput()) != SelectedLayout::PixelMajorRowStrip)
|
||||
return SelectedLayout::DenseNchw;
|
||||
if (!hasRowStripConsumer(reluPlan.getResult()))
|
||||
return SelectedLayout::DenseNchw;
|
||||
if (!allUsersCanHandleRowStrip(reluPlan.getResult(), layouts))
|
||||
return SelectedLayout::DenseNchw;
|
||||
return SelectedLayout::NchwRowStrip;
|
||||
return SelectedLayout::PixelMajorRowStrip;
|
||||
}
|
||||
|
||||
static SelectedLayout chooseBiasAddLayout(spatial::SpatBiasAddPlanOp biasAddPlan,
|
||||
llvm::DenseMap<Value, SelectedLayout>& layouts) {
|
||||
if (getSelectedLayout(layouts, biasAddPlan.getInput()) != SelectedLayout::NchwRowStrip)
|
||||
if (getSelectedLayout(layouts, biasAddPlan.getInput()) != SelectedLayout::PixelMajorRowStrip)
|
||||
return SelectedLayout::DenseNchw;
|
||||
auto resultType = dyn_cast<RankedTensorType>(biasAddPlan.getOutput().getType());
|
||||
if (!resultType || !isSupportedBiasAddValue(biasAddPlan.getBias(), resultType))
|
||||
@@ -113,11 +113,11 @@ static SelectedLayout chooseBiasAddLayout(spatial::SpatBiasAddPlanOp biasAddPlan
|
||||
return SelectedLayout::DenseNchw;
|
||||
if (!allUsersCanHandleRowStrip(biasAddPlan.getResult(), layouts))
|
||||
return SelectedLayout::DenseNchw;
|
||||
return SelectedLayout::NchwRowStrip;
|
||||
return SelectedLayout::PixelMajorRowStrip;
|
||||
}
|
||||
|
||||
static SelectedLayout chooseMaxPoolLayout(spatial::SpatMaxPool2DPlanOp maxPoolPlan) {
|
||||
return succeeded(canLowerMaxPoolPlanToRowStrip(maxPoolPlan)) ? SelectedLayout::NchwRowStrip
|
||||
return succeeded(canLowerMaxPoolPlanToRowStrip(maxPoolPlan)) ? SelectedLayout::PixelMajorRowStrip
|
||||
: SelectedLayout::DenseNchw;
|
||||
}
|
||||
|
||||
@@ -237,7 +237,7 @@ struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass,
|
||||
else
|
||||
continue;
|
||||
|
||||
if (getSelectedLayout(layouts, producedValue) != SelectedLayout::NchwRowStrip)
|
||||
if (getSelectedLayout(layouts, producedValue) != SelectedLayout::PixelMajorRowStrip)
|
||||
continue;
|
||||
|
||||
rewriter.setInsertionPointAfter(&op);
|
||||
|
||||
@@ -121,6 +121,63 @@ lowerMemRefCopyToPimCopy(memref::CopyOp copyOp,
|
||||
return success();
|
||||
}
|
||||
|
||||
static Value getForwardedInputConsumerOutput(OpOperand& use) {
|
||||
Operation* owner = use.getOwner();
|
||||
if (auto vmm = dyn_cast<pim::PimVMMOp>(owner))
|
||||
return &use == &vmm.getInputMutable() ? vmm.getOutputBuffer() : Value();
|
||||
if (isa<pim::PimVVAddOp,
|
||||
pim::PimVVSubOp,
|
||||
pim::PimVVMulOp,
|
||||
pim::PimVVMaxOp,
|
||||
pim::PimVVDMulOp>(owner))
|
||||
return use.getOperandNumber() < 2 ? owner->getOperand(2) : Value();
|
||||
return {};
|
||||
}
|
||||
|
||||
static void forwardSingleConsumerContiguousInputCopies(func::FuncOp funcOp) {
|
||||
SmallVector<memref::CopyOp> copies;
|
||||
funcOp.walk([&](memref::CopyOp copy) { copies.push_back(copy); });
|
||||
|
||||
for (memref::CopyOp copy : copies) {
|
||||
Value target = copy.getTarget();
|
||||
auto targetAlloc = target.getDefiningOp<memref::AllocOp>();
|
||||
if (!targetAlloc)
|
||||
continue;
|
||||
|
||||
OpOperand* consumerUse = nullptr;
|
||||
bool hasOtherUse = false;
|
||||
for (OpOperand& use : target.getUses()) {
|
||||
if (use.getOwner() == copy) {
|
||||
if (&use != ©.getTargetMutable())
|
||||
hasOtherUse = true;
|
||||
continue;
|
||||
}
|
||||
if (consumerUse)
|
||||
hasOtherUse = true;
|
||||
else
|
||||
consumerUse = &use;
|
||||
}
|
||||
if (hasOtherUse || !consumerUse)
|
||||
continue;
|
||||
|
||||
Value output = getForwardedInputConsumerOutput(*consumerUse);
|
||||
Value source = copy.getSource();
|
||||
if (!output || !isDeviceLocalPimAddress(source)
|
||||
|| (failed(resolveContiguousAddress(source)) && failed(compileContiguousAddressExpr(source))))
|
||||
continue;
|
||||
|
||||
FailureOr<Value> sourceBase = getPimAddressBase(source);
|
||||
FailureOr<Value> outputBase = getPimAddressBase(output);
|
||||
if (failed(sourceBase) || failed(outputBase) || *sourceBase == *outputBase)
|
||||
continue;
|
||||
|
||||
consumerUse->set(source);
|
||||
copy.erase();
|
||||
if (targetAlloc.use_empty())
|
||||
targetAlloc.erase();
|
||||
}
|
||||
}
|
||||
|
||||
enum class ExpectedPimCopyDirection { HostToDevice, DeviceToHost, DeviceToDevice };
|
||||
|
||||
static LogicalResult verifyPimCopyEndpoints(Operation* copy,
|
||||
@@ -276,6 +333,8 @@ void PimBufferizationPass::runOnOperation() {
|
||||
return;
|
||||
}
|
||||
|
||||
forwardSingleConsumerContiguousInputCopies(funcOp);
|
||||
|
||||
MLIRContext* ctx = moduleOp.getContext();
|
||||
PatternRewriter rewriter(ctx);
|
||||
|
||||
|
||||
@@ -406,7 +406,8 @@ LogicalResult SpatConcatOp::verify() {
|
||||
static bool isKnownLogicalLayout(StringRef layout) { return layout == "nchw"; }
|
||||
|
||||
static bool isKnownPhysicalLayout(StringRef layout) {
|
||||
return layout == "dense_nchw" || layout == "nchw_row_strip" || layout == "fragmented";
|
||||
return layout == "dense_nchw" || layout == "nchw_row_strip" || layout == "nhwc_row_strip"
|
||||
|| layout == "fragmented";
|
||||
}
|
||||
|
||||
static LogicalResult verifyPlanTensorTypes(Operation* op, Value input, Value output, StringRef kind) {
|
||||
|
||||
Reference in New Issue
Block a user