even faster on pimcomp models

2026-07-31 21:15:28 +02:00
parent 9ca1a0ed9f
commit f4a3b012cc
49 changed files with 1923 additions and 583 deletions
@@ -1256,7 +1256,7 @@ static Value buildPackedWeights(DenseElementsAttr wDenseAttr,
      for (int64_t kernelIndex = 0; kernelIndex < tiling.kernelElements; ++kernelIndex) {
        const int64_t kernelH = kernelIndex / wType.getDimSize(3);
        const int64_t kernelW = kernelIndex % wType.getDimSize(3);
-        const int64_t targetRow = localChannel * tiling.kernelElements + kernelIndex;
+        const int64_t targetRow = kernelIndex * tiling.channelsPerTile + localChannel;
        for (int64_t multiplierIndex = 0; multiplierIndex < tiling.outputMultiplier; ++multiplierIndex) {
          const int64_t globalOutChannel = globalChannel * tiling.outputMultiplier + multiplierIndex;
          const int64_t sourceFlatIndex =
@@ -1326,16 +1326,49 @@ static Value createInputTile(Value input,
  Value channelOffset = tiling.channelsPerTile == 1
                        ? channelTileIndex
                        : affineMulConst(rewriter, loc, channelTileIndex, tiling.channelsPerTile, anchorOp);
-  Value tile4D = createConvInputPatch(input,
-                                      inputTileType,
-                                      batchIndex,
-                                      channelOffset,
-                                      inputHeightOffset,
-                                      inputWidthOffset,
-                                      dilationHeight,
-                                      dilationWidth,
-                                      rewriter,
-                                      loc);
+  Value tile4D;
+  if (dilationHeight == 1 && dilationWidth == 1) {
+    SmallVector<OpFoldResult> offsets {batchIndex, inputHeightOffset, inputWidthOffset, channelOffset};
+    SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
+                                     rewriter.getIndexAttr(inputTileType.getDimSize(1)),
+                                     rewriter.getIndexAttr(inputTileType.getDimSize(2)),
+                                     rewriter.getIndexAttr(tiling.channelsPerTile)};
+    tile4D = tensor::ExtractSliceOp::create(
+      rewriter, loc, inputTileType, input, offsets, sizes, getUnitStrides(rewriter, 4));
+  }
+  else {
+    auto pixelType = RankedTensorType::get(
+      {1, 1, 1, tiling.channelsPerTile}, inputTileType.getElementType(), inputTileType.getEncoding());
+    tile4D = tensor::EmptyOp::create(rewriter, loc, inputTileType.getShape(), inputTileType.getElementType());
+    for (int64_t kernelH = 0; kernelH < inputTileType.getDimSize(1); ++kernelH)
+      for (int64_t kernelW = 0; kernelW < inputTileType.getDimSize(2); ++kernelW) {
+        Value sourceHeight = affineAddConst(rewriter, loc, inputHeightOffset, kernelH * dilationHeight, anchorOp);
+        Value sourceWidth = affineAddConst(rewriter, loc, inputWidthOffset, kernelW * dilationWidth, anchorOp);
+        SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
+                                        rewriter.getIndexAttr(1),
+                                        rewriter.getIndexAttr(1),
+                                        rewriter.getIndexAttr(tiling.channelsPerTile)};
+        Value pixel = tensor::ExtractSliceOp::create(
+          rewriter,
+          loc,
+          pixelType,
+          input,
+          SmallVector<OpFoldResult> {batchIndex, sourceHeight, sourceWidth, channelOffset},
+          sizes,
+          getUnitStrides(rewriter, 4));
+        tile4D = tensor::InsertSliceOp::create(
+          rewriter,
+          loc,
+          pixel,
+          tile4D,
+          SmallVector<OpFoldResult> {rewriter.getIndexAttr(0),
+                                     rewriter.getIndexAttr(kernelH),
+                                     rewriter.getIndexAttr(kernelW),
+                                     rewriter.getIndexAttr(0)},
+          sizes,
+          getUnitStrides(rewriter, 4));
+      }
+  }
  auto collapsedType = RankedTensorType::get({1, tiling.tileInputRows}, inputTileType.getElementType());
  return tensor::CollapseShapeOp::create(rewriter,
                                         loc,
@@ -1531,10 +1564,18 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter&
                                        state.padWidthEnd,
                                        rewriter,
                                        loc);
+  auto paddedInputType = cast<RankedTensorType>(paddedInput.getType());
+  auto channelLastInputType = RankedTensorType::get({paddedInputType.getDimSize(0),
+                                                      paddedInputType.getDimSize(2),
+                                                      paddedInputType.getDimSize(3),
+                                                      paddedInputType.getDimSize(1)},
+                                                     paddedInputType.getElementType());
+  Value channelLastInput = ONNXTransposeOp::create(
+    rewriter, loc, channelLastInputType, paddedInput, rewriter.getI64ArrayAttr({0, 2, 3, 1}));
  Value packedWeights = buildPackedWeights(wDenseAttr, state.wType, *tiling, rewriter, loc);

  Value expandedBias;
-  SmallVector<Value> batchInputs {paddedInput};
+  SmallVector<Value> batchInputs {channelLastInput};
  if (state.hasBias) {
    expandedBias = expandBiasIfNeeded(state.b, rewriter, loc);
    auto biasType = dyn_cast<RankedTensorType>(expandedBias.getType());
@@ -1553,9 +1594,8 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter&
    {1, static_cast<int64_t>(crossbarSize.getValue())}, state.outType.getElementType());
  auto piecesType = spatial::getGraphBatchPhysicalResultType(
    tiling->totalPatches * tiling->numChannelTiles, rowTileType);
-  auto paddedInputType = cast<RankedTensorType>(paddedInput.getType());
  auto inputTileType =
-    RankedTensorType::get({1, tiling->channelsPerTile, state.wType.getDimSize(2), state.wType.getDimSize(3)},
+    RankedTensorType::get({1, state.wType.getDimSize(2), state.wType.getDimSize(3), tiling->channelsPerTile},
                          paddedInputType.getElementType());
  SmallVector<Value> batchWeights;
  if (tiling->numChannelTiles == 1) {
@@ -1766,7 +1806,7 @@ static Value unpackRowsFromParallelGemm(Value packedRows,
 }

 static Value createWeightMatrix(
-  Value weights, const ConvGemmPlan& plan, PatternRewriter& rewriter, Location loc) {
+  Value weights, const ConvGemmPlan& plan, bool transpose, PatternRewriter& rewriter, Location loc) {
  auto buildWeightMatrix = [&](Value weight) -> Value {
    Value flattened = tensor::CollapseShapeOp::create(rewriter,
                                                      loc,
@@ -1776,6 +1816,8 @@ static Value createWeightMatrix(
                                                        {0},
                                                        {1, 2, 3}
    });
+    if (!transpose)
+      return flattened;
    return ONNXTransposeOp::create(rewriter, loc, plan.wTransType, flattened, rewriter.getI64ArrayAttr({1, 0}))
      .getResult();
  };
@@ -1783,8 +1825,9 @@ static Value createWeightMatrix(
  if (isCompileTimeComputable(weights))
    return buildWeightMatrix(weights);

+  RankedTensorType resultType = transpose ? plan.wTransType : plan.wFlatType;
  auto computeOp =
-    createSpatCompute<1>(rewriter, loc, TypeRange {plan.wTransType}, {}, ValueRange {weights}, [&](Value weight) {
+    createSpatCompute<1>(rewriter, loc, TypeRange {resultType}, {}, ValueRange {weights}, [&](Value weight) {
      spatial::SpatYieldOp::create(rewriter, loc, buildWeightMatrix(weight));
    });
  return computeOp.getResult(0);
@@ -1852,21 +1895,26 @@ static Value createPaddedPixelMajorWeightConstant(DenseElementsAttr sourceAttr,
                                                  const ConvLoweringState& state,
                                                  int64_t paddedK,
                                                  int64_t paddedC,
+                                                  int64_t packFactor,
                                                  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];
-        }
+  const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
+  for (int64_t copy = 0; copy < packFactor; ++copy)
+    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;
+            const int64_t packedRow = copy * patchSize + patchIndex;
+            const int64_t packedColumn = copy * state.numChannelsOut + outChannel;
+            paddedValues[packedRow * paddedC + packedColumn] = sourceValues[sourceFlatIndex];
+          }
  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(),
    DenseElementsAttr::get(paddedType, paddedValues), paddedType);
 }
@@ -2381,7 +2429,7 @@ static Value createStreamedConvRows(const ConvLoweringState& state,
  Value gemmBias = state.hasBias ? state.b : createZeroGemmBias(plan.gemmOutputRowsType, rewriter);
  Value packedBias = buildPackedBias(gemmBias, biasMatrix, biasDenseAttr, state, plan, rewriter, loc);
  Value gemmRows = ONNXGemmOp::create(rewriter, loc, plan.gemmOutputRowsType, inputRows,
-    packedWeights, packedBias, APFloat(1.0f), APFloat(1.0f), 0, 0).getY();
+    packedWeights, packedBias, APFloat(1.0f), APFloat(1.0f), 0, !wDenseAttr).getY();
  return maybeUnpackChunkRows(gemmRows, plan, rewriter, loc);
 }

@@ -2401,9 +2449,10 @@ static Value rewritePackedIm2ColConv(const ConvLoweringState& state,
  ConvGemmPlan plan =
    buildConvGemmPlan(state, static_cast<bool>(wDenseAttr), !state.hasBias || static_cast<bool>(biasDenseAttr), 0,
                      state.batchSize * state.outHeight * state.outWidth);
-  // Prepare weight matrix W for crossbar storage:
-  //   W: [Cout, Cin, KH, KW] -> [Cout, patchSize] -> [patchSize, Cout]
-  Value weightMatrix = createWeightMatrix(state.w, plan, rewriter, loc);
+  // Static weights use the crossbar [patchSize, Cout] layout. Runtime weights
+  // stay in ONNX's contiguous [Cout, patchSize] layout and Gemm consumes them
+  // through transB without materializing a transpose.
+  Value weightMatrix = createWeightMatrix(state.w, plan, static_cast<bool>(wDenseAttr), rewriter, loc);
  Value gemmInputRows = createIm2colRows(state, preparedInput, plan, rewriter, loc);
  Value gemmB = buildPackedWeights(wDenseAttr, weightMatrix, state, plan, rewriter, loc);
  Value gemmBias = createZeroGemmBias(plan.gemmOutputRowsType, rewriter);
@@ -2420,7 +2469,7 @@ static Value rewritePackedIm2ColConv(const ConvLoweringState& state,
                                      APFloat(1.0f),
                                      APFloat(1.0f),
                                      /*transA=*/0,
-                                      /*transB=*/0)
+                                      /*transB=*/!wDenseAttr)
                     .getY();

  return createCollectedConvOutput(ValueRange {gemmRows},
@@ -2452,7 +2501,7 @@ static Value rewriteStreamedConv(const ConvLoweringState& state,

  ConvGemmPlan seedPlan = buildConvGemmPlan(
    state, static_cast<bool>(wDenseAttr), !state.hasBias || static_cast<bool>(biasDenseAttr), 0, 1, forcedPackFactor);
-  Value weightMatrix = createWeightMatrix(state.w, seedPlan, rewriter, loc);
+  Value weightMatrix = createWeightMatrix(state.w, seedPlan, static_cast<bool>(wDenseAttr), rewriter, loc);
  Value collectedRows = createStreamedConvRows(state,
                                              preparedInput,
                                              weightMatrix,
@@ -2516,6 +2565,20 @@ static bool rowStripOutputChannelTileFitsOneCore(const ConvGeometry& geometry) {
         <= static_cast<int64_t>(crossbarCountInCore.getValue());
 }

+static int64_t chooseRowStripPixelPackFactor(const ConvLoweringState& state, int64_t xbarDim) {
+  const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
+  const int64_t baseWeightGroups = ceilIntegerDivide(patchSize, xbarDim)
+                                 * ceilIntegerDivide(state.numChannelsOut, xbarDim);
+  int64_t factor = std::min(state.outWidth, xbarDim / state.numChannelsOut);
+  while (factor > 1
+         && (state.outWidth % factor != 0
+             || ceilIntegerDivide(factor * patchSize, xbarDim)
+                  * ceilIntegerDivide(factor * state.numChannelsOut, xbarDim)
+                > baseWeightGroups))
+    --factor;
+  return std::max<int64_t>(factor, 1);
+}
+
 static bool canConsumePixelMajorRowStripFragments(const ConvLoweringState& state, StringRef& failureReason) {
  if (state.batchSize != 1) {
    failureReason = "batch_not_one";
@@ -2558,6 +2621,7 @@ static Value createZeroTensorConstant(RankedTensorType type, PatternRewriter& re
 }

 static FailureOr<Value> createBiasRowConstant(const ConvLoweringState& state,
+                                              int64_t packFactor,
                                              PatternRewriter& rewriter) {
  DenseElementsAttr denseAttr;
  if (!isSupportedBiasAddValue(state.b, state.outType, &denseAttr))
@@ -2566,10 +2630,14 @@ static FailureOr<Value> createBiasRowConstant(const ConvLoweringState& state,
  if (failed(channelValues))
    return failure();

-  auto biasType = RankedTensorType::get({1, state.numChannelsOut}, state.outType.getElementType());
+  SmallVector<Attribute> packedValues;
+  packedValues.reserve(packFactor * state.numChannelsOut);
+  for (int64_t copy = 0; copy < packFactor; ++copy)
+    packedValues.append(channelValues->begin(), channelValues->end());
+  auto biasType = RankedTensorType::get({1, packFactor * state.numChannelsOut}, state.outType.getElementType());
  return getOrCreateConstant(rewriter,
                             rewriter.getInsertionBlock()->getParentOp(),
-                             DenseElementsAttr::get(biasType, *channelValues),
+                             DenseElementsAttr::get(biasType, packedValues),
                             biasType);
 }

@@ -2825,39 +2893,47 @@ static FailureOr<Value> createConvInputWindow(Value input,
  Value initWindow = createZeroTensorConstant(paddedWindowType, rewriter);

  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
-  Value window = initWindow;
-  for (int64_t kernelRowIndex = 0; kernelRowIndex < state.wHeight; ++kernelRowIndex) {
-    Value kernelRow = getOrCreateIndexConstant(rewriter, anchorOp, kernelRowIndex);
+  Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
+  Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
+  Value cKernelRows = getOrCreateIndexConstant(rewriter, anchorOp, state.wHeight);
+  auto loop = buildNormalizedScfFor(
+    rewriter,
+    loc,
+    c0,
+    cKernelRows,
+    c1,
+    ValueRange {initWindow},
+    [&](OpBuilder&, Location rowLoc, Value kernelRow, ValueRange iterArgs,
+        SmallVectorImpl<Value>& yielded) -> LogicalResult {
    FailureOr<Value> sourceRow =
      denseInput
        ? FailureOr<Value>(
-            extractDenseConvWindowRow(input, sourceIndexTable, state, outputHeight, kernelRow, rewriter, loc))
-        : extractProjectedRowStripWindowRow(input, sourceIndexTable, state, outputHeight, kernelRow, rewriter, loc);
+            extractDenseConvWindowRow(input, sourceIndexTable, state, outputHeight, kernelRow, rewriter, rowLoc))
+        : extractProjectedRowStripWindowRow(input, sourceIndexTable, state, outputHeight, kernelRow, rewriter, rowLoc);
    if (failed(sourceRow))
      return failure();
    Value semanticRow = *sourceRow;
    if (state.padHeightBegin != 0 || state.padHeightEnd != 0) {
-      Value mask = extractProjectedRowStripWindowMask(*maskTable, state, outputHeight, kernelRow, rewriter, loc);
-      semanticRow = spatial::SpatVMulOp::create(rewriter, loc, fragmentType, semanticRow, mask).getResult();
+      Value mask = extractProjectedRowStripWindowMask(*maskTable, state, outputHeight, kernelRow, rewriter, rowLoc);
+      semanticRow = spatial::SpatVMulOp::create(rewriter, rowLoc, fragmentType, semanticRow, mask).getResult();
    }
-    Value paddedRow = createHorizontallyPaddedRowStripFragment(semanticRow, state, rewriter, loc);
-    window = tensor::InsertSliceOp::create(rewriter,
-                                           loc,
-                                           paddedRow,
-                                           window,
-                                           SmallVector<OpFoldResult> {rewriter.getIndexAttr(0),
-                                                                      rewriter.getIndexAttr(kernelRowIndex),
-                                                                      rewriter.getIndexAttr(0),
-                                                                      rewriter.getIndexAttr(0)},
-                                           SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
-                                                                      rewriter.getIndexAttr(1),
-                                                                      rewriter.getIndexAttr(
-                                                                        state.xWidth + state.padWidthBegin
-                                                                        + state.padWidthEnd),
-                                                                      rewriter.getIndexAttr(state.numChannelsIn)},
-                                           getUnitStrides(rewriter, 4));
-  }
-  return window;
+    Value paddedRow = createHorizontallyPaddedRowStripFragment(semanticRow, state, rewriter, rowLoc);
+    yielded.push_back(tensor::InsertSliceOp::create(
+      rewriter,
+      rowLoc,
+      paddedRow,
+      iterArgs.front(),
+      SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), kernelRow,
+                                 rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
+      SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1),
+                                 rewriter.getIndexAttr(state.xWidth + state.padWidthBegin
+                                                       + state.padWidthEnd),
+                                 rewriter.getIndexAttr(state.numChannelsIn)},
+      getUnitStrides(rewriter, 4)));
+    return success();
+  });
+  return failed(loop) ? FailureOr<Value>(failure())
+                      : FailureOr<Value>(loop->results.front());
 }

 static FailureOr<Value> createPixelMajorConvPatchRow(Value paddedWindow,
@@ -2878,20 +2954,92 @@ static FailureOr<Value> createPixelMajorConvPatchRow(Value paddedWindow,
                                   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);
+  Value patch;
+  if (state.dilationWidth == 1)
+    patch = tensor::ExtractSliceOp::create(
+      rewriter, loc, patchType, paddedWindow, offsets, sizes, getUnitStrides(rewriter, 4));
+  else {
+    auto columnType = RankedTensorType::get({1, state.wHeight, 1, state.numChannelsIn},
+                                            state.xType.getElementType(), state.xType.getEncoding());
+    patch = tensor::EmptyOp::create(rewriter, loc, patchType.getShape(), patchType.getElementType());
+    for (int64_t kernelColumn = 0; kernelColumn < state.wWidth; ++kernelColumn) {
+      Value sourceWidth =
+        affineAddConst(rewriter, loc, inputWidthOffset, kernelColumn * state.dilationWidth, anchorOp);
+      SmallVector<OpFoldResult> columnSizes {rewriter.getIndexAttr(1),
+                                             rewriter.getIndexAttr(state.wHeight),
+                                             rewriter.getIndexAttr(1),
+                                             rewriter.getIndexAttr(state.numChannelsIn)};
+      Value column = tensor::ExtractSliceOp::create(
+        rewriter, loc, columnType, paddedWindow,
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), sourceWidth,
+                                   rewriter.getIndexAttr(0)},
+        columnSizes, getUnitStrides(rewriter, 4));
+      patch = tensor::InsertSliceOp::create(
+        rewriter, loc, column, patch,
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0),
+                                   rewriter.getIndexAttr(kernelColumn), rewriter.getIndexAttr(0)},
+        columnSizes, getUnitStrides(rewriter, 4));
+    }
+  }
  return tensor::CollapseShapeOp::create(
           rewriter, loc, rowType, patch, SmallVector<ReassociationIndices> {{0}, {1, 2, 3}})
    .getResult();
 }

+static FailureOr<Value> createPackedPixelMajorConvPatchRow(Value paddedWindow,
+                                                           const ConvLoweringState& state,
+                                                           Value outputGroup,
+                                                           int64_t packFactor,
+                                                           PatternRewriter& rewriter,
+                                                           Location loc) {
+  if (packFactor == 1)
+    return createPixelMajorConvPatchRow(paddedWindow, state, outputGroup, rewriter, loc);
+
+  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+  const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
+  auto packedType = RankedTensorType::get(
+    {1, packFactor * patchSize}, state.xType.getElementType(), state.xType.getEncoding());
+  Value packed = tensor::EmptyOp::create(rewriter, loc, packedType.getShape(), packedType.getElementType());
+  Value outputStart = affineMulConst(rewriter, loc, outputGroup, packFactor, anchorOp);
+  Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
+  Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
+  Value cPackFactor = getOrCreateIndexConstant(rewriter, anchorOp, packFactor);
+  auto loop = buildNormalizedScfFor(
+    rewriter,
+    loc,
+    c0,
+    cPackFactor,
+    c1,
+    ValueRange {packed},
+    [&](OpBuilder&, Location copyLoc, Value copy, ValueRange iterArgs,
+        SmallVectorImpl<Value>& yielded) -> LogicalResult {
+      Value outputWidth = createOrFoldAffineApply(
+        rewriter, copyLoc, rewriter.getAffineDimExpr(0) + rewriter.getAffineDimExpr(1),
+        ValueRange {outputStart, copy}, anchorOp);
+      FailureOr<Value> patch =
+        createPixelMajorConvPatchRow(paddedWindow, state, outputWidth, rewriter, copyLoc);
+      if (failed(patch))
+        return failure();
+      Value packedOffset = affineMulConst(rewriter, copyLoc, copy, patchSize, anchorOp);
+      yielded.push_back(tensor::InsertSliceOp::create(
+        rewriter,
+        copyLoc,
+        *patch,
+        iterArgs.front(),
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), packedOffset},
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(patchSize)},
+        getUnitStrides(rewriter, 2)));
+      return success();
+    });
+  if (failed(loop))
+    return failure();
+  return loop->results.front();
+}
+
 static FailureOr<SmallVector<Value>> createConvInputTiles(Value paddedWindow,
                                                          const ConvLoweringState& state,
                                                          Value outputWidth,
+                                                          int64_t packFactor,
                                                          Value& partialInputScratch,
                                                          int64_t patchSize,
                                                          int64_t numKSlices,
@@ -2903,7 +3051,7 @@ static FailureOr<SmallVector<Value>> createConvInputTiles(Value paddedWindow,
  SmallVector<Value> inputTiles;
  inputTiles.reserve(numKSlices);

-  if (state.numChannelsIn % xbarDim == 0) {
+  if (packFactor == 1 && state.numChannelsIn % xbarDim == 0) {
    Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
    auto inputTileType = RankedTensorType::get(
      {1, 1, 1, xbarDim}, elementType, state.xType.getEncoding());
@@ -2938,8 +3086,8 @@ static FailureOr<SmallVector<Value>> createConvInputTiles(Value paddedWindow,
    return inputTiles;
  }

-  FailureOr<Value> patchRow =
-    createPixelMajorConvPatchRow(paddedWindow, state, outputWidth, rewriter, loc);
+  FailureOr<Value> patchRow = createPackedPixelMajorConvPatchRow(
+    paddedWindow, state, outputWidth, packFactor, rewriter, loc);
  if (failed(patchRow))
    return failure();

@@ -3081,17 +3229,18 @@ static FailureOr<Value> createRowStripConvOutput(const ConvLoweringState& state,
                                                 Value input,
                                                 Value paddedWeights,
                                                 Value bias,
+                                                 int64_t packFactor,
                                                 int64_t paddedK,
                                                 int64_t numKSlices,
                                                 int64_t xbarDim,
                                                 PatternRewriter& rewriter,
                                                 Location loc) {
  const int64_t laneCount = state.outHeight;
-  const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
+  const int64_t patchSize = packFactor * state.numChannelsIn * state.wHeight * state.wWidth;
  const bool hasPartialInputTile = patchSize % xbarDim != 0;
  auto elementType = state.outType.getElementType();
  auto partialInputScratchType = RankedTensorType::get({1, xbarDim}, elementType);
-  auto outputPixelType = RankedTensorType::get({1, 1, 1, state.numChannelsOut}, elementType);
+  auto outputPixelType = RankedTensorType::get({1, 1, packFactor, state.numChannelsOut}, elementType);
  auto fragmentType = getRowStripFragmentType(state.outType);
  auto storageType = getRowStripStorageType(state.outType);

@@ -3106,7 +3255,7 @@ static FailureOr<Value> createRowStripConvOutput(const ConvLoweringState& state,
      Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
      Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
      Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
-      Value cOutWidth = getOrCreateIndexConstant(rewriter, anchorOp, state.outWidth);
+      Value cOutWidth = getOrCreateIndexConstant(rewriter, anchorOp, state.outWidth / packFactor);
      FailureOr<Value> inputWindow =
        createConvInputWindow(args.inputs.front(), state, args.lane, rewriter, loc);
      if (failed(inputWindow))
@@ -3131,6 +3280,7 @@ static FailureOr<Value> createRowStripConvOutput(const ConvLoweringState& state,
          FailureOr<SmallVector<Value>> inputTiles = createConvInputTiles(*inputWindow,
                                                                          state,
                                                                          localColumn,
+                                                                          packFactor,
                                                                          partialInputScratch,
                                                                          patchSize,
                                                                          numKSlices,
@@ -3141,7 +3291,7 @@ static FailureOr<Value> createRowStripConvOutput(const ConvLoweringState& state,
            return failure();
          FailureOr<Value> output = createConvOutputRow(*inputTiles,
                                                        paddedK,
-                                                        state.numChannelsOut,
+                                                        packFactor * state.numChannelsOut,
                                                        args.weights.front(),
                                                        bias ? args.inputs[1] : Value(),
                                                        xbarDim,
@@ -3150,7 +3300,7 @@ static FailureOr<Value> createRowStripConvOutput(const ConvLoweringState& state,
          if (failed(output))
            return failure();
          Value outputPixel = tensor::ExpandShapeOp::create(
-            rewriter, pixelLoc, outputPixelType, *output, SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
+            rewriter, pixelLoc, outputPixelType, *output, SmallVector<ReassociationIndices> {{0, 1}, {2, 3}});
          Value next = tensor::InsertSliceOp::create(
            rewriter,
            pixelLoc,
@@ -3158,11 +3308,11 @@ static FailureOr<Value> createRowStripConvOutput(const ConvLoweringState& state,
            iterArgs.front(),
            SmallVector<OpFoldResult> {rewriter.getIndexAttr(0),
                                       rewriter.getIndexAttr(0),
-                                       localColumn,
+                                       affineMulConst(rewriter, pixelLoc, localColumn, packFactor, anchorOp),
                                       rewriter.getIndexAttr(0)},
            SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
                                       rewriter.getIndexAttr(1),
-                                       rewriter.getIndexAttr(1),
+                                       rewriter.getIndexAttr(packFactor),
                                       rewriter.getIndexAttr(state.numChannelsOut)},
            getUnitStrides(rewriter, 4));
          yielded.push_back(next);
@@ -3255,6 +3405,7 @@ static FailureOr<Value> createOutputChannelTiledRowStripConvOutput(const ConvLow
          FailureOr<SmallVector<Value>> inputTiles = createConvInputTiles(*inputWindow,
                                                                          state,
                                                                          widthIndex,
+                                                                          /*packFactor=*/1,
                                                                          partialInputScratch,
                                                                          patchSize,
                                                                          numKSlices,
@@ -3313,29 +3464,34 @@ static FailureOr<Value>
    return failure();

  const int64_t xbarDim = geometry.xbarSize;
-  const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
-  const int64_t numKSlices = ceilIntegerDivide(patchSize, xbarDim);
-  const int64_t paddedK = numKSlices * xbarDim;
+  const int64_t basePatchSize = state.numChannelsIn * state.wHeight * state.wWidth;
+  const int64_t baseNumKSlices = ceilIntegerDivide(basePatchSize, xbarDim);
+  const int64_t basePaddedK = baseNumKSlices * xbarDim;

  if (!rowStripOutputTileFitsOneCore(geometry)) {
    Value tiledWeights =
-      standard::createPaddedOutputChannelTiledWeightConstant(weightDenseAttr, state, paddedK, xbarDim, rewriter);
+      standard::createPaddedOutputChannelTiledWeightConstant(weightDenseAttr, state, basePaddedK, xbarDim, rewriter);
    return createOutputChannelTiledRowStripConvOutput(
-      state, state.x, tiledWeights, paddedK, numKSlices, xbarDim, rewriter, loc);
+      state, state.x, tiledWeights, basePaddedK, baseNumKSlices, xbarDim, rewriter, loc);
  }

-  const int64_t paddedOutputChannels = ceilIntegerDivide(state.numChannelsOut, xbarDim) * xbarDim;
-  Value paddedWeights =
-    standard::createPaddedPixelMajorWeightConstant(weightDenseAttr, state, paddedK, paddedOutputChannels, rewriter);
+  const int64_t packFactor = chooseRowStripPixelPackFactor(state, xbarDim);
+  const int64_t packedPatchSize = packFactor * basePatchSize;
+  const int64_t numKSlices = ceilIntegerDivide(packedPatchSize, xbarDim);
+  const int64_t paddedK = numKSlices * xbarDim;
+  const int64_t packedOutputChannels = packFactor * state.numChannelsOut;
+  const int64_t paddedOutputChannels = ceilIntegerDivide(packedOutputChannels, xbarDim) * xbarDim;
+  Value paddedWeights = standard::createPaddedPixelMajorWeightConstant(
+    weightDenseAttr, state, paddedK, paddedOutputChannels, packFactor, rewriter);
  FailureOr<Value> bias = failure();
  if (state.hasBias)
-    bias = createBiasRowConstant(state, rewriter);
+    bias = createBiasRowConstant(state, packFactor, rewriter);
  if (state.hasBias && failed(bias))
    return failure();

  return createRowStripConvOutput(
    state, state.x, paddedWeights, state.hasBias ? *bias : Value(),
-    paddedK, numKSlices, xbarDim, rewriter, loc);
+    packFactor, paddedK, numKSlices, xbarDim, rewriter, loc);
 }

 static FailureOr<Value> createConvOutputFromPixelMajorRowStripFragments(Value rowStripStorage,
@@ -3351,32 +3507,36 @@ static FailureOr<Value> createConvOutputFromPixelMajorRowStripFragments(Value ro

  ConvGeometry geometry = buildConvGeometry(state);
  const int64_t xbarDim = geometry.xbarSize;
-  const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
-  const int64_t numKSlices = ceilIntegerDivide(patchSize, xbarDim);
-  const int64_t paddedK = numKSlices * xbarDim;
+  const int64_t basePatchSize = state.numChannelsIn * state.wHeight * state.wWidth;
+  const int64_t baseNumKSlices = ceilIntegerDivide(basePatchSize, xbarDim);
+  const int64_t basePaddedK = baseNumKSlices * xbarDim;
  auto weightDenseAttr = getHostConstDenseElementsAttr(state.w);
  if (!weightDenseAttr)
    return failure();
  if (!rowStripOutputTileFitsOneCore(geometry)) {
    Value tiledWeights =
-      standard::createPaddedOutputChannelTiledWeightConstant(weightDenseAttr, state, paddedK, xbarDim, rewriter);
+      standard::createPaddedOutputChannelTiledWeightConstant(weightDenseAttr, state, basePaddedK, xbarDim, rewriter);
    return createOutputChannelTiledRowStripConvOutput(
-      state, rowStripStorage, tiledWeights, paddedK, numKSlices, xbarDim, rewriter, loc);
+      state, rowStripStorage, tiledWeights, basePaddedK, baseNumKSlices, xbarDim, rewriter, loc);
  }

-  const int64_t paddedOutputChannels =
-    ceilIntegerDivide(state.numChannelsOut, xbarDim) * xbarDim;
+  const int64_t packFactor = chooseRowStripPixelPackFactor(state, xbarDim);
+  const int64_t packedPatchSize = packFactor * basePatchSize;
+  const int64_t numKSlices = ceilIntegerDivide(packedPatchSize, xbarDim);
+  const int64_t paddedK = numKSlices * xbarDim;
+  const int64_t packedOutputChannels = packFactor * state.numChannelsOut;
+  const int64_t paddedOutputChannels = ceilIntegerDivide(packedOutputChannels, xbarDim) * xbarDim;
  Value paddedWeights = standard::createPaddedPixelMajorWeightConstant(
-    weightDenseAttr, state, paddedK, paddedOutputChannels, rewriter);
+    weightDenseAttr, state, paddedK, paddedOutputChannels, packFactor, rewriter);
  FailureOr<Value> bias = failure();
  if (state.hasBias)
-    bias = createBiasRowConstant(state, rewriter);
+    bias = createBiasRowConstant(state, packFactor, rewriter);
  if (state.hasBias && failed(bias))
    return failure();

  return createRowStripConvOutput(
    state, rowStripStorage, paddedWeights, state.hasBias ? *bias : Value(),
-    paddedK, numKSlices, xbarDim, rewriter, loc);
+    packFactor, paddedK, numKSlices, xbarDim, rewriter, loc);
 }

 static FailureOr<Value> createPointwiseOutputFromRowStripFragments(Value rowStripStorage,