Fix conv_relu_conv diamond shape

2026-07-06 11:22:39 +02:00
parent 83a54e28e4
commit 2bfc033af9
1 changed files with 192 additions and 127 deletions
@@ -2716,6 +2716,181 @@ static FailureOr<Value> createNchwRowStripConvPatchRow(Value paddedWindow,
    .getResult();
 }

+static FailureOr<Value> createPaddedConvOutputRow(Value patchRow,
+                                                  const ConvLoweringState& state,
+                                                  Value paddedWeights,
+                                                  Value paddedBias,
+                                                  int64_t paddedK,
+                                                  int64_t numKSlices,
+                                                  int64_t xbarDim,
+                                                  PatternRewriter& rewriter,
+                                                  Location loc) {
+  const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
+  auto elementType = state.outType.getElementType();
+  auto rowType = RankedTensorType::get({1, state.numChannelsOut}, elementType);
+  auto paddedRowType = RankedTensorType::get({1, xbarDim}, elementType);
+  auto paddedPatchRowType = RankedTensorType::get({1, paddedK}, elementType);
+  auto 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<OpFoldResult> aOffsets {rewriter.getIndexAttr(0), kOffset};
+    SmallVector<OpFoldResult> aSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)};
+    Value aTile = extractStaticSliceOrIdentity(
+      rewriter, pieceLoc, paddedPatchRow, paddedRowType, aOffsets, aSizes, getUnitStrides(rewriter, 2));
+    SmallVector<OpFoldResult> bOffsets {kOffset, rewriter.getIndexAttr(0)};
+    SmallVector<OpFoldResult> 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<Value>& reduceYielded) {
+        Value piece = createPiece(kSlice, reduceLoc);
+        reduceYielded.push_back(
+          spatial::SpatVAddOp::create(rewriter, reduceLoc, paddedRowType, reduceIterArgs.front(), piece).getResult());
+        return success();
+      });
+    if (failed(kLoop))
+      return failure();
+    rowResult = kLoop->results.front();
+  }
+
+  if (paddedBias)
+    rowResult = spatial::SpatVAddOp::create(rewriter, loc, paddedRowType, rowResult, paddedBias).getResult();
+  if (state.numChannelsOut == xbarDim)
+    return rowResult;
+
+  SmallVector<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
+  SmallVector<OpFoldResult> outputSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut)};
+  return tensor::ExtractSliceOp::create(
+           rewriter, loc, rowType, rowResult, outputOffsets, outputSizes, getUnitStrides(rewriter, 2))
+    .getResult();
+}
+
+static FailureOr<Value>
+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<Value> 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<Value>& 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<ReassociationIndices> {{0}, {1, 2, 3}});
+          FailureOr<Value> 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<ReassociationIndices> {{0}, {1, 2, 3}});
+          SmallVector<OpFoldResult> rowOffsets {
+            rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex};
+          SmallVector<OpFoldResult> 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<Value> createConvOutputFromNchwRowStripFragments(Value rowStripStorage,
                                                                  const ConvLoweringState& state,
                                                                  PatternRewriter& rewriter,
@@ -2734,11 +2909,7 @@ static FailureOr<Value> createConvOutputFromNchwRowStripFragments(Value rowStrip
  const int64_t numKSlices = ceilIntegerDivide(patchSize, xbarDim);
  const int64_t paddedK = numKSlices * xbarDim;
  auto elementType = state.outType.getElementType();
-  auto 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<Value> 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<Value> inputWindow = createNchwRowStripConvWindow(args.inputs.front(), state, args.lane, rewriter, loc);
      if (failed(inputWindow))
@@ -2782,58 +2951,22 @@ static FailureOr<Value> 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<Value>& reduceYielded) {
-              Value kOffset = arith::MulIOp::create(rewriter, reduceLoc, kSlice, cXbar);
-              SmallVector<OpFoldResult> aOffsets {rewriter.getIndexAttr(0), kOffset};
-              SmallVector<OpFoldResult> aSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)};
-              Value aTile = tensor::ExtractSliceOp::create(
-                rewriter, reduceLoc, paddedRowType, paddedRow, aOffsets, aSizes, getUnitStrides(rewriter, 2));
-              SmallVector<OpFoldResult> bOffsets {kOffset, rewriter.getIndexAttr(0)};
-              SmallVector<OpFoldResult> 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<Value> 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<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
-            SmallVector<OpFoldResult> 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<ReassociationIndices> {{0}, {1, 2, 3}});
          SmallVector<OpFoldResult> 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<Value> 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<bool>(wDenseAttr), !state.hasBias || static_cast<bool>(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<bool>(wDenseAttr), !state.hasBias || static_cast<bool>(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<DistributedTensorInfo> 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<Value> rows = createConvRowsForStrategy(rowState, decision, rewriter, planOp.getLoc());
-    if (failed(rows))
-      return failure();
-    FailureOr<Value> rowStripStorage = createRowStripStorageFromRows(*rows, state->outType, rewriter, planOp.getLoc());
+    FailureOr<Value> 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();
    }