Previous commit was broken in this one 9ms vs 7ms for vgg8 Arch-A

2026-07-28 12:09:59 +02:00
parent 87bd7b726d
commit 2b899b62a8
8 changed files with 518 additions and 210 deletions
@@ -2491,6 +2491,11 @@ static Value createZeroGemmBias(RankedTensorType gemmResultType, PatternRewriter
  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), zeroAttr, gemmResultType);
 }

+static bool rowStripOutputTileFitsOneCore(const ConvGeometry& geometry) {
+  return ceilIntegerDivide(geometry.k, geometry.xbarSize)
+         <= static_cast<int64_t>(crossbarCountInCore.getValue());
+}
+
 static bool canConsumePixelMajorRowStripFragments(const ConvLoweringState& state, StringRef& failureReason) {
  if (state.batchSize != 1) {
    failureReason = "batch_not_one";
@@ -2516,7 +2521,6 @@ static bool canConsumePixelMajorRowStripFragments(const ConvLoweringState& state
    failureReason = "dilation_not_one";
    return false;
  }
-  ConvGeometry geometry = buildConvGeometry(state);
  const bool pointwise = state.xHeight == 1 && state.xWidth == 1 && state.outHeight == 1 && state.outWidth == 1
                         && state.wHeight == 1 && state.wWidth == 1 && state.padHeightBegin == 0
                         && state.padHeightEnd == 0 && state.padWidthBegin == 0 && state.padWidthEnd == 0;
@@ -2547,12 +2551,12 @@ static bool canConsumePixelMajorRowStripFragments(const ConvLoweringState& state
    failureReason = "non_constant_weight";
    return false;
  }
-  if (state.hasBias && !isSupportedBiasAddValue(state.b, state.outType)) {
-    failureReason = "unsupported_bias";
+  if (!rowStripOutputTileFitsOneCore(buildConvGeometry(state))) {
+    failureReason = "output_row_does_not_fit_one_core";
    return false;
  }
-  if (geometry.c > geometry.xbarSize) {
-    failureReason = "output_channels_exceed_crossbar";
+  if (state.hasBias && !isSupportedBiasAddValue(state.b, state.outType)) {
+    failureReason = "unsupported_bias";
    return false;
  }
  return true;
@@ -2631,6 +2635,25 @@ static Value createRowStripWindowSourceRowTable(const ConvLoweringState& state,
  return getOrCreateConstant(rewriter, anchorOp, DenseElementsAttr::get(tableType, values), tableType);
 }

+static Value createRowStripWindowSourceSlotTable(const ConvLoweringState& state,
+                                                 int64_t tilesPerRow,
+                                                 PatternRewriter& rewriter) {
+  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+  auto tableType =
+    RankedTensorType::get({state.outHeight * state.wHeight * tilesPerRow}, rewriter.getIndexType());
+  SmallVector<Attribute> values;
+  values.reserve(tableType.getNumElements());
+  for (int64_t outputRow = 0; outputRow < state.outHeight; ++outputRow)
+    for (int64_t kernelRow = 0; kernelRow < state.wHeight; ++kernelRow) {
+      int64_t sourceRow =
+        outputRow * state.strideHeight + kernelRow * state.dilationHeight - state.padHeightBegin;
+      sourceRow = std::clamp(sourceRow, int64_t {0}, state.xHeight - 1);
+      for (int64_t tile = 0; tile < tilesPerRow; ++tile)
+        values.push_back(rewriter.getIndexAttr(sourceRow * tilesPerRow + tile));
+    }
+  return getOrCreateConstant(rewriter, anchorOp, DenseElementsAttr::get(tableType, values), tableType);
+}
+
 static Value createRowStripWindowTableIndex(Value outputHeight,
                                            Value kernelRow,
                                            const ConvLoweringState& state,
@@ -2644,16 +2667,72 @@ static Value createRowStripWindowTableIndex(Value outputHeight,
    rewriter, loc, outputRowExpr * state.wHeight + kernelRowExpr, ValueRange {outputHeight, kernelRow}, anchorOp);
 }

-static Value extractProjectedRowStripWindowRow(Value rowStripStorage,
-                                               Value sourceRowTable,
-                                               const ConvLoweringState& state,
-                                               Value outputHeight,
-                                               Value kernelRow,
-                                               PatternRewriter& rewriter,
-                                               Location loc) {
+static FailureOr<Value> extractProjectedRowStripWindowRow(Value rowStripStorage,
+                                                          Value sourceSlotTable,
+                                                          const ConvLoweringState& state,
+                                                          Value outputHeight,
+                                                          Value kernelRow,
+                                                          PatternRewriter& rewriter,
+                                                          Location loc) {
+  FailureOr<RowStripPhysicalValue> physical = describeRowStripPhysicalValue(rowStripStorage, state.xType);
+  if (failed(physical))
+    return failure();
  Value tableIndex = createRowStripWindowTableIndex(outputHeight, kernelRow, state, rewriter, loc);
-  Value sourceRow = tensor::ExtractOp::create(rewriter, loc, sourceRowTable, ValueRange {tableIndex}).getResult();
-  return extractRowStripFragment(rowStripStorage, state.xType, sourceRow, rewriter, loc);
+  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+  Value tileTableIndex =
+    affineMulConst(rewriter, loc, tableIndex, physical->tilesPerRow, anchorOp);
+  if (physical->tilesPerRow == 1) {
+    Value sourceSlot =
+      tensor::ExtractOp::create(rewriter, loc, sourceSlotTable, ValueRange {tileTableIndex}).getResult();
+    return extractGraphBatchPhysicalFragment(
+      rewriter, loc, rowStripStorage, sourceSlot, physical->fragmentType);
+  }
+
+  auto fullFragmentType = getRowStripFragmentType(state.xType);
+  Value fullFragment = tensor::EmptyOp::create(
+    rewriter, loc, fullFragmentType.getShape(), fullFragmentType.getElementType());
+  const int64_t tileChannels = physical->fragmentType.getDimSize(3);
+  for (int64_t tile = 0; tile < physical->tilesPerRow; ++tile) {
+    Value slotTableIndex = affineAddConst(rewriter, loc, tileTableIndex, tile, anchorOp);
+    Value tileSlot =
+      tensor::ExtractOp::create(rewriter, loc, sourceSlotTable, ValueRange {slotTableIndex}).getResult();
+    FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(
+      rewriter, loc, rowStripStorage, tileSlot, physical->fragmentType);
+    if (failed(fragment))
+      return failure();
+    const int64_t channelOffset = tile * tileChannels;
+    const int64_t validChannels = std::min(tileChannels, state.numChannelsIn - channelOffset);
+    auto validType = RankedTensorType::get(
+      {1, 1, state.xWidth, validChannels}, state.xType.getElementType(), state.xType.getEncoding());
+    Value validFragment = *fragment;
+    if (validChannels != tileChannels)
+      validFragment = tensor::ExtractSliceOp::create(
+        rewriter,
+        loc,
+        validType,
+        *fragment,
+        SmallVector<OpFoldResult>(4, rewriter.getIndexAttr(0)),
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
+                                   rewriter.getIndexAttr(1),
+                                   rewriter.getIndexAttr(state.xWidth),
+                                   rewriter.getIndexAttr(validChannels)},
+        getUnitStrides(rewriter, 4));
+    fullFragment = tensor::InsertSliceOp::create(
+      rewriter,
+      loc,
+      validFragment,
+      fullFragment,
+      SmallVector<OpFoldResult> {rewriter.getIndexAttr(0),
+                                 rewriter.getIndexAttr(0),
+                                 rewriter.getIndexAttr(0),
+                                 rewriter.getIndexAttr(channelOffset)},
+      SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
+                                 rewriter.getIndexAttr(1),
+                                 rewriter.getIndexAttr(state.xWidth),
+                                 rewriter.getIndexAttr(validChannels)},
+      getUnitStrides(rewriter, 4));
+  }
+  return fullFragment;
 }

 static Value extractDenseConvWindowRow(Value denseInput,
@@ -2739,13 +2818,17 @@ static FailureOr<Value> createConvInputWindow(Value input,
  auto fragmentType = getRowStripFragmentType(state.xType);
  auto inputType = dyn_cast<RankedTensorType>(input.getType());
  const bool denseInput = inputType == state.xType;
-  if (!denseInput && inputType != getRowStripStorageType(state.xType))
+  if (!denseInput && failed(describeRowStripPhysicalValue(input, state.xType)))
    return failure();
  auto paddedWindowType = RankedTensorType::get(
    {1, state.wHeight, state.xWidth + state.padWidthBegin + state.padWidthEnd, state.numChannelsIn},
    state.xType.getElementType(),
    state.xType.getEncoding());
-  Value sourceRowTable = createRowStripWindowSourceRowTable(state, rewriter);
+  FailureOr<RowStripPhysicalValue> physicalInput =
+    denseInput ? FailureOr<RowStripPhysicalValue>(failure()) : describeRowStripPhysicalValue(input, state.xType);
+  Value sourceIndexTable =
+    denseInput ? createRowStripWindowSourceRowTable(state, rewriter)
+               : createRowStripWindowSourceSlotTable(state, physicalInput->tilesPerRow, rewriter);
  FailureOr<Value> maskTable = createRowStripWindowMaskTable(state, rewriter);
  if (failed(maskTable))
    return failure();
@@ -2755,14 +2838,17 @@ static FailureOr<Value> createConvInputWindow(Value input,
  Value window = initWindow;
  for (int64_t kernelRowIndex = 0; kernelRowIndex < state.wHeight; ++kernelRowIndex) {
    Value kernelRow = getOrCreateIndexConstant(rewriter, anchorOp, kernelRowIndex);
-    Value sourceRow = denseInput
-                        ? extractDenseConvWindowRow(input, sourceRowTable, state, outputHeight, kernelRow, rewriter, loc)
-                        : extractProjectedRowStripWindowRow(
-                            input, sourceRowTable, state, outputHeight, kernelRow, rewriter, loc);
-    Value semanticRow = sourceRow;
+    FailureOr<Value> sourceRow =
+      denseInput
+        ? FailureOr<Value>(
+            extractDenseConvWindowRow(input, sourceIndexTable, state, outputHeight, kernelRow, rewriter, loc))
+        : extractProjectedRowStripWindowRow(input, sourceIndexTable, state, outputHeight, kernelRow, rewriter, loc);
+    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, sourceRow, mask).getResult();
+      semanticRow = spatial::SpatVMulOp::create(rewriter, loc, fragmentType, semanticRow, mask).getResult();
    }
    Value paddedRow = createHorizontallyPaddedRowStripFragment(semanticRow, state, rewriter, loc);
    window = tensor::InsertSliceOp::create(rewriter,
@@ -2814,7 +2900,7 @@ static FailureOr<Value> createPixelMajorConvPatchRow(Value paddedWindow,
 }

 static FailureOr<Value> createConvOutputTile(Value patchRow,
-                                             Value partialInputScratch,
+                                             Value& partialInputScratch,
                                             Value tileWeights,
                                             int64_t patchSize,
                                             int64_t numKSlices,
@@ -2853,7 +2939,7 @@ static FailureOr<Value> createConvOutputTile(Value patchRow,
        SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(kOffset)},
        SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(sliceSize)},
        getUnitStrides(rewriter, 2));
-      inputTile = tensor::InsertSliceOp::create(
+      partialInputScratch = tensor::InsertSliceOp::create(
        rewriter,
        loc,
        partial,
@@ -2861,6 +2947,7 @@ static FailureOr<Value> createConvOutputTile(Value patchRow,
        SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
        SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(sliceSize)},
        getUnitStrides(rewriter, 2));
+      inputTile = partialInputScratch;
    }
    SmallVector<OpFoldResult> bOffsets {
      rewriter.getIndexAttr(kOffset), rewriter.getIndexAttr(0)};
@@ -2878,7 +2965,7 @@ static FailureOr<Value> createConvOutputTile(Value patchRow,
 }

 static FailureOr<Value> createConvOutputRow(Value patchRow,
-                                            Value partialInputScratch,
+                                            Value& partialInputScratch,
                                            int64_t patchSize,
                                            int64_t paddedK,
                                            int64_t outputChannels,
@@ -2943,17 +3030,15 @@ static FailureOr<Value> createConvOutputRow(Value patchRow,

  SmallVector<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
  SmallVector<OpFoldResult> outputSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(outputChannels)};
-  return tensor::ExtractSliceOp::create(
-           rewriter, loc, rowType, paddedOutput, outputOffsets, outputSizes, getUnitStrides(rewriter, 2))
-    .getResult();
-}
-
-static bool rowStripOutputTileFitsOneCore(const ConvGeometry& geometry) {
-  return ceilIntegerDivide(geometry.k, geometry.xbarSize)
-         <= static_cast<int64_t>(crossbarCountInCore.getValue());
+  Value validRow = tensor::ExtractSliceOp::create(
+    rewriter, loc, rowType, paddedOutput, outputOffsets, outputSizes, getUnitStrides(rewriter, 2));
+  if (bias)
+    validRow = spatial::SpatVAddOp::create(rewriter, loc, rowType, validRow, bias).getResult();
+  return validRow;
 }

 static FailureOr<Value> createOutputChannelTiledRowStripConvOutput(const ConvLoweringState& state,
+                                                                   Value input,
                                                                   Value paddedWeights,
                                                                   int64_t paddedK,
                                                                   int64_t numKSlices,
@@ -2962,8 +3047,9 @@ static FailureOr<Value> createOutputChannelTiledRowStripConvOutput(const ConvLow
                                                                   Location loc) {
  const int64_t outputTileCount = ceilIntegerDivide(state.numChannelsOut, xbarDim);
  const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
+  const bool hasPartialInputTile = patchSize % xbarDim != 0;
  auto elementType = state.outType.getElementType();
-  auto paddedPatchRowType = RankedTensorType::get({1, paddedK}, elementType);
+  auto partialInputScratchType = RankedTensorType::get({1, xbarDim}, elementType);
  auto paddedRowType = RankedTensorType::get({1, xbarDim}, elementType);
  auto tilePixelType = RankedTensorType::get({1, 1, 1, xbarDim}, elementType);
  auto tileFragmentType = RankedTensorType::get({1, 1, state.outWidth, xbarDim}, elementType);
@@ -2975,101 +3061,102 @@ static FailureOr<Value> createOutputChannelTiledRowStripConvOutput(const ConvLow
    paddedBias = createPaddedBiasTileConstant(state, xbarDim, rewriter);
  if (state.hasBias && failed(paddedBias))
    return failure();
+
  auto tileBatch = createSpatComputeBatch(
-    rewriter, loc, TypeRange {tileStorageType}, laneCount, ValueRange {paddedWeights},
-    state.hasBias ? ValueRange {state.x, *paddedBias} : ValueRange {state.x},
+    rewriter,
+    loc,
+    TypeRange {tileStorageType},
+    laneCount,
+    ValueRange {paddedWeights},
+    state.hasBias ? ValueRange {input, *paddedBias} : ValueRange {input},
    [&](detail::SpatComputeBatchBodyArgs args) {
-        Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
-        Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
-        Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
-        Value cOutWidth = getOrCreateIndexConstant(rewriter, anchorOp, state.outWidth);
-        Value outputRow = affineFloorDivConst(rewriter, loc, args.lane, outputTileCount, anchorOp);
-        Value outputTile = affineModConst(rewriter, loc, args.lane, outputTileCount, anchorOp);
-        SmallVector<OpFoldResult> weightOffsets {
-          outputTile, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
-        SmallVector<OpFoldResult> weightSizes {
-          rewriter.getIndexAttr(1), rewriter.getIndexAttr(paddedK), rewriter.getIndexAttr(xbarDim)};
-        Value tileWeights = tensor::ExtractSliceOp::create(
-          rewriter, loc, tileWeightsType, args.weights.front(), weightOffsets, weightSizes, getUnitStrides(rewriter, 3));
-        FailureOr<Value> biasTile = failure();
-        if (state.hasBias)
-          biasTile = extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[1], outputTile, paddedRowType);
-        if (state.hasBias && failed(biasTile))
-          return failure();
-        FailureOr<Value> inputWindow =
-          createConvInputWindow(args.inputs.front(), state, outputRow, rewriter, loc);
-        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,
-          widthLoopInit,
-          [&](OpBuilder&,
-              Location widthLoc,
-              Value widthIndex,
-              ValueRange widthIterArgs,
-              SmallVectorImpl<Value>& widthYielded) {
-            FailureOr<Value> patchRow =
-              createPixelMajorConvPatchRow(*inputWindow, state, widthIndex, rewriter, widthLoc);
-            if (failed(patchRow))
-              return failure();
-            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> paddedOutputRow = createPaddedConvOutputTile(
-              paddedPatchRow, tileWeights, numKSlices, xbarDim, rewriter, widthLoc);
-            if (failed(paddedOutputRow))
-              return failure();
-            if (state.hasBias)
-              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}});
-            SmallVector<OpFoldResult> rowOffsets {
-              rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex, rewriter.getIndexAttr(0)};
-            SmallVector<OpFoldResult> rowSizes {rewriter.getIndexAttr(1),
-                                               rewriter.getIndexAttr(1),
-                                               rewriter.getIndexAttr(1),
-                                               rewriter.getIndexAttr(xbarDim)};
-            Value nextFragment = tensor::InsertSliceOp::create(rewriter,
-                                                                widthLoc,
-                                                                outputPixel,
-                                                                widthIterArgs.front(),
-                                                                rowOffsets,
-                                                                rowSizes,
-                                                                getUnitStrides(rewriter, 4));
-            widthYielded.push_back(nextFragment);
-            if (patchSize != paddedK)
-              widthYielded.push_back(paddedPatchRow);
-            return success();
-          });
-        if (failed(widthLoop))
-          return failure();
-        publishGraphBatchPhysicalFragment(
-          rewriter, loc, widthLoop->results.front(), args.outputs.front(), args.lane);
-        return success();
-      });
+      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 outputRow = affineFloorDivConst(rewriter, loc, args.lane, outputTileCount, anchorOp);
+      Value outputTile = affineModConst(rewriter, loc, args.lane, outputTileCount, anchorOp);
+      SmallVector<OpFoldResult> weightOffsets {
+        outputTile, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
+      SmallVector<OpFoldResult> weightSizes {
+        rewriter.getIndexAttr(1), rewriter.getIndexAttr(paddedK), rewriter.getIndexAttr(xbarDim)};
+      Value tileWeights = tensor::ExtractSliceOp::create(
+        rewriter, loc, tileWeightsType, args.weights.front(), weightOffsets, weightSizes, getUnitStrides(rewriter, 3));
+      FailureOr<Value> biasTile = failure();
+      if (state.hasBias)
+        biasTile = extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[1], outputTile, paddedRowType);
+      if (state.hasBias && failed(biasTile))
+        return failure();
+      FailureOr<Value> inputWindow =
+        createConvInputWindow(args.inputs.front(), state, outputRow, rewriter, loc);
+      if (failed(inputWindow))
+        return failure();
+      Value fragmentInit = tensor::EmptyOp::create(rewriter, loc, tileFragmentType.getShape(), elementType);
+      SmallVector<Value> widthLoopInit {fragmentInit};
+      if (hasPartialInputTile)
+        widthLoopInit.push_back(createZeroTensorConstant(partialInputScratchType, rewriter));
+      auto widthLoop = buildNormalizedScfFor(
+        rewriter,
+        loc,
+        c0,
+        cOutWidth,
+        c1,
+        widthLoopInit,
+        [&](OpBuilder&,
+            Location widthLoc,
+            Value widthIndex,
+            ValueRange widthIterArgs,
+            SmallVectorImpl<Value>& widthYielded) {
+          FailureOr<Value> patchRow =
+            createPixelMajorConvPatchRow(*inputWindow, state, widthIndex, rewriter, widthLoc);
+          if (failed(patchRow))
+            return failure();
+          Value partialInputScratch = hasPartialInputTile ? widthIterArgs[1] : Value();
+          FailureOr<Value> paddedOutputRow = createConvOutputTile(*patchRow,
+                                                                 partialInputScratch,
+                                                                 tileWeights,
+                                                                 patchSize,
+                                                                 numKSlices,
+                                                                 xbarDim,
+                                                                 rewriter,
+                                                                 widthLoc);
+          if (failed(paddedOutputRow))
+            return failure();
+          if (state.hasBias)
+            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}});
+          SmallVector<OpFoldResult> rowOffsets {
+            rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex, rewriter.getIndexAttr(0)};
+          SmallVector<OpFoldResult> rowSizes {rewriter.getIndexAttr(1),
+                                             rewriter.getIndexAttr(1),
+                                             rewriter.getIndexAttr(1),
+                                             rewriter.getIndexAttr(xbarDim)};
+          Value nextFragment = tensor::InsertSliceOp::create(rewriter,
+                                                              widthLoc,
+                                                              outputPixel,
+                                                              widthIterArgs.front(),
+                                                              rowOffsets,
+                                                              rowSizes,
+                                                              getUnitStrides(rewriter, 4));
+          widthYielded.push_back(nextFragment);
+          if (hasPartialInputTile)
+            widthYielded.push_back(partialInputScratch);
+          return success();
+        });
+      if (failed(widthLoop))
+        return failure();
+      publishGraphBatchPhysicalFragment(rewriter, loc, widthLoop->results.front(), args.outputs.front(), args.lane);
+      return success();
+    });
  if (failed(tileBatch))
    return failure();
  return tileBatch->getResult(0);
 }

 static FailureOr<Value>
-createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRewriter& rewriter, Location loc) {
+  createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRewriter& rewriter, Location loc) {
  ConvGeometry geometry = buildConvGeometry(state);
  if (state.group != 1 || state.batchSize != 1 || !rowStripOutputTileFitsOneCore(geometry))
    return failure();
@@ -3084,8 +3171,9 @@ createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRe
  const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
  const int64_t numKSlices = ceilIntegerDivide(patchSize, xbarDim);
  const int64_t paddedK = numKSlices * xbarDim;
+  const bool hasPartialInputTile = patchSize % xbarDim != 0;
  auto elementType = state.outType.getElementType();
-  auto paddedPatchRowType = RankedTensorType::get({1, paddedK}, elementType);
+  auto partialInputScratchType = RankedTensorType::get({1, xbarDim}, elementType);
  auto fragmentType = getRowStripFragmentType(state.outType);
  auto outputPixelType = RankedTensorType::get({1, 1, 1, state.numChannelsOut}, elementType);
  auto outputStorageType = getRowStripStorageType(state.outType);
@@ -3097,7 +3185,7 @@ createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRe
                              weightDenseAttr, state, paddedK, xbarDim, rewriter);
  if (state.numChannelsOut > xbarDim)
    return createOutputChannelTiledRowStripConvOutput(
-      state, paddedWeights, paddedK, numKSlices, xbarDim, rewriter, loc);
+      state, state.x, paddedWeights, paddedK, numKSlices, xbarDim, rewriter, loc);

  FailureOr<Value> bias = failure();
  if (state.hasBias)
@@ -3123,8 +3211,8 @@ createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRe
        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));
+      if (hasPartialInputTile)
+        widthLoopInit.push_back(createZeroTensorConstant(partialInputScratchType, rewriter));
      auto widthLoop = buildNormalizedScfFor(
        rewriter,
        loc,
@@ -3137,24 +3225,18 @@ createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRe
            createPixelMajorConvPatchRow(*inputWindow, state, widthIndex, rewriter, widthLoc);
          if (failed(patchRow))
            return failure();
-          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(),
-                                                                 numKSlices,
-                                                                 xbarDim,
-                                                                 rewriter,
-                                                                 widthLoc);
+          Value partialInputScratch = hasPartialInputTile ? widthIterArgs[1] : Value();
+          FailureOr<Value> outputRow = createConvOutputRow(*patchRow,
+                                                           partialInputScratch,
+                                                           patchSize,
+                                                           paddedK,
+                                                           state.numChannelsOut,
+                                                           args.weights.front(),
+                                                           state.hasBias ? args.inputs[1] : Value(),
+                                                           numKSlices,
+                                                           xbarDim,
+                                                           rewriter,
+                                                           widthLoc);
          if (failed(outputRow))
            return failure();

@@ -3171,8 +3253,8 @@ createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRe
          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);
+          if (hasPartialInputTile)
+            widthYielded.push_back(partialInputScratch);
          return success();
        });
      if (failed(widthLoop))
@@ -3190,8 +3272,7 @@ static FailureOr<Value> createConvOutputFromPixelMajorRowStripFragments(Value ro
                                                                        const ConvLoweringState& state,
                                                                        PatternRewriter& rewriter,
                                                                        Location loc) {
-  auto inputType = dyn_cast<RankedTensorType>(rowStripStorage.getType());
-  if (!inputType || inputType != getRowStripStorageType(state.xType))
+  if (failed(describeRowStripPhysicalValue(rowStripStorage, state.xType)))
    return failure();

  StringRef failureReason;
@@ -3203,15 +3284,22 @@ static FailureOr<Value> createConvOutputFromPixelMajorRowStripFragments(Value ro
  const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
  const int64_t numKSlices = ceilIntegerDivide(patchSize, xbarDim);
  const int64_t paddedK = numKSlices * xbarDim;
+  const bool hasPartialInputTile = patchSize % xbarDim != 0;
  auto elementType = state.outType.getElementType();
-  auto paddedPatchRowType = RankedTensorType::get({1, paddedK}, elementType);
+  auto partialInputScratchType = RankedTensorType::get({1, xbarDim}, elementType);
  auto outputPixelType = RankedTensorType::get({1, 1, 1, state.numChannelsOut}, elementType);
  auto outputStorageType = getRowStripStorageType(state.outType);
  auto weightDenseAttr = getHostConstDenseElementsAttr(state.w);
  if (!weightDenseAttr)
    return failure();
-  Value paddedWeights =
-    standard::createPaddedPixelMajorWeightConstant(weightDenseAttr, state, paddedK, xbarDim, rewriter);
+  Value paddedWeights = state.numChannelsOut <= xbarDim
+                          ? standard::createPaddedPixelMajorWeightConstant(
+                              weightDenseAttr, state, paddedK, xbarDim, rewriter)
+                          : standard::createPaddedOutputChannelTiledWeightConstant(
+                              weightDenseAttr, state, paddedK, xbarDim, rewriter);
+  if (state.numChannelsOut > xbarDim)
+    return createOutputChannelTiledRowStripConvOutput(
+      state, rowStripStorage, paddedWeights, paddedK, numKSlices, xbarDim, rewriter, loc);
  FailureOr<Value> bias = failure();
  if (state.hasBias)
    bias = createBiasRowConstant(state, rewriter);
@@ -3236,8 +3324,8 @@ static FailureOr<Value> createConvOutputFromPixelMajorRowStripFragments(Value ro
        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));
+      if (hasPartialInputTile)
+        widthLoopInit.push_back(createZeroTensorConstant(partialInputScratchType, rewriter));
      auto widthLoop = buildNormalizedScfFor(
        rewriter,
        loc,
@@ -3251,24 +3339,18 @@ static FailureOr<Value> createConvOutputFromPixelMajorRowStripFragments(Value ro
          if (failed(patchRow))
            return failure();

-          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(),
-                                                                 numKSlices,
-                                                                 xbarDim,
-                                                                 rewriter,
-                                                                 widthLoc);
+          Value partialInputScratch = hasPartialInputTile ? widthIterArgs[1] : Value();
+          FailureOr<Value> outputRow = createConvOutputRow(*patchRow,
+                                                           partialInputScratch,
+                                                           patchSize,
+                                                           paddedK,
+                                                           state.numChannelsOut,
+                                                           args.weights.front(),
+                                                           state.hasBias ? args.inputs[1] : Value(),
+                                                           numKSlices,
+                                                           xbarDim,
+                                                           rewriter,
+                                                           widthLoc);
          if (failed(outputRow))
            return failure();

@@ -3285,8 +3367,8 @@ static FailureOr<Value> createConvOutputFromPixelMajorRowStripFragments(Value ro
          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);
+          if (hasPartialInputTile)
+            widthYielded.push_back(partialInputScratch);
          return success();
        });
      if (failed(widthLoop))
@@ -4243,6 +4325,7 @@ LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp) {

  switch (decision.strategy) {
  case PimConvLoweringLegacy:
+  case PimConvLoweringDepthwise:
  case PimConvLoweringPackedIm2Col:
  case PimConvLoweringStreamedPatch:
  case PimConvLoweringOutputChannelTiled:
@@ -4250,7 +4333,6 @@ LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp) {
  case PimConvLoweringStreamedPacked:
    return success();
  case PimConvLoweringAuto:
-  case PimConvLoweringDepthwise:
  case PimConvLoweringInputKTiled:
    return failure();
  }