From 2bfc033af97cfc7b44a19cf1801393993a6b4ef7 Mon Sep 17 00:00:00 2001 From: ilgeco Date: Mon, 6 Jul 2026 11:22:39 +0200 Subject: [PATCH] Fix conv_relu_conv diamond shape --- .../ONNXToSpatial/Patterns/Math/Conv.cpp | 319 +++++++++++------- 1 file changed, 192 insertions(+), 127 deletions(-) diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp index 3f142b9..dc2995e 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp @@ -2716,6 +2716,181 @@ static FailureOr createNchwRowStripConvPatchRow(Value paddedWindow, .getResult(); } +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 paddedWeightTileType = RankedTensorType::get({xbarDim, xbarDim}, state.wType.getElementType()); + + Value paddedPatchRow = patchRow; + if (patchSize != paddedK) + paddedPatchRow = createZeroPaddedTensor( + paddedPatchRow, paddedPatchRowType, {0, 0}, {0, paddedK - patchSize}, rewriter, loc); + + 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)}; + SmallVector bSizes {rewriter.getIndexAttr(xbarDim), rewriter.getIndexAttr(xbarDim)}; + Value bTile = extractStaticSliceOrIdentity( + rewriter, pieceLoc, paddedWeights, paddedWeightTileType, bOffsets, bSizes, getUnitStrides(rewriter, 2)); + return spatial::SpatVMMOp::create(rewriter, pieceLoc, paddedRowType, bTile, aTile).getResult(); + }; + + Value rowResult = createPiece(c0, loc); + if (numKSlices > 1) { + auto kLoop = buildNormalizedScfFor( + rewriter, + loc, + c1, + cNumKSlices, + c1, + ValueRange {rowResult}, + [&](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(); + rowResult = kLoop->results.front(); + } + + 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(); +} + +static FailureOr +createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRewriter& rewriter, Location loc) { + ConvGeometry geometry = buildConvGeometry(state); + if (state.group != 1 || state.batchSize != 1 || geometry.c > geometry.xbarSize) + return failure(); + + auto weightDenseAttr = getHostConstDenseElementsAttr(state.w); + if (!weightDenseAttr) + return failure(); + if (state.hasBias && !isSupportedBiasAddValue(state.b, state.outType)) + return failure(); + + const int64_t xbarDim = geometry.xbarSize; + const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth; + const int64_t numKSlices = ceilIntegerDivide(patchSize, xbarDim); + const int64_t paddedK = numKSlices * xbarDim; + auto elementType = state.outType.getElementType(); + auto fragmentType = getRowStripFragmentType(state.outType); + auto outputPixelType = RankedTensorType::get({1, state.numChannelsOut, 1, 1}, elementType); + auto patchType = RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.wWidth}, state.xType.getElementType()); + auto patchRowType = RankedTensorType::get({1, patchSize}, state.xType.getElementType()); + auto outputStorageType = getRowStripStorageType(state.outType); + + PreparedConvInput preparedInput = standard::prepareInputForIm2Col(state, rewriter, loc); + Value paddedWeights = standard::createPaddedInputKTiledWeightConstant(weightDenseAttr, state, paddedK, xbarDim, rewriter); + FailureOr paddedBias = failure(); + if (state.hasBias) + paddedBias = createPaddedBiasRowConstant(state, xbarDim, rewriter); + if (state.hasBias && failed(paddedBias)) + return failure(); + + auto batchOp = createSpatComputeBatch( + rewriter, + loc, + TypeRange {outputStorageType}, + state.outHeight, + ValueRange {paddedWeights}, + state.hasBias ? ValueRange {preparedInput.value, *paddedBias} : ValueRange {preparedInput.value}, + [&](detail::SpatComputeBatchBodyArgs args) { + Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); + Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0); + Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1); + Value cOutWidth = getOrCreateIndexConstant(rewriter, anchorOp, state.outWidth); + Value inputHeightOffset = affineMulConst(rewriter, loc, args.lane, state.strideHeight, anchorOp); + Value fragmentInit = tensor::EmptyOp::create(rewriter, loc, fragmentType.getShape(), elementType); + auto widthLoop = buildNormalizedScfFor( + rewriter, + loc, + c0, + cOutWidth, + c1, + ValueRange {fragmentInit}, + [&](OpBuilder&, Location widthLoc, Value widthIndex, ValueRange widthIterArgs, SmallVectorImpl& widthYielded) { + Value inputWidthOffset = affineMulConst(rewriter, widthLoc, widthIndex, state.strideWidth, anchorOp); + Value patch = createConvInputPatch(args.inputs.front(), + patchType, + c0, + c0, + inputHeightOffset, + inputWidthOffset, + state.dilationHeight, + state.dilationWidth, + rewriter, + widthLoc); + Value patchRow = tensor::CollapseShapeOp::create( + rewriter, widthLoc, patchRowType, patch, SmallVector {{0}, {1, 2, 3}}); + FailureOr outputRow = createPaddedConvOutputRow(patchRow, + state, + args.weights.front(), + state.hasBias ? args.inputs[1] : Value(), + paddedK, + numKSlices, + xbarDim, + rewriter, + widthLoc); + if (failed(outputRow)) + return failure(); + + Value outputFragment = tensor::ExpandShapeOp::create(rewriter, + widthLoc, + outputPixelType, + *outputRow, + SmallVector {{0}, {1, 2, 3}}); + SmallVector rowOffsets { + rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex}; + SmallVector rowSizes { + rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut), rewriter.getIndexAttr(1), + rewriter.getIndexAttr(1)}; + Value nextFragment = tensor::InsertSliceOp::create( + rewriter, widthLoc, outputFragment, widthIterArgs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 4)); + widthYielded.push_back(nextFragment); + return success(); + }); + if (failed(widthLoop)) + return failure(); + + insertRowStripFragment(widthLoop->results.front(), args.outputs.front(), state.outType, args.lane, rewriter, loc); + return success(); + }); + if (failed(batchOp)) + return failure(); + return batchOp->getResult(0); +} + static FailureOr createConvOutputFromNchwRowStripFragments(Value rowStripStorage, const ConvLoweringState& state, PatternRewriter& rewriter, @@ -2734,11 +2909,7 @@ static FailureOr createConvOutputFromNchwRowStripFragments(Value rowStrip const int64_t numKSlices = ceilIntegerDivide(patchSize, xbarDim); const int64_t paddedK = numKSlices * xbarDim; auto elementType = state.outType.getElementType(); - auto rowType = RankedTensorType::get({1, state.numChannelsOut}, state.outType.getElementType()); auto outputPixelType = RankedTensorType::get({1, state.numChannelsOut, 1, 1}, elementType); - auto paddedRowType = RankedTensorType::get({1, xbarDim}, state.outType.getElementType()); - auto paddedPatchRowType = RankedTensorType::get({1, paddedK}, elementType, inputType.getEncoding()); - auto paddedWeightTileType = RankedTensorType::get({xbarDim, xbarDim}, state.wType.getElementType()); auto outputStorageType = getRowStripStorageType(state.outType); auto weightDenseAttr = getHostConstDenseElementsAttr(state.w); if (!weightDenseAttr) @@ -2761,9 +2932,7 @@ static FailureOr createConvOutputFromNchwRowStripFragments(Value rowStrip Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0); Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1); - Value cNumKSlices = getOrCreateIndexConstant(rewriter, anchorOp, numKSlices); Value cOutWidth = getOrCreateIndexConstant(rewriter, anchorOp, state.outWidth); - Value cXbar = getOrCreateIndexConstant(rewriter, anchorOp, xbarDim); auto fragmentType = getRowStripFragmentType(state.outType); FailureOr inputWindow = createNchwRowStripConvWindow(args.inputs.front(), state, args.lane, rewriter, loc); if (failed(inputWindow)) @@ -2782,58 +2951,22 @@ static FailureOr createConvOutputFromNchwRowStripFragments(Value rowStrip if (failed(patchRow)) return failure(); - Value paddedRow = *patchRow; - if (patchSize != paddedK) - paddedRow = createZeroPaddedTensor( - paddedRow, paddedPatchRowType, {0, 0}, {0, paddedK - patchSize}, rewriter, widthLoc); - - Value zeroRow = createZeroTensorConstant(paddedRowType, rewriter); - auto kLoop = buildNormalizedScfFor( - rewriter, - widthLoc, - c0, - cNumKSlices, - c1, - ValueRange {zeroRow}, - [&](OpBuilder&, Location reduceLoc, Value kSlice, ValueRange reduceIterArgs, SmallVectorImpl& reduceYielded) { - Value kOffset = arith::MulIOp::create(rewriter, reduceLoc, kSlice, cXbar); - SmallVector aOffsets {rewriter.getIndexAttr(0), kOffset}; - SmallVector aSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)}; - Value aTile = tensor::ExtractSliceOp::create( - rewriter, reduceLoc, paddedRowType, paddedRow, aOffsets, aSizes, getUnitStrides(rewriter, 2)); - SmallVector bOffsets {kOffset, rewriter.getIndexAttr(0)}; - SmallVector bSizes {rewriter.getIndexAttr(xbarDim), rewriter.getIndexAttr(xbarDim)}; - Value bTile = extractStaticSliceOrIdentity(rewriter, - reduceLoc, - args.weights.front(), - paddedWeightTileType, - bOffsets, - bSizes, - getUnitStrides(rewriter, 2)); - Value piece = spatial::SpatVMMOp::create(rewriter, reduceLoc, paddedRowType, bTile, aTile).getResult(); - reduceYielded.push_back( - spatial::SpatVAddOp::create(rewriter, reduceLoc, paddedRowType, reduceIterArgs.front(), piece).getResult()); - return success(); - }); - if (failed(kLoop)) + FailureOr outputRow = createPaddedConvOutputRow(*patchRow, + state, + args.weights.front(), + state.hasBias ? args.inputs[1] : Value(), + paddedK, + numKSlices, + xbarDim, + rewriter, + widthLoc); + if (failed(outputRow)) return failure(); - Value rowResult = kLoop->results.front(); - if (state.hasBias) - rowResult = spatial::SpatVAddOp::create(rewriter, widthLoc, paddedRowType, rowResult, args.inputs[1]).getResult(); - Value outputRow = rowResult; - if (state.numChannelsOut != xbarDim) { - SmallVector outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}; - SmallVector outputSizes { - rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut)}; - outputRow = tensor::ExtractSliceOp::create( - rewriter, widthLoc, rowType, rowResult, outputOffsets, outputSizes, getUnitStrides(rewriter, 2)); - } - Value outputFragment = tensor::ExpandShapeOp::create(rewriter, widthLoc, outputPixelType, - outputRow, + *outputRow, SmallVector {{0}, {1, 2, 3}}); SmallVector rowOffsets { rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex}; @@ -2915,72 +3048,6 @@ static Value createFragmentReciprocalConstant(const DistributedTensorStep& step, fragmentType); } -[[maybe_unused]] static FailureOr createConvRowsForStrategy(const ConvLoweringState& state, - const ConvLoweringDecision& decision, - PatternRewriter& rewriter, - Location loc) { - auto wDenseAttr = getHostConstDenseElementsAttr(state.w); - PreparedConvInput preparedInput = standard::prepareInputForIm2Col(state, rewriter, loc); - Value biasMatrix; - DenseElementsAttr biasDenseAttr; - if (state.hasBias) { - biasDenseAttr = getHostConstDenseElementsAttr(state.b); - biasMatrix = expandBiasIfNeeded(state.b, rewriter, loc); - } - - switch (decision.strategy) { - case PimConvLoweringLegacy: - case PimConvLoweringPackedIm2Col: { - standard::ConvGemmPlan plan = standard::buildConvGemmPlan( - state, static_cast(wDenseAttr), !state.hasBias || static_cast(biasDenseAttr), 0, - state.batchSize * state.outHeight * state.outWidth); - Value weightMatrix = standard::createWeightMatrix(state.w, plan, rewriter, loc); - Value gemmInputRows = standard::createIm2colRows(state, preparedInput, plan, rewriter, loc); - Value gemmB = standard::buildPackedWeights(wDenseAttr, weightMatrix, state, plan, rewriter, loc); - Value gemmBias = createZeroGemmBias(plan.gemmOutputRowsType, rewriter); - if (state.hasBias) - gemmBias = state.b; - Value gemmC = standard::buildPackedBias(gemmBias, biasMatrix, biasDenseAttr, state, plan, rewriter, loc); - Value gemmRows = ONNXGemmOp::create(rewriter, - loc, - plan.gemmOutputRowsType, - gemmInputRows, - gemmB, - gemmC, - APFloat(1.0f), - APFloat(1.0f), - /*transA=*/0, - /*transB=*/0) - .getY(); - return standard::maybeUnpackChunkRows(gemmRows, plan, rewriter, loc); - } - case PimConvLoweringStreamedPatch: - case PimConvLoweringOutputChannelTiled: - case PimConvLoweringTiled2D: - case PimConvLoweringStreamedPacked: { - standard::ConvGemmPlan seedPlan = standard::buildConvGemmPlan( - state, static_cast(wDenseAttr), !state.hasBias || static_cast(biasDenseAttr), 0, 1, - decision.strategy == PimConvLoweringStreamedPacked ? buildConvGeometry(state).pack : 1); - Value weightMatrix = standard::createWeightMatrix(state.w, seedPlan, rewriter, loc); - ConvGeometry geo = buildConvGeometry(state); - int64_t packFactor = decision.strategy == PimConvLoweringStreamedPacked ? geo.pack : 1; - uint64_t chunkPositions = chooseStreamChunkPositions(geo, packFactor); - return standard::createChunkedConvRows(state, - preparedInput, - weightMatrix, - biasMatrix, - wDenseAttr, - biasDenseAttr, - packFactor, - chunkPositions, - rewriter, - loc); - } - default: - return failure(); - } -} - [[maybe_unused]] static FailureOr applyDistributedPreservingStep(const DistributedTensorInfo& inputInfo, const DistributedTensorStep& step, PatternRewriter& rewriter, @@ -3741,6 +3808,8 @@ LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp) { return failure(); if (state->outType.getRank() != 4 || !state->outType.hasStaticShape()) return failure(); + if (!getHostConstDenseElementsAttr(state->w)) + return failure(); if (state->hasBias && !isSupportedBiasAddValue(state->b, state->outType)) return failure(); @@ -3752,6 +3821,8 @@ LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp) { analysis.barrierKind = DistributedConvBarrierKind::UnsupportedConsumer; analysis.barrierDetail = "selected row-strip layout"; ConvGeometry geometry = buildConvGeometry(*state); + if (geometry.c > geometry.xbarSize) + return failure(); ConvLoweringDecision decision = chooseConvLoweringStrategy(geometry, *requestedStrategy, analysis); if (decision.strategy == PimConvLoweringDepthwise && !depthwise::canUseStructuredRewrite(*state) && *requestedStrategy == PimConvLoweringAuto) { @@ -3830,21 +3901,15 @@ lowerSelectedConv2DPlan(spatial::SpatConv2DPlanOp planOp, const bool applyBiasAfterStorage = rowState.hasBias; Value originalBias = rowState.b; if (applyBiasAfterStorage) { - if (!isSupportedBiasAddValue(originalBias, rowState.outType)) - return planOp.emitOpError("selected row-strip Conv bias must be host-constant scalar/per-channel NCHW"), - failure(); rowState.b = Value(); rowState.hasBias = false; } - FailureOr rows = createConvRowsForStrategy(rowState, decision, rewriter, planOp.getLoc()); - if (failed(rows)) - return failure(); - FailureOr rowStripStorage = createRowStripStorageFromRows(*rows, state->outType, rewriter, planOp.getLoc()); + FailureOr rowStripStorage = createRowStripConvOutputFromDenseInput(rowState, rewriter, planOp.getLoc()); if (failed(rowStripStorage)) return planOp.emitOpError("failed to build row-strip fragment storage for the selected Conv plan"), failure(); if (applyBiasAfterStorage) { - rowStripStorage = applyRowStripBiasAdd(*rowStripStorage, rowState.outType, originalBias, rewriter, planOp.getLoc()); + rowStripStorage = applyRowStripBiasAdd(*rowStripStorage, state->outType, originalBias, rewriter, planOp.getLoc()); if (failed(rowStripStorage)) return planOp.emitOpError("failed to apply row-strip Conv bias per fragment"), failure(); }