finally fast googlenet with correct latency artifacts for fair comparison

src/PIM/Conversion/ONNXToSpatial/CMakeLists.txt

@@ -8,7 +8,6 @@ add_pim_library(OMONNXToSpatial
   ONNXToSpatialVerifier.cpp
   Patterns/Pre.cpp
   Patterns/Post.cpp
-  Patterns/GeneratedConversion.cpp
   Patterns/Math/Conv.cpp
   Patterns/Math/ConvGeometry.cpp
   Patterns/Math/Elementwise.cpp

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

@@ -19,9 +19,11 @@ FailureOr<RowStripPhysicalValue> describeRowStripPhysicalValue(Value storage, Ra
       || storageType.getRank() != 5 || logicalType.getRank() != 4 || logicalType.getDimSize(0) != 1
       || storageType.getElementType() != logicalType.getElementType()
       || storageType.getDimSize(1) != 1 || storageType.getDimSize(2) != 1
-      || storageType.getDimSize(3) != logicalType.getDimSize(3) || storageType.getDimSize(4) <= 0)
+      || storageType.getDimSize(3) != logicalType.getDimSize(3)
+      || storageType.getDimSize(4) <= 0)
     return failure();
-  const int64_t tilesPerRow = ceilIntegerDivide(logicalType.getDimSize(1), storageType.getDimSize(4));
+  const int64_t tilesPerRow =
+    ceilIntegerDivide(logicalType.getDimSize(1), storageType.getDimSize(4));
   if (storageType.getDimSize(0) != logicalType.getDimSize(2) * tilesPerRow)
     return failure();
   return RowStripPhysicalValue {storage, logicalType,
@@ -249,4 +251,111 @@ FailureOr<Value> applyRowStripBiasAdd(const RowStripPhysicalValue& value,
   return batchOp->getResult(0);
 }
 
+FailureOr<Value> applyRowStripAdd(const RowStripPhysicalValue& lhs,
+                                  const RowStripPhysicalValue& rhs,
+                                  PatternRewriter& rewriter,
+                                  Location loc) {
+  if (lhs.logicalType != rhs.logicalType || lhs.fragmentType != rhs.fragmentType
+      || lhs.storage.getType() != rhs.storage.getType() || lhs.tilesPerRow != rhs.tilesPerRow)
+    return failure();
+  auto storageType = cast<RankedTensorType>(lhs.storage.getType());
+  const int64_t laneCount = storageType.getDimSize(0);
+  auto batch = createSpatComputeBatch(
+    rewriter,
+    loc,
+    TypeRange {storageType},
+    laneCount,
+    {},
+    ValueRange {lhs.storage, rhs.storage},
+    [&](detail::SpatComputeBatchBodyArgs args) {
+      FailureOr<Value> lhsFragment =
+        extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[0], args.lane, lhs.fragmentType);
+      FailureOr<Value> rhsFragment =
+        extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[1], args.lane, rhs.fragmentType);
+      if (failed(lhsFragment) || failed(rhsFragment))
+        return failure();
+      Value added = spatial::SpatVAddOp::create(rewriter, loc, lhs.fragmentType, *lhsFragment, *rhsFragment);
+      publishGraphBatchPhysicalFragment(rewriter, loc, added, args.outputs.front(), args.lane);
+      return success();
+    });
+  if (failed(batch))
+    return failure();
+  return batch->getResult(0);
+}
+
+FailureOr<Value> applyRowStripConcat(ArrayRef<RowStripPhysicalValue> inputs,
+                                     RankedTensorType outputType,
+                                     PatternRewriter& rewriter,
+                                     Location loc) {
+  if (inputs.empty() || !outputType || !outputType.hasStaticShape() || outputType.getRank() != 4
+      || outputType.getDimSize(0) != 1)
+    return failure();
+  int64_t channels = 0;
+  for (const RowStripPhysicalValue& input : inputs) {
+    if (input.logicalType.getElementType() != outputType.getElementType()
+        || input.logicalType.getDimSize(0) != outputType.getDimSize(0)
+        || input.logicalType.getDimSize(2) != outputType.getDimSize(2)
+        || input.logicalType.getDimSize(3) != outputType.getDimSize(3))
+      return failure();
+    channels += input.logicalType.getDimSize(1);
+  }
+  if (channels != outputType.getDimSize(1))
+    return failure();
+
+  SmallVector<Value> storages;
+  llvm::transform(
+    inputs, std::back_inserter(storages), [](const RowStripPhysicalValue& input) { return input.storage; });
+  const int64_t tileWidth = outputType.getDimSize(3);
+  auto fragmentType = getRowStripFragmentType(outputType);
+  auto storageType = getRowStripStorageType(outputType);
+  auto batch = createSpatComputeBatch(
+    rewriter,
+    loc,
+    TypeRange {storageType},
+    outputType.getDimSize(2),
+    {},
+    storages,
+    [&](detail::SpatComputeBatchBodyArgs args) {
+      Operation* anchor = rewriter.getInsertionBlock()->getParentOp();
+      SmallVector<Value> fragments;
+      for (auto [inputIndex, input] : llvm::enumerate(inputs)) {
+        Value tileStart = affineMulConst(
+          rewriter, loc, args.lane, input.tilesPerRow, anchor);
+        for (int64_t tile = 0; tile < input.tilesPerRow; ++tile) {
+          Value slot =
+            affineAddConst(rewriter, loc, tileStart, tile, anchor);
+          FailureOr<Value> fragment =
+            extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[inputIndex], slot, input.fragmentType);
+          if (failed(fragment))
+            return failure();
+          int64_t channelOffset = tile * input.fragmentType.getDimSize(3);
+          int64_t validChannels =
+            std::min(input.fragmentType.getDimSize(3), input.logicalType.getDimSize(1) - channelOffset);
+          auto validType =
+            RankedTensorType::get(
+              {1, 1, tileWidth, validChannels},
+              outputType.getElementType());
+          MixedSliceGeometry slice;
+          slice.offsets.assign(4, rewriter.getIndexAttr(0));
+          slice.sizes = {rewriter.getIndexAttr(1),
+                         rewriter.getIndexAttr(1),
+                         rewriter.getIndexAttr(tileWidth),
+                         rewriter.getIndexAttr(validChannels)};
+          slice.strides.assign(4, rewriter.getIndexAttr(1));
+          Value valid = extractMixedSliceOrIdentity(rewriter, loc, *fragment, validType, slice);
+          if (!valid)
+            return failure();
+          fragments.push_back(valid);
+        }
+      }
+      Value concatenated =
+        spatial::SpatConcatOp::create(rewriter, loc, fragmentType, rewriter.getI64IntegerAttr(3), fragments);
+      publishGraphBatchPhysicalFragment(rewriter, loc, concatenated, args.outputs.front(), args.lane);
+      return success();
+    });
+  if (failed(batch))
+    return failure();
+  return batch->getResult(0);
+}
+
 } // namespace onnx_mlir

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

@@ -66,4 +66,14 @@ mlir::FailureOr<mlir::Value> applyRowStripBiasAdd(const RowStripPhysicalValue& v
                                                   mlir::PatternRewriter& rewriter,
                                                   mlir::Location loc);
 
+mlir::FailureOr<mlir::Value> applyRowStripAdd(const RowStripPhysicalValue& lhs,
+                                              const RowStripPhysicalValue& rhs,
+                                              mlir::PatternRewriter& rewriter,
+                                              mlir::Location loc);
+
+mlir::FailureOr<mlir::Value> applyRowStripConcat(llvm::ArrayRef<RowStripPhysicalValue> inputs,
+                                                 mlir::RankedTensorType outputType,
+                                                 mlir::PatternRewriter& rewriter,
+                                                 mlir::Location loc);
+
 } // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp

@@ -64,6 +64,22 @@ static FailureOr<Value> lowerRowStripBiasAdd(const RowStripPhysicalValue& input,
   return applyRowStripBiasAdd(input, planOp.getBias(), rewriter, planOp.getLoc());
 }
 
+static FailureOr<Value> lowerRowStripAdd(const RowStripPhysicalValue& lhs,
+                                         const RowStripPhysicalValue& rhs,
+                                         spatial::SpatAddPlanOp planOp,
+                                         PatternRewriter& rewriter) {
+  return applyRowStripAdd(lhs, rhs, rewriter, planOp.getLoc());
+}
+
+static FailureOr<Value> lowerRowStripConcat(ArrayRef<RowStripPhysicalValue> inputs,
+                                            spatial::SpatConcatPlanOp planOp,
+                                            PatternRewriter& rewriter) {
+  auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
+  if (!outputType)
+    return failure();
+  return applyRowStripConcat(inputs, outputType, rewriter, planOp.getLoc());
+}
+
 static FailureOr<Value>
 materializeRowStripToDense(const RowStripPhysicalValue& rowStripValue, Location loc, PatternRewriter& rewriter) {
   if (rowStripValue.logicalType.getRank() != 4 || !rowStripValue.logicalType.hasStaticShape())
@@ -122,6 +138,99 @@ static FailureOr<Value> lowerDenseBatchBiasAdd(Value input, Value bias, RankedTe
   return batch->getResult(0);
 }
 
+static LogicalResult lowerAddPlan(spatial::SpatAddPlanOp planOp,
+                                  llvm::DenseMap<Value, RowStripPhysicalValue>& rowStripValues,
+                                  llvm::SmallPtrSetImpl<Operation*>& eraseAfterLowering,
+                                  PatternRewriter& rewriter) {
+  FailureOr<RowStripPhysicalValue> lhs = getRowStripValue(rowStripValues, planOp.getLhs());
+  FailureOr<RowStripPhysicalValue> rhs = getRowStripValue(rowStripValues, planOp.getRhs());
+  if (succeeded(lhs) && succeeded(rhs)) {
+    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())
+      return planOp.emitOpError("row-strip add plan requires a row-strip blueprint result");
+
+    rewriter.setInsertionPoint(planOp);
+    FailureOr<Value> lowered = lowerRowStripAdd(*lhs, *rhs, planOp, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Spatial add plan");
+    auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
+    FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
+    if (failed(output))
+      return failure();
+    rowStripValues[blueprint.getResult()] = *output;
+    eraseAfterLowering.insert(planOp);
+    eraseAfterLowering.insert(blueprint);
+    return success();
+  }
+
+  rewriter.setInsertionPoint(planOp);
+  auto compute = createSpatCompute<2>(rewriter,
+                                      planOp.getLoc(),
+                                      planOp.getOutput().getType(),
+                                      {},
+                                      ValueRange {planOp.getLhs(), planOp.getRhs()},
+                                      [&](Value lhsValue, Value rhsValue) {
+                                        Value added = spatial::SpatVAddOp::create(
+                                          rewriter, planOp.getLoc(), planOp.getOutput().getType(), lhsValue, rhsValue);
+                                        spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added);
+                                      });
+  rewriter.replaceOp(planOp, compute.getResults());
+  return success();
+}
+
+static LogicalResult lowerConcatPlan(spatial::SpatConcatPlanOp planOp,
+                                     llvm::DenseMap<Value, RowStripPhysicalValue>& rowStripValues,
+                                     llvm::SmallPtrSetImpl<Operation*>& eraseAfterLowering,
+                                     PatternRewriter& rewriter) {
+  SmallVector<RowStripPhysicalValue> inputs;
+  for (Value input : planOp.getInputs()) {
+    FailureOr<RowStripPhysicalValue> physical = getRowStripValue(rowStripValues, input);
+    if (failed(physical)) {
+      inputs.clear();
+      break;
+    }
+    inputs.push_back(*physical);
+  }
+  if (inputs.size() == planOp.getInputs().size()) {
+    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())
+      return planOp.emitOpError("row-strip concat plan requires a row-strip blueprint result");
+    rewriter.setInsertionPoint(planOp);
+    FailureOr<Value> lowered = lowerRowStripConcat(inputs, planOp, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Spatial concat plan");
+    auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
+    FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
+    if (failed(output))
+      return failure();
+    rowStripValues[blueprint.getResult()] = *output;
+    eraseAfterLowering.insert(planOp);
+    eraseAfterLowering.insert(blueprint);
+    return success();
+  }
+
+  rewriter.setInsertionPoint(planOp);
+  auto compute = createSpatCompute(
+    rewriter,
+    planOp.getLoc(),
+    TypeRange {planOp.getOutput().getType()},
+    {},
+    planOp.getInputs(),
+    [&](ValueRange values) {
+      Value concatenated = spatial::SpatConcatOp::create(
+        rewriter, planOp.getLoc(), planOp.getOutput().getType(), rewriter.getI64IntegerAttr(planOp.getAxis()), values);
+      spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), concatenated);
+    });
+  rewriter.replaceOp(planOp, compute.getResults());
+  return success();
+}
+
 struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerSpatialPlansPass)
 
@@ -274,6 +383,40 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
         eraseAfterLowering.insert(blueprint);
         continue;
       }
+      if (auto planOp = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(&op)) {
+        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("selected global AveragePool plan requires a row-strip blueprint result");
+          signalPassFailure();
+          return;
+        }
+
+        FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
+        rewriter.setInsertionPoint(planOp);
+        std::optional<Value> physicalInput;
+        if (succeeded(input))
+          physicalInput = input->storage;
+        FailureOr<Value> lowered =
+          lowerSelectedGlobalAveragePoolPlan(planOp, physicalInput, rewriter);
+        if (failed(lowered)) {
+          planOp.emitOpError("failed to lower selected row-strip Spatial global AveragePool 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;
+      }
       if (auto planOp = dyn_cast<spatial::SpatBiasAddPlanOp>(&op)) {
         if (succeeded(getRowStripValue(rowStripValues, planOp.getInput()))) {
           auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
@@ -339,6 +482,20 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
         rewriter.replaceOp(planOp, computeOp.getResults());
         continue;
       }
+      if (auto planOp = dyn_cast<spatial::SpatAddPlanOp>(&op)) {
+        if (failed(lowerAddPlan(planOp, rowStripValues, eraseAfterLowering, rewriter))) {
+          signalPassFailure();
+          return;
+        }
+        continue;
+      }
+      if (auto planOp = dyn_cast<spatial::SpatConcatPlanOp>(&op)) {
+        if (failed(lowerConcatPlan(planOp, rowStripValues, eraseAfterLowering, rewriter))) {
+          signalPassFailure();
+          return;
+        }
+        continue;
+      }
       if (auto flattenOp = dyn_cast<spatial::SpatGraphCompute>(&op)) {
         if (flattenOp.getInputs().size() == 1) {
           FailureOr<RowStripPhysicalValue> input =
@@ -488,12 +645,15 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
           return;
         op->emitOpError("planning blueprint must not remain after LowerSpatialPlans");
         hasIllegalOps = true;
-      } else if (isa<spatial::SpatConv2DPlanOp,
-                     spatial::SpatBiasAddPlanOp,
-                     spatial::SpatReluPlanOp,
-                     spatial::SpatMaxPool2DPlanOp,
-                     spatial::SpatMaterializeLayoutOp>(op)
-          || op->getDialect()->getNamespace() == "onnx") {
+      }
+      else if (isa<spatial::SpatConv2DPlanOp,
+                   spatial::SpatBiasAddPlanOp,
+                   spatial::SpatAddPlanOp,
+                   spatial::SpatReluPlanOp,
+                   spatial::SpatMaxPool2DPlanOp,
+                   spatial::SpatGlobalAveragePoolPlanOp,
+                   spatial::SpatMaterializeLayoutOp>(op)
+               || op->getDialect()->getNamespace() == "onnx") {
         op->emitOpError("operation must not remain after LowerSpatialPlans");
         hasIllegalOps = true;
       }

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatial.td

@@ -60,11 +60,6 @@ def convAddToConvWithBiasRight : Pat<
 
 def replaceWithOperationOfValue : NativeCodeCall<"$0">;
 
-def removeLRN : Pat<
-  (ONNXLRNOp $A, $_, $_, $_, $_),
-  (replaceWithOperationOfValue $A)
->;
-
 def HaveSameStaticShape: Constraint<
   CPred<"onnx_mlir::haveSameStaticShape($0, $1)">,
   "Two tensors have the same static shape">;

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp

@@ -47,12 +47,17 @@ static void populateEmptyFunction(func::FuncOp funcOp) {
   SmallVector<spatial::SpatGraphComputeBatch> computeBatches(funcOp.getOps<spatial::SpatGraphComputeBatch>());
   SmallVector<spatial::SpatConv2DPlanOp> convPlans(funcOp.getOps<spatial::SpatConv2DPlanOp>());
   SmallVector<spatial::SpatBiasAddPlanOp> biasAddPlans(funcOp.getOps<spatial::SpatBiasAddPlanOp>());
+  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::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() || !reluPlans.empty()
-      || !maxPoolPlans.empty() || !blueprints.empty() || !materializers.empty()) {
+  if (!computes.empty() || !computeBatches.empty() || !convPlans.empty() || !biasAddPlans.empty() || !addPlans.empty()
+      || !concatPlans.empty() || !reluPlans.empty() || !maxPoolPlans.empty() || !blueprints.empty()
+      || !globalAveragePoolPlans.empty() || !materializers.empty()) {
     return;
   }
 

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp

@@ -146,8 +146,11 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter
             spatial::SpatGraphComputeBatch,
             spatial::SpatConv2DPlanOp,
             spatial::SpatBiasAddPlanOp,
+            spatial::SpatAddPlanOp,
+            spatial::SpatConcatPlanOp,
             spatial::SpatReluPlanOp,
             spatial::SpatMaxPool2DPlanOp,
+            spatial::SpatGlobalAveragePoolPlanOp,
             spatial::SpatBlueprintOp,
             spatial::SpatMaterializeLayoutOp>(&op)) {
       continue;

src/PIM/Conversion/ONNXToSpatial/Patterns.cpp

@@ -8,7 +8,6 @@ namespace onnx_mlir {
 void populatePrePatterns(RewritePatternSet& patterns, MLIRContext* ctx) { populateGeneratedPrePatterns(patterns, ctx); }
 
 void populateConversionPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
-  populateGeneratedConversionPatterns(patterns, ctx);
   populateElementwisePatterns(patterns, ctx);
   populateMatMulRewritePatterns(patterns, ctx);
   populateGemmPatterns(patterns, ctx);

src/PIM/Conversion/ONNXToSpatial/Patterns.hpp

@@ -13,7 +13,6 @@ void populateConversionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRCon
 void populatePostPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 
 void populateGeneratedPrePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
-void populateGeneratedConversionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateWeightPromotionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 
 void populateConvPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);

src/PIM/Conversion/ONNXToSpatial/Patterns/GeneratedConversion.cpp

@@ -1,18 +0,0 @@
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
-
-using namespace mlir;
-
-namespace onnx_mlir {
-
-namespace {
-
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatial.hpp.inc"
-
-} // namespace
-
-void populateGeneratedConversionPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
-  patterns.add<removeLRN>(ctx);
-}
-
-} // namespace onnx_mlir

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

@@ -1242,8 +1242,9 @@ static Value buildPackedWeights(DenseElementsAttr wDenseAttr,
                                 const Tiling& tiling,
                                 PatternRewriter& rewriter,
                                 Location loc) {
+  const int64_t paddedOutputChannels = static_cast<int64_t>(crossbarSize.getValue());
   auto packedWeightType = RankedTensorType::get(
-    {tiling.numChannelTiles, tiling.tileInputRows, tiling.tileOutputChannels}, wType.getElementType());
+    {tiling.numChannelTiles, tiling.tileInputRows, paddedOutputChannels}, wType.getElementType());
   SmallVector<Attribute> packedValues(packedWeightType.getNumElements(),
                                       cast<Attribute>(rewriter.getZeroAttr(wType.getElementType())));
   SmallVector<Attribute> sourceValues(wDenseAttr.getValues<Attribute>());
@@ -1262,7 +1263,7 @@ static Value buildPackedWeights(DenseElementsAttr wDenseAttr,
             ((globalOutChannel * wType.getDimSize(1) * wType.getDimSize(2)) + kernelH) * wType.getDimSize(3) + kernelW;
           const int64_t targetCol = localChannel * tiling.outputMultiplier + multiplierIndex;
           const int64_t targetFlatIndex =
-            ((tileIndex * tiling.tileInputRows) + targetRow) * tiling.tileOutputChannels + targetCol;
+            ((tileIndex * tiling.tileInputRows) + targetRow) * paddedOutputChannels + targetCol;
           packedValues[targetFlatIndex] = sourceValues[sourceFlatIndex];
         }
       }
@@ -1353,11 +1354,12 @@ static Value createWeightTile(Value packedWeights,
                               PatternRewriter& rewriter,
                               Location loc) {
   SmallVector<OpFoldResult> offsets {channelTileIndex, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
+  const int64_t paddedOutputChannels = static_cast<int64_t>(crossbarSize.getValue());
   SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
                                    rewriter.getIndexAttr(tiling.tileInputRows),
-                                   rewriter.getIndexAttr(tiling.tileOutputChannels)};
+                                   rewriter.getIndexAttr(paddedOutputChannels)};
   auto collapsedType =
-    RankedTensorType::get({tiling.tileInputRows, tiling.tileOutputChannels}, packedWeightType.getElementType());
+    RankedTensorType::get({tiling.tileInputRows, paddedOutputChannels}, packedWeightType.getElementType());
   return extractMixedSliceOrIdentity(
     rewriter, loc, packedWeights, collapsedType,
     {offsets, sizes, getUnitStrides(rewriter, 3)});
@@ -1547,6 +1549,8 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter&
   auto gemmOutType =
     RankedTensorType::get({tiling->totalPatches, state.outType.getDimSize(1)}, state.outType.getElementType());
   auto rowTileType = RankedTensorType::get({1, tiling->tileOutputChannels}, state.outType.getElementType());
+  auto paddedRowTileType = RankedTensorType::get(
+    {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());
@@ -1617,7 +1621,17 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter&
                                             *tiling,
                                             rewriter,
                                             loc);
-      Value rowTile = spatial::SpatVMMOp::create(rewriter, loc, rowTileType, weightTile, inputTile).getResult();
+      Value paddedRowTile =
+        spatial::SpatVMMOp::create(rewriter, loc, paddedRowTileType, weightTile, inputTile).getResult();
+      Value rowTile = tensor::ExtractSliceOp::create(
+        rewriter,
+        loc,
+        rowTileType,
+        paddedRowTile,
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
+                                  rewriter.getIndexAttr(tiling->tileOutputChannels)},
+        getUnitStrides(rewriter, 2));
       if (args.inputs.size() > 1) {
         Value biasArg = pickInputByRank(/*rank=*/2);
         if (!biasArg) {
@@ -2493,6 +2507,7 @@ static Value createZeroGemmBias(RankedTensorType gemmResultType, PatternRewriter
 
 static bool rowStripOutputTileFitsOneCore(const ConvGeometry& geometry) {
   return ceilIntegerDivide(geometry.k, geometry.xbarSize)
+           * ceilIntegerDivide(geometry.c, geometry.xbarSize)
          <= static_cast<int64_t>(crossbarCountInCore.getValue());
 }
 
@@ -2521,28 +2536,6 @@ static bool canConsumePixelMajorRowStripFragments(const ConvLoweringState& state
     failureReason = "dilation_not_one";
     return false;
   }
-  const bool pointwise = state.xHeight == 1 && state.xWidth == 1 && state.outHeight == 1 && state.outWidth == 1
-                         && state.wHeight == 1 && state.wWidth == 1 && state.padHeightBegin == 0
-                         && state.padHeightEnd == 0 && state.padWidthBegin == 0 && state.padWidthEnd == 0;
-  if (pointwise) {
-    if (!getHostConstDenseElementsAttr(state.w)) {
-      failureReason = "non_constant_weight";
-      return false;
-    }
-    if (state.hasBias && !isSupportedBiasAddValue(state.b, state.outType)) {
-      failureReason = "unsupported_bias";
-      return false;
-    }
-    return true;
-  }
-  if (state.wHeight != 3 || state.wWidth != 3) {
-    failureReason = "kernel_not_3x3";
-    return false;
-  }
-  if (state.padHeightBegin != 1 || state.padHeightEnd != 1 || state.padWidthBegin != 1 || state.padWidthEnd != 1) {
-    failureReason = "padding_not_1";
-    return false;
-  }
   if (state.outHeight != state.xHeight || state.outWidth != state.xWidth) {
     failureReason = "not_same_spatial_shape";
     return false;
@@ -3043,57 +3036,154 @@ static FailureOr<Value> createConvOutputRow(ValueRange inputTiles,
                                             Location loc) {
   auto elementType = cast<RankedTensorType>(inputTiles.front().getType()).getElementType();
   auto rowType = RankedTensorType::get({1, outputChannels}, elementType);
-  auto tileWeightsType =
-    RankedTensorType::get({paddedK, xbarDim},
-                          cast<RankedTensorType>(paddedWeights.getType()).getElementType());
   const int64_t outputTileCount = ceilIntegerDivide(outputChannels, xbarDim);
-
-  auto getTileWeights = [&](int64_t outputTile) {
-    if (outputTileCount == 1)
-      return paddedWeights;
-    SmallVector<OpFoldResult> offsets {
-      rewriter.getIndexAttr(outputTile), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
-    SmallVector<OpFoldResult> sizes {
-      rewriter.getIndexAttr(1), rewriter.getIndexAttr(paddedK), rewriter.getIndexAttr(xbarDim)};
-    return extractStaticSliceOrIdentity(
-      rewriter, loc, paddedWeights, tileWeightsType, offsets, sizes, getUnitStrides(rewriter, 3));
-  };
-
-  if (outputTileCount == 1) {
-    FailureOr<Value> rowResult = createConvOutputTile(
-      inputTiles, getTileWeights(0), outputChannels, xbarDim, rewriter, loc);
-    if (failed(rowResult))
-      return failure();
-    Value validRow = *rowResult;
-    if (bias)
-      validRow = spatial::SpatVAddOp::create(rewriter, loc, rowType, validRow, bias).getResult();
-    return validRow;
-  }
-
   const int64_t paddedOutputChannels = outputTileCount * xbarDim;
   auto paddedOutputType = RankedTensorType::get({1, paddedOutputChannels}, elementType);
-  Value paddedOutput = tensor::EmptyOp::create(rewriter, loc, paddedOutputType.getShape(), elementType);
-  for (int64_t outputTile = 0; outputTile < outputTileCount; ++outputTile) {
-    FailureOr<Value> tileResult = createConvOutputTile(
-      inputTiles, getTileWeights(outputTile), xbarDim, xbarDim, rewriter, loc);
-    if (failed(tileResult))
-      return failure();
-    SmallVector<OpFoldResult> tileOffsets {
-      rewriter.getIndexAttr(0), rewriter.getIndexAttr(outputTile * xbarDim)};
-    SmallVector<OpFoldResult> tileSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)};
-    paddedOutput = tensor::InsertSliceOp::create(
-      rewriter, loc, *tileResult, paddedOutput, tileOffsets, tileSizes, getUnitStrides(rewriter, 2));
+  auto weightSliceType = RankedTensorType::get(
+    {xbarDim, paddedOutputChannels},
+    cast<RankedTensorType>(paddedWeights.getType()).getElementType());
+
+  Value paddedOutput;
+  for (auto [kSlice, inputTile] : llvm::enumerate(inputTiles)) {
+    const int64_t kOffset = static_cast<int64_t>(kSlice) * xbarDim;
+    Value weightSlice = extractStaticSliceOrIdentity(
+      rewriter,
+      loc,
+      paddedWeights,
+      weightSliceType,
+      SmallVector<OpFoldResult> {rewriter.getIndexAttr(kOffset), rewriter.getIndexAttr(0)},
+      SmallVector<OpFoldResult> {rewriter.getIndexAttr(xbarDim),
+                                rewriter.getIndexAttr(paddedOutputChannels)},
+      getUnitStrides(rewriter, 2));
+    Value piece =
+      spatial::SpatVMMOp::create(rewriter, loc, paddedOutputType, weightSlice, inputTile).getResult();
+    paddedOutput = paddedOutput
+                     ? spatial::SpatVAddOp::create(
+                         rewriter, loc, paddedOutputType, paddedOutput, piece).getResult()
+                     : piece;
   }
 
-  SmallVector<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
-  SmallVector<OpFoldResult> outputSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(outputChannels)};
-  Value validRow = tensor::ExtractSliceOp::create(
-    rewriter, loc, rowType, paddedOutput, outputOffsets, outputSizes, getUnitStrides(rewriter, 2));
+  Value validRow = outputChannels == paddedOutputChannels
+                     ? paddedOutput
+                     : tensor::ExtractSliceOp::create(
+                         rewriter,
+                         loc,
+                         rowType,
+                         paddedOutput,
+                         SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
+                         SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
+                                                   rewriter.getIndexAttr(outputChannels)},
+                         getUnitStrides(rewriter, 2))
+                         .getResult();
   if (bias)
     validRow = spatial::SpatVAddOp::create(rewriter, loc, rowType, validRow, bias).getResult();
   return validRow;
 }
 
+static FailureOr<Value> createRowStripConvOutput(const ConvLoweringState& state,
+                                                 Value input,
+                                                 Value paddedWeights,
+                                                 Value bias,
+                                                 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 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 fragmentType = getRowStripFragmentType(state.outType);
+  auto storageType = getRowStripStorageType(state.outType);
+
+  auto batch = createSpatComputeBatch(
+    rewriter,
+    loc,
+    TypeRange {storageType},
+    laneCount,
+    ValueRange {paddedWeights},
+    bias ? ValueRange {input, bias} : ValueRange {input},
+    [&](detail::SpatComputeBatchBodyArgs args) {
+      Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+      Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
+      Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
+      Value cOutWidth = getOrCreateIndexConstant(rewriter, anchorOp, state.outWidth);
+      FailureOr<Value> inputWindow =
+        createConvInputWindow(args.inputs.front(), state, args.lane, rewriter, loc);
+      if (failed(inputWindow))
+        return failure();
+      Value fragmentInit = tensor::EmptyOp::create(rewriter, loc, fragmentType.getShape(), elementType);
+      SmallVector<Value> loopInit {fragmentInit};
+      if (hasPartialInputTile)
+        loopInit.push_back(createZeroTensorConstant(partialInputScratchType, rewriter));
+      auto loop = buildNormalizedScfFor(
+        rewriter,
+        loc,
+        c0,
+        cOutWidth,
+        c1,
+        loopInit,
+        [&](OpBuilder&,
+            Location pixelLoc,
+            Value localColumn,
+            ValueRange iterArgs,
+            SmallVectorImpl<Value>& yielded) {
+          Value partialInputScratch = hasPartialInputTile ? iterArgs[1] : Value();
+          FailureOr<SmallVector<Value>> inputTiles = createConvInputTiles(*inputWindow,
+                                                                          state,
+                                                                          localColumn,
+                                                                          partialInputScratch,
+                                                                          patchSize,
+                                                                          numKSlices,
+                                                                          xbarDim,
+                                                                          rewriter,
+                                                                          pixelLoc);
+          if (failed(inputTiles))
+            return failure();
+          FailureOr<Value> output = createConvOutputRow(*inputTiles,
+                                                        paddedK,
+                                                        state.numChannelsOut,
+                                                        args.weights.front(),
+                                                        bias ? args.inputs[1] : Value(),
+                                                        xbarDim,
+                                                        rewriter,
+                                                        pixelLoc);
+          if (failed(output))
+            return failure();
+          Value outputPixel = tensor::ExpandShapeOp::create(
+            rewriter, pixelLoc, outputPixelType, *output, SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
+          Value next = tensor::InsertSliceOp::create(
+            rewriter,
+            pixelLoc,
+            outputPixel,
+            iterArgs.front(),
+            SmallVector<OpFoldResult> {rewriter.getIndexAttr(0),
+                                       rewriter.getIndexAttr(0),
+                                       localColumn,
+                                       rewriter.getIndexAttr(0)},
+            SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
+                                       rewriter.getIndexAttr(1),
+                                       rewriter.getIndexAttr(1),
+                                       rewriter.getIndexAttr(state.numChannelsOut)},
+            getUnitStrides(rewriter, 4));
+          yielded.push_back(next);
+          if (hasPartialInputTile)
+            yielded.push_back(partialInputScratch);
+          return success();
+        });
+      if (failed(loop))
+        return failure();
+      publishGraphBatchPhysicalFragment(
+        rewriter, loc, loop->results.front(), args.outputs.front(), args.lane);
+      return success();
+    });
+  if (failed(batch))
+    return failure();
+  return batch->getResult(0);
+}
+
 static FailureOr<Value> createOutputChannelTiledRowStripConvOutput(const ConvLoweringState& state,
                                                                    Value input,
                                                                    Value paddedWeights,
@@ -3229,21 +3319,11 @@ static FailureOr<Value>
   const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
   const int64_t numKSlices = ceilIntegerDivide(patchSize, xbarDim);
   const int64_t paddedK = numKSlices * xbarDim;
-  const bool hasPartialInputTile = patchSize % xbarDim != 0;
-  auto elementType = state.outType.getElementType();
-  auto partialInputScratchType = RankedTensorType::get({1, xbarDim}, elementType);
-  auto fragmentType = getRowStripFragmentType(state.outType);
-  auto outputPixelType = RankedTensorType::get({1, 1, 1, state.numChannelsOut}, elementType);
-  auto outputStorageType = getRowStripStorageType(state.outType);
 
-  Value paddedWeights = state.numChannelsOut <= xbarDim
-                          ? standard::createPaddedPixelMajorWeightConstant(
-                              weightDenseAttr, state, paddedK, xbarDim, rewriter)
-                          : standard::createPaddedOutputChannelTiledWeightConstant(
-                              weightDenseAttr, state, paddedK, xbarDim, rewriter);
-  if (state.numChannelsOut > xbarDim)
-    return createOutputChannelTiledRowStripConvOutput(
-      state, state.x, paddedWeights, paddedK, numKSlices, xbarDim, rewriter, loc);
+  const int64_t paddedOutputChannels =
+    ceilIntegerDivide(state.numChannelsOut, xbarDim) * xbarDim;
+  Value paddedWeights = standard::createPaddedPixelMajorWeightConstant(
+    weightDenseAttr, state, paddedK, paddedOutputChannels, rewriter);
 
   FailureOr<Value> bias = failure();
   if (state.hasBias)
@@ -3251,83 +3331,9 @@ static FailureOr<Value>
   if (state.hasBias && failed(bias))
     return failure();
 
-  auto batchOp = createSpatComputeBatch(
-    rewriter,
-    loc,
-    TypeRange {outputStorageType},
-    state.outHeight,
-    ValueRange {paddedWeights},
-    state.hasBias ? ValueRange {state.x, *bias} : ValueRange {state.x},
-    [&](detail::SpatComputeBatchBodyArgs args) {
-      Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
-      Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
-      Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
-      Value cOutWidth = getOrCreateIndexConstant(rewriter, anchorOp, state.outWidth);
-      FailureOr<Value> inputWindow =
-        createConvInputWindow(args.inputs.front(), state, args.lane, rewriter, loc);
-      if (failed(inputWindow))
-        return failure();
-      Value fragmentInit = tensor::EmptyOp::create(rewriter, loc, fragmentType.getShape(), elementType);
-      SmallVector<Value> widthLoopInit {fragmentInit};
-      if (hasPartialInputTile)
-        widthLoopInit.push_back(createZeroTensorConstant(partialInputScratchType, rewriter));
-      auto widthLoop = buildNormalizedScfFor(
-        rewriter,
-        loc,
-        c0,
-        cOutWidth,
-        c1,
-        widthLoopInit,
-        [&](OpBuilder&, Location widthLoc, Value widthIndex, ValueRange widthIterArgs, SmallVectorImpl<Value>& widthYielded) {
-          Value partialInputScratch = hasPartialInputTile ? widthIterArgs[1] : Value();
-          FailureOr<SmallVector<Value>> inputTiles = createConvInputTiles(*inputWindow,
-                                                                          state,
-                                                                          widthIndex,
-                                                                          partialInputScratch,
-                                                                          patchSize,
-                                                                          numKSlices,
-                                                                          xbarDim,
-                                                                          rewriter,
-                                                                          widthLoc);
-          if (failed(inputTiles))
-            return failure();
-          FailureOr<Value> outputRow = createConvOutputRow(*inputTiles,
-                                                           paddedK,
-                                                           state.numChannelsOut,
-                                                           args.weights.front(),
-                                                           state.hasBias ? args.inputs[1] : Value(),
-                                                           xbarDim,
-                                                           rewriter,
-                                                           widthLoc);
-          if (failed(outputRow))
-            return failure();
-
-          Value outputFragment = tensor::ExpandShapeOp::create(rewriter,
-                                                               widthLoc,
-                                                               outputPixelType,
-                                                               *outputRow,
-                                                               SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
-          SmallVector<OpFoldResult> rowOffsets {
-            rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex, rewriter.getIndexAttr(0)};
-          SmallVector<OpFoldResult> rowSizes {
-            rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1),
-            rewriter.getIndexAttr(state.numChannelsOut)};
-          Value nextFragment = tensor::InsertSliceOp::create(
-            rewriter, widthLoc, outputFragment, widthIterArgs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 4));
-          widthYielded.push_back(nextFragment);
-          if (hasPartialInputTile)
-            widthYielded.push_back(partialInputScratch);
-          return success();
-        });
-      if (failed(widthLoop))
-        return failure();
-
-      insertRowStripFragment(widthLoop->results.front(), args.outputs.front(), state.outType, args.lane, rewriter, loc);
-      return success();
-    });
-  if (failed(batchOp))
-    return failure();
-  return batchOp->getResult(0);
+  return createRowStripConvOutput(
+    state, state.x, paddedWeights, state.hasBias ? *bias : Value(),
+    paddedK, numKSlices, xbarDim, rewriter, loc);
 }
 
 static FailureOr<Value> createConvOutputFromPixelMajorRowStripFragments(Value rowStripStorage,
@@ -3346,105 +3352,22 @@ static FailureOr<Value> createConvOutputFromPixelMajorRowStripFragments(Value ro
   const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
   const int64_t numKSlices = ceilIntegerDivide(patchSize, xbarDim);
   const int64_t paddedK = numKSlices * xbarDim;
-  const bool hasPartialInputTile = patchSize % xbarDim != 0;
-  auto elementType = state.outType.getElementType();
-  auto partialInputScratchType = RankedTensorType::get({1, xbarDim}, elementType);
-  auto outputPixelType = RankedTensorType::get({1, 1, 1, state.numChannelsOut}, elementType);
-  auto outputStorageType = getRowStripStorageType(state.outType);
   auto weightDenseAttr = getHostConstDenseElementsAttr(state.w);
   if (!weightDenseAttr)
     return failure();
-  Value paddedWeights = state.numChannelsOut <= xbarDim
-                          ? standard::createPaddedPixelMajorWeightConstant(
-                              weightDenseAttr, state, paddedK, xbarDim, rewriter)
-                          : standard::createPaddedOutputChannelTiledWeightConstant(
-                              weightDenseAttr, state, paddedK, xbarDim, rewriter);
-  if (state.numChannelsOut > xbarDim)
-    return createOutputChannelTiledRowStripConvOutput(
-      state, rowStripStorage, paddedWeights, paddedK, numKSlices, xbarDim, rewriter, loc);
+  const int64_t paddedOutputChannels =
+    ceilIntegerDivide(state.numChannelsOut, xbarDim) * xbarDim;
+  Value paddedWeights = standard::createPaddedPixelMajorWeightConstant(
+    weightDenseAttr, state, paddedK, paddedOutputChannels, rewriter);
   FailureOr<Value> bias = failure();
   if (state.hasBias)
     bias = createBiasRowConstant(state, rewriter);
   if (state.hasBias && failed(bias))
     return failure();
 
-  auto batchOp = createSpatComputeBatch(
-    rewriter,
-    loc,
-    TypeRange {outputStorageType},
-    state.outHeight,
-    ValueRange {paddedWeights},
-    state.hasBias ? ValueRange {rowStripStorage, *bias} : ValueRange {rowStripStorage},
-    [&](detail::SpatComputeBatchBodyArgs args) {
-      Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
-      Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
-      Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
-      Value cOutWidth = getOrCreateIndexConstant(rewriter, anchorOp, state.outWidth);
-      auto fragmentType = getRowStripFragmentType(state.outType);
-      FailureOr<Value> inputWindow = createConvInputWindow(args.inputs.front(), state, args.lane, rewriter, loc);
-      if (failed(inputWindow))
-        return failure();
-      Value fragmentInit = tensor::EmptyOp::create(rewriter, loc, fragmentType.getShape(), elementType);
-      SmallVector<Value> widthLoopInit {fragmentInit};
-      if (hasPartialInputTile)
-        widthLoopInit.push_back(createZeroTensorConstant(partialInputScratchType, rewriter));
-      auto widthLoop = buildNormalizedScfFor(
-        rewriter,
-        loc,
-        c0,
-        cOutWidth,
-        c1,
-        widthLoopInit,
-        [&](OpBuilder&, Location widthLoc, Value widthIndex, ValueRange widthIterArgs, SmallVectorImpl<Value>& widthYielded) {
-          Value partialInputScratch = hasPartialInputTile ? widthIterArgs[1] : Value();
-          FailureOr<SmallVector<Value>> inputTiles = createConvInputTiles(*inputWindow,
-                                                                          state,
-                                                                          widthIndex,
-                                                                          partialInputScratch,
-                                                                          patchSize,
-                                                                          numKSlices,
-                                                                          xbarDim,
-                                                                          rewriter,
-                                                                          widthLoc);
-          if (failed(inputTiles))
-            return failure();
-          FailureOr<Value> outputRow = createConvOutputRow(*inputTiles,
-                                                           paddedK,
-                                                           state.numChannelsOut,
-                                                           args.weights.front(),
-                                                           state.hasBias ? args.inputs[1] : Value(),
-                                                           xbarDim,
-                                                           rewriter,
-                                                           widthLoc);
-          if (failed(outputRow))
-            return failure();
-
-          Value outputFragment = tensor::ExpandShapeOp::create(rewriter,
-                                                               widthLoc,
-                                                               outputPixelType,
-                                                               *outputRow,
-                                                               SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
-          SmallVector<OpFoldResult> rowOffsets {
-            rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex, rewriter.getIndexAttr(0)};
-          SmallVector<OpFoldResult> rowSizes {
-            rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1),
-            rewriter.getIndexAttr(state.numChannelsOut)};
-          Value nextFragment = tensor::InsertSliceOp::create(
-            rewriter, widthLoc, outputFragment, widthIterArgs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 4));
-          widthYielded.push_back(nextFragment);
-          if (hasPartialInputTile)
-            widthYielded.push_back(partialInputScratch);
-          return success();
-        });
-      if (failed(widthLoop))
-        return failure();
-
-      insertRowStripFragment(widthLoop->results.front(), args.outputs.front(), state.outType, args.lane, rewriter, loc);
-      return success();
-    });
-  if (failed(batchOp))
-    return failure();
-  return batchOp->getResult(0);
+  return createRowStripConvOutput(
+    state, rowStripStorage, paddedWeights, state.hasBias ? *bias : Value(),
+    paddedK, numKSlices, xbarDim, rewriter, loc);
 }
 
 static FailureOr<Value> createPointwiseOutputFromRowStripFragments(Value rowStripStorage,

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

@@ -193,6 +193,13 @@ struct AddToSpatialCompute : OpConversionPattern<ONNXAddOp> {
       return success();
     }
 
+    if (resultType.getRank() == 4 && adaptor.getA().getType() == resultType && adaptor.getB().getType() == resultType) {
+      auto plan = spatial::SpatAddPlanOp::create(
+        rewriter, op.getLoc(), resultType, adaptor.getA(), adaptor.getB(), rewriter.getStringAttr("nchw"));
+      rewriter.replaceOp(op, plan.getResult());
+      return success();
+    }
+
     auto lhs = prepareElementwiseOperand(adaptor.getA(), resultType, rewriter, op.getLoc());
     if (failed(lhs))
       return failure();

src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Pool.cpp

@@ -246,6 +246,15 @@ struct PoolToSpatialComputeBase : public OpConversionPattern<PoolOp> {
         return success();
       }
     }
+    else if (batchSize == 1 && outputHeight == 1 && outputWidth == 1
+             && kernelHeight == inputHeight && kernelWidth == inputWidth
+             && dilationHeight == 1 && dilationWidth == 1 && padTop == 0
+             && padLeft == 0 && padBottom == 0 && padRight == 0) {
+      auto plan = spatial::SpatGlobalAveragePoolPlanOp::create(
+        rewriter, loc, outType, x, rewriter.getStringAttr("nchw"));
+      rewriter.replaceOp(poolOp, plan.getResult());
+      return success();
+    }
 
     const int64_t xbarSize = static_cast<int64_t>(crossbarSize.getValue());
     const int64_t channelTileCount = (channels + xbarSize - 1) / xbarSize;
@@ -676,6 +685,132 @@ FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
   return batch->getResult(0);
 }
 
+LogicalResult canLowerGlobalAveragePoolPlanToRowStrip(spatial::SpatGlobalAveragePoolPlanOp planOp) {
+  auto inputType = dyn_cast<RankedTensorType>(planOp.getInput().getType());
+  auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
+  if (!inputType || !outputType || !inputType.hasStaticShape() || !outputType.hasStaticShape())
+    return failure();
+  if (inputType.getRank() != 4 || outputType.getRank() != 4 || inputType.getDimSize(0) != 1
+      || outputType.getDimSize(0) != 1 || inputType.getDimSize(1) != outputType.getDimSize(1)
+      || outputType.getDimSize(2) != 1 || outputType.getDimSize(3) != 1)
+    return failure();
+  return success();
+}
+
+FailureOr<Value> lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp,
+                                                    std::optional<Value> rowStripInput,
+                                                    PatternRewriter& rewriter) {
+  if (failed(canLowerGlobalAveragePoolPlanToRowStrip(planOp)))
+    return failure();
+
+  Location loc = planOp.getLoc();
+  auto inputType = cast<RankedTensorType>(planOp.getInput().getType());
+  auto outputType = cast<RankedTensorType>(planOp.getOutput().getType());
+  auto elementType = dyn_cast<FloatType>(inputType.getElementType());
+  if (!elementType)
+    return failure();
+
+  Value input = rowStripInput.value_or(planOp.getInput());
+  auto actualInputType = dyn_cast<RankedTensorType>(input.getType());
+  FailureOr<RowStripPhysicalValue> physicalValue = describeRowStripPhysicalValue(input, inputType);
+  const bool physicalInput = succeeded(physicalValue);
+  if (!physicalInput && actualInputType != inputType)
+    return failure();
+
+  const int64_t height = inputType.getDimSize(2);
+  const int64_t width = inputType.getDimSize(3);
+  const int64_t channels = inputType.getDimSize(1);
+  const int64_t tilesPerRow = physicalInput ? physicalValue->tilesPerRow : 1;
+  auto inputFragmentType =
+    physicalInput ? physicalValue->fragmentType : getRowStripFragmentType(inputType);
+  auto nchwInputFragmentType = RankedTensorType::get(
+    {1, channels, 1, width}, inputType.getElementType(), inputType.getEncoding());
+  auto outputFragmentType = RankedTensorType::get(
+    {1, 1, 1, inputFragmentType.getDimSize(3)}, elementType, outputType.getEncoding());
+  auto outputStorageType =
+    spatial::getGraphBatchPhysicalResultType(tilesPerRow, outputFragmentType);
+  auto zero = getOrCreateConstant(
+    rewriter, rewriter.getInsertionBlock()->getParentOp(), rewriter.getZeroAttr(outputFragmentType), outputFragmentType);
+  auto scaleAttr = DenseElementsAttr::get(
+    outputFragmentType, rewriter.getFloatAttr(elementType, 1.0 / static_cast<double>(height * width)));
+  auto scale = getOrCreateConstant(
+    rewriter, rewriter.getInsertionBlock()->getParentOp(), scaleAttr, outputFragmentType);
+
+  auto batch = createSpatComputeBatch(
+    rewriter,
+    loc,
+    TypeRange {outputStorageType},
+    tilesPerRow,
+    ValueRange {zero, scale},
+    ValueRange {input},
+    [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
+      Value reduced = args.weights[0];
+      for (int64_t row = 0; row < height; ++row) {
+        Value fragment;
+        if (physicalInput) {
+          Value sourceSlot = args.lane;
+          if (row != 0)
+            sourceSlot = arith::AddIOp::create(
+              rewriter,
+              loc,
+              sourceSlot,
+              getOrCreateIndexConstant(
+                rewriter, rewriter.getInsertionBlock()->getParentOp(), row * tilesPerRow));
+          FailureOr<Value> physicalFragment = extractGraphBatchPhysicalFragment(
+            rewriter, loc, args.inputs.front(), sourceSlot, inputFragmentType);
+          if (failed(physicalFragment))
+            return failure();
+          fragment = *physicalFragment;
+        }
+        else {
+          Value nchw = tensor::ExtractSliceOp::create(
+            rewriter,
+            loc,
+            nchwInputFragmentType,
+            args.inputs.front(),
+            SmallVector<OpFoldResult> {rewriter.getIndexAttr(0),
+                                      rewriter.getIndexAttr(0),
+                                      rewriter.getIndexAttr(row),
+                                      rewriter.getIndexAttr(0)},
+            SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
+                                      rewriter.getIndexAttr(channels),
+                                      rewriter.getIndexAttr(1),
+                                      rewriter.getIndexAttr(width)},
+            getUnitStrides(rewriter, 4));
+          fragment = ONNXTransposeOp::create(
+            rewriter, loc, inputFragmentType, nchw, rewriter.getI64ArrayAttr({0, 2, 3, 1}));
+        }
+        for (int64_t column = 0; column < width; ++column) {
+          Value point = tensor::ExtractSliceOp::create(
+            rewriter,
+            loc,
+            outputFragmentType,
+            fragment,
+            SmallVector<OpFoldResult> {rewriter.getIndexAttr(0),
+                                      rewriter.getIndexAttr(0),
+                                      rewriter.getIndexAttr(column),
+                                      rewriter.getIndexAttr(0)},
+            SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
+                                      rewriter.getIndexAttr(1),
+                                      rewriter.getIndexAttr(1),
+                                      rewriter.getIndexAttr(inputFragmentType.getDimSize(3))},
+            getUnitStrides(rewriter, 4));
+          point = materializeTileTensor(rewriter, loc, point);
+          reduced = spatial::SpatVAddOp::create(
+            rewriter, loc, outputFragmentType, reduced, point);
+        }
+      }
+      reduced = spatial::SpatVMulOp::create(
+        rewriter, loc, outputFragmentType, reduced, args.weights[1]);
+      publishGraphBatchPhysicalFragment(
+        rewriter, loc, reduced, args.outputs.front(), args.lane);
+      return success();
+    });
+  if (failed(batch))
+    return failure();
+  return batch->getResult(0);
+}
+
 void populatePoolPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
   patterns.insert<PoolToSpatialCompute<ONNXMaxPoolSingleOutOp>>(ctx);
   patterns.insert<PoolToSpatialCompute<ONNXAveragePoolOp>>(ctx);

src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Concat.cpp

@@ -25,6 +25,17 @@ struct Concat : public OpConversionPattern<ONNXConcatOp> {
       return success();
     }
 
+    auto resultType = dyn_cast<RankedTensorType>(maxpoolOp.getResult().getType());
+    if (axis == 1 && resultType && resultType.hasStaticShape() && resultType.getRank() == 4
+        && llvm::all_of(inputs, [](Value input) {
+             auto type = dyn_cast<RankedTensorType>(input.getType());
+             return type && type.hasStaticShape() && type.getRank() == 4;
+           })) {
+      rewriter.replaceOpWithNewOp<spatial::SpatConcatPlanOp>(
+        maxpoolOp, resultType, inputs, rewriter.getI64IntegerAttr(axis), rewriter.getStringAttr("nchw"));
+      return success();
+    }
+
     auto computeOp = createSpatCompute(
       rewriter, maxpoolOp.getLoc(), TypeRange {maxpoolOp.getResult().getType()}, {}, inputs, [&](ValueRange args) {
         spatial::SpatYieldOp::create(

src/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp

@@ -27,6 +27,14 @@ lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
                            std::optional<mlir::Value> rowStripInput,
                            mlir::PatternRewriter& rewriter);
 
+mlir::LogicalResult
+canLowerGlobalAveragePoolPlanToRowStrip(spatial::SpatGlobalAveragePoolPlanOp planOp);
+
+mlir::FailureOr<mlir::Value>
+lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp,
+                                   std::optional<mlir::Value> rowStripInput,
+                                   mlir::PatternRewriter& rewriter);
+
 mlir::LogicalResult canLowerFlattenFromRowStrip(spatial::SpatGraphCompute flattenOp);
 
 mlir::LogicalResult lowerFlattenFromRowStrip(const RowStripPhysicalValue& input,

src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp

@@ -36,10 +36,16 @@ static bool usesSelectedRowStrip(Operation* user, llvm::DenseMap<Value, Selected
     return getSelectedLayout(layouts, reluPlan.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))
+    return getSelectedLayout(layouts, addPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
+  if (auto concatPlan = dyn_cast<spatial::SpatConcatPlanOp>(user))
+    return getSelectedLayout(layouts, concatPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
   if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(user))
     return getSelectedLayout(layouts, convPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
   if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(user))
     return getSelectedLayout(layouts, maxPoolPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
+  if (auto averagePoolPlan = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(user))
+    return getSelectedLayout(layouts, averagePoolPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
   if (auto flattenCompute = dyn_cast<spatial::SpatGraphCompute>(user))
     return succeeded(canLowerFlattenFromRowStrip(flattenCompute));
   return false;
@@ -62,10 +68,16 @@ static bool canConsumeRowStripAsUser(Operation* user) {
     auto resultType = dyn_cast<RankedTensorType>(biasAddPlan.getOutput().getType());
     return resultType && isSupportedBiasAddValue(biasAddPlan.getBias(), resultType);
   }
+  if (isa<spatial::SpatAddPlanOp>(user))
+    return true;
+  if (isa<spatial::SpatConcatPlanOp>(user))
+    return true;
   if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(user))
     return succeeded(canConsumeAndProduceRowStrip(convPlan));
   if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(user))
     return succeeded(canLowerMaxPoolPlanToRowStrip(maxPoolPlan));
+  if (auto averagePoolPlan = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(user))
+    return succeeded(canLowerGlobalAveragePoolPlanToRowStrip(averagePoolPlan));
   return false;
 }
 
@@ -118,11 +130,38 @@ static SelectedLayout chooseBiasAddLayout(spatial::SpatBiasAddPlanOp biasAddPlan
   return SelectedLayout::PixelMajorRowStrip;
 }
 
+static SelectedLayout chooseAddLayout(spatial::SpatAddPlanOp addPlan, llvm::DenseMap<Value, SelectedLayout>& layouts) {
+  if (getSelectedLayout(layouts, addPlan.getLhs()) != SelectedLayout::PixelMajorRowStrip
+      || getSelectedLayout(layouts, addPlan.getRhs()) != SelectedLayout::PixelMajorRowStrip)
+    return SelectedLayout::DenseNchw;
+  if (!allUsersCanHandleRowStrip(addPlan.getResult(), layouts))
+    return SelectedLayout::DenseNchw;
+  return SelectedLayout::PixelMajorRowStrip;
+}
+
+static SelectedLayout chooseConcatLayout(spatial::SpatConcatPlanOp concatPlan,
+                                         llvm::DenseMap<Value, SelectedLayout>& layouts) {
+  if (llvm::any_of(concatPlan.getInputs(), [&](Value input) {
+        return getSelectedLayout(layouts, input) != SelectedLayout::PixelMajorRowStrip;
+      }))
+    return SelectedLayout::DenseNchw;
+  if (!allUsersCanHandleRowStrip(concatPlan.getResult(), layouts))
+    return SelectedLayout::DenseNchw;
+  return SelectedLayout::PixelMajorRowStrip;
+}
+
 static SelectedLayout chooseMaxPoolLayout(spatial::SpatMaxPool2DPlanOp maxPoolPlan) {
   return succeeded(canLowerMaxPoolPlanToRowStrip(maxPoolPlan)) ? SelectedLayout::PixelMajorRowStrip
                                                                : SelectedLayout::DenseNchw;
 }
 
+static SelectedLayout chooseGlobalAveragePoolLayout(
+  spatial::SpatGlobalAveragePoolPlanOp averagePoolPlan) {
+  return succeeded(canLowerGlobalAveragePoolPlanToRowStrip(averagePoolPlan))
+           ? SelectedLayout::PixelMajorRowStrip
+           : SelectedLayout::DenseNchw;
+}
+
 static spatial::SpatBlueprintOp insertRowStripBlueprint(IRRewriter& rewriter, Value value) {
   auto outputType = cast<RankedTensorType>(value.getType());
   auto [offsets, sizes] = buildRowStripMetadata(outputType);
@@ -215,6 +254,22 @@ struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass,
           }
           continue;
         }
+        if (auto addPlan = dyn_cast<spatial::SpatAddPlanOp>(&op)) {
+          SelectedLayout selected = chooseAddLayout(addPlan, layouts);
+          if (layouts[addPlan.getResult()] != selected) {
+            layouts[addPlan.getResult()] = selected;
+            changed = true;
+          }
+          continue;
+        }
+        if (auto concatPlan = dyn_cast<spatial::SpatConcatPlanOp>(&op)) {
+          SelectedLayout selected = chooseConcatLayout(concatPlan, layouts);
+          if (layouts[concatPlan.getResult()] != selected) {
+            layouts[concatPlan.getResult()] = selected;
+            changed = true;
+          }
+          continue;
+        }
         if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(&op)) {
           SelectedLayout selected = chooseMaxPoolLayout(maxPoolPlan);
           if (layouts[maxPoolPlan.getResult()] != selected) {
@@ -223,6 +278,14 @@ struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass,
           }
           continue;
         }
+        if (auto averagePoolPlan = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(&op)) {
+          SelectedLayout selected = chooseGlobalAveragePoolLayout(averagePoolPlan);
+          if (layouts[averagePoolPlan.getResult()] != selected) {
+            layouts[averagePoolPlan.getResult()] = selected;
+            changed = true;
+          }
+          continue;
+        }
       }
     }
 
@@ -232,10 +295,16 @@ struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass,
         producedValue = convPlan.getResult();
       else if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(&op))
         producedValue = biasAddPlan.getResult();
+      else if (auto addPlan = dyn_cast<spatial::SpatAddPlanOp>(&op))
+        producedValue = addPlan.getResult();
+      else if (auto concatPlan = dyn_cast<spatial::SpatConcatPlanOp>(&op))
+        producedValue = concatPlan.getResult();
       else if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(&op))
         producedValue = reluPlan.getResult();
       else if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(&op))
         producedValue = maxPoolPlan.getResult();
+      else if (auto averagePoolPlan = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(&op))
+        producedValue = averagePoolPlan.getResult();
       else
         continue;
 

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp

@@ -264,7 +264,13 @@ LogicalResult raptor::SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func:
     auto outputType = cast<RankedTensorType>(vmmOp.getOutput().getType());
     ArrayRef<int64_t> outputShape = outputType.getShape();
     assert(isHVectorShape(outputShape) && "expected a horizontal vector output");
-    assert(outputShape[1] <= static_cast<int64_t>(crossbarSize) && "output width must fit in one crossbar");
+    auto weightType = cast<RankedTensorType>(vmmOp.getWeight().getType());
+    const int64_t xbarDim = static_cast<int64_t>(crossbarSize);
+    const int64_t paddedOutputWidth = ceilIntegerDivide(outputShape[1], xbarDim) * xbarDim;
+    assert(weightType.getRank() == 2 && weightType.getDimSize(1) == paddedOutputWidth
+           && "expected VMM weight width to match the padded output width");
+    assert(paddedOutputWidth / xbarDim <= static_cast<int64_t>(crossbarCountInCore)
+           && "output width must fit in one core");
 
     rewriter.setInsertionPoint(vmmOp);
     auto paddedInput = padHVectorInputToCrossbarSize(rewriter, vmmOp.getLoc(), vmmOp.getInput());
@@ -273,8 +279,8 @@ LogicalResult raptor::SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func:
       return WalkResult::interrupt();
     }
     auto paddedOutputType = RankedTensorType::get(
-      {outputShape[0], static_cast<int64_t>(crossbarSize)}, outputType.getElementType(), outputType.getEncoding());
-    Value paddedOutputBuffer = outputShape[1] == static_cast<int64_t>(crossbarSize)
+      {outputShape[0], paddedOutputWidth}, outputType.getElementType(), outputType.getEncoding());
+    Value paddedOutputBuffer = outputShape[1] == paddedOutputWidth
                                ? vmmOp.getOutputBuffer()
                                : createEmptyTensorFromShaped(rewriter, vmmOp.getLoc(), paddedOutputType).getResult();
     vmmOp.getInputMutable().assign(*paddedInput);
@@ -282,7 +288,7 @@ LogicalResult raptor::SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func:
 
     vmmOp.getOutput().setType(paddedOutputType);
 
-    if (outputShape[1] == static_cast<int64_t>(crossbarSize))
+    if (outputShape[1] == paddedOutputWidth)
       return WalkResult::advance();
 
     SmallVector<OpFoldResult> offsets = {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};