finally fast googlenet with correct latency artifacts for fair comparison

src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp

@@ -1242,8 +1242,9 @@ static Value buildPackedWeights(DenseElementsAttr wDenseAttr,
                                 const Tiling& tiling,
                                 PatternRewriter& rewriter,
                                 Location loc) {
+  const int64_t paddedOutputChannels = static_cast<int64_t>(crossbarSize.getValue());
   auto packedWeightType = RankedTensorType::get(
-    {tiling.numChannelTiles, tiling.tileInputRows, tiling.tileOutputChannels}, wType.getElementType());
+    {tiling.numChannelTiles, tiling.tileInputRows, paddedOutputChannels}, wType.getElementType());
   SmallVector<Attribute> packedValues(packedWeightType.getNumElements(),
                                       cast<Attribute>(rewriter.getZeroAttr(wType.getElementType())));
   SmallVector<Attribute> sourceValues(wDenseAttr.getValues<Attribute>());
@@ -1262,7 +1263,7 @@ static Value buildPackedWeights(DenseElementsAttr wDenseAttr,
             ((globalOutChannel * wType.getDimSize(1) * wType.getDimSize(2)) + kernelH) * wType.getDimSize(3) + kernelW;
           const int64_t targetCol = localChannel * tiling.outputMultiplier + multiplierIndex;
           const int64_t targetFlatIndex =
-            ((tileIndex * tiling.tileInputRows) + targetRow) * tiling.tileOutputChannels + targetCol;
+            ((tileIndex * tiling.tileInputRows) + targetRow) * paddedOutputChannels + targetCol;
           packedValues[targetFlatIndex] = sourceValues[sourceFlatIndex];
         }
       }
@@ -1353,11 +1354,12 @@ static Value createWeightTile(Value packedWeights,
                               PatternRewriter& rewriter,
                               Location loc) {
   SmallVector<OpFoldResult> offsets {channelTileIndex, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
+  const int64_t paddedOutputChannels = static_cast<int64_t>(crossbarSize.getValue());
   SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
                                    rewriter.getIndexAttr(tiling.tileInputRows),
-                                   rewriter.getIndexAttr(tiling.tileOutputChannels)};
+                                   rewriter.getIndexAttr(paddedOutputChannels)};
   auto collapsedType =
-    RankedTensorType::get({tiling.tileInputRows, tiling.tileOutputChannels}, packedWeightType.getElementType());
+    RankedTensorType::get({tiling.tileInputRows, paddedOutputChannels}, packedWeightType.getElementType());
   return extractMixedSliceOrIdentity(
     rewriter, loc, packedWeights, collapsedType,
     {offsets, sizes, getUnitStrides(rewriter, 3)});
@@ -1547,6 +1549,8 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter&
   auto gemmOutType =
     RankedTensorType::get({tiling->totalPatches, state.outType.getDimSize(1)}, state.outType.getElementType());
   auto rowTileType = RankedTensorType::get({1, tiling->tileOutputChannels}, state.outType.getElementType());
+  auto paddedRowTileType = RankedTensorType::get(
+    {1, static_cast<int64_t>(crossbarSize.getValue())}, state.outType.getElementType());
   auto piecesType = spatial::getGraphBatchPhysicalResultType(
     tiling->totalPatches * tiling->numChannelTiles, rowTileType);
   auto paddedInputType = cast<RankedTensorType>(paddedInput.getType());
@@ -1617,7 +1621,17 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter&
                                             *tiling,
                                             rewriter,
                                             loc);
-      Value rowTile = spatial::SpatVMMOp::create(rewriter, loc, rowTileType, weightTile, inputTile).getResult();
+      Value paddedRowTile =
+        spatial::SpatVMMOp::create(rewriter, loc, paddedRowTileType, weightTile, inputTile).getResult();
+      Value rowTile = tensor::ExtractSliceOp::create(
+        rewriter,
+        loc,
+        rowTileType,
+        paddedRowTile,
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
+                                  rewriter.getIndexAttr(tiling->tileOutputChannels)},
+        getUnitStrides(rewriter, 2));
       if (args.inputs.size() > 1) {
         Value biasArg = pickInputByRank(/*rank=*/2);
         if (!biasArg) {
@@ -2493,6 +2507,7 @@ static Value createZeroGemmBias(RankedTensorType gemmResultType, PatternRewriter
 
 static bool rowStripOutputTileFitsOneCore(const ConvGeometry& geometry) {
   return ceilIntegerDivide(geometry.k, geometry.xbarSize)
+           * ceilIntegerDivide(geometry.c, geometry.xbarSize)
          <= static_cast<int64_t>(crossbarCountInCore.getValue());
 }
 
@@ -2521,28 +2536,6 @@ static bool canConsumePixelMajorRowStripFragments(const ConvLoweringState& state
     failureReason = "dilation_not_one";
     return false;
   }
-  const bool pointwise = state.xHeight == 1 && state.xWidth == 1 && state.outHeight == 1 && state.outWidth == 1
-                         && state.wHeight == 1 && state.wWidth == 1 && state.padHeightBegin == 0
-                         && state.padHeightEnd == 0 && state.padWidthBegin == 0 && state.padWidthEnd == 0;
-  if (pointwise) {
-    if (!getHostConstDenseElementsAttr(state.w)) {
-      failureReason = "non_constant_weight";
-      return false;
-    }
-    if (state.hasBias && !isSupportedBiasAddValue(state.b, state.outType)) {
-      failureReason = "unsupported_bias";
-      return false;
-    }
-    return true;
-  }
-  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;
@@ -3043,57 +3036,154 @@ static FailureOr<Value> createConvOutputRow(ValueRange inputTiles,
                                             Location loc) {
   auto elementType = cast<RankedTensorType>(inputTiles.front().getType()).getElementType();
   auto rowType = RankedTensorType::get({1, outputChannels}, elementType);
-  auto tileWeightsType =
-    RankedTensorType::get({paddedK, xbarDim},
-                          cast<RankedTensorType>(paddedWeights.getType()).getElementType());
   const int64_t outputTileCount = ceilIntegerDivide(outputChannels, xbarDim);
-
-  auto getTileWeights = [&](int64_t outputTile) {
-    if (outputTileCount == 1)
-      return paddedWeights;
-    SmallVector<OpFoldResult> offsets {
-      rewriter.getIndexAttr(outputTile), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
-    SmallVector<OpFoldResult> sizes {
-      rewriter.getIndexAttr(1), rewriter.getIndexAttr(paddedK), rewriter.getIndexAttr(xbarDim)};
-    return extractStaticSliceOrIdentity(
-      rewriter, loc, paddedWeights, tileWeightsType, offsets, sizes, getUnitStrides(rewriter, 3));
-  };
-
-  if (outputTileCount == 1) {
-    FailureOr<Value> rowResult = createConvOutputTile(
-      inputTiles, getTileWeights(0), outputChannels, xbarDim, rewriter, loc);
-    if (failed(rowResult))
-      return failure();
-    Value validRow = *rowResult;
-    if (bias)
-      validRow = spatial::SpatVAddOp::create(rewriter, loc, rowType, validRow, bias).getResult();
-    return validRow;
-  }
-
   const int64_t paddedOutputChannels = outputTileCount * xbarDim;
   auto paddedOutputType = RankedTensorType::get({1, paddedOutputChannels}, elementType);
-  Value paddedOutput = tensor::EmptyOp::create(rewriter, loc, paddedOutputType.getShape(), elementType);
-  for (int64_t outputTile = 0; outputTile < outputTileCount; ++outputTile) {
-    FailureOr<Value> tileResult = createConvOutputTile(
-      inputTiles, getTileWeights(outputTile), xbarDim, xbarDim, rewriter, loc);
-    if (failed(tileResult))
-      return failure();
-    SmallVector<OpFoldResult> tileOffsets {
-      rewriter.getIndexAttr(0), rewriter.getIndexAttr(outputTile * xbarDim)};
-    SmallVector<OpFoldResult> tileSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)};
-    paddedOutput = tensor::InsertSliceOp::create(
-      rewriter, loc, *tileResult, paddedOutput, tileOffsets, tileSizes, getUnitStrides(rewriter, 2));
+  auto weightSliceType = RankedTensorType::get(
+    {xbarDim, paddedOutputChannels},
+    cast<RankedTensorType>(paddedWeights.getType()).getElementType());
+
+  Value paddedOutput;
+  for (auto [kSlice, inputTile] : llvm::enumerate(inputTiles)) {
+    const int64_t kOffset = static_cast<int64_t>(kSlice) * xbarDim;
+    Value weightSlice = extractStaticSliceOrIdentity(
+      rewriter,
+      loc,
+      paddedWeights,
+      weightSliceType,
+      SmallVector<OpFoldResult> {rewriter.getIndexAttr(kOffset), rewriter.getIndexAttr(0)},
+      SmallVector<OpFoldResult> {rewriter.getIndexAttr(xbarDim),
+                                rewriter.getIndexAttr(paddedOutputChannels)},
+      getUnitStrides(rewriter, 2));
+    Value piece =
+      spatial::SpatVMMOp::create(rewriter, loc, paddedOutputType, weightSlice, inputTile).getResult();
+    paddedOutput = paddedOutput
+                     ? spatial::SpatVAddOp::create(
+                         rewriter, loc, paddedOutputType, paddedOutput, piece).getResult()
+                     : piece;
   }
 
-  SmallVector<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
-  SmallVector<OpFoldResult> outputSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(outputChannels)};
-  Value validRow = tensor::ExtractSliceOp::create(
-    rewriter, loc, rowType, paddedOutput, outputOffsets, outputSizes, getUnitStrides(rewriter, 2));
+  Value validRow = outputChannels == paddedOutputChannels
+                     ? paddedOutput
+                     : tensor::ExtractSliceOp::create(
+                         rewriter,
+                         loc,
+                         rowType,
+                         paddedOutput,
+                         SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
+                         SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
+                                                   rewriter.getIndexAttr(outputChannels)},
+                         getUnitStrides(rewriter, 2))
+                         .getResult();
   if (bias)
     validRow = spatial::SpatVAddOp::create(rewriter, loc, rowType, validRow, bias).getResult();
   return validRow;
 }
 
+static FailureOr<Value> createRowStripConvOutput(const ConvLoweringState& state,
+                                                 Value input,
+                                                 Value paddedWeights,
+                                                 Value bias,
+                                                 int64_t paddedK,
+                                                 int64_t numKSlices,
+                                                 int64_t xbarDim,
+                                                 PatternRewriter& rewriter,
+                                                 Location loc) {
+  const int64_t laneCount = state.outHeight;
+  const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
+  const bool hasPartialInputTile = patchSize % xbarDim != 0;
+  auto elementType = state.outType.getElementType();
+  auto partialInputScratchType = RankedTensorType::get({1, xbarDim}, elementType);
+  auto outputPixelType = RankedTensorType::get({1, 1, 1, state.numChannelsOut}, elementType);
+  auto fragmentType = getRowStripFragmentType(state.outType);
+  auto storageType = getRowStripStorageType(state.outType);
+
+  auto batch = createSpatComputeBatch(
+    rewriter,
+    loc,
+    TypeRange {storageType},
+    laneCount,
+    ValueRange {paddedWeights},
+    bias ? ValueRange {input, bias} : ValueRange {input},
+    [&](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);
+      FailureOr<Value> inputWindow =
+        createConvInputWindow(args.inputs.front(), state, args.lane, rewriter, loc);
+      if (failed(inputWindow))
+        return failure();
+      Value fragmentInit = tensor::EmptyOp::create(rewriter, loc, fragmentType.getShape(), elementType);
+      SmallVector<Value> loopInit {fragmentInit};
+      if (hasPartialInputTile)
+        loopInit.push_back(createZeroTensorConstant(partialInputScratchType, rewriter));
+      auto loop = buildNormalizedScfFor(
+        rewriter,
+        loc,
+        c0,
+        cOutWidth,
+        c1,
+        loopInit,
+        [&](OpBuilder&,
+            Location pixelLoc,
+            Value localColumn,
+            ValueRange iterArgs,
+            SmallVectorImpl<Value>& yielded) {
+          Value partialInputScratch = hasPartialInputTile ? iterArgs[1] : Value();
+          FailureOr<SmallVector<Value>> inputTiles = createConvInputTiles(*inputWindow,
+                                                                          state,
+                                                                          localColumn,
+                                                                          partialInputScratch,
+                                                                          patchSize,
+                                                                          numKSlices,
+                                                                          xbarDim,
+                                                                          rewriter,
+                                                                          pixelLoc);
+          if (failed(inputTiles))
+            return failure();
+          FailureOr<Value> output = createConvOutputRow(*inputTiles,
+                                                        paddedK,
+                                                        state.numChannelsOut,
+                                                        args.weights.front(),
+                                                        bias ? args.inputs[1] : Value(),
+                                                        xbarDim,
+                                                        rewriter,
+                                                        pixelLoc);
+          if (failed(output))
+            return failure();
+          Value outputPixel = tensor::ExpandShapeOp::create(
+            rewriter, pixelLoc, outputPixelType, *output, SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
+          Value next = tensor::InsertSliceOp::create(
+            rewriter,
+            pixelLoc,
+            outputPixel,
+            iterArgs.front(),
+            SmallVector<OpFoldResult> {rewriter.getIndexAttr(0),
+                                       rewriter.getIndexAttr(0),
+                                       localColumn,
+                                       rewriter.getIndexAttr(0)},
+            SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
+                                       rewriter.getIndexAttr(1),
+                                       rewriter.getIndexAttr(1),
+                                       rewriter.getIndexAttr(state.numChannelsOut)},
+            getUnitStrides(rewriter, 4));
+          yielded.push_back(next);
+          if (hasPartialInputTile)
+            yielded.push_back(partialInputScratch);
+          return success();
+        });
+      if (failed(loop))
+        return failure();
+      publishGraphBatchPhysicalFragment(
+        rewriter, loc, loop->results.front(), args.outputs.front(), args.lane);
+      return success();
+    });
+  if (failed(batch))
+    return failure();
+  return batch->getResult(0);
+}
+
 static FailureOr<Value> createOutputChannelTiledRowStripConvOutput(const ConvLoweringState& state,
                                                                    Value input,
                                                                    Value paddedWeights,
@@ -3229,21 +3319,11 @@ static FailureOr<Value>
   const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
   const int64_t numKSlices = ceilIntegerDivide(patchSize, xbarDim);
   const int64_t paddedK = numKSlices * xbarDim;
-  const bool hasPartialInputTile = patchSize % xbarDim != 0;
-  auto elementType = state.outType.getElementType();
-  auto partialInputScratchType = RankedTensorType::get({1, xbarDim}, elementType);
-  auto fragmentType = getRowStripFragmentType(state.outType);
-  auto outputPixelType = RankedTensorType::get({1, 1, 1, state.numChannelsOut}, elementType);
-  auto outputStorageType = getRowStripStorageType(state.outType);
 
-  Value paddedWeights = state.numChannelsOut <= xbarDim
-                          ? standard::createPaddedPixelMajorWeightConstant(
-                              weightDenseAttr, state, paddedK, xbarDim, rewriter)
-                          : standard::createPaddedOutputChannelTiledWeightConstant(
-                              weightDenseAttr, state, paddedK, xbarDim, rewriter);
-  if (state.numChannelsOut > xbarDim)
-    return createOutputChannelTiledRowStripConvOutput(
-      state, state.x, paddedWeights, paddedK, numKSlices, xbarDim, rewriter, loc);
+  const int64_t paddedOutputChannels =
+    ceilIntegerDivide(state.numChannelsOut, xbarDim) * xbarDim;
+  Value paddedWeights = standard::createPaddedPixelMajorWeightConstant(
+    weightDenseAttr, state, paddedK, paddedOutputChannels, rewriter);
 
   FailureOr<Value> bias = failure();
   if (state.hasBias)
@@ -3251,83 +3331,9 @@ static FailureOr<Value>
   if (state.hasBias && failed(bias))
     return failure();
 
-  auto batchOp = createSpatComputeBatch(
-    rewriter,
-    loc,
-    TypeRange {outputStorageType},
-    state.outHeight,
-    ValueRange {paddedWeights},
-    state.hasBias ? ValueRange {state.x, *bias} : ValueRange {state.x},
-    [&](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);
-      FailureOr<Value> inputWindow =
-        createConvInputWindow(args.inputs.front(), state, args.lane, rewriter, loc);
-      if (failed(inputWindow))
-        return failure();
-      Value fragmentInit = tensor::EmptyOp::create(rewriter, loc, fragmentType.getShape(), elementType);
-      SmallVector<Value> widthLoopInit {fragmentInit};
-      if (hasPartialInputTile)
-        widthLoopInit.push_back(createZeroTensorConstant(partialInputScratchType, rewriter));
-      auto widthLoop = buildNormalizedScfFor(
-        rewriter,
-        loc,
-        c0,
-        cOutWidth,
-        c1,
-        widthLoopInit,
-        [&](OpBuilder&, Location widthLoc, Value widthIndex, ValueRange widthIterArgs, SmallVectorImpl<Value>& widthYielded) {
-          Value partialInputScratch = hasPartialInputTile ? widthIterArgs[1] : Value();
-          FailureOr<SmallVector<Value>> inputTiles = createConvInputTiles(*inputWindow,
-                                                                          state,
-                                                                          widthIndex,
-                                                                          partialInputScratch,
-                                                                          patchSize,
-                                                                          numKSlices,
-                                                                          xbarDim,
-                                                                          rewriter,
-                                                                          widthLoc);
-          if (failed(inputTiles))
-            return failure();
-          FailureOr<Value> outputRow = createConvOutputRow(*inputTiles,
-                                                           paddedK,
-                                                           state.numChannelsOut,
-                                                           args.weights.front(),
-                                                           state.hasBias ? args.inputs[1] : Value(),
-                                                           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), widthIndex, rewriter.getIndexAttr(0)};
-          SmallVector<OpFoldResult> rowSizes {
-            rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1),
-            rewriter.getIndexAttr(state.numChannelsOut)};
-          Value nextFragment = tensor::InsertSliceOp::create(
-            rewriter, widthLoc, outputFragment, widthIterArgs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 4));
-          widthYielded.push_back(nextFragment);
-          if (hasPartialInputTile)
-            widthYielded.push_back(partialInputScratch);
-          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);
+  return createRowStripConvOutput(
+    state, state.x, paddedWeights, state.hasBias ? *bias : Value(),
+    paddedK, numKSlices, xbarDim, rewriter, loc);
 }
 
 static FailureOr<Value> createConvOutputFromPixelMajorRowStripFragments(Value rowStripStorage,
@@ -3346,105 +3352,22 @@ static FailureOr<Value> createConvOutputFromPixelMajorRowStripFragments(Value ro
   const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
   const int64_t numKSlices = ceilIntegerDivide(patchSize, xbarDim);
   const int64_t paddedK = numKSlices * xbarDim;
-  const bool hasPartialInputTile = patchSize % xbarDim != 0;
-  auto elementType = state.outType.getElementType();
-  auto partialInputScratchType = RankedTensorType::get({1, xbarDim}, elementType);
-  auto outputPixelType = RankedTensorType::get({1, 1, 1, state.numChannelsOut}, elementType);
-  auto outputStorageType = getRowStripStorageType(state.outType);
   auto weightDenseAttr = getHostConstDenseElementsAttr(state.w);
   if (!weightDenseAttr)
     return failure();
-  Value paddedWeights = state.numChannelsOut <= xbarDim
-                          ? standard::createPaddedPixelMajorWeightConstant(
-                              weightDenseAttr, state, paddedK, xbarDim, rewriter)
-                          : standard::createPaddedOutputChannelTiledWeightConstant(
-                              weightDenseAttr, state, paddedK, xbarDim, rewriter);
-  if (state.numChannelsOut > xbarDim)
-    return createOutputChannelTiledRowStripConvOutput(
-      state, rowStripStorage, paddedWeights, paddedK, numKSlices, xbarDim, rewriter, loc);
+  const int64_t paddedOutputChannels =
+    ceilIntegerDivide(state.numChannelsOut, xbarDim) * xbarDim;
+  Value paddedWeights = standard::createPaddedPixelMajorWeightConstant(
+    weightDenseAttr, state, paddedK, paddedOutputChannels, rewriter);
   FailureOr<Value> bias = failure();
   if (state.hasBias)
     bias = createBiasRowConstant(state, rewriter);
   if (state.hasBias && failed(bias))
     return failure();
 
-  auto batchOp = createSpatComputeBatch(
-    rewriter,
-    loc,
-    TypeRange {outputStorageType},
-    state.outHeight,
-    ValueRange {paddedWeights},
-    state.hasBias ? ValueRange {rowStripStorage, *bias} : ValueRange {rowStripStorage},
-    [&](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);
-      auto fragmentType = getRowStripFragmentType(state.outType);
-      FailureOr<Value> inputWindow = createConvInputWindow(args.inputs.front(), state, args.lane, rewriter, loc);
-      if (failed(inputWindow))
-        return failure();
-      Value fragmentInit = tensor::EmptyOp::create(rewriter, loc, fragmentType.getShape(), elementType);
-      SmallVector<Value> widthLoopInit {fragmentInit};
-      if (hasPartialInputTile)
-        widthLoopInit.push_back(createZeroTensorConstant(partialInputScratchType, rewriter));
-      auto widthLoop = buildNormalizedScfFor(
-        rewriter,
-        loc,
-        c0,
-        cOutWidth,
-        c1,
-        widthLoopInit,
-        [&](OpBuilder&, Location widthLoc, Value widthIndex, ValueRange widthIterArgs, SmallVectorImpl<Value>& widthYielded) {
-          Value partialInputScratch = hasPartialInputTile ? widthIterArgs[1] : Value();
-          FailureOr<SmallVector<Value>> inputTiles = createConvInputTiles(*inputWindow,
-                                                                          state,
-                                                                          widthIndex,
-                                                                          partialInputScratch,
-                                                                          patchSize,
-                                                                          numKSlices,
-                                                                          xbarDim,
-                                                                          rewriter,
-                                                                          widthLoc);
-          if (failed(inputTiles))
-            return failure();
-          FailureOr<Value> outputRow = createConvOutputRow(*inputTiles,
-                                                           paddedK,
-                                                           state.numChannelsOut,
-                                                           args.weights.front(),
-                                                           state.hasBias ? args.inputs[1] : Value(),
-                                                           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), widthIndex, rewriter.getIndexAttr(0)};
-          SmallVector<OpFoldResult> rowSizes {
-            rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1),
-            rewriter.getIndexAttr(state.numChannelsOut)};
-          Value nextFragment = tensor::InsertSliceOp::create(
-            rewriter, widthLoc, outputFragment, widthIterArgs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 4));
-          widthYielded.push_back(nextFragment);
-          if (hasPartialInputTile)
-            widthYielded.push_back(partialInputScratch);
-          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);
+  return createRowStripConvOutput(
+    state, rowStripStorage, paddedWeights, state.hasBias ? *bias : Value(),
+    paddedK, numKSlices, xbarDim, rewriter, loc);
 }
 
 static FailureOr<Value> createPointwiseOutputFromRowStripFragments(Value rowStripStorage,

src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Elementwise.cpp

@@ -193,6 +193,13 @@ struct AddToSpatialCompute : OpConversionPattern<ONNXAddOp> {
       return success();
     }
 
+    if (resultType.getRank() == 4 && adaptor.getA().getType() == resultType && adaptor.getB().getType() == resultType) {
+      auto plan = spatial::SpatAddPlanOp::create(
+        rewriter, op.getLoc(), resultType, adaptor.getA(), adaptor.getB(), rewriter.getStringAttr("nchw"));
+      rewriter.replaceOp(op, plan.getResult());
+      return success();
+    }
+
     auto lhs = prepareElementwiseOperand(adaptor.getA(), resultType, rewriter, op.getLoc());
     if (failed(lhs))
       return failure();