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,
@@ -16,6 +16,28 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace {

+struct SiluToSpatialPlan : OpRewritePattern<ONNXMulOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ONNXMulOp mulOp, PatternRewriter& rewriter) const override {
+    ONNXSigmoidOp sigmoidOp = mulOp.getA().getDefiningOp<ONNXSigmoidOp>();
+    Value input = mulOp.getB();
+    if (!sigmoidOp) {
+      sigmoidOp = mulOp.getB().getDefiningOp<ONNXSigmoidOp>();
+      input = mulOp.getA();
+    }
+    if (!sigmoidOp || sigmoidOp.getX() != input || !sigmoidOp->hasOneUse()
+        || sigmoidOp.getResult().getType() != input.getType() || mulOp.getResult().getType() != input.getType())
+      return failure();
+
+    auto plan = spatial::SpatSiluPlanOp::create(
+      rewriter, mulOp.getLoc(), mulOp.getResult().getType(), input, rewriter.getStringAttr("nchw"));
+    rewriter.replaceOp(mulOp, plan.getResult());
+    rewriter.eraseOp(sigmoidOp);
+    return success();
+  }
+};
+
 static DenseElementsAttr getDenseConstantAttr(Value value) {
  if (auto constantOp = value.getDefiningOp<arith::ConstantOp>())
    return dyn_cast<DenseElementsAttr>(constantOp.getValue());
@@ -219,6 +241,10 @@ struct AddToSpatialCompute : OpConversionPattern<ONNXAddOp> {

 } // namespace

+void populateElementwiseFusionPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
+  patterns.add<SiluToSpatialPlan>(ctx);
+}
+
 void populateElementwisePatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
  patterns.add<AddToSpatialCompute>(ctx);
  patterns.add<BinaryElementwiseToSpatialCompute<ONNXSubOp, spatial::SpatVSubOp>>(ctx);
@@ -259,17 +259,6 @@ static FailureOr<spatial::SpatComputeBatch> createVmmBatch(Value a,
  return *batchOp;
 }

-static Value
-createDynamicGemmBatchRow(Value lane, int64_t numOutCols, ConversionPatternRewriter& rewriter, Location loc) {
-  if (numOutCols == 1)
-    return lane;
-
-  MLIRContext* context = rewriter.getContext();
-  AffineExpr d0 = getAffineDimExpr(0, context);
-  return createOrFoldAffineApply(
-    rewriter, loc, d0.floorDiv(numOutCols), ValueRange {lane}, rewriter.getInsertionBlock()->getParentOp());
-}
-
 static Value extractDynamicGemmBColumn(
  Value matrix, Value column, RankedTensorType vectorType, ConversionPatternRewriter& rewriter, Location loc) {
  SmallVector<OpFoldResult> offsets {rewriter.getIndexAttr(0), column};
@@ -373,33 +362,56 @@ static FailureOr<spatial::SpatComputeBatch> createVvdmulBatch(Value a,
                                                              Value b,
                                                              RankedTensorType aType,
                                                              RankedTensorType bType,
-                                                              RankedTensorType scalarPiecesType,
+                                                              RankedTensorType columnPiecesType,
                                                              RankedTensorType outType,
+                                                              bool transposeB,
                                                              ConversionPatternRewriter& rewriter,
                                                              Location loc) {
  const int64_t numOutRows = outType.getDimSize(0);
  const int64_t numOutCols = outType.getDimSize(1);
  const int64_t reductionSize = aType.getDimSize(1);
-  const int64_t laneCount = numOutRows * numOutCols;
+  auto vectorType = RankedTensorType::get({1, reductionSize}, aType.getElementType());
+  auto scalarType = RankedTensorType::get({1, 1}, outType.getElementType());
+  auto columnType = RankedTensorType::get({numOutRows, 1}, outType.getElementType());
  auto batchOp = createSpatComputeBatch(
    rewriter,
    loc,
-    TypeRange {scalarPiecesType},
-    laneCount,
+    TypeRange {columnPiecesType},
+    numOutCols,
    ValueRange {},
    ValueRange {a, b},
    [&](detail::SpatComputeBatchBodyArgs args) {
-      Value row = createDynamicGemmBatchRow(args.lane, numOutCols, rewriter, loc);
-      Value column =
-        onnx_mlir::affineModConst(rewriter, loc, args.lane, numOutCols, rewriter.getInsertionBlock()->getParentOp());
-
-      auto vectorType = RankedTensorType::get({1, reductionSize}, aType.getElementType());
-      auto scalarType = RankedTensorType::get({1, 1}, outType.getElementType());
-      Value aVector = extractDynamicGemmRowVector(args.inputs[0], row, vectorType, rewriter, loc);
-      Value bVector = extractDynamicGemmBColumn(args.inputs[1], column, vectorType, rewriter, loc);
-      Value scalar = spatial::SpatVVDMulOp::create(rewriter, loc, scalarType, aVector, bVector).getResult();
-
-      publishGraphBatchPhysicalFragment(rewriter, loc, scalar, args.outputs.front(), args.lane);
+      Value bVector = transposeB
+                        ? extractDynamicGemmRowVector(args.inputs[1], args.lane, vectorType, rewriter, loc)
+                        : extractDynamicGemmBColumn(args.inputs[1], args.lane, vectorType, rewriter, loc);
+      Value columnInit = tensor::EmptyOp::create(rewriter, loc, columnType.getShape(), columnType.getElementType());
+      Value c0 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
+      Value c1 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 1);
+      Value cNumOutRows =
+        getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), numOutRows);
+      auto loop = buildNormalizedScfFor(
+        rewriter,
+        loc,
+        c0,
+        cNumOutRows,
+        c1,
+        ValueRange {columnInit},
+        [&](OpBuilder&, Location nestedLoc, Value row, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
+          Value aVector = extractDynamicGemmRowVector(args.inputs[0], row, vectorType, rewriter, nestedLoc);
+          Value scalar = spatial::SpatVVDMulOp::create(rewriter, nestedLoc, scalarType, aVector, bVector).getResult();
+          Value next = tensor::InsertSliceOp::create(rewriter,
+                                                     nestedLoc,
+                                                     scalar,
+                                                     iterArgs.front(),
+                                                     SmallVector<OpFoldResult> {row, rewriter.getIndexAttr(0)},
+                                                     SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
+                                                                               rewriter.getIndexAttr(1)},
+                                                     getUnitStrides(rewriter, 2));
+          yielded.push_back(next);
+          return success();
+        });
+      assert(succeeded(loop) && "dynamic Gemm row loop construction must succeed");
+      publishGraphBatchPhysicalFragment(rewriter, loc, loop->results.front(), args.outputs.front(), args.lane);
    });
  if (failed(batchOp))
    return failure();
@@ -415,7 +427,7 @@ static FailureOr<spatial::SpatCompute> createDynamicGemmOutputCompute(Value scal
                                                                      float beta,
                                                                      ConversionPatternRewriter& rewriter,
                                                                      Location loc) {
-  const int64_t laneCount = scalarPiecesType.getDimSize(0);
+  const int64_t numOutRows = outType.getDimSize(0);
  const int64_t numOutCols = outType.getDimSize(1);
  SmallVector<Value> inputs {scalarPieces};
  if (bias)
@@ -428,43 +440,62 @@ static FailureOr<spatial::SpatCompute> createDynamicGemmOutputCompute(Value scal
    Value outputInit = tensor::EmptyOp::create(rewriter, loc, outType.getShape(), outType.getElementType()).getResult();
    Value c0 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
    Value c1 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 1);
-    Value cLaneCount = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), laneCount);
+    Value cNumOutCols = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), numOutCols);
+    auto columnType = RankedTensorType::get({numOutRows, 1}, outType.getElementType());
    auto loop = buildNormalizedScfFor(
      rewriter,
      loc,
      c0,
-      cLaneCount,
+      cNumOutCols,
      c1,
      ValueRange {outputInit},
-      [&](OpBuilder&, Location nestedLoc, Value lane, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
+      [&](OpBuilder&, Location nestedLoc, Value column, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
        Value outputAcc = iterArgs.front();
-        Value row = createDynamicGemmBatchRow(lane, numOutCols, rewriter, nestedLoc);
-        Value column =
-          onnx_mlir::affineModConst(rewriter, nestedLoc, lane, numOutCols, rewriter.getInsertionBlock()->getParentOp());
-        SmallVector<OpFoldResult> scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-        SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-        FailureOr<Value> scalar = extractGraphBatchPhysicalFragment(rewriter, nestedLoc, pieces, lane, scalarType);
-        if (failed(scalar))
+        FailureOr<Value> columnPiece =
+          extractGraphBatchPhysicalFragment(rewriter, nestedLoc, pieces, column, columnType);
+        if (failed(columnPiece))
          return failure();
-        if (alpha != 1.0f) {
-          Value alphaTensor = createScalarTensorConstant(scalarType, alpha, rewriter, nestedLoc);
-          *scalar = spatial::SpatVMulOp::create(rewriter, nestedLoc, scalarType, *scalar, alphaTensor).getResult();
-        }
-        if (biasArg) {
-          Value biasScalar =
-            createBroadcastedBiasScalar(biasArg, biasType, row, column, scalarType, rewriter, nestedLoc);
-          if (beta != 1.0f) {
-            Value betaTensor = createScalarTensorConstant(scalarType, beta, rewriter, nestedLoc);
-            biasScalar =
-              spatial::SpatVMulOp::create(rewriter, nestedLoc, scalarType, biasScalar, betaTensor).getResult();
-          }
-          *scalar = spatial::SpatVAddOp::create(rewriter, nestedLoc, scalarType, *scalar, biasScalar).getResult();
-        }
-        SmallVector<OpFoldResult> outputOffsets {row, column};
-        Value outputNext =
-          tensor::InsertSliceOp::create(rewriter, nestedLoc, *scalar, outputAcc, outputOffsets, scalarSizes, unitStrides)
-            .getResult();
-        yielded.push_back(outputNext);
+        Value cNumOutRows =
+          getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), numOutRows);
+        auto rowLoop = buildNormalizedScfFor(
+          rewriter,
+          nestedLoc,
+          c0,
+          cNumOutRows,
+          c1,
+          ValueRange {outputAcc},
+          [&](OpBuilder&, Location rowLoc, Value row, ValueRange rowIterArgs, SmallVectorImpl<Value>& rowYielded) {
+            SmallVector<OpFoldResult> scalarOffsets {row, rewriter.getIndexAttr(0)};
+            SmallVector<OpFoldResult> scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
+            SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
+            Value scalar = tensor::ExtractSliceOp::create(
+              rewriter, rowLoc, scalarType, *columnPiece, scalarOffsets, scalarSizes, unitStrides);
+            if (alpha != 1.0f) {
+              Value alphaTensor = createScalarTensorConstant(scalarType, alpha, rewriter, rowLoc);
+              scalar = spatial::SpatVMulOp::create(rewriter, rowLoc, scalarType, scalar, alphaTensor).getResult();
+            }
+            if (biasArg) {
+              Value biasScalar = createBroadcastedBiasScalar(biasArg, biasType, row, column, scalarType, rewriter, rowLoc);
+              if (beta != 1.0f) {
+                Value betaTensor = createScalarTensorConstant(scalarType, beta, rewriter, rowLoc);
+                biasScalar =
+                  spatial::SpatVMulOp::create(rewriter, rowLoc, scalarType, biasScalar, betaTensor).getResult();
+              }
+              scalar = spatial::SpatVAddOp::create(rewriter, rowLoc, scalarType, scalar, biasScalar).getResult();
+            }
+            Value next = tensor::InsertSliceOp::create(rewriter,
+                                                       rowLoc,
+                                                       scalar,
+                                                       rowIterArgs.front(),
+                                                       SmallVector<OpFoldResult> {row, column},
+                                                       scalarSizes,
+                                                       unitStrides);
+            rowYielded.push_back(next);
+            return success();
+          });
+        if (failed(rowLoop))
+          return failure();
+        yielded.push_back(rowLoop->results.front());
        return success();
      });
    if (failed(loop))
@@ -660,16 +691,10 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
    aType = transposedType;
  }

-  if (gemmOpAdaptor.getTransB()) {
-    auto bShape = bType.getShape();
-    auto transposedType = RankedTensorType::get({bShape[1], bShape[0]}, bType.getElementType(), bType.getEncoding());
-    b = ONNXTransposeOp::create(rewriter, loc, transposedType, b, rewriter.getI64ArrayAttr({1, 0})).getResult();
-    bType = transposedType;
-  }
-
  const int64_t numOutRows = outType.getDimSize(0);
  const int64_t numOutCols = outType.getDimSize(1);
  const int64_t reductionSize = aType.getDimSize(1);
+  const bool transposeB = gemmOpAdaptor.getTransB();

  if (!isCompileTimeComputable(b)) {
    bool hasC = hasGemmBias(c);
@@ -690,8 +715,9 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
      biasType = *verifiedBiasType;
    }

-    if (aType.getDimSize(0) != numOutRows || bType.getDimSize(0) != reductionSize
-        || bType.getDimSize(1) != numOutCols) {
+    const int64_t bReductionSize = bType.getDimSize(transposeB ? 1 : 0);
+    const int64_t bOutputColumns = bType.getDimSize(transposeB ? 0 : 1);
+    if (aType.getDimSize(0) != numOutRows || bReductionSize != reductionSize || bOutputColumns != numOutCols) {
      gemmOp.emitOpError("has inconsistent A, B, and output shapes");
      return failure();
    }
@@ -702,8 +728,9 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
      return failure();
    }

-    auto scalarPiecesType = spatial::getGraphBatchPhysicalResultType(laneCount64, RankedTensorType::get({1, 1}, outType.getElementType()));
-    auto batchOp = createVvdmulBatch(a, b, aType, bType, scalarPiecesType, outType, rewriter, loc);
+    auto columnType = RankedTensorType::get({numOutRows, 1}, outType.getElementType());
+    auto scalarPiecesType = spatial::getGraphBatchPhysicalResultType(numOutCols, columnType);
+    auto batchOp = createVvdmulBatch(a, b, aType, bType, scalarPiecesType, outType, transposeB, rewriter, loc);
    if (failed(batchOp))
      return failure();
    auto outputCompute = createDynamicGemmOutputCompute(
@@ -714,6 +741,13 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
    return success();
  }

+  if (transposeB) {
+    auto bShape = bType.getShape();
+    auto transposedType = RankedTensorType::get({bShape[1], bShape[0]}, bType.getElementType(), bType.getEncoding());
+    b = ONNXTransposeOp::create(rewriter, loc, transposedType, b, rewriter.getI64ArrayAttr({1, 0})).getResult();
+    bType = transposedType;
+  }
+
  auto scaledB = materializeScaledConstantTensor(b, gemmOpAdaptor.getAlpha().convertToFloat(), rewriter, loc);
  if (failed(scaledB)) {
    gemmOp.emitOpError("requires constant Gemm input B when alpha is not 1.0");
@@ -389,41 +389,6 @@ static FailureOr<spatial::SpatComputeBatch> createBatchedVmmBatch(Value a,
  return *batchOp;
 }

-static Value extractDynamicBatchedBColumn(Value matrix,
-                                          ArrayRef<int64_t> sourceBatchShape,
-                                          ArrayRef<int64_t> outputBatchShape,
-                                          Value outputBatchIndex,
-                                          Value column,
-                                          RankedTensorType vectorType,
-                                          PatternRewriter& rewriter,
-                                          Location loc) {
-  auto columnSliceType = RankedTensorType::get({1, vectorType.getDimSize(1), 1}, vectorType.getElementType());
-  Value sourceBatchIndex =
-    mapOutputBatchIndexToSourceBatchIndex(outputBatchIndex, sourceBatchShape, outputBatchShape, rewriter, loc);
-  SmallVector<OpFoldResult> offsets {OpFoldResult(sourceBatchIndex), rewriter.getIndexAttr(0), column};
-  SmallVector<OpFoldResult> sizes {
-    rewriter.getIndexAttr(1), rewriter.getIndexAttr(vectorType.getDimSize(1)), rewriter.getIndexAttr(1)};
-  SmallVector<OpFoldResult> strides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-  Value columnSlice = tensor::ExtractSliceOp::create(rewriter, loc, columnSliceType, matrix, offsets, sizes, strides);
-  auto collapsedType = RankedTensorType::get({vectorType.getDimSize(1)}, vectorType.getElementType());
-  Value collapsed = tensor::CollapseShapeOp::create(rewriter,
-                                                    loc,
-                                                    collapsedType,
-                                                    columnSlice,
-                                                    SmallVector<ReassociationIndices> {
-                                                      {0, 1, 2}
-  })
-                      .getResult();
-  return tensor::ExpandShapeOp::create(rewriter,
-                                       loc,
-                                       vectorType,
-                                       collapsed,
-                                       SmallVector<ReassociationIndices> {
-                                         {0, 1}
-  })
-    .getResult();
-}
-
 static Value extractDynamicBatchedRowVector(Value matrix,
                                            ArrayRef<int64_t> sourceBatchShape,
                                            ArrayRef<int64_t> outputBatchShape,
@@ -449,7 +414,7 @@ static FailureOr<spatial::SpatComputeBatch> createBatchedVvdmulBatch(Value a,
                                                                     ArrayRef<int64_t> outputBatchShape,
                                                                     RankedTensorType aType,
                                                                     RankedTensorType bType,
-                                                                     RankedTensorType scalarPiecesType,
+                                                                     RankedTensorType columnPiecesType,
                                                                     RankedTensorType outType,
                                                                     PatternRewriter& rewriter,
                                                                     Location loc) {
@@ -457,29 +422,52 @@ static FailureOr<spatial::SpatComputeBatch> createBatchedVvdmulBatch(Value a,
  const int64_t numOutRows = outType.getDimSize(1);
  const int64_t numOutCols = outType.getDimSize(2);
  const int64_t reductionSize = aType.getDimSize(2);
-  const int64_t laneCount = numBatches * numOutRows * numOutCols;
+  const int64_t laneCount = numBatches * numOutCols;
+  auto vectorType = RankedTensorType::get({1, reductionSize}, aType.getElementType());
+  auto scalarType = RankedTensorType::get({1, 1}, outType.getElementType());
+  auto columnType = RankedTensorType::get({numOutRows, 1}, outType.getElementType());
  auto batchOp = createSpatComputeBatch(
    rewriter,
    loc,
-    TypeRange {scalarPiecesType},
+    TypeRange {columnPiecesType},
    laneCount,
    ValueRange {},
    ValueRange {a, b},
    [&](detail::SpatComputeBatchBodyArgs args) {
      Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
-      Value batch = affineFloorDivConst(rewriter, loc, args.lane, numOutRows * numOutCols, anchorOp);
-      Value batchLane = affineModConst(rewriter, loc, args.lane, numOutRows * numOutCols, anchorOp);
-      Value row = affineFloorDivConst(rewriter, loc, batchLane, numOutCols, anchorOp);
-      Value column = affineModConst(rewriter, loc, batchLane, numOutCols, anchorOp);
-
-      auto vectorType = RankedTensorType::get({1, reductionSize}, aType.getElementType());
-      auto scalarType = RankedTensorType::get({1, 1}, outType.getElementType());
-      Value aVector = extractDynamicBatchedRowVector(
-        args.inputs[0], aBatchShape, outputBatchShape, batch, row, vectorType, rewriter, loc);
-      Value bVector = extractDynamicBatchedBColumn(
+      Value batch = affineFloorDivConst(rewriter, loc, args.lane, numOutCols, anchorOp);
+      Value column = affineModConst(rewriter, loc, args.lane, numOutCols, anchorOp);
+      Value bVector = extractDynamicBatchedRowVector(
        args.inputs[1], bBatchShape, outputBatchShape, batch, column, vectorType, rewriter, loc);
-      Value scalar = spatial::SpatVVDMulOp::create(rewriter, loc, scalarType, aVector, bVector).getResult();
-      publishGraphBatchPhysicalFragment(rewriter, loc, scalar, args.outputs.front(), args.lane);
+      Value columnInit = tensor::EmptyOp::create(rewriter, loc, columnType.getShape(), columnType.getElementType());
+      Value c0 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
+      Value c1 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 1);
+      Value cNumOutRows =
+        getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), numOutRows);
+      auto loop = buildNormalizedScfFor(
+        rewriter,
+        loc,
+        c0,
+        cNumOutRows,
+        c1,
+        ValueRange {columnInit},
+        [&](OpBuilder&, Location nestedLoc, Value row, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
+          Value aVector = extractDynamicBatchedRowVector(
+            args.inputs[0], aBatchShape, outputBatchShape, batch, row, vectorType, rewriter, nestedLoc);
+          Value scalar = spatial::SpatVVDMulOp::create(rewriter, nestedLoc, scalarType, aVector, bVector).getResult();
+          Value next = tensor::InsertSliceOp::create(rewriter,
+                                                     nestedLoc,
+                                                     scalar,
+                                                     iterArgs.front(),
+                                                     SmallVector<OpFoldResult> {row, rewriter.getIndexAttr(0)},
+                                                     SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
+                                                                               rewriter.getIndexAttr(1)},
+                                                     getUnitStrides(rewriter, 2));
+          yielded.push_back(next);
+          return success();
+        });
+      assert(succeeded(loop) && "dynamic MatMul row loop construction must succeed");
+      publishGraphBatchPhysicalFragment(rewriter, loc, loop->results.front(), args.outputs.front(), args.lane);
    });
  if (failed(batchOp))
    return failure();
@@ -492,9 +480,8 @@ static FailureOr<Value> createBatchedDynamicOutputCompute(Value scalarPieces,
                                                          PatternRewriter& rewriter,
                                                          Location loc) {
  const int64_t laneCount = scalarPiecesType.getDimSize(0);
-  const int64_t numOutRows = outType.getDimSize(1);
  const int64_t numOutCols = outType.getDimSize(2);
-  auto scalarType = RankedTensorType::get({1, 1}, outType.getElementType());
+  auto columnType = RankedTensorType::get({outType.getDimSize(1), 1}, outType.getElementType());

  auto computeOp = createSpatCompute<1>(
    rewriter, loc, TypeRange {outType}, {}, ValueRange {scalarPieces}, [&](Value pieces) -> LogicalResult {
@@ -513,19 +500,18 @@ static FailureOr<Value> createBatchedDynamicOutputCompute(Value scalarPieces,
        [&](OpBuilder&, Location nestedLoc, Value lane, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
          Value outputAcc = iterArgs.front();
          Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
-          Value batch = affineFloorDivConst(rewriter, nestedLoc, lane, numOutRows * numOutCols, anchorOp);
-          Value batchLane = affineModConst(rewriter, nestedLoc, lane, numOutRows * numOutCols, anchorOp);
-          Value row = affineFloorDivConst(rewriter, nestedLoc, batchLane, numOutCols, anchorOp);
-          Value column = affineModConst(rewriter, nestedLoc, batchLane, numOutCols, anchorOp);
-          FailureOr<Value> scalar = extractGraphBatchPhysicalFragment(rewriter, nestedLoc, pieces, lane, scalarType);
-          if (failed(scalar))
+          Value batch = affineFloorDivConst(rewriter, nestedLoc, lane, numOutCols, anchorOp);
+          Value column = affineModConst(rewriter, nestedLoc, lane, numOutCols, anchorOp);
+          FailureOr<Value> columnPiece =
+            extractGraphBatchPhysicalFragment(rewriter, nestedLoc, pieces, lane, columnType);
+          if (failed(columnPiece))
            return failure();
-          SmallVector<OpFoldResult> outputOffsets {batch, row, column};
+          SmallVector<OpFoldResult> outputOffsets {batch, rewriter.getIndexAttr(0), column};
          SmallVector<OpFoldResult> outputSizes = {
-            rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
+            rewriter.getIndexAttr(1), rewriter.getIndexAttr(outType.getDimSize(1)), rewriter.getIndexAttr(1)};
          Value next =
            tensor::InsertSliceOp::create(
-              rewriter, nestedLoc, *scalar, outputAcc, outputOffsets, outputSizes, getUnitStrides(rewriter, 3))
+              rewriter, nestedLoc, *columnPiece, outputAcc, outputOffsets, outputSizes, getUnitStrides(rewriter, 3))
              .getResult();
          yielded.push_back(next);
          return success();
@@ -1011,12 +997,14 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
        return success();
      }
    }
-    const int64_t laneCount = plan.batch * plan.m * plan.n;
+    const int64_t laneCount = plan.batch * plan.n;
+    auto columnType = RankedTensorType::get({plan.m, 1}, shapeInfo->outType.getElementType());
    auto scalarPiecesType = spatial::getGraphBatchPhysicalResultType(
-      laneCount, RankedTensorType::get({1, 1}, shapeInfo->outType.getElementType()));
+      laneCount, columnType);
+    Value transposedRhs = transposeLastTwoDims(plan.rhs, rewriter, loc);
    auto batchOp = createBatchedVvdmulBatch(plan.lhs,
                                            plan.lhsBatchShape,
-                                            plan.rhs,
+                                            transposedRhs,
                                            plan.rhsBatchShape,
                                            plan.outputBatchShape,
                                            plan.lhsType,