diff --git a/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp b/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp index e198b77..99b9762 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp @@ -9,6 +9,8 @@ #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" #include "src/Dialect/ONNX/ONNXOps.hpp" +#include + using namespace mlir; namespace onnx_mlir { @@ -255,10 +257,18 @@ FailureOr applyRowStripAdd(const RowStripPhysicalValue& lhs, const RowStripPhysicalValue& rhs, PatternRewriter& rewriter, Location loc) { - if (lhs.logicalType != rhs.logicalType || lhs.fragmentType != rhs.fragmentType - || lhs.storage.getType() != rhs.storage.getType() || lhs.tilesPerRow != rhs.tilesPerRow) + if (lhs.logicalType != rhs.logicalType) return failure(); - auto storageType = cast(lhs.storage.getType()); + + const int64_t fragmentChannels = + std::gcd(lhs.fragmentType.getDimSize(3), rhs.fragmentType.getDimSize(3)); + auto fragmentType = RankedTensorType::get( + {1, 1, lhs.logicalType.getDimSize(3), fragmentChannels}, + lhs.logicalType.getElementType(), + lhs.logicalType.getEncoding()); + const int64_t tilesPerRow = ceilIntegerDivide(lhs.logicalType.getDimSize(1), fragmentChannels); + auto storageType = + spatial::getGraphBatchPhysicalResultType(lhs.logicalType.getDimSize(2) * tilesPerRow, fragmentType); const int64_t laneCount = storageType.getDimSize(0); auto batch = createSpatComputeBatch( rewriter, @@ -268,13 +278,43 @@ FailureOr applyRowStripAdd(const RowStripPhysicalValue& lhs, {}, ValueRange {lhs.storage, rhs.storage}, [&](detail::SpatComputeBatchBodyArgs args) { - FailureOr lhsFragment = - extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[0], args.lane, lhs.fragmentType); - FailureOr rhsFragment = - extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[1], args.lane, rhs.fragmentType); + auto extractFragment = [&](Value storage, const RowStripPhysicalValue& input) -> FailureOr { + Operation* anchor = rewriter.getInsertionBlock()->getParentOp(); + MLIRContext* context = rewriter.getContext(); + AffineExpr lane = getAffineDimExpr(0, context); + AffineExpr outputTile = lane % tilesPerRow; + AffineExpr channelOffset = outputTile * fragmentChannels; + Value sourceSlot = createOrFoldAffineApply( + rewriter, + loc, + lane.floorDiv(tilesPerRow) * input.tilesPerRow + + channelOffset.floorDiv(input.fragmentType.getDimSize(3)), + ValueRange {args.lane}, + anchor); + FailureOr source = + extractGraphBatchPhysicalFragment(rewriter, loc, storage, sourceSlot, input.fragmentType); + if (failed(source) || input.fragmentType == fragmentType) + return source; + Value sourceOffset = createOrFoldAffineApply( + rewriter, loc, channelOffset % input.fragmentType.getDimSize(3), ValueRange {args.lane}, anchor); + return tensor::ExtractSliceOp::create( + rewriter, + loc, + fragmentType, + *source, + SmallVector { + rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), sourceOffset}, + SmallVector {rewriter.getIndexAttr(1), + rewriter.getIndexAttr(1), + rewriter.getIndexAttr(lhs.logicalType.getDimSize(3)), + rewriter.getIndexAttr(fragmentChannels)}, + getUnitStrides(rewriter, 4)).getResult(); + }; + FailureOr lhsFragment = extractFragment(args.inputs[0], lhs); + FailureOr rhsFragment = extractFragment(args.inputs[1], rhs); if (failed(lhsFragment) || failed(rhsFragment)) return failure(); - Value added = spatial::SpatVAddOp::create(rewriter, loc, lhs.fragmentType, *lhsFragment, *rhsFragment); + Value added = spatial::SpatVAddOp::create(rewriter, loc, fragmentType, *lhsFragment, *rhsFragment); publishGraphBatchPhysicalFragment(rewriter, loc, added, args.outputs.front(), args.lane); return success(); }); diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp index 5fcdb35..847288d 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp @@ -2511,6 +2511,11 @@ static bool rowStripOutputTileFitsOneCore(const ConvGeometry& geometry) { <= static_cast(crossbarCountInCore.getValue()); } +static bool rowStripOutputChannelTileFitsOneCore(const ConvGeometry& geometry) { + return ceilIntegerDivide(geometry.k, geometry.xbarSize) + <= static_cast(crossbarCountInCore.getValue()); +} + static bool canConsumePixelMajorRowStripFragments(const ConvLoweringState& state, StringRef& failureReason) { if (state.batchSize != 1) { failureReason = "batch_not_one"; @@ -2528,24 +2533,16 @@ static bool canConsumePixelMajorRowStripFragments(const ConvLoweringState& state failureReason = "non_float_input"; return false; } - if (state.strideHeight != 1 || state.strideWidth != 1) { - failureReason = "stride_not_one"; - return false; - } if (state.dilationHeight != 1 || state.dilationWidth != 1) { failureReason = "dilation_not_one"; return false; } - if (state.outHeight != state.xHeight || state.outWidth != state.xWidth) { - failureReason = "not_same_spatial_shape"; - return false; - } if (!getHostConstDenseElementsAttr(state.w)) { failureReason = "non_constant_weight"; return false; } - if (!rowStripOutputTileFitsOneCore(buildConvGeometry(state))) { - failureReason = "output_row_does_not_fit_one_core"; + if (!rowStripOutputChannelTileFitsOneCore(buildConvGeometry(state))) { + failureReason = "output_channel_tile_does_not_fit_one_core"; return false; } if (state.hasBias && !isSupportedBiasAddValue(state.b, state.outType)) { @@ -3306,7 +3303,7 @@ static FailureOr createOutputChannelTiledRowStripConvOutput(const ConvLow static FailureOr createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRewriter& rewriter, Location loc) { ConvGeometry geometry = buildConvGeometry(state); - if (state.group != 1 || state.batchSize != 1 || !rowStripOutputTileFitsOneCore(geometry)) + if (state.group != 1 || state.batchSize != 1 || !rowStripOutputChannelTileFitsOneCore(geometry)) return failure(); auto weightDenseAttr = getHostConstDenseElementsAttr(state.w); @@ -3320,11 +3317,16 @@ static FailureOr const int64_t numKSlices = ceilIntegerDivide(patchSize, xbarDim); const int64_t paddedK = numKSlices * xbarDim; - const int64_t paddedOutputChannels = - ceilIntegerDivide(state.numChannelsOut, xbarDim) * xbarDim; - Value paddedWeights = standard::createPaddedPixelMajorWeightConstant( - weightDenseAttr, state, paddedK, paddedOutputChannels, rewriter); + if (!rowStripOutputTileFitsOneCore(geometry)) { + Value tiledWeights = + standard::createPaddedOutputChannelTiledWeightConstant(weightDenseAttr, state, paddedK, xbarDim, rewriter); + return createOutputChannelTiledRowStripConvOutput( + state, state.x, tiledWeights, paddedK, numKSlices, xbarDim, rewriter, loc); + } + const int64_t paddedOutputChannels = ceilIntegerDivide(state.numChannelsOut, xbarDim) * xbarDim; + Value paddedWeights = + standard::createPaddedPixelMajorWeightConstant(weightDenseAttr, state, paddedK, paddedOutputChannels, rewriter); FailureOr bias = failure(); if (state.hasBias) bias = createBiasRowConstant(state, rewriter); @@ -3355,6 +3357,13 @@ static FailureOr createConvOutputFromPixelMajorRowStripFragments(Value ro auto weightDenseAttr = getHostConstDenseElementsAttr(state.w); if (!weightDenseAttr) return failure(); + if (!rowStripOutputTileFitsOneCore(geometry)) { + Value tiledWeights = + standard::createPaddedOutputChannelTiledWeightConstant(weightDenseAttr, state, paddedK, xbarDim, rewriter); + return createOutputChannelTiledRowStripConvOutput( + state, rowStripStorage, tiledWeights, paddedK, numKSlices, xbarDim, rewriter, loc); + } + const int64_t paddedOutputChannels = ceilIntegerDivide(state.numChannelsOut, xbarDim) * xbarDim; Value paddedWeights = standard::createPaddedPixelMajorWeightConstant( @@ -4298,7 +4307,7 @@ LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp) { analysis.barrierKind = DistributedConvBarrierKind::UnsupportedConsumer; analysis.barrierDetail = "selected row-strip layout"; ConvGeometry geometry = buildConvGeometry(*state); - if (!rowStripOutputTileFitsOneCore(geometry)) + if (!rowStripOutputChannelTileFitsOneCore(geometry)) return failure(); ConvLoweringDecision decision = chooseConvLoweringStrategy(geometry, *requestedStrategy, analysis); if (decision.strategy == PimConvLoweringDepthwise && !depthwise::canUseStructuredRewrite(*state) diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.cpp index 98bf1e1..c47f6a4 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.cpp @@ -501,7 +501,9 @@ static Value extractReductionPiece(Value partialPiecesArg, SmallVector unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)}; SmallVector pieceSizes {rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())}; SmallVector pieceOffsets { - createPartialGroupOffset(hSlice, kSlice, numKSlices, numOutRows, rewriter, loc), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}; + createPartialGroupOffset(hSlice, kSlice, numKSlices, numOutRows, rewriter, loc), + rewriter.getIndexAttr(0), + rewriter.getIndexAttr(0)}; return extractMixedSliceOrIdentity( rewriter, loc, partialPiecesArg, pieceType, {pieceOffsets, pieceSizes, unitStrides}); @@ -535,94 +537,72 @@ static Value reducePartialPiecesForHSlice(Value partialPiecesArg, return activePieces.front(); } -static FailureOr createReductionCompute(Value partialPieces, - Value bias, - RankedTensorType partialPiecesType, - RankedTensorType outType, - RankedTensorType paddedOutType, - int64_t numKSlices, - ConversionPatternRewriter& rewriter, - Location loc) { +static FailureOr createReductionOutput(Value partialPieces, + Value bias, + RankedTensorType partialPiecesType, + RankedTensorType outType, + RankedTensorType paddedOutType, + int64_t numKSlices, + ConversionPatternRewriter& rewriter, + Location loc) { + const int64_t numOutRows = outType.getDimSize(0); + const int64_t numOutHSlices = ceilIntegerDivide(outType.getDimSize(1), crossbarSize.getValue()); + auto pieceType = RankedTensorType::get({numOutRows, static_cast(crossbarSize.getValue())}, + partialPiecesType.getElementType()); + + if (bias && cast(bias.getType()) != paddedOutType) + bias = createZeroPaddedTensor(bias, paddedOutType, rewriter, loc); SmallVector inputs {partialPieces}; if (bias) inputs.push_back(bias); - auto computeOp = - createSpatCompute(rewriter, loc, TypeRange {outType}, {}, inputs, [&](ValueRange blockArgs) -> LogicalResult { - Value partialPiecesArg = blockArgs[0]; - Value biasArg = bias ? blockArgs[1] : Value(); - if (biasArg && cast(biasArg.getType()) != paddedOutType) - biasArg = createZeroPaddedTensor(biasArg, paddedOutType, rewriter, loc); - - const int64_t numOutRows = outType.getDimSize(0); - const int64_t numOutHSlices = ceilIntegerDivide(outType.getDimSize(1), crossbarSize.getValue()); - auto pieceType = RankedTensorType::get({numOutRows, static_cast(crossbarSize.getValue())}, - partialPiecesType.getElementType()); - - Value outputInit = - tensor::EmptyOp::create(rewriter, loc, paddedOutType.getShape(), paddedOutType.getElementType()).getResult(); - SmallVector unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)}; - SmallVector pieceSizes {rewriter.getIndexAttr(numOutRows), - rewriter.getIndexAttr(crossbarSize.getValue())}; - - auto buildOutputSlice = [&](Value outputAcc, Value hSlice) -> Value { - Value reduced = - reducePartialPiecesForHSlice(partialPiecesArg, hSlice, pieceType, numKSlices, numOutRows, rewriter, loc); - Value hOffset = onnx_mlir::affineMulConst( - rewriter, loc, hSlice, crossbarSize.getValue(), rewriter.getInsertionBlock()->getParentOp()); - if (biasArg) { - SmallVector biasOffsets {rewriter.getIndexAttr(0), hOffset}; - Value biasSlice = - tensor::ExtractSliceOp::create(rewriter, loc, pieceType, biasArg, biasOffsets, pieceSizes, unitStrides) - .getResult(); - reduced = spatial::SpatVAddOp::create(rewriter, loc, pieceType, reduced, biasSlice).getResult(); - } - - SmallVector outputOffsets {rewriter.getIndexAttr(0), hOffset}; - return tensor::InsertSliceOp::create(rewriter, loc, reduced, outputAcc, outputOffsets, pieceSizes, unitStrides) - .getResult(); - }; - - Value paddedOutput = outputInit; - if (numOutHSlices == 1) { - Value hSlice = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0); - paddedOutput = buildOutputSlice(outputInit, hSlice); - } - else { - Value c0 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0); - Value c1 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 1); - Value cOutHSlices = - getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), numOutHSlices); - auto hLoop = buildNormalizedScfFor( - rewriter, - loc, - c0, - cOutHSlices, - c1, - ValueRange {outputInit}, - [&](OpBuilder&, Location, Value hSlice, ValueRange iterArgs, SmallVectorImpl& yielded) { - yielded.push_back(buildOutputSlice(iterArgs.front(), hSlice)); - return success(); - }); - if (failed(hLoop)) - return failure(); - paddedOutput = hLoop->results.front(); - } - - Value result = paddedOutput; - if (paddedOutType != outType) { - SmallVector outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}; - SmallVector outputSizes {rewriter.getIndexAttr(outType.getDimSize(0)), - rewriter.getIndexAttr(outType.getDimSize(1))}; - result = - tensor::ExtractSliceOp::create(rewriter, loc, outType, paddedOutput, outputOffsets, outputSizes, unitStrides) + SmallVector outputSlices; + outputSlices.reserve(numOutHSlices); + for (int64_t hSlice = 0; hSlice < numOutHSlices; ++hSlice) { + const int64_t columnOffset = hSlice * crossbarSize.getValue(); + const int64_t columns = + std::min(static_cast(crossbarSize.getValue()), outType.getDimSize(1) - columnOffset); + auto outputSliceType = RankedTensorType::get({numOutRows, columns}, outType.getElementType()); + auto computeOp = createSpatCompute( + rewriter, loc, TypeRange {outputSliceType}, {}, inputs, [&](ValueRange blockArgs) -> LogicalResult { + Value hSliceValue = + getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), hSlice); + Value reduced = reducePartialPiecesForHSlice( + blockArgs[0], hSliceValue, pieceType, numKSlices, numOutRows, rewriter, loc); + if (bias) { + SmallVector biasOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(columnOffset)}; + SmallVector pieceSizes {rewriter.getIndexAttr(numOutRows), + rewriter.getIndexAttr(crossbarSize.getValue())}; + SmallVector unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)}; + Value biasSlice = + tensor::ExtractSliceOp::create(rewriter, loc, pieceType, blockArgs[1], biasOffsets, pieceSizes, unitStrides) .getResult(); + reduced = spatial::SpatVAddOp::create(rewriter, loc, pieceType, reduced, biasSlice).getResult(); } - spatial::SpatYieldOp::create(rewriter, loc, result); + + if (outputSliceType != pieceType) { + SmallVector offsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}; + SmallVector sizes {rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(columns)}; + reduced = tensor::ExtractSliceOp::create( + rewriter, loc, outputSliceType, reduced, offsets, sizes, getUnitStrides(rewriter, 2)); + } + spatial::SpatYieldOp::create(rewriter, loc, reduced); return success(); }); + if (failed(computeOp)) + return failure(); + outputSlices.push_back(computeOp->getResult(0)); + } + if (outputSlices.size() == 1) + return outputSlices.front(); - return computeOp; + auto concatCompute = + createSpatCompute(rewriter, loc, TypeRange {outType}, {}, outputSlices, [&](ValueRange blockArgs) { + Value result = + spatial::SpatConcatOp::create(rewriter, loc, outType, rewriter.getI64IntegerAttr(1), blockArgs).getOutput(); + spatial::SpatYieldOp::create(rewriter, loc, result); + }); + return concatCompute.getResult(0); } struct GemmToSpatialComputes : OpConversionPattern { @@ -800,12 +780,12 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp, createVmmBatch(a, b, aType, paddedBType, partialPiecesType, numOutRows, numKSlices, numOutHSlices, rewriter, loc); if (failed(batchOp)) return failure(); - auto reductionCompute = createReductionCompute( + auto reductionOutput = createReductionOutput( batchOp->getResult(0), bias, partialPiecesType, outType, paddedOutType, numKSlices, rewriter, loc); - if (failed(reductionCompute)) + if (failed(reductionOutput)) return failure(); - rewriter.replaceOp(gemmOp, reductionCompute->getResults()); + rewriter.replaceOp(gemmOp, *reductionOutput); return success(); } diff --git a/validation/networks/pimcomp_models/results.csv b/validation/networks/pimcomp_models/results.csv index 440abe6..878a2a2 100644 --- a/validation/networks/pimcomp_models/results.csv +++ b/validation/networks/pimcomp_models/results.csv @@ -1,5 +1,5 @@ model,raptor_latency_ms,pimcomp_latency_ms,raptor_energy_pj,pimcomp_energy_pj,faster_compiler,speedup -vgg8,2.451664,7.985074,647608825.120001,1597904071.120000,raptor,3.26 -resnet18,113.483357,58.853733,23204566467.119949,13983168748.119974,pimcomp,1.93 -resnet34,120.655292,91.607980,27614808203.679939,22722922369.680016,pimcomp,1.32 -googlenet,22.936609,62.923463,7798109358.239990,14547526780.240000,raptor,2.74 +vgg8,2.463417,7.985074,649524564.120001,1597904071.120000,raptor,3.24 +resnet18,31.449811,58.853733,9335506856.119984,13983168748.119974,raptor,1.87 +resnet34,58.071522,91.607980,17027195415.679953,22722922369.680016,raptor,1.58 +googlenet,24.247113,62.923463,8019605936.239990,14547526780.240000,raptor,2.60