meno diamantini

2026-07-06 10:12:20 +02:00
parent cc9b025a35
commit 83a54e28e4
26 changed files with 3756 additions and 1278 deletions
@@ -23,7 +23,9 @@
 #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
 #include "src/Accelerators/PIM/Common/Support/ReportUtils.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns/Math/ConvGeometry.hpp"
@@ -221,10 +223,6 @@ struct DistributedConvReportTotals {
 };

 static Value createZeroGemmBias(RankedTensorType gemmResultType, PatternRewriter& rewriter);
-static FailureOr<Value> createRowStripPackedRows(Value rows,
-                                                 const ConvLoweringState& state,
-                                                 PatternRewriter& rewriter,
-                                                 Location loc);

 static FailureOr<ConvLoweringState> analyzeConvLoweringState(ONNXConvOp convOp, Value x, Value w, Value b);

@@ -467,59 +465,6 @@ classifyDistributedBinaryConsumer(Operation* user,
  return std::nullopt;
 }

-static bool covers(RowInterval acquired, RowInterval needed) {
-  return acquired.begin <= needed.begin && acquired.end >= needed.end;
-}
-
-static bool canConsumeRowStripHwcInput(const ConvLoweringState& state, StringRef& failureReason) {
-  if (state.batchSize != 1) {
-    failureReason = "unsupported_batch";
-    return false;
-  }
-  if (state.group != 1) {
-    failureReason = "unsupported_groups";
-    return false;
-  }
-  if (isDepthwiseConv(state.group, state.numChannelsIn, state.numChannelsOut, state.numChannelsInPerGroup)) {
-    failureReason = "unsupported_depthwise";
-    return false;
-  }
-  if (state.strideHeight != 1 || state.strideWidth != 1) {
-    failureReason = "unsupported_stride";
-    return false;
-  }
-  if (state.dilationHeight != 1 || state.dilationWidth != 1) {
-    failureReason = "unsupported_dilation";
-    return false;
-  }
-  if (state.padHeightBegin != state.padHeightEnd || state.padWidthBegin != state.padWidthEnd) {
-    failureReason = "unsupported_padding";
-    return false;
-  }
-  if (state.padHeightBegin != 1 || state.padWidthBegin != 1) {
-    failureReason = "unsupported_padding";
-    return false;
-  }
-  if (state.wHeight != 3 || state.wWidth != 3) {
-    failureReason = "unsupported_kernel";
-    return false;
-  }
-  if (state.outHeight != state.xHeight || state.outWidth != state.xWidth) {
-    failureReason = "unsupported_output_shape";
-    return false;
-  }
-  if (!getHostConstDenseElementsAttr(state.w)) {
-    failureReason = "non_constant_weights";
-    return false;
-  }
-  if (state.hasBias && !getHostConstDenseElementsAttr(state.b)) {
-    failureReason = "non_constant_bias";
-    return false;
-  }
-  failureReason = "";
-  return true;
-}
-
 static std::string stringifyDistributedTensorOpKind(DistributedTensorOpKind kind) {
  switch (kind) {
  case DistributedTensorOpKind::Relu: return "Relu";
@@ -2493,21 +2438,15 @@ static Value rewriteStreamedConv(const ConvLoweringState& state,

 } // namespace standard

-static RankedTensorType getRowStripFragmentType(RankedTensorType tensorType, int64_t width) {
-  return RankedTensorType::get(
-    {tensorType.getDimSize(0), tensorType.getDimSize(1), 1, width}, tensorType.getElementType(), tensorType.getEncoding());
-}
-
 static SmallVector<DistributedFragmentInfo, 8> buildRowStripFragments(RankedTensorType tensorType) {
  SmallVector<DistributedFragmentInfo, 8> fragments;
-  const int64_t height = tensorType.getDimSize(2);
-  const int64_t width = tensorType.getDimSize(3);
-  const int64_t channels = tensorType.getDimSize(1);
-  fragments.reserve(height);
-  for (int64_t row = 0; row < height; ++row) {
+  auto [offsets, sizes] = buildRowStripMetadata(tensorType);
+  const int64_t rank = tensorType.getRank();
+  fragments.reserve(offsets.size() / rank);
+  for (int64_t row = 0; row < static_cast<int64_t>(offsets.size() / rank); ++row) {
    fragments.push_back(DistributedFragmentInfo {
-      {0, 0, row, 0},
-      {1, channels, 1, width},
+      {offsets.begin() + row * rank, offsets.begin() + (row + 1) * rank},
+      {sizes.begin() + row * rank, sizes.begin() + (row + 1) * rank},
      {1, 1, 1, 1},
      row,
    });
@@ -2527,101 +2466,266 @@ static DistributedTensorInfo makeDistributedTensorInfo(Value storage, RankedTens
  return info;
 }

-static Value createPerChannelConstantFragment(DenseElementsAttr denseAttr,
-                                              RankedTensorType fragmentType,
-                                              PatternRewriter& rewriter) {
-  auto denseType = cast<RankedTensorType>(denseAttr.getType());
-  SmallVector<Attribute> channelValues;
-  channelValues.reserve(fragmentType.getDimSize(1));
-  SmallVector<Attribute> flattened(denseAttr.getValues<Attribute>());
-  if (denseType.getRank() == 1) {
-    channelValues = flattened;
-  }
-  else if (denseType.getRank() == 2) {
-    channelValues = flattened;
-  }
-  else {
-    for (int64_t channel = 0; channel < denseType.getDimSize(1); ++channel)
-      channelValues.push_back(flattened[channel]);
-  }
-
-  SmallVector<Attribute> values;
-  values.reserve(fragmentType.getNumElements());
-  for (int64_t n = 0; n < fragmentType.getDimSize(0); ++n)
-    for (int64_t channel = 0; channel < fragmentType.getDimSize(1); ++channel)
-      for (int64_t h = 0; h < fragmentType.getDimSize(2); ++h)
-        for (int64_t w = 0; w < fragmentType.getDimSize(3); ++w)
-          values.push_back(channelValues[channel]);
-
-  auto attr = DenseElementsAttr::get(fragmentType, values);
-  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), attr, fragmentType);
-}
-
 static Value createZeroGemmBias(RankedTensorType gemmResultType, PatternRewriter& rewriter) {
  auto zeroAttr = DenseElementsAttr::get(gemmResultType, rewriter.getZeroAttr(gemmResultType.getElementType()));
  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), zeroAttr, gemmResultType);
 }

-static bool canDirectLowerRowStripConv(const ConvLoweringState& state, StringRef& failureReason) {
-  if (!canConsumeRowStripHwcInput(state, failureReason))
-    return false;
-
-  ConvGeometry geometry = buildConvGeometry(state);
-  if (state.numChannelsOut > geometry.xbarSize) {
-    failureReason = "unsupported_output_channels";
+static bool canConsumeNchwRowStripFragments(const ConvLoweringState& state, StringRef& failureReason) {
+  if (state.batchSize != 1) {
+    failureReason = "batch_not_one";
+    return false;
+  }
+  if (state.group != 1) {
+    failureReason = "grouped_conv";
+    return false;
+  }
+  if (!state.xType.hasStaticShape() || !state.wType.hasStaticShape() || !state.outType.hasStaticShape()) {
+    failureReason = "dynamic_shape";
+    return false;
+  }
+  if (!isa<FloatType>(state.xType.getElementType())) {
+    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.wHeight != 3 || state.wWidth != 3) {
+    failureReason = "kernel_not_3x3";
+    return false;
+  }
+  if (state.padHeightBegin != 1 || state.padHeightEnd != 1 || state.padWidthBegin != 1 || state.padWidthEnd != 1) {
+    failureReason = "padding_not_1";
+    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 (state.hasBias && !isSupportedBiasAddValue(state.b, state.outType)) {
+    failureReason = "unsupported_bias";
+    return false;
+  }
+  ConvGeometry geometry = buildConvGeometry(state);
+  if (geometry.c > geometry.xbarSize) {
+    failureReason = "output_channels_exceed_crossbar";
    return false;
  }
-
-  failureReason = "";
  return true;
 }

-static FailureOr<Value> createRowStripPackedRows(Value rows,
-                                                 const ConvLoweringState& state,
-                                                 PatternRewriter& rewriter,
-                                                 Location loc) {
-  auto rowsType = dyn_cast<RankedTensorType>(rows.getType());
-  if (!rowsType || !rowsType.hasStaticShape() || rowsType.getRank() != 2)
-    return failure();
-
-  if (state.batchSize != 1)
-    return failure();
-  if (state.outType.getRank() != 4 || !state.outType.hasStaticShape())
-    return failure();
-
-  const int64_t outHeight = state.outType.getDimSize(2);
-  const int64_t outWidth = state.outType.getDimSize(3);
-  const int64_t outChannels = state.outType.getDimSize(1);
-  if (rowsType.getDimSize(0) != outHeight * outWidth || rowsType.getDimSize(1) != outChannels)
-    return failure();
-
-  auto packedType = RankedTensorType::get({outHeight, outWidth, outChannels}, rowsType.getElementType(), rowsType.getEncoding());
-  auto packedRows =
-    createSpatCompute<1>(rewriter, loc, TypeRange {packedType}, {}, rows, [&](Value rowValues) {
-      Value packed = tensor::ExpandShapeOp::create(
-        rewriter, loc, packedType, rowValues, SmallVector<ReassociationIndices> {{0, 1}, {2}});
-      spatial::SpatYieldOp::create(rewriter, loc, packed);
-    });
-  return packedRows.getResult(0);
+static Value createZeroTensorConstant(RankedTensorType type, PatternRewriter& rewriter) {
+  auto zeroAttr = DenseElementsAttr::get(type, rewriter.getZeroAttr(type.getElementType()));
+  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), zeroAttr, type);
 }

-static FailureOr<Value> createConvOutputFromRowStripHwc(Value inputHwc,
-                                                        const ConvLoweringState& state,
-                                                        PatternRewriter& rewriter,
-                                                        Location loc) {
-  auto inputType = dyn_cast<RankedTensorType>(inputHwc.getType());
-  if (!inputType || !inputType.hasStaticShape() || inputType.getRank() != 3)
+static FailureOr<Value> createPaddedBiasRowConstant(const ConvLoweringState& state,
+                                                    int64_t paddedChannels,
+                                                    PatternRewriter& rewriter) {
+  DenseElementsAttr denseAttr;
+  if (!isSupportedBiasAddValue(state.b, state.outType, &denseAttr))
    return failure();
-  if (inputType.getDimSize(0) != state.xHeight || inputType.getDimSize(1) != state.xWidth
-      || inputType.getDimSize(2) != state.numChannelsIn)
+  FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, state.outType);
+  if (failed(channelValues))
+    return failure();
+
+  auto biasType = RankedTensorType::get({1, paddedChannels}, state.outType.getElementType());
+  SmallVector<Attribute> values(biasType.getNumElements(), cast<Attribute>(rewriter.getZeroAttr(biasType.getElementType())));
+  for (int64_t channel = 0; channel < state.numChannelsOut; ++channel)
+    values[channel] = (*channelValues)[channel];
+  auto biasAttr = DenseElementsAttr::get(biasType, values);
+  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), biasAttr, biasType);
+}
+
+static Value createHorizontallyPaddedRowStripFragment(Value fragment,
+                                                      const ConvLoweringState& state,
+                                                      PatternRewriter& rewriter,
+                                                      Location loc) {
+  auto paddedType = RankedTensorType::get({1, state.numChannelsIn, 1, state.xWidth + 2},
+                                          state.xType.getElementType(),
+                                          state.xType.getEncoding());
+  return createZeroPaddedTensor(fragment, paddedType, {0, 0, 0, 1}, {0, 0, 0, 1}, rewriter, loc);
+}
+
+static Value createRowStripWindowSourceRowTable(const ConvLoweringState& state, PatternRewriter& rewriter) {
+  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+  auto tableType = RankedTensorType::get({state.outHeight * state.wHeight}, rewriter.getIndexType());
+  SmallVector<Attribute> values;
+  values.reserve(tableType.getNumElements());
+  for (int64_t outputRow = 0; outputRow < state.outHeight; ++outputRow) {
+    for (int64_t kernelRow = 0; kernelRow < state.wHeight; ++kernelRow) {
+      int64_t sourceRow = outputRow + kernelRow - state.padHeightBegin;
+      sourceRow = std::clamp(sourceRow, int64_t {0}, state.xHeight - 1);
+      values.push_back(rewriter.getIndexAttr(sourceRow));
+    }
+  }
+
+  return getOrCreateConstant(rewriter, anchorOp, DenseElementsAttr::get(tableType, values), tableType);
+}
+
+static Value createRowStripWindowTableIndex(Value outputHeight,
+                                            Value kernelRow,
+                                            const ConvLoweringState& state,
+                                            PatternRewriter& rewriter,
+                                            Location loc) {
+  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+  MLIRContext* ctx = rewriter.getContext();
+  AffineExpr outputRowExpr = getAffineDimExpr(0, ctx);
+  AffineExpr kernelRowExpr = getAffineDimExpr(1, ctx);
+  return createOrFoldAffineApply(
+    rewriter, loc, outputRowExpr * state.wHeight + kernelRowExpr, ValueRange {outputHeight, kernelRow}, anchorOp);
+}
+
+static Value extractProjectedRowStripWindowRow(Value rowStripStorage,
+                                               Value sourceRowTable,
+                                               const ConvLoweringState& state,
+                                               Value outputHeight,
+                                               Value kernelRow,
+                                               PatternRewriter& rewriter,
+                                               Location loc) {
+  Value tableIndex = createRowStripWindowTableIndex(outputHeight, kernelRow, state, rewriter, loc);
+  Value sourceRow = tensor::ExtractOp::create(rewriter, loc, sourceRowTable, ValueRange {tableIndex}).getResult();
+  return extractRowStripFragment(rowStripStorage, state.xType, sourceRow, rewriter, loc);
+}
+
+static FailureOr<Value> createRowStripWindowMaskTable(const ConvLoweringState& state, PatternRewriter& rewriter) {
+  auto elementType = state.xType.getElementType();
+  auto floatType = dyn_cast<FloatType>(elementType);
+  if (!floatType)
+    return failure();
+
+  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+  auto tableType = RankedTensorType::get({state.outHeight * state.wHeight, state.numChannelsIn, 1, state.xWidth},
+                                         elementType,
+                                         state.xType.getEncoding());
+  Attribute zero = rewriter.getZeroAttr(elementType);
+  Attribute one = rewriter.getFloatAttr(floatType, 1.0);
+  SmallVector<Attribute> values;
+  values.reserve(tableType.getNumElements());
+  for (int64_t outputRow = 0; outputRow < state.outHeight; ++outputRow) {
+    for (int64_t kernelRow = 0; kernelRow < state.wHeight; ++kernelRow) {
+      int64_t sourceRow = outputRow + kernelRow - 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)
+          values.push_back(value);
+    }
+  }
+
+  return getOrCreateConstant(rewriter, anchorOp, DenseElementsAttr::get(tableType, values), tableType);
+}
+
+static Value extractProjectedRowStripWindowMask(Value maskTable,
+                                                const ConvLoweringState& state,
+                                                Value outputHeight,
+                                                Value kernelRow,
+                                                PatternRewriter& rewriter,
+                                                Location loc) {
+  Value tableIndex = createRowStripWindowTableIndex(outputHeight, kernelRow, state, rewriter, loc);
+  auto fragmentType = getRowStripFragmentType(state.xType);
+  SmallVector<OpFoldResult> offsets {
+    tableIndex, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
+  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
+                                   rewriter.getIndexAttr(state.numChannelsIn),
+                                   rewriter.getIndexAttr(1),
+                                   rewriter.getIndexAttr(state.xWidth)};
+  return tensor::ExtractSliceOp::create(rewriter,
+                                        loc,
+                                        fragmentType,
+                                        maskTable,
+                                        offsets,
+                                        sizes,
+                                        getUnitStrides(rewriter, 4));
+}
+
+static FailureOr<Value> createNchwRowStripConvWindow(Value rowStripStorage,
+                                                     const ConvLoweringState& state,
+                                                     Value outputHeight,
+                                                     PatternRewriter& rewriter,
+                                                     Location loc) {
+  auto fragmentType = getRowStripFragmentType(state.xType);
+  auto paddedWindowType = RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.xWidth + 2},
+                                                state.xType.getElementType(),
+                                                state.xType.getEncoding());
+  Value sourceRowTable = createRowStripWindowSourceRowTable(state, rewriter);
+  FailureOr<Value> maskTable = createRowStripWindowMaskTable(state, rewriter);
+  if (failed(maskTable))
+    return failure();
+  Value initWindow = createZeroTensorConstant(paddedWindowType, rewriter);
+
+  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+  Value window = initWindow;
+  for (int64_t kernelRowIndex = 0; kernelRowIndex < state.wHeight; ++kernelRowIndex) {
+    Value kernelRow = getOrCreateIndexConstant(rewriter, anchorOp, kernelRowIndex);
+    Value sourceRow =
+      extractProjectedRowStripWindowRow(rowStripStorage, sourceRowTable, state, outputHeight, kernelRow, rewriter, loc);
+    Value mask = extractProjectedRowStripWindowMask(*maskTable, state, outputHeight, kernelRow, rewriter, loc);
+    Value semanticRow = spatial::SpatVMulOp::create(rewriter, loc, fragmentType, sourceRow, mask).getResult();
+    Value paddedRow = createHorizontallyPaddedRowStripFragment(semanticRow, state, rewriter, loc);
+    window = tensor::InsertSliceOp::create(rewriter,
+                                           loc,
+                                           paddedRow,
+                                           window,
+                                           SmallVector<OpFoldResult> {rewriter.getIndexAttr(0),
+                                                                      rewriter.getIndexAttr(0),
+                                                                      rewriter.getIndexAttr(kernelRowIndex),
+                                                                      rewriter.getIndexAttr(0)},
+                                           SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
+                                                                      rewriter.getIndexAttr(state.numChannelsIn),
+                                                                      rewriter.getIndexAttr(1),
+                                                                      rewriter.getIndexAttr(state.xWidth + 2)},
+                                           getUnitStrides(rewriter, 4));
+  }
+  return window;
+}
+
+static FailureOr<Value> createNchwRowStripConvPatchRow(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},
+                                         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 patch = createConvInputPatch(paddedWindow,
+                                     patchType,
+                                     c0,
+                                     c0,
+                                     c0,
+                                     outputWidth,
+                                     state.dilationHeight,
+                                     state.dilationWidth,
+                                     rewriter,
+                                     loc);
+  return tensor::CollapseShapeOp::create(
+           rewriter, loc, rowType, patch, SmallVector<ReassociationIndices> {{0}, {1, 2, 3}})
+    .getResult();
+}
+
+static FailureOr<Value> createConvOutputFromNchwRowStripFragments(Value rowStripStorage,
+                                                                  const ConvLoweringState& state,
+                                                                  PatternRewriter& rewriter,
+                                                                  Location loc) {
+  auto inputType = dyn_cast<RankedTensorType>(rowStripStorage.getType());
+  if (!inputType || inputType != getRowStripStorageType(state.xType))
    return failure();

  StringRef failureReason;
-  if (!canDirectLowerRowStripConv(state, failureReason))
-    return failure();
-
-  ConvRowDemand demand = buildConvRowDemand(RowInterval {0, state.outHeight}, state);
-  if (!covers(demand.acquiredInputRows, demand.neededInputRows))
+  if (!canConsumeNchwRowStripFragments(state, failureReason))
    return failure();

  ConvGeometry geometry = buildConvGeometry(state);
@@ -2629,121 +2733,61 @@ static FailureOr<Value> createConvOutputFromRowStripHwc(Value inputHwc,
  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 = inputType.getElementType();
-  auto paddedInputType = RankedTensorType::get({state.xHeight + state.padHeightBegin + state.padHeightEnd,
-                                                state.xWidth + state.padWidthBegin + state.padWidthEnd,
-                                                state.numChannelsIn},
-                                               elementType,
-                                               inputType.getEncoding());
-  auto paddedPatchType =
-    RankedTensorType::get({state.wHeight, state.wWidth, 1}, elementType, inputType.getEncoding());
-  auto flatPatchType = RankedTensorType::get({state.wHeight * state.wWidth}, elementType, inputType.getEncoding());
-  auto rowChunkType = RankedTensorType::get({1, state.wHeight * state.wWidth}, elementType, inputType.getEncoding());
+  auto elementType = state.outType.getElementType();
  auto rowType = RankedTensorType::get({1, state.numChannelsOut}, state.outType.getElementType());
-  auto packedOutputType =
-    RankedTensorType::get({state.outHeight, state.outWidth, state.numChannelsOut}, state.outType.getElementType());
-  auto packedOutputSliceType =
-    RankedTensorType::get({1, 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)
    return failure();
  Value paddedWeights = standard::createPaddedInputKTiledWeightConstant(weightDenseAttr, state, paddedK, xbarDim, rewriter);
-
-  Value paddedBias;
-  if (state.hasBias) {
-    Value biasMatrix = expandBiasIfNeeded(state.b, rewriter, loc);
-    auto biasMatrixType = cast<RankedTensorType>(biasMatrix.getType());
-    auto paddedBiasType = RankedTensorType::get({1, xbarDim}, state.outType.getElementType());
-    if (auto biasDenseAttr = getHostConstDenseElementsAttr(state.b))
-      paddedBias = standard::createPaddedConstantMatrix(biasDenseAttr, biasMatrixType, paddedBiasType, rewriter);
-    else
-      paddedBias = materializeOrComputeUnary(
-        biasMatrix, paddedBiasType, rewriter, loc, [&](Value biasValue) {
-          return standard::createPaddedConvMatrix(biasValue, biasMatrixType, paddedBiasType, rewriter, loc);
-        });
-  }
-
-  auto paddedInputOp =
-    createSpatCompute<1>(rewriter, loc, TypeRange {paddedInputType}, {}, inputHwc, [&](Value hwcInputArg) {
-      Value paddedInput = createZeroPaddedTensor(hwcInputArg,
-                                                 paddedInputType,
-                                                 {state.padHeightBegin, state.padWidthBegin, 0},
-                                                 {state.padHeightEnd, state.padWidthEnd, 0},
-                                                 rewriter,
-                                                 loc);
-      spatial::SpatYieldOp::create(rewriter, loc, paddedInput);
-    });
-
-  SmallVector<Value> batchInputs {paddedInputOp.getResult(0)};
+  FailureOr<Value> paddedBias = failure();
  if (state.hasBias)
-    batchInputs.push_back(paddedBias);
+    paddedBias = createPaddedBiasRowConstant(state, xbarDim, rewriter);
+  if (state.hasBias && failed(paddedBias))
+    return failure();
+
  auto batchOp = createSpatComputeBatch(
    rewriter,
    loc,
-    TypeRange {packedOutputType},
+    TypeRange {outputStorageType},
    state.outHeight,
    ValueRange {paddedWeights},
-    batchInputs,
+    state.hasBias ? ValueRange {rowStripStorage, *paddedBias} : ValueRange {rowStripStorage},
    [&](detail::SpatComputeBatchBodyArgs args) {
      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 cNumChannels = getOrCreateIndexConstant(rewriter, anchorOp, state.numChannelsIn);
-      Value localHeightOffset = args.lane;
-      Value packedRowInit =
-        tensor::EmptyOp::create(rewriter, loc, ArrayRef<int64_t> {1, state.outWidth, state.numChannelsOut}, elementType);
+      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))
+        return failure();
+      Value fragmentInit = tensor::EmptyOp::create(rewriter, loc, fragmentType.getShape(), elementType);
      auto widthLoop = buildNormalizedScfFor(
        rewriter,
        loc,
        c0,
        cOutWidth,
        c1,
-        ValueRange {packedRowInit},
+        ValueRange {fragmentInit},
        [&](OpBuilder&, Location widthLoc, Value widthIndex, ValueRange widthIterArgs, SmallVectorImpl<Value>& widthYielded) {
-          Value localWidthOffset = widthIndex;
-          Value rowInit = tensor::EmptyOp::create(rewriter, widthLoc, ArrayRef<int64_t> {1, patchSize}, elementType);
-          auto rowLoop = buildNormalizedScfFor(
-            rewriter,
-            widthLoc,
-            c0,
-            cNumChannels,
-            c1,
-            ValueRange {rowInit},
-            [&](OpBuilder&, Location rowLoc, Value channel, ValueRange rowIterArgs, SmallVectorImpl<Value>& rowYielded) {
-              SmallVector<OpFoldResult> patchOffsets {localHeightOffset, localWidthOffset, channel};
-              SmallVector<OpFoldResult> patchSizes {
-                rewriter.getIndexAttr(state.wHeight), rewriter.getIndexAttr(state.wWidth), rewriter.getIndexAttr(1)};
-              Value channelPatch = tensor::ExtractSliceOp::create(
-                rewriter, rowLoc, paddedPatchType, args.inputs.front(), patchOffsets, patchSizes, getUnitStrides(rewriter, 3));
-              Value flatPatch = tensor::CollapseShapeOp::create(
-                rewriter, rowLoc, flatPatchType, channelPatch, SmallVector<ReassociationIndices> {{0, 1, 2}});
-              Value rowChunk = tensor::ExpandShapeOp::create(
-                rewriter, rowLoc, rowChunkType, flatPatch, SmallVector<ReassociationIndices> {{0, 1}});
-              Value flatOffset = affineMulConst(
-                rewriter, rowLoc, channel, state.wHeight * state.wWidth, anchorOp);
-              SmallVector<OpFoldResult> rowOffsets {rewriter.getIndexAttr(0), flatOffset};
-              SmallVector<OpFoldResult> rowSizes {
-                rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.wHeight * state.wWidth)};
-              Value nextRow = tensor::InsertSliceOp::create(
-                rewriter, rowLoc, rowChunk, rowIterArgs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 2));
-              rowYielded.push_back(nextRow);
-              return success();
-            });
-          if (failed(rowLoop))
+          FailureOr<Value> patchRow =
+            createNchwRowStripConvPatchRow(*inputWindow, state, widthIndex, rewriter, widthLoc);
+          if (failed(patchRow))
            return failure();

-          Value paddedRow = rowLoop->results.front();
+          Value paddedRow = *patchRow;
          if (patchSize != paddedK)
            paddedRow = createZeroPaddedTensor(
              paddedRow, paddedPatchRowType, {0, 0}, {0, paddedK - patchSize}, rewriter, widthLoc);

-          auto zeroAttr = DenseElementsAttr::get(paddedRowType, rewriter.getZeroAttr(state.outType.getElementType()));
-          Value zeroRow = getOrCreateConstant(rewriter, anchorOp, zeroAttr, paddedRowType);
+          Value zeroRow = createZeroTensorConstant(paddedRowType, rewriter);
          auto kLoop = buildNormalizedScfFor(
            rewriter,
            widthLoc,
@@ -2752,7 +2796,7 @@ static FailureOr<Value> createConvOutputFromRowStripHwc(Value inputHwc,
            c1,
            ValueRange {zeroRow},
            [&](OpBuilder&, Location reduceLoc, Value kSlice, ValueRange reduceIterArgs, SmallVectorImpl<Value>& reduceYielded) {
-              Value kOffset = affineMulConst(rewriter, reduceLoc, kSlice, xbarDim, anchorOp);
+              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(
@@ -2776,9 +2820,7 @@ static FailureOr<Value> createConvOutputFromRowStripHwc(Value inputHwc,

          Value rowResult = kLoop->results.front();
          if (state.hasBias)
-            rowResult =
-              spatial::SpatVAddOp::create(rewriter, widthLoc, paddedRowType, rowResult, args.inputs[1]).getResult();
-
+            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)};
@@ -2790,25 +2832,23 @@ static FailureOr<Value> createConvOutputFromRowStripHwc(Value inputHwc,

          Value outputFragment = tensor::ExpandShapeOp::create(rewriter,
                                                               widthLoc,
-                                                               packedOutputSliceType,
+                                                               outputPixelType,
                                                               outputRow,
-                                                               SmallVector<ReassociationIndices> {{0}, {1, 2}});
-          SmallVector<OpFoldResult> rowOffsets {rewriter.getIndexAttr(0), widthIndex, rewriter.getIndexAttr(0)};
+                                                               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(1), rewriter.getIndexAttr(state.numChannelsOut)};
-          Value nextPackedRow = tensor::InsertSliceOp::create(
-            rewriter, widthLoc, outputFragment, widthIterArgs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 3));
-          widthYielded.push_back(nextPackedRow);
+            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();

-      SmallVector<OpFoldResult> batchOffsets {args.lane, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> batchSizes {
-        rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.outWidth), rewriter.getIndexAttr(state.numChannelsOut)};
-      createParallelInsertSliceIntoBatchOutput(
-        rewriter, loc, widthLoop->results.front(), args.outputs.front(), batchOffsets, batchSizes, getUnitStrides(rewriter, 3));
+      insertRowStripFragment(widthLoop->results.front(), args.outputs.front(), state.outType, args.lane, rewriter, loc);
      return success();
    });
  if (failed(batchOp))
@@ -2816,19 +2856,22 @@ static FailureOr<Value> createConvOutputFromRowStripHwc(Value inputHwc,
  return batchOp->getResult(0);
 }

-static FailureOr<Value> createConvRowsFromRowStripInput(const ConvLoweringState& state,
-                                                        [[maybe_unused]] const ConvLoweringDecision& decision,
-                                                        Value rowStripInput,
-                                                        PatternRewriter& rewriter,
-                                                        Location loc) {
-  return createConvOutputFromRowStripHwc(rowStripInput, state, rewriter, loc);
+static FailureOr<Value> createConvOutputFromRowStripInput(const ConvLoweringState& state,
+                                                          [[maybe_unused]] const ConvLoweringDecision& decision,
+                                                          Value rowStripInput,
+                                                          PatternRewriter& rewriter,
+                                                          Location loc) {
+  return createConvOutputFromNchwRowStripFragments(rowStripInput, state, rewriter, loc);
 }

 static Value createFragmentConstant(const DistributedTensorStep& step,
                                    RankedTensorType fragmentType,
                                    PatternRewriter& rewriter) {
-  if (step.constantKind == DistributedTensorConstantKind::PerChannel)
-    return createPerChannelConstantFragment(step.constantAttr, fragmentType, rewriter);
+  if (step.constantKind == DistributedTensorConstantKind::PerChannel) {
+    FailureOr<Value> constant = createPerChannelConstantFragment(step.constantAttr, fragmentType, rewriter);
+    assert(succeeded(constant) && "distributed per-channel constants are classified before lowering");
+    return *constant;
+  }

  Attribute splatValue = step.constantAttr.getSplatValue<Attribute>();
  return getOrCreateConstant(rewriter,
@@ -2938,67 +2981,22 @@ static Value createFragmentReciprocalConstant(const DistributedTensorStep& step,
  }
 }

-[[maybe_unused]] static FailureOr<DistributedTensorInfo> createDistributedTensorFromRows(Value rows,
-                                                                                         RankedTensorType logicalType,
-                                                                                         PatternRewriter& rewriter,
-                                                                                         Location loc) {
-  const int64_t width = logicalType.getDimSize(3);
-  const int64_t height = logicalType.getDimSize(2);
-  auto rowsType = cast<RankedTensorType>(rows.getType());
-  auto rowSliceType =
-    RankedTensorType::get({width, logicalType.getDimSize(1)}, logicalType.getElementType(), rowsType.getEncoding());
-  auto channelWidthType =
-    RankedTensorType::get({logicalType.getDimSize(1), width}, logicalType.getElementType(), rowsType.getEncoding());
-  auto fragmentType = getRowStripFragmentType(logicalType, width);
-  auto batchOp = createSpatComputeBatch(
-    rewriter, loc, TypeRange {logicalType}, height, {}, ValueRange {rows}, [&](detail::SpatComputeBatchBodyArgs args) {
-      Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
-      Value rowStart = affineMulConst(rewriter, loc, args.lane, width, anchorOp);
-      SmallVector<OpFoldResult> rowOffsets {rowStart, rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> rowSizes {rewriter.getIndexAttr(width), rewriter.getIndexAttr(logicalType.getDimSize(1))};
-      Value rowSlice = tensor::ExtractSliceOp::create(
-        rewriter, loc, rowSliceType, args.inputs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 2));
-      Value channelWidth = ONNXTransposeOp::create(
-        rewriter, loc, channelWidthType, rowSlice, rewriter.getI64ArrayAttr({1, 0})).getResult();
-      Value fragment = tensor::ExpandShapeOp::create(rewriter,
-                                                     loc,
-                                                     fragmentType,
-                                                     channelWidth,
-                                                     SmallVector<ReassociationIndices> {{0, 1}, {2, 3}});
-      SmallVector<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), args.lane,
-                                               rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> outputSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(logicalType.getDimSize(1)),
-                                             rewriter.getIndexAttr(1), rewriter.getIndexAttr(width)};
-      createParallelInsertSliceIntoBatchOutput(
-        rewriter, loc, fragment, args.outputs.front(), outputOffsets, outputSizes, getUnitStrides(rewriter, 4));
-      return success();
-    });
-  if (failed(batchOp))
-    return failure();
-  return makeDistributedTensorInfo(batchOp->getResult(0), logicalType);
-}
-
 [[maybe_unused]] static FailureOr<DistributedTensorInfo> applyDistributedPreservingStep(const DistributedTensorInfo& inputInfo,
                                                                                        const DistributedTensorStep& step,
                                                                                        PatternRewriter& rewriter,
                                                                                        Location loc) {
  auto logicalType = inputInfo.logicalType;
-  const int64_t width = logicalType.getDimSize(3);
-  auto fragmentType = getRowStripFragmentType(logicalType, width);
+  auto fragmentType = getRowStripFragmentType(logicalType);
+  auto storageType = getRowStripStorageType(logicalType);
  auto batchOp = createSpatComputeBatch(rewriter,
                                        loc,
-                                        TypeRange {logicalType},
+                                        TypeRange {storageType},
                                        inputInfo.laneCount,
                                        {},
                                        ValueRange {inputInfo.storage},
                                        [&](detail::SpatComputeBatchBodyArgs args) {
-                                          SmallVector<OpFoldResult> offsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0),
-                                                                             args.lane, rewriter.getIndexAttr(0)};
-                                          SmallVector<OpFoldResult> sizes {
-                                            rewriter.getIndexAttr(1), rewriter.getIndexAttr(logicalType.getDimSize(1)),
-                                            rewriter.getIndexAttr(1), rewriter.getIndexAttr(width)};
-                                          Value fragment = tensor::ExtractSliceOp::create(
-                                            rewriter, loc, fragmentType, args.inputs.front(), offsets, sizes, getUnitStrides(rewriter, 4));
+                                          Value fragment =
+                                            extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc);
                                          switch (step.kind) {
                                          case DistributedTensorOpKind::Relu:
                                            fragment = spatial::SpatReluOp::create(rewriter, loc, fragmentType, fragment).getResult();
@@ -3034,8 +3032,8 @@ static Value createFragmentReciprocalConstant(const DistributedTensorStep& step,
                                          case DistributedTensorOpKind::Conv:
                                            return failure();
                                          }
-                                          createParallelInsertSliceIntoBatchOutput(
-                                            rewriter, loc, fragment, args.outputs.front(), offsets, sizes, getUnitStrides(rewriter, 4));
+                                          insertRowStripFragment(
+                                            fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
                                          return success();
                                        });
  if (failed(batchOp))
@@ -3743,6 +3741,8 @@ LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp) {
    return failure();
  if (state->outType.getRank() != 4 || !state->outType.hasStaticShape())
    return failure();
+  if (state->hasBias && !isSupportedBiasAddValue(state->b, state->outType))
+    return failure();

  FailureOr<PimConvLoweringType> requestedStrategy = resolveRequestedConvLoweringStrategy(planOp.getOperation());
  if (failed(requestedStrategy))
@@ -3786,7 +3786,7 @@ LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp) {
    return failure();

  StringRef failureReason;
-  return canDirectLowerRowStripConv(*state, failureReason) ? success() : failure();
+  return canConsumeNchwRowStripFragments(*state, failureReason) ? success() : failure();
 }

 FailureOr<Value>
@@ -3822,17 +3822,33 @@ lowerSelectedConv2DPlan(spatial::SpatConv2DPlanOp planOp,
    if (rowStripInput) {
      if (failed(canConsumeAndProduceRowStrip(planOp)))
        return planOp.emitOpError("selected row-strip input/output layout is not supported for this Conv plan"), failure();
-      return createConvRowsFromRowStripInput(*state, decision, *rowStripInput, rewriter, planOp.getLoc());
+      return createConvOutputFromRowStripInput(*state, decision, *rowStripInput, rewriter, planOp.getLoc());
    }
    if (failed(canLowerConvPlanToRowStrip(planOp)))
      return planOp.emitOpError("selected row-strip layout is not supported for this Conv plan"), failure();
-    FailureOr<Value> rows = createConvRowsForStrategy(*state, decision, rewriter, planOp.getLoc());
+    ConvLoweringState rowState = *state;
+    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> packedRows = createRowStripPackedRows(*rows, *state, rewriter, planOp.getLoc());
-    if (failed(packedRows))
-      return planOp.emitOpError("failed to pack Conv rows into the selected row-strip physical layout"), failure();
-    return *packedRows;
+    FailureOr<Value> rowStripStorage = createRowStripStorageFromRows(*rows, state->outType, 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());
+      if (failed(rowStripStorage))
+        return planOp.emitOpError("failed to apply row-strip Conv bias per fragment"), failure();
+    }
+    return *rowStripStorage;
  }

  if (decision.strategy == PimConvLoweringDepthwise)
@@ -5,7 +5,7 @@

 #include "llvm/ADT/SmallVector.h"

-#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
@@ -47,38 +47,28 @@ static FailureOr<Value> materializeBroadcastedConstantTensor(Value value,
    return failure();

  const int64_t rankOffset = static_cast<int64_t>(resultShape.size() - sourceShape.size());
-  for (int64_t i = 0; i < static_cast<int64_t>(resultShape.size()); ++i) {
-    const int64_t sourceIndex = i - rankOffset;
-    const int64_t sourceDim = sourceIndex < 0 ? 1 : sourceShape[sourceIndex];
-    const int64_t resultDim = resultShape[i];
-    if (sourceDim != 1 && sourceDim != resultDim)
-      return failure();
-  }
-
-  SmallVector<Attribute> sourceValues(denseAttr.getValues<Attribute>());
  SmallVector<int64_t> sourceStrides = computeRowMajorStrides(sourceShape);
  SmallVector<int64_t> resultStrides = computeRowMajorStrides(resultShape);
-
+  SmallVector<Attribute> sourceValues(denseAttr.getValues<Attribute>());
  SmallVector<Attribute> resultValues;
  resultValues.reserve(resultType.getNumElements());

  for (int64_t flatIndex = 0; flatIndex < resultType.getNumElements(); ++flatIndex) {
    int64_t remaining = flatIndex;
    int64_t sourceFlatIndex = 0;
-
    for (int64_t i = 0; i < static_cast<int64_t>(resultShape.size()); ++i) {
      const int64_t resultIndex = resultStrides.empty() ? 0 : remaining / resultStrides[i];
      remaining = resultStrides.empty() ? 0 : remaining % resultStrides[i];
-
      const int64_t sourceIndex = i - rankOffset;
      if (sourceIndex < 0)
        continue;
-
      const int64_t sourceDim = sourceShape[sourceIndex];
+      const int64_t resultDim = resultShape[i];
+      if (sourceDim != 1 && sourceDim != resultDim)
+        return failure();
      const int64_t mappedIndex = sourceDim == 1 ? 0 : resultIndex;
      sourceFlatIndex += mappedIndex * sourceStrides[sourceIndex];
    }
-
    resultValues.push_back(sourceValues[sourceFlatIndex]);
  }

@@ -106,7 +96,7 @@ static FailureOr<Value> materializeReciprocalTensor(Value value,
  if (failed(broadcastedValue))
    return failure();

-  auto denseAttr = dyn_cast<DenseFPElementsAttr>(getDenseConstantAttr(*broadcastedValue));
+  auto denseAttr = dyn_cast<DenseFPElementsAttr>(getHostConstDenseElementsAttr(*broadcastedValue));
  if (!denseAttr)
    return failure();

@@ -185,10 +175,45 @@ struct DivToSpatialCompute : OpConversionPattern<ONNXDivOp> {
  }
 };

+struct AddToSpatialCompute : OpConversionPattern<ONNXAddOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(ONNXAddOp op, ONNXAddOpAdaptor adaptor, ConversionPatternRewriter& rewriter) const override {
+    auto resultType = dyn_cast<RankedTensorType>(op.getResult().getType());
+    if (!resultType || !resultType.hasStaticShape())
+      return failure();
+
+    FailureOr<BiasAddPlanCandidate> candidate =
+      classifyBiasAddPlanCandidate(adaptor.getA(), adaptor.getB(), resultType);
+    if (succeeded(candidate)) {
+      auto plan = spatial::SpatBiasAddPlanOp::create(
+        rewriter, op.getLoc(), resultType, candidate->data, candidate->bias, rewriter.getStringAttr("nchw"));
+      rewriter.replaceOp(op, plan.getResult());
+      return success();
+    }
+
+    auto lhs = prepareElementwiseOperand(adaptor.getA(), resultType, rewriter, op.getLoc());
+    if (failed(lhs))
+      return failure();
+    auto rhs = prepareElementwiseOperand(adaptor.getB(), resultType, rewriter, op.getLoc());
+    if (failed(rhs))
+      return failure();
+
+    auto computeOp =
+      createSpatCompute<2>(rewriter, op.getLoc(), resultType, {}, ValueRange {*lhs, *rhs}, [&](Value x, Value y) {
+        auto loweredOp = spatial::SpatVAddOp::create(rewriter, op.getLoc(), resultType, x, y);
+        spatial::SpatYieldOp::create(rewriter, op.getLoc(), loweredOp.getResult());
+      });
+    rewriter.replaceOp(op, computeOp);
+    return success();
+  }
+};
+
 } // namespace

 void populateElementwisePatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
-  patterns.add<BinaryElementwiseToSpatialCompute<ONNXAddOp, spatial::SpatVAddOp>>(ctx);
+  patterns.add<AddToSpatialCompute>(ctx);
  patterns.add<BinaryElementwiseToSpatialCompute<ONNXSubOp, spatial::SpatVSubOp>>(ctx);
  patterns.add<BinaryElementwiseToSpatialCompute<ONNXMulOp, spatial::SpatVMulOp>>(ctx);
  patterns.add<DivToSpatialCompute>(ctx);