relu_conv_relu Faster on Arch-A

2026-07-28 11:16:28 +02:00
parent e76c278d1c
commit e4d8fba579
9 changed files with 446 additions and 257 deletions
@@ -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,9 +2563,8 @@ static Value createZeroTensorConstant(RankedTensorType type, PatternRewriter& re
  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), zeroAttr, type);
 }

-static FailureOr<Value> createPaddedBiasRowConstant(const ConvLoweringState& state,
-                                                    int64_t paddedChannels,
-                                                    PatternRewriter& rewriter) {
+static FailureOr<Value> createBiasRowConstant(const ConvLoweringState& state,
+                                              PatternRewriter& rewriter) {
  DenseElementsAttr denseAttr;
  if (!isSupportedBiasAddValue(state.b, state.outType, &denseAttr))
    return failure();
@@ -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());
-  SmallVector<Attribute> values(biasType.getNumElements(), cast<Attribute>(rewriter.getZeroAttr(biasType.getElementType())));
-  for (int64_t channel = 0; channel < state.numChannelsOut; ++channel)
-    values[channel] = (*channelValues)[channel];
-  auto biasAttr = DenseElementsAttr::get(biasType, values);
-  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), biasAttr, biasType);
+  auto biasType = RankedTensorType::get({1, state.numChannelsOut}, state.outType.getElementType());
+  return getOrCreateConstant(rewriter,
+                             rewriter.getInsertionBlock()->getParentOp(),
+                             DenseElementsAttr::get(biasType, *channelValues),
+                             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, 0, state.padWidthEnd},
+                                {0, 0, state.padWidthBegin, 0},
+                                {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(
-    rewriter, loc, fragmentType, denseInput, offsets, sizes, getUnitStrides(rewriter, 4));
+  Value nchw = tensor::ExtractSliceOp::create(
+    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 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,
-                                                       const ConvLoweringState& state,
-                                                       Value outputWidth,
-                                                       PatternRewriter& rewriter,
-                                                       Location loc) {
+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,
-                                     patchType,
-                                     c0,
-                                     c0,
-                                     c0,
-                                     inputWidthOffset,
-                                     state.dilationHeight,
-                                     state.dilationWidth,
-                                     rewriter,
-                                     loc);
+  SmallVector<OpFoldResult> offsets {
+    rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), inputWidthOffset, rewriter.getIndexAttr(0)};
+  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
+                                   rewriter.getIndexAttr(state.wHeight),
+                                   rewriter.getIndexAttr(state.wWidth),
+                                   rewriter.getIndexAttr(state.numChannelsIn)};
+  SmallVector<OpFoldResult> strides {rewriter.getIndexAttr(1),
+                                     rewriter.getIndexAttr(1),
+                                     rewriter.getIndexAttr(state.dilationWidth),
+                                     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,
-                                                   Value tileWeights,
-                                                   int64_t numKSlices,
-                                                   int64_t xbarDim,
-                                                   PatternRewriter& rewriter,
-                                                   Location loc) {
-  auto elementType = cast<RankedTensorType>(paddedPatchRow.getType()).getElementType();
+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>(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 c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
-  Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
-  Value cNumKSlices = getOrCreateIndexConstant(rewriter, anchorOp, numKSlices);
-  Value cXbar = getOrCreateIndexConstant(rewriter, anchorOp, xbarDim);
-  auto createPiece = [&](Value kSlice, Location pieceLoc) -> Value {
-    Value kOffset = arith::MulIOp::create(rewriter, pieceLoc, kSlice, cXbar);
-    SmallVector<OpFoldResult> aOffsets {rewriter.getIndexAttr(0), kOffset};
-    SmallVector<OpFoldResult> aSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)};
-    Value aTile = extractStaticSliceOrIdentity(
-      rewriter, pieceLoc, paddedPatchRow, paddedRowType, aOffsets, aSizes, getUnitStrides(rewriter, 2));
-    SmallVector<OpFoldResult> bOffsets {kOffset, rewriter.getIndexAttr(0)};
+  Value tileResult;
+  for (int64_t kSlice = 0; kSlice < numKSlices; ++kSlice) {
+    const int64_t kOffset = kSlice * xbarDim;
+    const int64_t sliceSize = std::min(xbarDim, patchSize - kOffset);
+    Value inputTile;
+    if (sliceSize == xbarDim) {
+      inputTile = extractStaticSliceOrIdentity(
+        rewriter,
+        loc,
+        patchRow,
+        paddedRowType,
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(kOffset)},
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)},
+        getUnitStrides(rewriter, 2));
+    }
+    else {
+      if (!partialInputScratch)
+        return failure();
+      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, 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,
-    cNumKSlices,
-    c1,
-    ValueRange {tileResult},
-    [&](OpBuilder&, Location reduceLoc, Value kSlice, ValueRange reduceIterArgs, SmallVectorImpl<Value>& reduceYielded) {
-      Value piece = createPiece(kSlice, reduceLoc);
-      reduceYielded.push_back(
-        spatial::SpatVAddOp::create(rewriter, reduceLoc, paddedRowType, reduceIterArgs.front(), piece).getResult());
-      return success();
-    });
-  if (failed(kLoop))
-    return failure();
-  return kLoop->results.front();
+      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,
-                                                  const ConvLoweringState& state,
-                                                  Value paddedWeights,
-                                                  Value paddedBias,
-                                                  int64_t paddedK,
-                                                  int64_t numKSlices,
-                                                  int64_t xbarDim,
-                                                  PatternRewriter& rewriter,
-                                                  Location loc) {
-  const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
-  auto elementType = state.outType.getElementType();
-  auto rowType = RankedTensorType::get({1, state.numChannelsOut}, elementType);
-  auto paddedRowType = RankedTensorType::get({1, xbarDim}, elementType);
-  auto paddedPatchRowType = RankedTensorType::get({1, paddedK}, elementType);
-  auto tileWeightsType = RankedTensorType::get({paddedK, xbarDim}, state.wType.getElementType());
-  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);
+static FailureOr<Value> createConvOutputRow(Value patchRow,
+                                            Value partialInputScratch,
+                                            int64_t patchSize,
+                                            int64_t paddedK,
+                                            int64_t outputChannels,
+                                            Value paddedWeights,
+                                            Value bias,
+                                            int64_t numKSlices,
+                                            int64_t xbarDim,
+                                            PatternRewriter& rewriter,
+                                            Location loc) {
+  auto elementType = cast<RankedTensorType>(patchRow.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)
@@ -2871,29 +2907,31 @@ static FailureOr<Value> createPaddedConvOutputRow(Value patchRow,
  };

  if (outputTileCount == 1) {
-    FailureOr<Value> rowResult = createPaddedConvOutputTile(
-      paddedPatchRow, getTileWeights(0), numKSlices, xbarDim, rewriter, loc);
+    FailureOr<Value> rowResult = createConvOutputTile(
+      patchRow, partialInputScratch, getTileWeights(0), patchSize, numKSlices, xbarDim, rewriter, loc);
    if (failed(rowResult))
      return failure();
-    if (paddedBias)
-      rowResult = spatial::SpatVAddOp::create(rewriter, loc, paddedRowType, *rowResult, paddedBias).getResult();
-    if (state.numChannelsOut == xbarDim)
-      return *rowResult;
-
-    SmallVector<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
-    SmallVector<OpFoldResult> outputSizes {
-      rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut)};
-    return tensor::ExtractSliceOp::create(
-             rewriter, loc, rowType, *rowResult, outputOffsets, outputSizes, getUnitStrides(rewriter, 2))
-      .getResult();
+    Value validRow = *rowResult;
+    if (outputChannels != xbarDim)
+      validRow = tensor::ExtractSliceOp::create(
+        rewriter,
+        loc,
+        rowType,
+        validRow,
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(outputChannels)},
+        getUnitStrides(rewriter, 2));
+    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(
-      paddedPatchRow, getTileWeights(outputTile), numKSlices, xbarDim, rewriter, loc);
+    FailureOr<Value> tileResult = createConvOutputTile(
+      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 tileFragmentType = RankedTensorType::get({1, xbarDim, 1, state.outWidth}, elementType);
+  auto tilePixelType = RankedTensorType::get({1, 1, 1, xbarDim}, 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, paddedPatchRowType, {0, 0}, {0, paddedK - patchSize}, rewriter, widthLoc);
+              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> 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);
-  if (state.hasBias && failed(paddedBias))
+    bias = createBiasRowConstant(state, rewriter);
+  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,
-                                                                 state,
+          Value paddedPatchRow = *patchRow;
+          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,16 +3186,16 @@ createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRe
  return batchOp->getResult(0);
 }

-static FailureOr<Value> createConvOutputFromNchwRowStripFragments(Value rowStripStorage,
-                                                                  const ConvLoweringState& state,
-                                                                  PatternRewriter& rewriter,
-                                                                  Location loc) {
+static FailureOr<Value> createConvOutputFromPixelMajorRowStripFragments(Value rowStripStorage,
+                                                                        const ConvLoweringState& state,
+                                                                        PatternRewriter& rewriter,
+                                                                        Location loc) {
  auto inputType = dyn_cast<RankedTensorType>(rowStripStorage.getType());
  if (!inputType || inputType != getRowStripStorageType(state.xType))
    return failure();

  StringRef failureReason;
-  if (!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);
-  FailureOr<Value> paddedBias = failure();
+  Value paddedWeights =
+    standard::createPaddedPixelMajorWeightConstant(weightDenseAttr, state, paddedK, xbarDim, rewriter);
+  FailureOr<Value> bias = failure();
  if (state.hasBias)
-    paddedBias = createPaddedBiasRowConstant(state, xbarDim, rewriter);
-  if (state.hasBias && failed(paddedBias))
+    bias = createBiasRowConstant(state, rewriter);
+  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,
-                                                                 state,
+          Value paddedPatchRow = *patchRow;
+          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(2); ++h)
-          for (int64_t w = 0; w < fragmentType.getDimSize(3); ++w) {
+      for (int64_t h = 0; h < fragmentType.getDimSize(1); ++h)
+        for (int64_t w = 0; w < fragmentType.getDimSize(2); ++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>