relu_conv_relu Faster on Arch-A

src/PIM/Common/IR/AddressAnalysis.cpp

@@ -58,8 +58,18 @@ static mlir::Value resolveForYieldedAliasToInit(mlir::scf::ForOp forOp,
 mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnowledge* knowledge) {
   value = resolveAlias(value, knowledge);
 
-  if (mlir::isa<mlir::BlockArgument>(value))
+  if (auto blockArgument = mlir::dyn_cast<mlir::BlockArgument>(value)) {
+    auto forOp = mlir::dyn_cast_or_null<mlir::scf::ForOp>(blockArgument.getOwner()->getParentOp());
+    if (forOp && blockArgument.getArgNumber() > 0) {
+      const unsigned iterArgIndex = blockArgument.getArgNumber() - 1;
+      auto yieldOp = mlir::dyn_cast<mlir::scf::YieldOp>(forOp.getBody()->getTerminator());
+      if (iterArgIndex < forOp.getInitArgs().size() && yieldOp
+          && iterArgIndex < yieldOp.getNumOperands()
+          && resolveAlias(yieldOp.getOperand(iterArgIndex), knowledge) == blockArgument)
+        return resolveLoopCarriedAliasImpl(forOp.getInitArgs()[iterArgIndex], knowledge);
+    }
     return value;
+  }
 
   mlir::Operation* definingOp = value.getDefiningOp();
   if (!definingOp)

src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp

@@ -18,10 +18,10 @@ FailureOr<RowStripPhysicalValue> describeRowStripPhysicalValue(Value storage, Ra
   if (!storageType || !storageType.hasStaticShape() || !logicalType || !logicalType.hasStaticShape()
       || storageType.getRank() != 5 || logicalType.getRank() != 4 || logicalType.getDimSize(0) != 1
       || storageType.getElementType() != logicalType.getElementType()
-      || storageType.getDimSize(1) != 1 || storageType.getDimSize(2) <= 0
+      || storageType.getDimSize(1) != 1 || storageType.getDimSize(2) != 1
-      || storageType.getDimSize(3) != 1 || storageType.getDimSize(4) != logicalType.getDimSize(3))
+      || storageType.getDimSize(3) != logicalType.getDimSize(3) || storageType.getDimSize(4) <= 0)
     return failure();
-  const int64_t tilesPerRow = ceilIntegerDivide(logicalType.getDimSize(1), storageType.getDimSize(2));
+  const int64_t tilesPerRow = ceilIntegerDivide(logicalType.getDimSize(1), storageType.getDimSize(4));
   if (storageType.getDimSize(0) != logicalType.getDimSize(2) * tilesPerRow)
     return failure();
   return RowStripPhysicalValue {storage, logicalType,
@@ -30,7 +30,8 @@ FailureOr<RowStripPhysicalValue> describeRowStripPhysicalValue(Value storage, Ra
 }
 
 RankedTensorType getRowStripFragmentType(RankedTensorType logicalType) {
-  return RankedTensorType::get({logicalType.getDimSize(0), logicalType.getDimSize(1), 1, logicalType.getDimSize(3)},
+  return RankedTensorType::get({logicalType.getDimSize(0), 1, logicalType.getDimSize(3),
+                                logicalType.getDimSize(1)},
                                logicalType.getElementType(),
                                logicalType.getEncoding());
 }
@@ -78,16 +79,18 @@ void insertRowStripFragment(Value fragment,
 FailureOr<Value> createPerChannelConstantFragment(DenseElementsAttr denseAttr,
                                                   RankedTensorType fragmentType,
                                                   PatternRewriter& rewriter) {
-  FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, fragmentType);
+  auto logicalType = RankedTensorType::get(
+    {1, fragmentType.getDimSize(3), 1, 1}, fragmentType.getElementType());
+  FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, logicalType);
   if (failed(channelValues))
     return failure();
 
   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(1); ++h)
-      for (int64_t h = 0; h < fragmentType.getDimSize(2); ++h)
+      for (int64_t w = 0; w < fragmentType.getDimSize(2); ++w)
-        for (int64_t w = 0; w < fragmentType.getDimSize(3); ++w)
+        for (int64_t channel = 0; channel < fragmentType.getDimSize(3); ++channel)
           values.push_back((*channelValues)[channel]);
 
   auto attr = DenseElementsAttr::get(fragmentType, values);
@@ -117,7 +120,6 @@ FailureOr<Value> createRowStripStorageFromRows(Value rows,
     return failure();
 
   auto rowSliceType = RankedTensorType::get({width, channels}, logicalType.getElementType(), rowsType.getEncoding());
-  auto channelWidthType = RankedTensorType::get({channels, width}, logicalType.getElementType(), rowsType.getEncoding());
   auto fragmentType = getRowStripFragmentType(logicalType);
   auto storageType = getRowStripStorageType(logicalType);
   auto batchOp = createSpatComputeBatch(
@@ -128,10 +130,8 @@ FailureOr<Value> createRowStripStorageFromRows(Value rows,
       SmallVector<OpFoldResult> rowSizes {rewriter.getIndexAttr(width), rewriter.getIndexAttr(channels)};
       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}});
+        rewriter, loc, fragmentType, rowSlice, SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
       insertRowStripFragment(fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
       return success();
     });
@@ -143,8 +143,28 @@ FailureOr<Value> createRowStripStorageFromRows(Value rows,
 FailureOr<Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& value,
                                                  PatternRewriter& rewriter,
                                                  Location loc) {
+  const int64_t laneCount = cast<RankedTensorType>(value.storage.getType()).getDimSize(0);
+  const int64_t tileChannels = value.fragmentType.getDimSize(3);
+  auto nchwFragmentType = RankedTensorType::get(
+    {1, tileChannels, 1, value.logicalType.getDimSize(3)}, value.fragmentType.getElementType(),
+    value.fragmentType.getEncoding());
+  auto nchwStorageType = spatial::getGraphBatchPhysicalResultType(laneCount, nchwFragmentType);
+  auto transposed = createSpatComputeBatch(
+    rewriter, loc, TypeRange {nchwStorageType}, laneCount, {}, ValueRange {value.storage},
+    [&](detail::SpatComputeBatchBodyArgs args) {
+      FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(
+        rewriter, loc, args.inputs.front(), args.lane, value.fragmentType);
+      if (failed(fragment))
+        return failure();
+      Value nchw = ONNXTransposeOp::create(
+        rewriter, loc, nchwFragmentType, *fragment, rewriter.getI64ArrayAttr({0, 3, 1, 2}));
+      publishGraphBatchPhysicalFragment(rewriter, loc, nchw, args.outputs.front(), args.lane);
+      return success();
+    });
+  if (failed(transposed))
+    return failure();
+
   SmallVector<FragmentAssemblyEntry> entries;
-  const int64_t tileChannels = value.fragmentType.getDimSize(1);
   for (int64_t row = 0; row < value.logicalType.getDimSize(2); ++row)
     for (int64_t tile = 0; tile < value.tilesPerRow; ++tile) {
       const int64_t channelOffset = tile * tileChannels;
@@ -152,7 +172,7 @@ FailureOr<Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& va
         {1, std::min(tileChannels, value.logicalType.getDimSize(1) - channelOffset), 1,
          value.logicalType.getDimSize(3)}});
     }
-  return createFragmentAssemblyBlueprint(value.storage, value.logicalType, entries, "nchw_row_strip",
+  return createFragmentAssemblyBlueprint(transposed->getResult(0), value.logicalType, entries, "nhwc_row_strip",
                                          kRowStripIndexMap, rewriter, loc);
 }
 
@@ -193,11 +213,12 @@ FailureOr<Value> applyRowStripBiasAdd(const RowStripPhysicalValue& value,
   auto biasStorageType = spatial::getGraphBatchPhysicalResultType(value.tilesPerRow, value.fragmentType);
   SmallVector<Attribute> biasValues(
     biasStorageType.getNumElements(), cast<Attribute>(rewriter.getZeroAttr(value.fragmentType.getElementType())));
-  const int64_t tileChannels = value.fragmentType.getDimSize(1);
+  const int64_t tileChannels = value.fragmentType.getDimSize(3);
-  const int64_t width = value.fragmentType.getDimSize(3);
+  const int64_t width = value.fragmentType.getDimSize(2);
   for (int64_t channel = 0; channel < value.logicalType.getDimSize(1); ++channel)
     for (int64_t w = 0; w < width; ++w)
-      biasValues[((channel / tileChannels) * tileChannels + channel % tileChannels) * width + w] =
+      biasValues[(channel / tileChannels) * width * tileChannels + w * tileChannels
+                 + channel % tileChannels] =
         (*channelValues)[channel];
   Value biasStorage = getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(),
     DenseElementsAttr::get(biasStorageType, biasValues), biasStorageType);

src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp

@@ -6,7 +6,7 @@
 
 namespace onnx_mlir {
 
-inline constexpr llvm::StringLiteral kRowStripIndexMap = "nchw_row_strip_fragments";
+inline constexpr llvm::StringLiteral kRowStripIndexMap = "nhwc_row_strip_fragments";
 
 struct RowStripPhysicalValue {
   mlir::Value storage;

src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp

@@ -30,7 +30,7 @@ namespace onnx_mlir {
 namespace {
 
 static constexpr StringLiteral kDenseLayout = "dense_nchw";
-static constexpr StringLiteral kRowStripLayout = "nchw_row_strip";
+static constexpr StringLiteral kRowStripLayout = "nhwc_row_strip";
 
 static FailureOr<RowStripPhysicalValue> getRowStripValue(llvm::DenseMap<Value, RowStripPhysicalValue>& rowStripValues,
                                                          Value value) {
@@ -251,20 +251,8 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
         FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
         rewriter.setInsertionPoint(planOp);
         std::optional<Value> physicalInput;
-        if (succeeded(input)) {
+        if (succeeded(input))
-          if (input->tilesPerRow == 1) {
           physicalInput = input->storage;
-          }
-          else {
-            FailureOr<Value> denseInput = materializeRowStripToDense(*input, planOp.getLoc(), rewriter);
-            if (failed(denseInput)) {
-              planOp.emitOpError("failed to materialize tiled row-strip input for MaxPool");
-              signalPassFailure();
-              return;
-            }
-            planOp.getInputMutable().assign(*denseInput);
-          }
-        }
         FailureOr<Value> lowered = lowerSelectedMaxPool2DPlan(
           planOp, physicalInput, rewriter);
         if (failed(lowered)) {

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

@@ -1834,6 +1834,29 @@ static Value createPaddedInputKTiledWeightConstant(DenseElementsAttr sourceAttr,
   return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), paddedAttr, paddedType);
 }
 
+static Value createPaddedPixelMajorWeightConstant(DenseElementsAttr sourceAttr,
+                                                  const ConvLoweringState& state,
+                                                  int64_t paddedK,
+                                                  int64_t paddedC,
+                                                  PatternRewriter& rewriter) {
+  auto paddedType = RankedTensorType::get({paddedK, paddedC}, state.wType.getElementType());
+  SmallVector<Attribute> sourceValues(sourceAttr.getValues<Attribute>());
+  SmallVector<Attribute> paddedValues(
+    paddedType.getNumElements(), cast<Attribute>(rewriter.getZeroAttr(paddedType.getElementType())));
+  for (int64_t outChannel = 0; outChannel < state.numChannelsOut; ++outChannel)
+    for (int64_t kernelH = 0; kernelH < state.wHeight; ++kernelH)
+      for (int64_t kernelW = 0; kernelW < state.wWidth; ++kernelW)
+        for (int64_t inChannel = 0; inChannel < state.numChannelsIn; ++inChannel) {
+          const int64_t sourceFlatIndex =
+            (((outChannel * state.numChannelsIn) + inChannel) * state.wHeight + kernelH) * state.wWidth + kernelW;
+          const int64_t patchIndex =
+            ((kernelH * state.wWidth) + kernelW) * state.numChannelsIn + inChannel;
+          paddedValues[patchIndex * paddedC + outChannel] = sourceValues[sourceFlatIndex];
+        }
+  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(),
+    DenseElementsAttr::get(paddedType, paddedValues), paddedType);
+}
+
 static Value createPaddedOutputChannelTiledWeightConstant(DenseElementsAttr sourceAttr,
                                                           const ConvLoweringState& state,
                                                           int64_t paddedK,
@@ -1853,7 +1876,8 @@ static Value createPaddedOutputChannelTiledWeightConstant(DenseElementsAttr sour
         for (int64_t kernelW = 0; kernelW < state.wWidth; ++kernelW) {
           const int64_t sourceFlatIndex =
             (((outChannel * state.numChannelsIn) + inChannel) * state.wHeight + kernelH) * state.wWidth + kernelW;
-          const int64_t patchIndex = ((inChannel * state.wHeight) + kernelH) * state.wWidth + kernelW;
+          const int64_t patchIndex =
+            ((kernelH * state.wWidth) + kernelW) * state.numChannelsIn + inChannel;
           const int64_t destinationFlatIndex =
             ((outputTile * paddedK) + patchIndex) * xbarDim + tileChannel;
           paddedValues[destinationFlatIndex] = sourceValues[sourceFlatIndex];
@@ -2467,7 +2491,7 @@ static Value createZeroGemmBias(RankedTensorType gemmResultType, PatternRewriter
   return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), zeroAttr, gemmResultType);
 }
 
-static bool canConsumeNchwRowStripFragments(const ConvLoweringState& state, StringRef& failureReason) {
+static bool canConsumePixelMajorRowStripFragments(const ConvLoweringState& state, StringRef& failureReason) {
   if (state.batchSize != 1) {
     failureReason = "batch_not_one";
     return false;
@@ -2539,8 +2563,7 @@ static Value createZeroTensorConstant(RankedTensorType type, PatternRewriter& re
   return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), zeroAttr, type);
 }
 
-static FailureOr<Value> createPaddedBiasRowConstant(const ConvLoweringState& state,
+static FailureOr<Value> createBiasRowConstant(const ConvLoweringState& state,
-                                                    int64_t paddedChannels,
                                               PatternRewriter& rewriter) {
   DenseElementsAttr denseAttr;
   if (!isSupportedBiasAddValue(state.b, state.outType, &denseAttr))
@@ -2549,12 +2572,11 @@ static FailureOr<Value> createPaddedBiasRowConstant(const ConvLoweringState& sta
   if (failed(channelValues))
     return failure();
 
-  auto biasType = RankedTensorType::get({1, paddedChannels}, state.outType.getElementType());
+  auto biasType = RankedTensorType::get({1, state.numChannelsOut}, state.outType.getElementType());
-  SmallVector<Attribute> values(biasType.getNumElements(), cast<Attribute>(rewriter.getZeroAttr(biasType.getElementType())));
+  return getOrCreateConstant(rewriter,
-  for (int64_t channel = 0; channel < state.numChannelsOut; ++channel)
+                             rewriter.getInsertionBlock()->getParentOp(),
-    values[channel] = (*channelValues)[channel];
+                             DenseElementsAttr::get(biasType, *channelValues),
-  auto biasAttr = DenseElementsAttr::get(biasType, values);
+                             biasType);
-  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), biasAttr, biasType);
 }
 
 static FailureOr<Value> createPaddedBiasTileConstant(const ConvLoweringState& state,
@@ -2581,13 +2603,13 @@ static Value createHorizontallyPaddedRowStripFragment(Value fragment,
                                                       PatternRewriter& rewriter,
                                                       Location loc) {
   auto paddedType = RankedTensorType::get(
-    {1, state.numChannelsIn, 1, state.xWidth + state.padWidthBegin + state.padWidthEnd},
+    {1, 1, state.xWidth + state.padWidthBegin + state.padWidthEnd, state.numChannelsIn},
     state.xType.getElementType(),
     state.xType.getEncoding());
   return createZeroPaddedTensor(fragment,
                                 paddedType,
-                                {0, 0, 0, state.padWidthBegin},
+                                {0, 0, state.padWidthBegin, 0},
-                                {0, 0, 0, state.padWidthEnd},
+                                {0, 0, state.padWidthEnd, 0},
                                 rewriter,
                                 loc);
 }
@@ -2643,6 +2665,8 @@ static Value extractDenseConvWindowRow(Value denseInput,
                                        Location loc) {
   Value tableIndex = createRowStripWindowTableIndex(outputHeight, kernelRow, state, rewriter, loc);
   Value sourceRow = tensor::ExtractOp::create(rewriter, loc, sourceRowTable, ValueRange {tableIndex}).getResult();
+  auto nchwType = RankedTensorType::get(
+    {1, state.numChannelsIn, 1, state.xWidth}, state.xType.getElementType(), state.xType.getEncoding());
   auto fragmentType = getRowStripFragmentType(state.xType);
   SmallVector<OpFoldResult> offsets {
     rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), sourceRow, rewriter.getIndexAttr(0)};
@@ -2650,8 +2674,10 @@ static Value extractDenseConvWindowRow(Value denseInput,
                                    rewriter.getIndexAttr(state.numChannelsIn),
                                    rewriter.getIndexAttr(1),
                                    rewriter.getIndexAttr(state.xWidth)};
-  return tensor::ExtractSliceOp::create(
+  Value nchw = tensor::ExtractSliceOp::create(
-    rewriter, loc, fragmentType, denseInput, offsets, sizes, getUnitStrides(rewriter, 4));
+    rewriter, loc, nchwType, denseInput, offsets, sizes, getUnitStrides(rewriter, 4));
+  return ONNXTransposeOp::create(
+    rewriter, loc, fragmentType, nchw, rewriter.getI64ArrayAttr({0, 2, 3, 1}));
 }
 
 static FailureOr<Value> createRowStripWindowMaskTable(const ConvLoweringState& state, PatternRewriter& rewriter) {
@@ -2661,7 +2687,7 @@ static FailureOr<Value> createRowStripWindowMaskTable(const ConvLoweringState& s
     return failure();
 
   Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
-  auto tableType = RankedTensorType::get({state.outHeight * state.wHeight, state.numChannelsIn, 1, state.xWidth},
+  auto tableType = RankedTensorType::get({state.outHeight * state.wHeight, 1, state.xWidth, state.numChannelsIn},
                                          elementType,
                                          state.xType.getEncoding());
   Attribute zero = rewriter.getZeroAttr(elementType);
@@ -2673,8 +2699,8 @@ static FailureOr<Value> createRowStripWindowMaskTable(const ConvLoweringState& s
       int64_t sourceRow =
         outputRow * state.strideHeight + kernelRow * state.dilationHeight - 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)
+        for (int64_t channel = 0; channel < state.numChannelsIn; ++channel)
           values.push_back(value);
     }
   }
@@ -2693,9 +2719,9 @@ static Value extractProjectedRowStripWindowMask(Value maskTable,
   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)};
+                                   rewriter.getIndexAttr(state.xWidth),
+                                   rewriter.getIndexAttr(state.numChannelsIn)};
   return tensor::ExtractSliceOp::create(rewriter,
                                         loc,
                                         fragmentType,
@@ -2716,7 +2742,7 @@ static FailureOr<Value> createConvInputWindow(Value input,
   if (!denseInput && inputType != getRowStripStorageType(state.xType))
     return failure();
   auto paddedWindowType = RankedTensorType::get(
-    {1, state.numChannelsIn, state.wHeight, state.xWidth + state.padWidthBegin + state.padWidthEnd},
+    {1, state.wHeight, state.xWidth + state.padWidthBegin + state.padWidthEnd, state.numChannelsIn},
     state.xType.getElementType(),
     state.xType.getEncoding());
   Value sourceRowTable = createRowStripWindowSourceRowTable(state, rewriter);
@@ -2744,120 +2770,130 @@ static FailureOr<Value> createConvInputWindow(Value input,
                                            paddedRow,
                                            window,
                                            SmallVector<OpFoldResult> {rewriter.getIndexAttr(0),
-                                                                      rewriter.getIndexAttr(0),
                                                                       rewriter.getIndexAttr(kernelRowIndex),
+                                                                      rewriter.getIndexAttr(0),
                                                                       rewriter.getIndexAttr(0)},
                                            SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
-                                                                      rewriter.getIndexAttr(state.numChannelsIn),
                                                                       rewriter.getIndexAttr(1),
                                                                       rewriter.getIndexAttr(
                                                                         state.xWidth + state.padWidthBegin
-                                                                        + state.padWidthEnd)},
+                                                                        + state.padWidthEnd),
+                                                                      rewriter.getIndexAttr(state.numChannelsIn)},
                                            getUnitStrides(rewriter, 4));
   }
   return window;
 }
 
-static FailureOr<Value> createNchwRowStripConvPatchRow(Value paddedWindow,
+static FailureOr<Value> createPixelMajorConvPatchRow(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},
+  auto patchType = RankedTensorType::get({1, state.wHeight, state.wWidth, state.numChannelsIn},
                                          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 inputWidthOffset = affineMulConst(rewriter, loc, outputWidth, state.strideWidth, anchorOp);
-  Value patch = createConvInputPatch(paddedWindow,
+  SmallVector<OpFoldResult> offsets {
-                                     patchType,
+    rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), inputWidthOffset, rewriter.getIndexAttr(0)};
-                                     c0,
+  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
-                                     c0,
+                                   rewriter.getIndexAttr(state.wHeight),
-                                     c0,
+                                   rewriter.getIndexAttr(state.wWidth),
-                                     inputWidthOffset,
+                                   rewriter.getIndexAttr(state.numChannelsIn)};
-                                     state.dilationHeight,
+  SmallVector<OpFoldResult> strides {rewriter.getIndexAttr(1),
-                                     state.dilationWidth,
+                                     rewriter.getIndexAttr(1),
-                                     rewriter,
+                                     rewriter.getIndexAttr(state.dilationWidth),
-                                     loc);
+                                     rewriter.getIndexAttr(1)};
+  Value patch = tensor::ExtractSliceOp::create(
+    rewriter, loc, patchType, paddedWindow, offsets, sizes, strides);
   return tensor::CollapseShapeOp::create(
            rewriter, loc, rowType, patch, SmallVector<ReassociationIndices> {{0}, {1, 2, 3}})
     .getResult();
 }
 
-static FailureOr<Value> createPaddedConvOutputTile(Value paddedPatchRow,
+static FailureOr<Value> createConvOutputTile(Value patchRow,
+                                             Value partialInputScratch,
                                              Value tileWeights,
+                                             int64_t patchSize,
                                              int64_t numKSlices,
                                              int64_t xbarDim,
                                              PatternRewriter& rewriter,
                                              Location loc) {
-  auto elementType = cast<RankedTensorType>(paddedPatchRow.getType()).getElementType();
+  auto elementType = cast<RankedTensorType>(patchRow.getType()).getElementType();
   auto paddedRowType = RankedTensorType::get({1, xbarDim}, elementType);
   auto weightElementType = cast<RankedTensorType>(tileWeights.getType()).getElementType();
   auto paddedWeightTileType = RankedTensorType::get({xbarDim, xbarDim}, weightElementType);
 
-  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+  Value tileResult;
-  Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
+  for (int64_t kSlice = 0; kSlice < numKSlices; ++kSlice) {
-  Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
+    const int64_t kOffset = kSlice * xbarDim;
-  Value cNumKSlices = getOrCreateIndexConstant(rewriter, anchorOp, numKSlices);
+    const int64_t sliceSize = std::min(xbarDim, patchSize - kOffset);
-  Value cXbar = getOrCreateIndexConstant(rewriter, anchorOp, xbarDim);
+    Value inputTile;
-  auto createPiece = [&](Value kSlice, Location pieceLoc) -> Value {
+    if (sliceSize == xbarDim) {
-    Value kOffset = arith::MulIOp::create(rewriter, pieceLoc, kSlice, cXbar);
+      inputTile = extractStaticSliceOrIdentity(
-    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, tileWeights, paddedWeightTileType, bOffsets, bSizes, getUnitStrides(rewriter, 2));
-    return spatial::SpatVMMOp::create(rewriter, pieceLoc, paddedRowType, bTile, aTile).getResult();
-  };
-
-  Value tileResult = createPiece(c0, loc);
-  if (numKSlices == 1)
-    return tileResult;
-
-  auto kLoop = buildNormalizedScfFor(
         rewriter,
         loc,
-    c1,
+        patchRow,
-    cNumKSlices,
+        paddedRowType,
-    c1,
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(kOffset)},
-    ValueRange {tileResult},
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)},
-    [&](OpBuilder&, Location reduceLoc, Value kSlice, ValueRange reduceIterArgs, SmallVectorImpl<Value>& reduceYielded) {
+        getUnitStrides(rewriter, 2));
-      Value piece = createPiece(kSlice, reduceLoc);
+    }
-      reduceYielded.push_back(
+    else {
-        spatial::SpatVAddOp::create(rewriter, reduceLoc, paddedRowType, reduceIterArgs.front(), piece).getResult());
+      if (!partialInputScratch)
-      return success();
-    });
-  if (failed(kLoop))
         return failure();
-  return kLoop->results.front();
+      auto partialType = RankedTensorType::get({1, sliceSize}, elementType);
+      Value partial = extractStaticSliceOrIdentity(
+        rewriter,
+        loc,
+        patchRow,
+        partialType,
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(kOffset)},
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(sliceSize)},
+        getUnitStrides(rewriter, 2));
+      inputTile = tensor::InsertSliceOp::create(
+        rewriter,
+        loc,
+        partial,
+        partialInputScratch,
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(sliceSize)},
+        getUnitStrides(rewriter, 2));
+    }
+    SmallVector<OpFoldResult> bOffsets {
+      rewriter.getIndexAttr(kOffset), rewriter.getIndexAttr(0)};
+    SmallVector<OpFoldResult> bSizes {rewriter.getIndexAttr(xbarDim), rewriter.getIndexAttr(xbarDim)};
+    Value bTile = extractStaticSliceOrIdentity(
+      rewriter, loc, tileWeights, paddedWeightTileType, bOffsets, bSizes, getUnitStrides(rewriter, 2));
+    Value piece = spatial::SpatVMMOp::create(
+      rewriter, loc, paddedRowType, bTile, inputTile).getResult();
+    tileResult = tileResult
+                   ? spatial::SpatVAddOp::create(
+                       rewriter, loc, paddedRowType, tileResult, piece).getResult()
+                   : piece;
+  }
+  return tileResult;
 }
 
-static FailureOr<Value> createPaddedConvOutputRow(Value patchRow,
+static FailureOr<Value> createConvOutputRow(Value patchRow,
-                                                  const ConvLoweringState& state,
+                                            Value partialInputScratch,
-                                                  Value paddedWeights,
+                                            int64_t patchSize,
-                                                  Value paddedBias,
                                             int64_t paddedK,
+                                            int64_t outputChannels,
+                                            Value paddedWeights,
+                                            Value bias,
                                             int64_t numKSlices,
                                             int64_t xbarDim,
                                             PatternRewriter& rewriter,
                                             Location loc) {
-  const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
+  auto elementType = cast<RankedTensorType>(patchRow.getType()).getElementType();
-  auto elementType = state.outType.getElementType();
+  auto rowType = RankedTensorType::get({1, outputChannels}, elementType);
-  auto rowType = RankedTensorType::get({1, state.numChannelsOut}, elementType);
+  auto tileWeightsType =
-  auto paddedRowType = RankedTensorType::get({1, xbarDim}, elementType);
+    RankedTensorType::get({paddedK, xbarDim},
-  auto paddedPatchRowType = RankedTensorType::get({1, paddedK}, elementType);
+                          cast<RankedTensorType>(paddedWeights.getType()).getElementType());
-  auto tileWeightsType = RankedTensorType::get({paddedK, xbarDim}, state.wType.getElementType());
+  const int64_t outputTileCount = ceilIntegerDivide(outputChannels, xbarDim);
-  const int64_t outputTileCount = ceilIntegerDivide(state.numChannelsOut, xbarDim);
-
-  Value paddedPatchRow = patchRow;
-  if (patchSize != paddedK)
-    paddedPatchRow = createZeroPaddedTensor(
-      paddedPatchRow, paddedPatchRowType, {0, 0}, {0, paddedK - patchSize}, rewriter, loc);
 
   auto getTileWeights = [&](int64_t outputTile) {
     if (outputTileCount == 1)
@@ -2871,29 +2907,31 @@ static FailureOr<Value> createPaddedConvOutputRow(Value patchRow,
   };
 
   if (outputTileCount == 1) {
-    FailureOr<Value> rowResult = createPaddedConvOutputTile(
+    FailureOr<Value> rowResult = createConvOutputTile(
-      paddedPatchRow, getTileWeights(0), numKSlices, xbarDim, rewriter, loc);
+      patchRow, partialInputScratch, getTileWeights(0), patchSize, numKSlices, xbarDim, rewriter, loc);
     if (failed(rowResult))
       return failure();
-    if (paddedBias)
+    Value validRow = *rowResult;
-      rowResult = spatial::SpatVAddOp::create(rewriter, loc, paddedRowType, *rowResult, paddedBias).getResult();
+    if (outputChannels != xbarDim)
-    if (state.numChannelsOut == xbarDim)
+      validRow = tensor::ExtractSliceOp::create(
-      return *rowResult;
+        rewriter,
-
+        loc,
-    SmallVector<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
+        rowType,
-    SmallVector<OpFoldResult> outputSizes {
+        validRow,
-      rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut)};
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
-    return tensor::ExtractSliceOp::create(
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(outputChannels)},
-             rewriter, loc, rowType, *rowResult, outputOffsets, outputSizes, getUnitStrides(rewriter, 2))
+        getUnitStrides(rewriter, 2));
-      .getResult();
+    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 = createPaddedConvOutputTile(
+    FailureOr<Value> tileResult = createConvOutputTile(
-      paddedPatchRow, getTileWeights(outputTile), numKSlices, xbarDim, rewriter, loc);
+      patchRow, partialInputScratch, getTileWeights(outputTile), patchSize, numKSlices, xbarDim, rewriter, loc);
     if (failed(tileResult))
       return failure();
     SmallVector<OpFoldResult> tileOffsets {
@@ -2904,7 +2942,7 @@ static FailureOr<Value> createPaddedConvOutputRow(Value patchRow,
   }
 
   SmallVector<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
-  SmallVector<OpFoldResult> outputSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut)};
+  SmallVector<OpFoldResult> outputSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(outputChannels)};
   return tensor::ExtractSliceOp::create(
            rewriter, loc, rowType, paddedOutput, outputOffsets, outputSizes, getUnitStrides(rewriter, 2))
     .getResult();
@@ -2927,8 +2965,8 @@ static FailureOr<Value> createOutputChannelTiledRowStripConvOutput(const ConvLow
   auto elementType = state.outType.getElementType();
   auto paddedPatchRowType = RankedTensorType::get({1, paddedK}, elementType);
   auto paddedRowType = RankedTensorType::get({1, xbarDim}, elementType);
-  auto tilePixelType = RankedTensorType::get({1, xbarDim, 1, 1}, elementType);
+  auto tilePixelType = RankedTensorType::get({1, 1, 1, xbarDim}, elementType);
-  auto tileFragmentType = RankedTensorType::get({1, xbarDim, 1, state.outWidth}, elementType);
+  auto tileFragmentType = RankedTensorType::get({1, 1, state.outWidth, xbarDim}, elementType);
   auto tileWeightsType = RankedTensorType::get({paddedK, xbarDim}, state.wType.getElementType());
   const int64_t laneCount = state.outHeight * outputTileCount;
   auto tileStorageType = spatial::getGraphBatchPhysicalResultType(laneCount, tileFragmentType);
@@ -2963,26 +3001,35 @@ static FailureOr<Value> createOutputChannelTiledRowStripConvOutput(const ConvLow
         if (failed(inputWindow))
           return failure();
         Value fragmentInit = tensor::EmptyOp::create(rewriter, loc, tileFragmentType.getShape(), elementType);
+        SmallVector<Value> widthLoopInit {fragmentInit};
+        if (patchSize != paddedK)
+          widthLoopInit.push_back(createZeroTensorConstant(paddedPatchRowType, rewriter));
         auto widthLoop = buildNormalizedScfFor(
           rewriter,
           loc,
           c0,
           cOutWidth,
           c1,
-          ValueRange {fragmentInit},
+          widthLoopInit,
           [&](OpBuilder&,
               Location widthLoc,
               Value widthIndex,
               ValueRange widthIterArgs,
               SmallVectorImpl<Value>& widthYielded) {
             FailureOr<Value> patchRow =
-              createNchwRowStripConvPatchRow(*inputWindow, state, widthIndex, rewriter, widthLoc);
+              createPixelMajorConvPatchRow(*inputWindow, state, widthIndex, rewriter, widthLoc);
             if (failed(patchRow))
               return failure();
             Value paddedPatchRow = *patchRow;
             if (patchSize != paddedK)
-              paddedPatchRow = createZeroPaddedTensor(
+              paddedPatchRow = tensor::InsertSliceOp::create(
-                paddedPatchRow, paddedPatchRowType, {0, 0}, {0, paddedK - patchSize}, rewriter, widthLoc);
+                rewriter,
+                widthLoc,
+                paddedPatchRow,
+                widthIterArgs[1],
+                SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
+                SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(patchSize)},
+                getUnitStrides(rewriter, 2));
             FailureOr<Value> paddedOutputRow = createPaddedConvOutputTile(
               paddedPatchRow, tileWeights, numKSlices, xbarDim, rewriter, widthLoc);
             if (failed(paddedOutputRow))
@@ -2991,13 +3038,13 @@ static FailureOr<Value> createOutputChannelTiledRowStripConvOutput(const ConvLow
               paddedOutputRow = spatial::SpatVAddOp::create(
                 rewriter, widthLoc, paddedRowType, *paddedOutputRow, *biasTile).getResult();
             Value outputPixel = tensor::ExpandShapeOp::create(
-              rewriter, widthLoc, tilePixelType, *paddedOutputRow, SmallVector<ReassociationIndices> {{0}, {1, 2, 3}});
+              rewriter, widthLoc, tilePixelType, *paddedOutputRow, SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
             SmallVector<OpFoldResult> rowOffsets {
-              rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex};
+              rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex, rewriter.getIndexAttr(0)};
             SmallVector<OpFoldResult> rowSizes {rewriter.getIndexAttr(1),
-                                               rewriter.getIndexAttr(xbarDim),
                                                rewriter.getIndexAttr(1),
-                                               rewriter.getIndexAttr(1)};
+                                               rewriter.getIndexAttr(1),
+                                               rewriter.getIndexAttr(xbarDim)};
             Value nextFragment = tensor::InsertSliceOp::create(rewriter,
                                                                 widthLoc,
                                                                 outputPixel,
@@ -3006,6 +3053,8 @@ static FailureOr<Value> createOutputChannelTiledRowStripConvOutput(const ConvLow
                                                                 rowSizes,
                                                                 getUnitStrides(rewriter, 4));
             widthYielded.push_back(nextFragment);
+            if (patchSize != paddedK)
+              widthYielded.push_back(paddedPatchRow);
             return success();
           });
         if (failed(widthLoop))
@@ -3036,12 +3085,13 @@ createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRe
   const int64_t numKSlices = ceilIntegerDivide(patchSize, xbarDim);
   const int64_t paddedK = numKSlices * xbarDim;
   auto elementType = state.outType.getElementType();
+  auto paddedPatchRowType = RankedTensorType::get({1, paddedK}, elementType);
   auto fragmentType = getRowStripFragmentType(state.outType);
-  auto outputPixelType = RankedTensorType::get({1, state.numChannelsOut, 1, 1}, elementType);
+  auto outputPixelType = RankedTensorType::get({1, 1, 1, state.numChannelsOut}, elementType);
   auto outputStorageType = getRowStripStorageType(state.outType);
 
   Value paddedWeights = state.numChannelsOut <= xbarDim
-                          ? standard::createPaddedInputKTiledWeightConstant(
+                          ? standard::createPaddedPixelMajorWeightConstant(
                               weightDenseAttr, state, paddedK, xbarDim, rewriter)
                           : standard::createPaddedOutputChannelTiledWeightConstant(
                               weightDenseAttr, state, paddedK, xbarDim, rewriter);
@@ -3049,10 +3099,10 @@ createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRe
     return createOutputChannelTiledRowStripConvOutput(
       state, paddedWeights, paddedK, numKSlices, xbarDim, rewriter, loc);
 
-  FailureOr<Value> paddedBias = failure();
+  FailureOr<Value> bias = failure();
   if (state.hasBias)
-    paddedBias = createPaddedBiasRowConstant(state, xbarDim, rewriter);
+    bias = createBiasRowConstant(state, rewriter);
-  if (state.hasBias && failed(paddedBias))
+  if (state.hasBias && failed(bias))
     return failure();
 
   auto batchOp = createSpatComputeBatch(
@@ -3061,7 +3111,7 @@ createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRe
     TypeRange {outputStorageType},
     state.outHeight,
     ValueRange {paddedWeights},
-    state.hasBias ? ValueRange {state.x, *paddedBias} : ValueRange {state.x},
+    state.hasBias ? ValueRange {state.x, *bias} : ValueRange {state.x},
     [&](detail::SpatComputeBatchBodyArgs args) {
       Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
       Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
@@ -3072,23 +3122,35 @@ createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRe
       if (failed(inputWindow))
         return failure();
       Value fragmentInit = tensor::EmptyOp::create(rewriter, loc, fragmentType.getShape(), elementType);
+      SmallVector<Value> widthLoopInit {fragmentInit};
+      if (patchSize != paddedK)
+        widthLoopInit.push_back(createZeroTensorConstant(paddedPatchRowType, rewriter));
       auto widthLoop = buildNormalizedScfFor(
         rewriter,
         loc,
         c0,
         cOutWidth,
         c1,
-        ValueRange {fragmentInit},
+        widthLoopInit,
         [&](OpBuilder&, Location widthLoc, Value widthIndex, ValueRange widthIterArgs, SmallVectorImpl<Value>& widthYielded) {
           FailureOr<Value> patchRow =
-            createNchwRowStripConvPatchRow(*inputWindow, state, widthIndex, rewriter, widthLoc);
+            createPixelMajorConvPatchRow(*inputWindow, state, widthIndex, rewriter, widthLoc);
           if (failed(patchRow))
             return failure();
-          FailureOr<Value> outputRow = createPaddedConvOutputRow(*patchRow,
+          Value paddedPatchRow = *patchRow;
-                                                                 state,
+          if (patchSize != paddedK)
+            paddedPatchRow = tensor::InsertSliceOp::create(
+              rewriter,
+              widthLoc,
+              paddedPatchRow,
+              widthIterArgs[1],
+              SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
+              SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(patchSize)},
+              getUnitStrides(rewriter, 2));
+          FailureOr<Value> outputRow = createPaddedConvOutputRow(paddedPatchRow,
+                                                                 state.numChannelsOut,
                                                                  args.weights.front(),
                                                                  state.hasBias ? args.inputs[1] : Value(),
-                                                                 paddedK,
                                                                  numKSlices,
                                                                  xbarDim,
                                                                  rewriter,
@@ -3100,15 +3162,17 @@ createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRe
                                                                widthLoc,
                                                                outputPixelType,
                                                                *outputRow,
-                                                               SmallVector<ReassociationIndices> {{0}, {1, 2, 3}});
+                                                               SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
           SmallVector<OpFoldResult> rowOffsets {
-            rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex};
+            rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex, rewriter.getIndexAttr(0)};
           SmallVector<OpFoldResult> rowSizes {
-            rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut), rewriter.getIndexAttr(1),
+            rewriter.getIndexAttr(1), 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 (patchSize != paddedK)
+            widthYielded.push_back(paddedPatchRow);
           return success();
         });
       if (failed(widthLoop))
@@ -3122,7 +3186,7 @@ createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRe
   return batchOp->getResult(0);
 }
 
-static FailureOr<Value> createConvOutputFromNchwRowStripFragments(Value rowStripStorage,
+static FailureOr<Value> createConvOutputFromPixelMajorRowStripFragments(Value rowStripStorage,
                                                                         const ConvLoweringState& state,
                                                                         PatternRewriter& rewriter,
                                                                         Location loc) {
@@ -3131,7 +3195,7 @@ static FailureOr<Value> createConvOutputFromNchwRowStripFragments(Value rowStrip
     return failure();
 
   StringRef failureReason;
-  if (!canConsumeNchwRowStripFragments(state, failureReason))
+  if (!canConsumePixelMajorRowStripFragments(state, failureReason))
     return failure();
 
   ConvGeometry geometry = buildConvGeometry(state);
@@ -3140,16 +3204,18 @@ 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 outputPixelType = RankedTensorType::get({1, state.numChannelsOut, 1, 1}, elementType);
+  auto paddedPatchRowType = RankedTensorType::get({1, paddedK}, 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 = standard::createPaddedInputKTiledWeightConstant(weightDenseAttr, state, paddedK, xbarDim, rewriter);
+  Value paddedWeights =
-  FailureOr<Value> paddedBias = failure();
+    standard::createPaddedPixelMajorWeightConstant(weightDenseAttr, state, paddedK, xbarDim, rewriter);
+  FailureOr<Value> bias = failure();
   if (state.hasBias)
-    paddedBias = createPaddedBiasRowConstant(state, xbarDim, rewriter);
+    bias = createBiasRowConstant(state, rewriter);
-  if (state.hasBias && failed(paddedBias))
+  if (state.hasBias && failed(bias))
     return failure();
 
   auto batchOp = createSpatComputeBatch(
@@ -3158,7 +3224,7 @@ static FailureOr<Value> createConvOutputFromNchwRowStripFragments(Value rowStrip
     TypeRange {outputStorageType},
     state.outHeight,
     ValueRange {paddedWeights},
-    state.hasBias ? ValueRange {rowStripStorage, *paddedBias} : ValueRange {rowStripStorage},
+    state.hasBias ? ValueRange {rowStripStorage, *bias} : ValueRange {rowStripStorage},
     [&](detail::SpatComputeBatchBodyArgs args) {
       Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
       Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
@@ -3169,24 +3235,36 @@ static FailureOr<Value> createConvOutputFromNchwRowStripFragments(Value rowStrip
       if (failed(inputWindow))
         return failure();
       Value fragmentInit = tensor::EmptyOp::create(rewriter, loc, fragmentType.getShape(), elementType);
+      SmallVector<Value> widthLoopInit {fragmentInit};
+      if (patchSize != paddedK)
+        widthLoopInit.push_back(createZeroTensorConstant(paddedPatchRowType, rewriter));
       auto widthLoop = buildNormalizedScfFor(
         rewriter,
         loc,
         c0,
         cOutWidth,
         c1,
-        ValueRange {fragmentInit},
+        widthLoopInit,
         [&](OpBuilder&, Location widthLoc, Value widthIndex, ValueRange widthIterArgs, SmallVectorImpl<Value>& widthYielded) {
           FailureOr<Value> patchRow =
-            createNchwRowStripConvPatchRow(*inputWindow, state, widthIndex, rewriter, widthLoc);
+            createPixelMajorConvPatchRow(*inputWindow, state, widthIndex, rewriter, widthLoc);
           if (failed(patchRow))
             return failure();
 
-          FailureOr<Value> outputRow = createPaddedConvOutputRow(*patchRow,
+          Value paddedPatchRow = *patchRow;
-                                                                 state,
+          if (patchSize != paddedK)
+            paddedPatchRow = tensor::InsertSliceOp::create(
+              rewriter,
+              widthLoc,
+              paddedPatchRow,
+              widthIterArgs[1],
+              SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
+              SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(patchSize)},
+              getUnitStrides(rewriter, 2));
+          FailureOr<Value> outputRow = createPaddedConvOutputRow(paddedPatchRow,
+                                                                 state.numChannelsOut,
                                                                  args.weights.front(),
                                                                  state.hasBias ? args.inputs[1] : Value(),
-                                                                 paddedK,
                                                                  numKSlices,
                                                                  xbarDim,
                                                                  rewriter,
@@ -3198,15 +3276,17 @@ static FailureOr<Value> createConvOutputFromNchwRowStripFragments(Value rowStrip
                                                                widthLoc,
                                                                outputPixelType,
                                                                *outputRow,
-                                                               SmallVector<ReassociationIndices> {{0}, {1, 2, 3}});
+                                                               SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
           SmallVector<OpFoldResult> rowOffsets {
-            rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex};
+            rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex, rewriter.getIndexAttr(0)};
           SmallVector<OpFoldResult> rowSizes {
-            rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut), rewriter.getIndexAttr(1),
+            rewriter.getIndexAttr(1), 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 (patchSize != paddedK)
+            widthYielded.push_back(paddedPatchRow);
           return success();
         });
       if (failed(widthLoop))
@@ -3228,7 +3308,7 @@ static FailureOr<Value> createPointwiseOutputFromRowStripFragments(Value rowStri
   if (failed(input)) return failure();
   ConvGeometry geometry = buildConvGeometry(state);
   const int64_t xbarDim = geometry.xbarSize;
-  const int64_t inputFragmentChannels = input->fragmentType.getDimSize(1);
+  const int64_t inputFragmentChannels = input->fragmentType.getDimSize(3);
   if (inputFragmentChannels % xbarDim != 0 || state.numChannelsIn % xbarDim != 0)
     return failure();
   auto weightDenseAttr = getHostConstDenseElementsAttr(state.w);
@@ -3241,7 +3321,7 @@ static FailureOr<Value> createPointwiseOutputFromRowStripFragments(Value rowStri
   auto inputRowType = RankedTensorType::get({1, inputFragmentChannels}, elementType);
   auto weightTileType = RankedTensorType::get({state.numChannelsIn, xbarDim}, state.wType.getElementType());
   auto weightSliceType = RankedTensorType::get({xbarDim, xbarDim}, state.wType.getElementType());
-  auto outputFragmentType = RankedTensorType::get({1, xbarDim, 1, 1}, elementType);
+  auto outputFragmentType = RankedTensorType::get({1, 1, 1, xbarDim}, elementType);
   auto outputStorageType = spatial::getGraphBatchPhysicalResultType(outputTileCount, outputFragmentType);
   Value paddedWeights = standard::createPaddedOutputChannelTiledWeightConstant(
     weightDenseAttr, state, state.numChannelsIn, xbarDim, rewriter);
@@ -3272,7 +3352,7 @@ static FailureOr<Value> createPointwiseOutputFromRowStripFragments(Value rowStri
           rewriter, pieceLoc, args.inputs.front(), sourceSlot, input->fragmentType);
         if (failed(fragment)) return failure();
         Value inputRow = tensor::CollapseShapeOp::create(rewriter, pieceLoc, inputRowType, *fragment,
-          SmallVector<ReassociationIndices> {{0}, {1, 2, 3}});
+          SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
         Value inputSlice = tensor::ExtractSliceOp::create(rewriter, pieceLoc, paddedRowType, inputRow,
           SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), sourceOffset},
           SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)},
@@ -3305,7 +3385,7 @@ static FailureOr<Value> createPointwiseOutputFromRowStripFragments(Value rowStri
         result = spatial::SpatVAddOp::create(rewriter, loc, paddedRowType, *result, *bias).getResult();
       }
       Value fragment = tensor::ExpandShapeOp::create(rewriter, loc, outputFragmentType, *result,
-        SmallVector<ReassociationIndices> {{0}, {1, 2, 3}});
+        SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
       publishGraphBatchPhysicalFragment(rewriter, loc, fragment, args.outputs.front(), args.lane);
       return success();
     });
@@ -3320,7 +3400,7 @@ static FailureOr<Value> createConvOutputFromRowStripInput(const ConvLoweringStat
                                                           Location loc) {
   if (state.xHeight == 1 && state.xWidth == 1 && state.wHeight == 1 && state.wWidth == 1)
     return createPointwiseOutputFromRowStripFragments(rowStripInput, state, rewriter, loc);
-  return createConvOutputFromNchwRowStripFragments(rowStripInput, state, rewriter, loc);
+  return createConvOutputFromPixelMajorRowStripFragments(rowStripInput, state, rewriter, loc);
 }
 
 static Value createFragmentConstant(const DistributedTensorStep& step,
@@ -3350,9 +3430,9 @@ static Value createFragmentReciprocalConstant(const DistributedTensorStep& step,
       channelValues.push_back(value);
     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(1); ++h)
-        for (int64_t h = 0; h < fragmentType.getDimSize(2); ++h)
+        for (int64_t w = 0; w < fragmentType.getDimSize(2); ++w)
-          for (int64_t w = 0; w < fragmentType.getDimSize(3); ++w) {
+          for (int64_t channel = 0; channel < fragmentType.getDimSize(3); ++channel) {
             APFloat reciprocal = channelValues[channel];
             APFloat one(reciprocal.getSemantics(), 1);
             [[maybe_unused]] APFloat::opStatus status = one.divide(reciprocal, APFloat::rmNearestTiesToEven);
@@ -4183,7 +4263,7 @@ LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp) {
     return failure();
 
   StringRef failureReason;
-  return canConsumeNchwRowStripFragments(*state, failureReason) ? success() : failure();
+  return canConsumePixelMajorRowStripFragments(*state, failureReason) ? success() : failure();
 }
 
 FailureOr<Value>

src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Pool.cpp

@@ -481,9 +481,12 @@ FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
   const int64_t kernelWidth = planOp.getKernelShape()[1];
   Value input = rowStripInput.value_or(planOp.getInput());
   auto actualInputType = dyn_cast<RankedTensorType>(input.getType());
-  const bool physicalInput = actualInputType == getRowStripStorageType(inputType);
+  FailureOr<RowStripPhysicalValue> physicalValue = describeRowStripPhysicalValue(input, inputType);
+  const bool physicalInput = succeeded(physicalValue);
   if (!physicalInput && actualInputType != inputType)
     return failure();
+  const int64_t tilesPerRow = physicalInput ? physicalValue->tilesPerRow : 1;
+  const int64_t tileChannels = physicalInput ? physicalValue->fragmentType.getDimSize(3) : channels;
 
   Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
   Value rowTable = createClampedPoolIndexTable(rewriter,
@@ -502,51 +505,78 @@ FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
                                                   planOp.getDilations()[1],
                                                   planOp.getPads()[1],
                                                   inputWidth);
-  auto inputFragmentType = getRowStripFragmentType(inputType);
+  auto inputFragmentType =
-  auto outputFragmentType = getRowStripFragmentType(outputType);
+    physicalInput ? physicalValue->fragmentType : getRowStripFragmentType(inputType);
-  auto outputStorageType = getRowStripStorageType(outputType);
+  auto nchwInputFragmentType = RankedTensorType::get(
-  auto tileType = RankedTensorType::get({1, channels, 1, 1}, outputType.getElementType());
+    {1, channels, 1, inputWidth}, inputType.getElementType(), inputType.getEncoding());
+  auto outputFragmentType = RankedTensorType::get(
+    {1, 1, outputWidth, tileChannels}, outputType.getElementType(), outputType.getEncoding());
+  auto outputStorageType =
+    spatial::getGraphBatchPhysicalResultType(outputHeight * tilesPerRow, outputFragmentType);
+  auto tileType = RankedTensorType::get({1, 1, 1, tileChannels}, outputType.getElementType());
   auto batch = createSpatComputeBatch(
     rewriter,
     loc,
     TypeRange {outputStorageType},
-    outputHeight,
+    outputHeight * tilesPerRow,
     {},
     ValueRange {input},
     [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
       SmallVector<Value> inputRows;
       inputRows.reserve(kernelHeight);
+      Value outputRow = tilesPerRow == 1
+                          ? args.lane
+                          : affineFloorDivConst(rewriter, loc, args.lane, tilesPerRow, anchorOp);
+      Value channelTile = tilesPerRow == 1
+                            ? getOrCreateIndexConstant(rewriter, anchorOp, 0)
+                            : affineModConst(rewriter, loc, args.lane, tilesPerRow, anchorOp);
       for (int64_t kernelRow = 0; kernelRow < kernelHeight; ++kernelRow) {
         Value sourceRow =
-          extractPoolIndex(rewriter, loc, anchorOp, rowTable, args.lane, kernelRow, kernelHeight);
+          extractPoolIndex(rewriter, loc, anchorOp, rowTable, outputRow, kernelRow, kernelHeight);
         if (physicalInput) {
-          inputRows.push_back(
+          Value sourceSlot = sourceRow;
-            extractRowStripFragment(args.inputs.front(), inputType, sourceRow, rewriter, loc));
+          if (tilesPerRow != 1) {
+            sourceSlot = arith::AddIOp::create(
+              rewriter,
+              loc,
+              arith::MulIOp::create(rewriter,
+                                    loc,
+                                    sourceRow,
+                                    getOrCreateIndexConstant(rewriter, anchorOp, tilesPerRow)),
+              channelTile);
+          }
+          FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(
+            rewriter, loc, args.inputs.front(), sourceSlot, inputFragmentType);
+          if (failed(fragment))
+            return failure();
+          inputRows.push_back(*fragment);
         }
         else {
           SmallVector<OpFoldResult> offsets {
             rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), sourceRow, rewriter.getIndexAttr(0)};
-          inputRows.push_back(tensor::ExtractSliceOp::create(rewriter,
+          Value nchw = tensor::ExtractSliceOp::create(rewriter,
                                                      loc,
-                                                            inputFragmentType,
+                                                     nchwInputFragmentType,
                                                      args.inputs.front(),
                                                      offsets,
                                                      SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
                                                                                 rewriter.getIndexAttr(channels),
                                                                                 rewriter.getIndexAttr(1),
                                                                                 rewriter.getIndexAttr(inputWidth)},
-                                                            getUnitStrides(rewriter, 4)));
+                                                     getUnitStrides(rewriter, 4));
+          inputRows.push_back(ONNXTransposeOp::create(
+            rewriter, loc, inputFragmentType, nchw, rewriter.getI64ArrayAttr({0, 2, 3, 1})));
         }
       }
 
       auto windowType = RankedTensorType::get(
-        {1, channels, kernelHeight, inputWidth}, inputType.getElementType(), inputType.getEncoding());
+        {1, kernelHeight, inputWidth, tileChannels}, inputType.getElementType(), inputType.getEncoding());
       Value window = tensor::EmptyOp::create(
         rewriter, loc, windowType.getShape(), windowType.getElementType());
       for (int64_t kernelRow = 0; kernelRow < kernelHeight; ++kernelRow) {
         SmallVector<OpFoldResult> offsets {rewriter.getIndexAttr(0),
-                                           rewriter.getIndexAttr(0),
                                            rewriter.getIndexAttr(kernelRow),
+                                           rewriter.getIndexAttr(0),
                                            rewriter.getIndexAttr(0)};
         window = tensor::InsertSliceOp::create(rewriter,
                                                loc,
@@ -554,9 +584,9 @@ FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
                                                window,
                                                offsets,
                                                SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
-                                                                          rewriter.getIndexAttr(channels),
                                                                           rewriter.getIndexAttr(1),
-                                                                          rewriter.getIndexAttr(inputWidth)},
+                                                                          rewriter.getIndexAttr(inputWidth),
+                                                                          rewriter.getIndexAttr(tileChannels)},
                                                getUnitStrides(rewriter, 4));
       }
 
@@ -585,43 +615,43 @@ FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
                                                     kernelColumn,
                                                     kernelWidth);
               SmallVector<OpFoldResult> offsets {
-                rewriter.getIndexAttr(0),
                 rewriter.getIndexAttr(0),
                 rewriter.getIndexAttr(kernelRow),
-                sourceColumn};
+                sourceColumn,
+                rewriter.getIndexAttr(0)};
               Value point = tensor::ExtractSliceOp::create(rewriter,
                                                            nestedLoc,
                                                            tileType,
                                                            window,
                                                            offsets,
                                                            SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
-                                                                                      rewriter.getIndexAttr(channels),
                                                                                       rewriter.getIndexAttr(1),
-                                                                                      rewriter.getIndexAttr(1)},
+                                                                                      rewriter.getIndexAttr(1),
+                                                                                      rewriter.getIndexAttr(tileChannels)},
                                                            getUnitStrides(rewriter, 4));
               reduced = reduced ? spatial::SpatVMaxOp::create(rewriter, nestedLoc, tileType, reduced, point).getResult()
                                 : materializeTileTensor(rewriter, nestedLoc, point);
             }
           }
           SmallVector<OpFoldResult> outputOffsets {
-            rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), outputColumn};
+            rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), outputColumn, rewriter.getIndexAttr(0)};
           Value updated = tensor::InsertSliceOp::create(rewriter,
                                                         nestedLoc,
                                                         reduced,
                                                         iterArgs.front(),
                                                         outputOffsets,
                                                         SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
-                                                                                   rewriter.getIndexAttr(channels),
                                                                                    rewriter.getIndexAttr(1),
-                                                                                   rewriter.getIndexAttr(1)},
+                                                                                   rewriter.getIndexAttr(1),
+                                                                                   rewriter.getIndexAttr(tileChannels)},
                                                         getUnitStrides(rewriter, 4));
           yielded.push_back(updated);
           return success();
         });
       if (failed(outputLoop))
         return failure();
-      insertRowStripFragment(
+      publishGraphBatchPhysicalFragment(
-        outputLoop->results.front(), args.outputs.front(), outputType, args.lane, rewriter, loc);
+        rewriter, loc, outputLoop->results.front(), args.outputs.front(), args.lane);
       return success();
     });
   if (failed(batch))

src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp

@@ -19,11 +19,11 @@ namespace {
 
 static constexpr StringLiteral kLogicalLayout = "nchw";
 static constexpr StringLiteral kDenseLayout = "dense_nchw";
-static constexpr StringLiteral kRowStripLayout = "nchw_row_strip";
+static constexpr StringLiteral kRowStripLayout = "nhwc_row_strip";
 
 enum class SelectedLayout {
   DenseNchw,
-  NchwRowStrip,
+  PixelMajorRowStrip,
 };
 
 static SelectedLayout getSelectedLayout(llvm::DenseMap<Value, SelectedLayout>& layouts, Value value) {
@@ -33,13 +33,13 @@ static SelectedLayout getSelectedLayout(llvm::DenseMap<Value, SelectedLayout>& l
 
 static bool usesSelectedRowStrip(Operation* user, llvm::DenseMap<Value, SelectedLayout>& layouts) {
   if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(user))
-    return getSelectedLayout(layouts, reluPlan.getResult()) == SelectedLayout::NchwRowStrip;
+    return getSelectedLayout(layouts, reluPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
   if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(user))
-    return getSelectedLayout(layouts, biasAddPlan.getResult()) == SelectedLayout::NchwRowStrip;
+    return getSelectedLayout(layouts, biasAddPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
   if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(user))
-    return getSelectedLayout(layouts, convPlan.getResult()) == SelectedLayout::NchwRowStrip;
+    return getSelectedLayout(layouts, convPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
   if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(user))
-    return getSelectedLayout(layouts, maxPoolPlan.getResult()) == SelectedLayout::NchwRowStrip;
+    return getSelectedLayout(layouts, maxPoolPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
   return false;
 }
 
@@ -77,7 +77,7 @@ static bool hasRowStripConsumer(Value value) {
 static bool canSelectConvRowStrip(spatial::SpatConv2DPlanOp convPlan,
                                   llvm::DenseMap<Value, SelectedLayout>& layouts) {
   SelectedLayout inputLayout = getSelectedLayout(layouts, convPlan.getInput());
-  if (inputLayout == SelectedLayout::NchwRowStrip)
+  if (inputLayout == SelectedLayout::PixelMajorRowStrip)
     return succeeded(canConsumeAndProduceRowStrip(convPlan));
   return succeeded(canLowerConvPlanToRowStrip(convPlan));
 }
@@ -88,23 +88,23 @@ static SelectedLayout chooseConvLayout(spatial::SpatConv2DPlanOp convPlan,
     return SelectedLayout::DenseNchw;
   if (!allUsersCanHandleRowStrip(convPlan.getResult(), layouts))
     return SelectedLayout::DenseNchw;
-  return SelectedLayout::NchwRowStrip;
+  return SelectedLayout::PixelMajorRowStrip;
 }
 
 static SelectedLayout chooseReluLayout(spatial::SpatReluPlanOp reluPlan,
                                        llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  if (getSelectedLayout(layouts, reluPlan.getInput()) != SelectedLayout::NchwRowStrip)
+  if (getSelectedLayout(layouts, reluPlan.getInput()) != SelectedLayout::PixelMajorRowStrip)
     return SelectedLayout::DenseNchw;
   if (!hasRowStripConsumer(reluPlan.getResult()))
     return SelectedLayout::DenseNchw;
   if (!allUsersCanHandleRowStrip(reluPlan.getResult(), layouts))
     return SelectedLayout::DenseNchw;
-  return SelectedLayout::NchwRowStrip;
+  return SelectedLayout::PixelMajorRowStrip;
 }
 
 static SelectedLayout chooseBiasAddLayout(spatial::SpatBiasAddPlanOp biasAddPlan,
                                           llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  if (getSelectedLayout(layouts, biasAddPlan.getInput()) != SelectedLayout::NchwRowStrip)
+  if (getSelectedLayout(layouts, biasAddPlan.getInput()) != SelectedLayout::PixelMajorRowStrip)
     return SelectedLayout::DenseNchw;
   auto resultType = dyn_cast<RankedTensorType>(biasAddPlan.getOutput().getType());
   if (!resultType || !isSupportedBiasAddValue(biasAddPlan.getBias(), resultType))
@@ -113,11 +113,11 @@ static SelectedLayout chooseBiasAddLayout(spatial::SpatBiasAddPlanOp biasAddPlan
     return SelectedLayout::DenseNchw;
   if (!allUsersCanHandleRowStrip(biasAddPlan.getResult(), layouts))
     return SelectedLayout::DenseNchw;
-  return SelectedLayout::NchwRowStrip;
+  return SelectedLayout::PixelMajorRowStrip;
 }
 
 static SelectedLayout chooseMaxPoolLayout(spatial::SpatMaxPool2DPlanOp maxPoolPlan) {
-  return succeeded(canLowerMaxPoolPlanToRowStrip(maxPoolPlan)) ? SelectedLayout::NchwRowStrip
+  return succeeded(canLowerMaxPoolPlanToRowStrip(maxPoolPlan)) ? SelectedLayout::PixelMajorRowStrip
                                                                : SelectedLayout::DenseNchw;
 }
 
@@ -237,7 +237,7 @@ struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass,
       else
         continue;
 
-      if (getSelectedLayout(layouts, producedValue) != SelectedLayout::NchwRowStrip)
+      if (getSelectedLayout(layouts, producedValue) != SelectedLayout::PixelMajorRowStrip)
         continue;
 
       rewriter.setInsertionPointAfter(&op);

src/PIM/Dialect/Pim/Transforms/Bufferization/PimBufferizationPass.cpp

@@ -121,6 +121,63 @@ lowerMemRefCopyToPimCopy(memref::CopyOp copyOp,
   return success();
 }
 
+static Value getForwardedInputConsumerOutput(OpOperand& use) {
+  Operation* owner = use.getOwner();
+  if (auto vmm = dyn_cast<pim::PimVMMOp>(owner))
+    return &use == &vmm.getInputMutable() ? vmm.getOutputBuffer() : Value();
+  if (isa<pim::PimVVAddOp,
+          pim::PimVVSubOp,
+          pim::PimVVMulOp,
+          pim::PimVVMaxOp,
+          pim::PimVVDMulOp>(owner))
+    return use.getOperandNumber() < 2 ? owner->getOperand(2) : Value();
+  return {};
+}
+
+static void forwardSingleConsumerContiguousInputCopies(func::FuncOp funcOp) {
+  SmallVector<memref::CopyOp> copies;
+  funcOp.walk([&](memref::CopyOp copy) { copies.push_back(copy); });
+
+  for (memref::CopyOp copy : copies) {
+    Value target = copy.getTarget();
+    auto targetAlloc = target.getDefiningOp<memref::AllocOp>();
+    if (!targetAlloc)
+      continue;
+
+    OpOperand* consumerUse = nullptr;
+    bool hasOtherUse = false;
+    for (OpOperand& use : target.getUses()) {
+      if (use.getOwner() == copy) {
+        if (&use != &copy.getTargetMutable())
+          hasOtherUse = true;
+        continue;
+      }
+      if (consumerUse)
+        hasOtherUse = true;
+      else
+        consumerUse = &use;
+    }
+    if (hasOtherUse || !consumerUse)
+      continue;
+
+    Value output = getForwardedInputConsumerOutput(*consumerUse);
+    Value source = copy.getSource();
+    if (!output || !isDeviceLocalPimAddress(source)
+        || (failed(resolveContiguousAddress(source)) && failed(compileContiguousAddressExpr(source))))
+      continue;
+
+    FailureOr<Value> sourceBase = getPimAddressBase(source);
+    FailureOr<Value> outputBase = getPimAddressBase(output);
+    if (failed(sourceBase) || failed(outputBase) || *sourceBase == *outputBase)
+      continue;
+
+    consumerUse->set(source);
+    copy.erase();
+    if (targetAlloc.use_empty())
+      targetAlloc.erase();
+  }
+}
+
 enum class ExpectedPimCopyDirection { HostToDevice, DeviceToHost, DeviceToDevice };
 
 static LogicalResult verifyPimCopyEndpoints(Operation* copy,
@@ -276,6 +333,8 @@ void PimBufferizationPass::runOnOperation() {
     return;
   }
 
+  forwardSingleConsumerContiguousInputCopies(funcOp);
+
   MLIRContext* ctx = moduleOp.getContext();
   PatternRewriter rewriter(ctx);
 

src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp

@@ -406,7 +406,8 @@ LogicalResult SpatConcatOp::verify() {
 static bool isKnownLogicalLayout(StringRef layout) { return layout == "nchw"; }
 
 static bool isKnownPhysicalLayout(StringRef layout) {
-  return layout == "dense_nchw" || layout == "nchw_row_strip" || layout == "fragmented";
+  return layout == "dense_nchw" || layout == "nchw_row_strip" || layout == "nhwc_row_strip"
+         || layout == "fragmented";
 }
 
 static LogicalResult verifyPlanTensorTypes(Operation* op, Value input, Value output, StringRef kind) {