diff --git a/src/PIM/Common/IR/AddressAnalysis.cpp b/src/PIM/Common/IR/AddressAnalysis.cpp index e090765..1fb13bc 100644 --- a/src/PIM/Common/IR/AddressAnalysis.cpp +++ b/src/PIM/Common/IR/AddressAnalysis.cpp @@ -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(value)) + if (auto blockArgument = mlir::dyn_cast(value)) { + auto forOp = mlir::dyn_cast_or_null(blockArgument.getOwner()->getParentOp()); + if (forOp && blockArgument.getArgNumber() > 0) { + const unsigned iterArgIndex = blockArgument.getArgNumber() - 1; + auto yieldOp = mlir::dyn_cast(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) diff --git a/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp b/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp index 396e925..ccaf398 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp @@ -18,10 +18,10 @@ FailureOr 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 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 createPerChannelConstantFragment(DenseElementsAttr denseAttr, RankedTensorType fragmentType, PatternRewriter& rewriter) { - FailureOr> channelValues = getBiasChannelValues(denseAttr, fragmentType); + auto logicalType = RankedTensorType::get( + {1, fragmentType.getDimSize(3), 1, 1}, fragmentType.getElementType()); + FailureOr> channelValues = getBiasChannelValues(denseAttr, logicalType); if (failed(channelValues)) return failure(); SmallVector 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 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 createRowStripStorageFromRows(Value rows, SmallVector 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 {{0, 1}, {2, 3}}); + rewriter, loc, fragmentType, rowSlice, SmallVector {{0, 1, 2}, {3}}); insertRowStripFragment(fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc); return success(); }); @@ -143,8 +143,28 @@ FailureOr createRowStripStorageFromRows(Value rows, FailureOr createRowStripAssemblyBlueprint(const RowStripPhysicalValue& value, PatternRewriter& rewriter, Location loc) { + const int64_t laneCount = cast(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 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 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 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 applyRowStripBiasAdd(const RowStripPhysicalValue& value, auto biasStorageType = spatial::getGraphBatchPhysicalResultType(value.tilesPerRow, value.fragmentType); SmallVector biasValues( biasStorageType.getNumElements(), cast(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); diff --git a/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp b/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp index 2e74af6..e8c24a6 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp +++ b/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp @@ -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; diff --git a/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp b/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp index 539cb9e..2264a7e 100644 --- a/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp @@ -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 getRowStripValue(llvm::DenseMap& rowStripValues, Value value) { @@ -251,20 +251,8 @@ struct LowerSpatialPlansPass final : PassWrapper input = getRowStripValue(rowStripValues, planOp.getInput()); rewriter.setInsertionPoint(planOp); std::optional physicalInput; - if (succeeded(input)) { - if (input->tilesPerRow == 1) { - physicalInput = input->storage; - } - else { - FailureOr 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 lowered = lowerSelectedMaxPool2DPlan( planOp, physicalInput, rewriter); if (failed(lowered)) { diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp index 5dc5dd2..9a67ebf 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp @@ -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 sourceValues(sourceAttr.getValues()); + SmallVector paddedValues( + paddedType.getNumElements(), cast(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 createPaddedBiasRowConstant(const ConvLoweringState& state, - int64_t paddedChannels, - PatternRewriter& rewriter) { +static FailureOr createBiasRowConstant(const ConvLoweringState& state, + PatternRewriter& rewriter) { DenseElementsAttr denseAttr; if (!isSupportedBiasAddValue(state.b, state.outType, &denseAttr)) return failure(); @@ -2549,12 +2572,11 @@ static FailureOr createPaddedBiasRowConstant(const ConvLoweringState& sta if (failed(channelValues)) return failure(); - auto biasType = RankedTensorType::get({1, paddedChannels}, state.outType.getElementType()); - SmallVector values(biasType.getNumElements(), cast(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 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 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 createRowStripWindowMaskTable(const ConvLoweringState& state, PatternRewriter& rewriter) { @@ -2661,7 +2687,7 @@ static FailureOr 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 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 offsets { tableIndex, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}; SmallVector 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 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 createConvInputWindow(Value input, paddedRow, window, SmallVector {rewriter.getIndexAttr(0), - rewriter.getIndexAttr(0), rewriter.getIndexAttr(kernelRowIndex), + rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}, SmallVector {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 createNchwRowStripConvPatchRow(Value paddedWindow, - const ConvLoweringState& state, - Value outputWidth, - PatternRewriter& rewriter, - Location loc) { +static FailureOr 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 offsets { + rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), inputWidthOffset, rewriter.getIndexAttr(0)}; + SmallVector sizes {rewriter.getIndexAttr(1), + rewriter.getIndexAttr(state.wHeight), + rewriter.getIndexAttr(state.wWidth), + rewriter.getIndexAttr(state.numChannelsIn)}; + SmallVector 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 {{0}, {1, 2, 3}}) .getResult(); } -static FailureOr createPaddedConvOutputTile(Value paddedPatchRow, - Value tileWeights, - int64_t numKSlices, - int64_t xbarDim, - PatternRewriter& rewriter, - Location loc) { - auto elementType = cast(paddedPatchRow.getType()).getElementType(); +static FailureOr createConvOutputTile(Value patchRow, + Value partialInputScratch, + Value tileWeights, + int64_t patchSize, + int64_t numKSlices, + int64_t xbarDim, + PatternRewriter& rewriter, + Location loc) { + auto elementType = cast(patchRow.getType()).getElementType(); auto paddedRowType = RankedTensorType::get({1, xbarDim}, elementType); auto weightElementType = cast(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 aOffsets {rewriter.getIndexAttr(0), kOffset}; - SmallVector aSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)}; - Value aTile = extractStaticSliceOrIdentity( - rewriter, pieceLoc, paddedPatchRow, paddedRowType, aOffsets, aSizes, getUnitStrides(rewriter, 2)); - SmallVector 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 {rewriter.getIndexAttr(0), rewriter.getIndexAttr(kOffset)}, + SmallVector {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 {rewriter.getIndexAttr(0), rewriter.getIndexAttr(kOffset)}, + SmallVector {rewriter.getIndexAttr(1), rewriter.getIndexAttr(sliceSize)}, + getUnitStrides(rewriter, 2)); + inputTile = tensor::InsertSliceOp::create( + rewriter, + loc, + partial, + partialInputScratch, + SmallVector {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}, + SmallVector {rewriter.getIndexAttr(1), rewriter.getIndexAttr(sliceSize)}, + getUnitStrides(rewriter, 2)); + } + SmallVector bOffsets { + rewriter.getIndexAttr(kOffset), rewriter.getIndexAttr(0)}; SmallVector 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& 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 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 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(patchRow.getType()).getElementType(); + auto rowType = RankedTensorType::get({1, outputChannels}, elementType); + auto tileWeightsType = + RankedTensorType::get({paddedK, xbarDim}, + cast(paddedWeights.getType()).getElementType()); + const int64_t outputTileCount = ceilIntegerDivide(outputChannels, xbarDim); auto getTileWeights = [&](int64_t outputTile) { if (outputTileCount == 1) @@ -2871,29 +2907,31 @@ static FailureOr createPaddedConvOutputRow(Value patchRow, }; if (outputTileCount == 1) { - FailureOr rowResult = createPaddedConvOutputTile( - paddedPatchRow, getTileWeights(0), numKSlices, xbarDim, rewriter, loc); + FailureOr 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 outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}; - SmallVector 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 {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}, + SmallVector {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 tileResult = createPaddedConvOutputTile( - paddedPatchRow, getTileWeights(outputTile), numKSlices, xbarDim, rewriter, loc); + FailureOr tileResult = createConvOutputTile( + patchRow, partialInputScratch, getTileWeights(outputTile), patchSize, numKSlices, xbarDim, rewriter, loc); if (failed(tileResult)) return failure(); SmallVector tileOffsets { @@ -2904,7 +2942,7 @@ static FailureOr createPaddedConvOutputRow(Value patchRow, } SmallVector outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}; - SmallVector outputSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut)}; + SmallVector 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 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 createOutputChannelTiledRowStripConvOutput(const ConvLow if (failed(inputWindow)) return failure(); Value fragmentInit = tensor::EmptyOp::create(rewriter, loc, tileFragmentType.getShape(), elementType); + SmallVector 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& widthYielded) { FailureOr 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 {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}, + SmallVector {rewriter.getIndexAttr(1), rewriter.getIndexAttr(patchSize)}, + getUnitStrides(rewriter, 2)); FailureOr paddedOutputRow = createPaddedConvOutputTile( paddedPatchRow, tileWeights, numKSlices, xbarDim, rewriter, widthLoc); if (failed(paddedOutputRow)) @@ -2991,13 +3038,13 @@ static FailureOr createOutputChannelTiledRowStripConvOutput(const ConvLow paddedOutputRow = spatial::SpatVAddOp::create( rewriter, widthLoc, paddedRowType, *paddedOutputRow, *biasTile).getResult(); Value outputPixel = tensor::ExpandShapeOp::create( - rewriter, widthLoc, tilePixelType, *paddedOutputRow, SmallVector {{0}, {1, 2, 3}}); + rewriter, widthLoc, tilePixelType, *paddedOutputRow, SmallVector {{0, 1, 2}, {3}}); SmallVector rowOffsets { - rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex}; + rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex, rewriter.getIndexAttr(0)}; SmallVector 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 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 paddedBias = failure(); + FailureOr 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 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& widthYielded) { FailureOr patchRow = - createNchwRowStripConvPatchRow(*inputWindow, state, widthIndex, rewriter, widthLoc); + createPixelMajorConvPatchRow(*inputWindow, state, widthIndex, rewriter, widthLoc); if (failed(patchRow)) return failure(); - FailureOr outputRow = createPaddedConvOutputRow(*patchRow, - state, + Value paddedPatchRow = *patchRow; + if (patchSize != paddedK) + paddedPatchRow = tensor::InsertSliceOp::create( + rewriter, + widthLoc, + paddedPatchRow, + widthIterArgs[1], + SmallVector {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}, + SmallVector {rewriter.getIndexAttr(1), rewriter.getIndexAttr(patchSize)}, + getUnitStrides(rewriter, 2)); + FailureOr 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 {{0}, {1, 2, 3}}); + SmallVector {{0, 1, 2}, {3}}); SmallVector rowOffsets { - rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex}; + rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex, rewriter.getIndexAttr(0)}; SmallVector 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 createConvOutputFromNchwRowStripFragments(Value rowStripStorage, - const ConvLoweringState& state, - PatternRewriter& rewriter, - Location loc) { +static FailureOr createConvOutputFromPixelMajorRowStripFragments(Value rowStripStorage, + const ConvLoweringState& state, + PatternRewriter& rewriter, + Location loc) { auto inputType = dyn_cast(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 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 paddedBias = failure(); + Value paddedWeights = + standard::createPaddedPixelMajorWeightConstant(weightDenseAttr, state, paddedK, xbarDim, rewriter); + FailureOr 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 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 createConvOutputFromNchwRowStripFragments(Value rowStrip if (failed(inputWindow)) return failure(); Value fragmentInit = tensor::EmptyOp::create(rewriter, loc, fragmentType.getShape(), elementType); + SmallVector 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& widthYielded) { FailureOr patchRow = - createNchwRowStripConvPatchRow(*inputWindow, state, widthIndex, rewriter, widthLoc); + createPixelMajorConvPatchRow(*inputWindow, state, widthIndex, rewriter, widthLoc); if (failed(patchRow)) return failure(); - FailureOr outputRow = createPaddedConvOutputRow(*patchRow, - state, + Value paddedPatchRow = *patchRow; + if (patchSize != paddedK) + paddedPatchRow = tensor::InsertSliceOp::create( + rewriter, + widthLoc, + paddedPatchRow, + widthIterArgs[1], + SmallVector {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}, + SmallVector {rewriter.getIndexAttr(1), rewriter.getIndexAttr(patchSize)}, + getUnitStrides(rewriter, 2)); + FailureOr outputRow = createPaddedConvOutputRow(paddedPatchRow, + state.numChannelsOut, args.weights.front(), state.hasBias ? args.inputs[1] : Value(), - paddedK, numKSlices, xbarDim, rewriter, @@ -3198,15 +3276,17 @@ static FailureOr createConvOutputFromNchwRowStripFragments(Value rowStrip widthLoc, outputPixelType, *outputRow, - SmallVector {{0}, {1, 2, 3}}); + SmallVector {{0, 1, 2}, {3}}); SmallVector rowOffsets { - rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex}; + rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex, rewriter.getIndexAttr(0)}; SmallVector 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 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 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 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 {{0}, {1, 2, 3}}); + SmallVector {{0, 1, 2}, {3}}); Value inputSlice = tensor::ExtractSliceOp::create(rewriter, pieceLoc, paddedRowType, inputRow, SmallVector {rewriter.getIndexAttr(0), sourceOffset}, SmallVector {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)}, @@ -3305,7 +3385,7 @@ static FailureOr createPointwiseOutputFromRowStripFragments(Value rowStri result = spatial::SpatVAddOp::create(rewriter, loc, paddedRowType, *result, *bias).getResult(); } Value fragment = tensor::ExpandShapeOp::create(rewriter, loc, outputFragmentType, *result, - SmallVector {{0}, {1, 2, 3}}); + SmallVector {{0, 1, 2}, {3}}); publishGraphBatchPhysicalFragment(rewriter, loc, fragment, args.outputs.front(), args.lane); return success(); }); @@ -3320,7 +3400,7 @@ static FailureOr 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 diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Pool.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Pool.cpp index abce781..3c41c2e 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Pool.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Pool.cpp @@ -481,9 +481,12 @@ FailureOr lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp, const int64_t kernelWidth = planOp.getKernelShape()[1]; Value input = rowStripInput.value_or(planOp.getInput()); auto actualInputType = dyn_cast(input.getType()); - const bool physicalInput = actualInputType == getRowStripStorageType(inputType); + FailureOr 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 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 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 fragment = extractGraphBatchPhysicalFragment( + rewriter, loc, args.inputs.front(), sourceSlot, inputFragmentType); + if (failed(fragment)) + return failure(); + inputRows.push_back(*fragment); } else { SmallVector 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 {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 {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 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 lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp, window, offsets, SmallVector {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 lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp, kernelColumn, kernelWidth); SmallVector 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 {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 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 {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)) diff --git a/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp b/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp index 5c8df36..1e2443b 100644 --- a/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp @@ -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& layouts, Value value) { @@ -33,13 +33,13 @@ static SelectedLayout getSelectedLayout(llvm::DenseMap& l static bool usesSelectedRowStrip(Operation* user, llvm::DenseMap& layouts) { if (auto reluPlan = dyn_cast(user)) - return getSelectedLayout(layouts, reluPlan.getResult()) == SelectedLayout::NchwRowStrip; + return getSelectedLayout(layouts, reluPlan.getResult()) == SelectedLayout::PixelMajorRowStrip; if (auto biasAddPlan = dyn_cast(user)) - return getSelectedLayout(layouts, biasAddPlan.getResult()) == SelectedLayout::NchwRowStrip; + return getSelectedLayout(layouts, biasAddPlan.getResult()) == SelectedLayout::PixelMajorRowStrip; if (auto convPlan = dyn_cast(user)) - return getSelectedLayout(layouts, convPlan.getResult()) == SelectedLayout::NchwRowStrip; + return getSelectedLayout(layouts, convPlan.getResult()) == SelectedLayout::PixelMajorRowStrip; if (auto maxPoolPlan = dyn_cast(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& 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& 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& layouts) { - if (getSelectedLayout(layouts, biasAddPlan.getInput()) != SelectedLayout::NchwRowStrip) + if (getSelectedLayout(layouts, biasAddPlan.getInput()) != SelectedLayout::PixelMajorRowStrip) return SelectedLayout::DenseNchw; auto resultType = dyn_cast(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(owner)) + return &use == &vmm.getInputMutable() ? vmm.getOutputBuffer() : Value(); + if (isa(owner)) + return use.getOperandNumber() < 2 ? owner->getOperand(2) : Value(); + return {}; +} + +static void forwardSingleConsumerContiguousInputCopies(func::FuncOp funcOp) { + SmallVector copies; + funcOp.walk([&](memref::CopyOp copy) { copies.push_back(copy); }); + + for (memref::CopyOp copy : copies) { + Value target = copy.getTarget(); + auto targetAlloc = target.getDefiningOp(); + 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 sourceBase = getPimAddressBase(source); + FailureOr 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); diff --git a/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp b/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp index 2766058..1e92546 100644 --- a/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp +++ b/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp @@ -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) {