even faster on pimcomp models

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

@@ -180,7 +180,11 @@ FailureOr<Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& va
                                          kRowStripIndexMap, rewriter, loc);
 }
 
-FailureOr<Value> applyRowStripRelu(const RowStripPhysicalValue& value, PatternRewriter& rewriter, Location loc) {
+template <typename BuildActivation>
+static FailureOr<Value> applyRowStripActivation(const RowStripPhysicalValue& value,
+                                                PatternRewriter& rewriter,
+                                                Location loc,
+                                                BuildActivation buildActivation) {
   auto storageType = cast<RankedTensorType>(value.storage.getType());
   const int64_t laneCount = storageType.getDimSize(0);
   auto batchOp = createSpatComputeBatch(rewriter,
@@ -193,10 +197,9 @@ FailureOr<Value> applyRowStripRelu(const RowStripPhysicalValue& value, PatternRe
                                           FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(
                                             rewriter, loc, args.inputs.front(), args.lane, value.fragmentType);
                                           if (failed(fragment)) return failure();
-                                          Value relu = spatial::SpatReluOp::create(
-                                            rewriter, loc, value.fragmentType, *fragment).getResult();
+                                          Value result = buildActivation(*fragment);
                                           publishGraphBatchPhysicalFragment(
-                                            rewriter, loc, relu, args.outputs.front(), args.lane);
+                                            rewriter, loc, result, args.outputs.front(), args.lane);
                                           return success();
                                         });
   if (failed(batchOp))
@@ -204,6 +207,19 @@ FailureOr<Value> applyRowStripRelu(const RowStripPhysicalValue& value, PatternRe
   return batchOp->getResult(0);
 }
 
+FailureOr<Value> applyRowStripRelu(const RowStripPhysicalValue& value, PatternRewriter& rewriter, Location loc) {
+  return applyRowStripActivation(value, rewriter, loc, [&](Value fragment) {
+    return spatial::SpatReluOp::create(rewriter, loc, value.fragmentType, fragment).getResult();
+  });
+}
+
+FailureOr<Value> applyRowStripSilu(const RowStripPhysicalValue& value, PatternRewriter& rewriter, Location loc) {
+  return applyRowStripActivation(value, rewriter, loc, [&](Value fragment) {
+    Value sigmoid = spatial::SpatSigmoidOp::create(rewriter, loc, value.fragmentType, fragment).getResult();
+    return spatial::SpatVMulOp::create(rewriter, loc, value.fragmentType, fragment, sigmoid).getResult();
+  });
+}
+
 FailureOr<Value> applyRowStripBiasAdd(const RowStripPhysicalValue& value,
                                      Value bias,
                                      PatternRewriter& rewriter,

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

@@ -61,6 +61,10 @@ mlir::FailureOr<mlir::Value> applyRowStripRelu(const RowStripPhysicalValue& valu
                                                mlir::PatternRewriter& rewriter,
                                                mlir::Location loc);
 
+mlir::FailureOr<mlir::Value> applyRowStripSilu(const RowStripPhysicalValue& value,
+                                               mlir::PatternRewriter& rewriter,
+                                               mlir::Location loc);
+
 mlir::FailureOr<mlir::Value> applyRowStripBiasAdd(const RowStripPhysicalValue& value,
                                                   mlir::Value bias,
                                                   mlir::PatternRewriter& rewriter,

src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp

@@ -58,6 +58,11 @@ lowerRowStripRelu(const RowStripPhysicalValue& input, spatial::SpatReluPlanOp pl
   return applyRowStripRelu(input, rewriter, planOp.getLoc());
 }
 
+static FailureOr<Value>
+lowerRowStripSilu(const RowStripPhysicalValue& input, spatial::SpatSiluPlanOp planOp, PatternRewriter& rewriter) {
+  return applyRowStripSilu(input, rewriter, planOp.getLoc());
+}
+
 static FailureOr<Value> lowerRowStripBiasAdd(const RowStripPhysicalValue& input,
                                              spatial::SpatBiasAddPlanOp planOp,
                                              PatternRewriter& rewriter) {
@@ -349,6 +354,51 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
         rewriter.replaceOp(planOp, computeOp.getResults());
         continue;
       }
+
+      if (auto planOp = dyn_cast<spatial::SpatSiluPlanOp>(&op)) {
+        if (succeeded(getRowStripValue(rowStripValues, planOp.getInput()))) {
+          auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
+            auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
+            return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
+          });
+          if (outputBlueprint == planOp.getResult().getUsers().end()) {
+            planOp.emitOpError("row-strip SiLU plan requires a row-strip blueprint result");
+            signalPassFailure();
+            return;
+          }
+
+          FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
+          rewriter.setInsertionPoint(planOp);
+          FailureOr<Value> lowered = lowerRowStripSilu(*input, planOp, rewriter);
+          if (failed(lowered)) {
+            planOp.emitOpError("failed to lower selected row-strip Spatial SiLU plan");
+            signalPassFailure();
+            return;
+          }
+          auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
+          FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
+          if (failed(output)) {
+            signalPassFailure();
+            return;
+          }
+          rowStripValues[blueprint.getResult()] = *output;
+          eraseAfterLowering.insert(planOp);
+          eraseAfterLowering.insert(blueprint);
+          continue;
+        }
+
+        rewriter.setInsertionPoint(planOp);
+        auto computeOp = createSpatCompute<1>(
+          rewriter, planOp.getLoc(), planOp.getOutput().getType(), {}, planOp.getInput(), [&](Value x) {
+            Value sigmoid = spatial::SpatSigmoidOp::create(
+              rewriter, planOp.getLoc(), planOp.getOutput().getType(), x).getResult();
+            Value silu = spatial::SpatVMulOp::create(
+              rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, sigmoid).getResult();
+            spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), silu);
+          });
+        rewriter.replaceOp(planOp, computeOp.getResults());
+        continue;
+      }
       if (auto planOp = dyn_cast<spatial::SpatMaxPool2DPlanOp>(&op)) {
         auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
           auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
@@ -650,6 +700,7 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
                    spatial::SpatBiasAddPlanOp,
                    spatial::SpatAddPlanOp,
                    spatial::SpatReluPlanOp,
+                   spatial::SpatSiluPlanOp,
                    spatial::SpatMaxPool2DPlanOp,
                    spatial::SpatGlobalAveragePoolPlanOp,
                    spatial::SpatMaterializeLayoutOp>(op)

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp

@@ -8,6 +8,7 @@
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/Passes.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 #include "llvm/ADT/SmallVector.h"
 
@@ -50,13 +51,14 @@ static void populateEmptyFunction(func::FuncOp funcOp) {
   SmallVector<spatial::SpatAddPlanOp> addPlans(funcOp.getOps<spatial::SpatAddPlanOp>());
   SmallVector<spatial::SpatConcatPlanOp> concatPlans(funcOp.getOps<spatial::SpatConcatPlanOp>());
   SmallVector<spatial::SpatReluPlanOp> reluPlans(funcOp.getOps<spatial::SpatReluPlanOp>());
+  SmallVector<spatial::SpatSiluPlanOp> siluPlans(funcOp.getOps<spatial::SpatSiluPlanOp>());
   SmallVector<spatial::SpatMaxPool2DPlanOp> maxPoolPlans(funcOp.getOps<spatial::SpatMaxPool2DPlanOp>());
   SmallVector<spatial::SpatGlobalAveragePoolPlanOp> globalAveragePoolPlans(
     funcOp.getOps<spatial::SpatGlobalAveragePoolPlanOp>());
   SmallVector<spatial::SpatBlueprintOp> blueprints(funcOp.getOps<spatial::SpatBlueprintOp>());
   SmallVector<spatial::SpatMaterializeLayoutOp> materializers(funcOp.getOps<spatial::SpatMaterializeLayoutOp>());
   if (!computes.empty() || !computeBatches.empty() || !convPlans.empty() || !biasAddPlans.empty() || !addPlans.empty()
-      || !concatPlans.empty() || !reluPlans.empty() || !maxPoolPlans.empty() || !blueprints.empty()
+      || !concatPlans.empty() || !reluPlans.empty() || !siluPlans.empty() || !maxPoolPlans.empty() || !blueprints.empty()
       || !globalAveragePoolPlans.empty() || !materializers.empty()) {
     return;
   }
@@ -121,6 +123,14 @@ void ONNXToSpatialPass::runOnOperation() {
     return;
   }
 
+  RewritePatternSet fusionPatterns(ctx);
+  populateElementwiseFusionPatterns(fusionPatterns, ctx);
+  if (failed(applyPatternsGreedily(moduleOp, std::move(fusionPatterns)))) {
+    moduleOp.emitError("failed to fuse layout-aware ONNX elementwise patterns");
+    signalPassFailure();
+    return;
+  }
+
   auto entryFunc = getPimEntryFunc(moduleOp);
   if (failed(entryFunc)) {
     moduleOp.emitError("failed to locate the PIM entry function during ONNX-to-Spatial lowering");

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp

@@ -149,6 +149,7 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter
             spatial::SpatAddPlanOp,
             spatial::SpatConcatPlanOp,
             spatial::SpatReluPlanOp,
+            spatial::SpatSiluPlanOp,
             spatial::SpatMaxPool2DPlanOp,
             spatial::SpatGlobalAveragePoolPlanOp,
             spatial::SpatBlueprintOp,

src/PIM/Conversion/ONNXToSpatial/Patterns.hpp

@@ -17,6 +17,7 @@ void populateWeightPromotionPatterns(mlir::RewritePatternSet& patterns, mlir::ML
 
 void populateConvPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateElementwisePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+void populateElementwiseFusionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateGemmPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateMatMulRewritePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populatePoolPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);

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

@@ -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,

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

@@ -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);

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

@@ -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");

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

@@ -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,

src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp

@@ -34,6 +34,8 @@ static SelectedLayout getSelectedLayout(llvm::DenseMap<Value, SelectedLayout>& l
 static bool usesSelectedRowStrip(Operation* user, llvm::DenseMap<Value, SelectedLayout>& layouts) {
   if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(user))
     return getSelectedLayout(layouts, reluPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
+  if (auto siluPlan = dyn_cast<spatial::SpatSiluPlanOp>(user))
+    return getSelectedLayout(layouts, siluPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
   if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(user))
     return getSelectedLayout(layouts, biasAddPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
   if (auto addPlan = dyn_cast<spatial::SpatAddPlanOp>(user))
@@ -62,7 +64,7 @@ static bool allUsersCanHandleRowStrip(Value value, llvm::DenseMap<Value, Selecte
 }
 
 static bool canConsumeRowStripAsUser(Operation* user) {
-  if (isa<spatial::SpatReluPlanOp>(user))
+  if (isa<spatial::SpatReluPlanOp, spatial::SpatSiluPlanOp>(user))
     return true;
   if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(user)) {
     auto resultType = dyn_cast<RankedTensorType>(biasAddPlan.getOutput().getType());
@@ -105,13 +107,12 @@ static SelectedLayout chooseConvLayout(spatial::SpatConv2DPlanOp convPlan,
   return SelectedLayout::PixelMajorRowStrip;
 }
 
-static SelectedLayout chooseReluLayout(spatial::SpatReluPlanOp reluPlan,
-                                       llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  if (getSelectedLayout(layouts, reluPlan.getInput()) != SelectedLayout::PixelMajorRowStrip)
+static SelectedLayout chooseActivationLayout(Value input,
+                                             Value result,
+                                             llvm::DenseMap<Value, SelectedLayout>& layouts) {
+  if (getSelectedLayout(layouts, input) != SelectedLayout::PixelMajorRowStrip)
     return SelectedLayout::DenseNchw;
-  if (!hasRowStripConsumer(reluPlan.getResult()))
-    return SelectedLayout::DenseNchw;
-  if (!allUsersCanHandleRowStrip(reluPlan.getResult(), layouts))
+  if (!allUsersCanHandleRowStrip(result, layouts))
     return SelectedLayout::DenseNchw;
   return SelectedLayout::PixelMajorRowStrip;
 }
@@ -239,13 +240,21 @@ struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass,
           continue;
         }
         if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(&op)) {
-          SelectedLayout selected = chooseReluLayout(reluPlan, layouts);
+          SelectedLayout selected = chooseActivationLayout(reluPlan.getInput(), reluPlan.getResult(), layouts);
           if (layouts[reluPlan.getResult()] != selected) {
             layouts[reluPlan.getResult()] = selected;
             changed = true;
           }
           continue;
         }
+        if (auto siluPlan = dyn_cast<spatial::SpatSiluPlanOp>(&op)) {
+          SelectedLayout selected = chooseActivationLayout(siluPlan.getInput(), siluPlan.getResult(), layouts);
+          if (layouts[siluPlan.getResult()] != selected) {
+            layouts[siluPlan.getResult()] = selected;
+            changed = true;
+          }
+          continue;
+        }
         if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(&op)) {
           SelectedLayout selected = chooseBiasAddLayout(biasAddPlan, layouts);
           if (layouts[biasAddPlan.getResult()] != selected) {
@@ -301,6 +310,8 @@ struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass,
         producedValue = concatPlan.getResult();
       else if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(&op))
         producedValue = reluPlan.getResult();
+      else if (auto siluPlan = dyn_cast<spatial::SpatSiluPlanOp>(&op))
+        producedValue = siluPlan.getResult();
       else if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(&op))
         producedValue = maxPoolPlan.getResult();
       else if (auto averagePoolPlan = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(&op))