Merge branch 'TestRottoConDeadLock' of chef.heaplab.deib.polimi.it:nnicolosi/Raptor into TestRottoConDeadLock

.gitignore

@@ -11,4 +11,6 @@ build_*
 compile.sh
 pimcomp_utils/*
 
+*.zip
+**/*.zip
 **/__pycache__/

src/PIM/Common/IR/TensorSliceUtils.cpp

@@ -81,8 +81,17 @@ Value extractMixedSliceOrIdentity(OpBuilder &rewriter,
 
 Value insertMixedSlice(OpBuilder &builder, Location loc, Value source,
                        Value dest, const MixedSliceGeometry &geometry) {
+  SmallVector<OpFoldResult> sizes(geometry.sizes);
+  auto sourceType = dyn_cast<RankedTensorType>(source.getType());
+  auto destType = dyn_cast<RankedTensorType>(dest.getType());
+  if (sourceType && destType && sourceType.hasStaticShape()
+      && sourceType.getRank() == destType.getRank()) {
+    sizes.clear();
+    for (int64_t dimension : sourceType.getShape())
+      sizes.push_back(builder.getIndexAttr(dimension));
+  }
   return tensor::InsertSliceOp::create(builder, loc, source, dest,
-                                       geometry.offsets, geometry.sizes,
+                                       geometry.offsets, sizes,
                                        geometry.strides);
 }
 

src/PIM/Compiler/PimCodeGen.cpp

@@ -1248,7 +1248,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
     }
 
     for (auto [slot, fileName] : llvm::enumerate(weightFiles)) {
-      xbarsPerGroup.push_back(static_cast<int64_t>(slot));
+      xbarsPerGroup.push_back(1);
       std::string sourcePath = outputDirPath + "/weights/" + fileName;
       std::string targetPath = coreWeightsDirPath + "/crossbar_" + std::to_string(slot) + ".bin";
       sys::fs::remove(targetPath);

src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp

@@ -242,6 +242,37 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
         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);
+          return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
+        });
+        if (outputBlueprint == planOp.getResult().getUsers().end()) {
+          planOp.emitOpError("selected MaxPool plan requires a row-strip blueprint result");
+          signalPassFailure();
+          return;
+        }
+
+        FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
+        rewriter.setInsertionPoint(planOp);
+        FailureOr<Value> lowered = lowerSelectedMaxPool2DPlan(
+          planOp, succeeded(input) ? std::optional<Value> {input->storage} : std::nullopt, rewriter);
+        if (failed(lowered)) {
+          planOp.emitOpError("failed to lower selected row-strip Spatial MaxPool 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) {
@@ -444,6 +475,7 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
       } else if (isa<spatial::SpatConv2DPlanOp,
                      spatial::SpatBiasAddPlanOp,
                      spatial::SpatReluPlanOp,
+                     spatial::SpatMaxPool2DPlanOp,
                      spatial::SpatMaterializeLayoutOp>(op)
           || op->getDialect()->getNamespace() == "onnx") {
         op->emitOpError("operation must not remain after LowerSpatialPlans");

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp

@@ -48,10 +48,11 @@ static void populateEmptyFunction(func::FuncOp funcOp) {
   SmallVector<spatial::SpatConv2DPlanOp> convPlans(funcOp.getOps<spatial::SpatConv2DPlanOp>());
   SmallVector<spatial::SpatBiasAddPlanOp> biasAddPlans(funcOp.getOps<spatial::SpatBiasAddPlanOp>());
   SmallVector<spatial::SpatReluPlanOp> reluPlans(funcOp.getOps<spatial::SpatReluPlanOp>());
+  SmallVector<spatial::SpatMaxPool2DPlanOp> maxPoolPlans(funcOp.getOps<spatial::SpatMaxPool2DPlanOp>());
   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()
-      || !blueprints.empty() || !materializers.empty()) {
+      || !maxPoolPlans.empty() || !blueprints.empty() || !materializers.empty()) {
     return;
   }
 

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp

@@ -147,6 +147,7 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter
             spatial::SpatConv2DPlanOp,
             spatial::SpatBiasAddPlanOp,
             spatial::SpatReluPlanOp,
+            spatial::SpatMaxPool2DPlanOp,
             spatial::SpatBlueprintOp,
             spatial::SpatMaterializeLayoutOp>(&op)) {
       continue;

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

@@ -1834,6 +1834,37 @@ static Value createPaddedInputKTiledWeightConstant(DenseElementsAttr sourceAttr,
   return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), paddedAttr, paddedType);
 }
 
+static Value createPaddedOutputChannelTiledWeightConstant(DenseElementsAttr sourceAttr,
+                                                          const ConvLoweringState& state,
+                                                          int64_t paddedK,
+                                                          int64_t xbarDim,
+                                                          PatternRewriter& rewriter) {
+  const int64_t outputTileCount = ceilIntegerDivide(state.numChannelsOut, xbarDim);
+  auto paddedType =
+    RankedTensorType::get({outputTileCount, paddedK, xbarDim}, 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) {
+    const int64_t outputTile = outChannel / xbarDim;
+    const int64_t tileChannel = outChannel % xbarDim;
+    for (int64_t inChannel = 0; inChannel < state.numChannelsIn; ++inChannel) {
+      for (int64_t kernelH = 0; kernelH < state.wHeight; ++kernelH) {
+        for (int64_t kernelW = 0; kernelW < state.wWidth; ++kernelW) {
+          const int64_t sourceFlatIndex =
+            (((outChannel * state.numChannelsIn) + inChannel) * state.wHeight + kernelH) * state.wWidth + kernelW;
+          const int64_t patchIndex = ((inChannel * state.wHeight) + kernelH) * state.wWidth + kernelW;
+          const int64_t destinationFlatIndex =
+            ((outputTile * paddedK) + patchIndex) * xbarDim + tileChannel;
+          paddedValues[destinationFlatIndex] = sourceValues[sourceFlatIndex];
+        }
+      }
+    }
+  }
+  auto paddedAttr = DenseElementsAttr::get(paddedType, paddedValues);
+  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), paddedAttr, paddedType);
+}
+
 static FailureOr<Value> rewriteInputKTiledConv(const ConvLoweringState& state,
                                                ArrayRef<DistributedTensorStep> distributedConsumers,
                                                PatternRewriter& rewriter,
@@ -2548,10 +2579,16 @@ static Value createHorizontallyPaddedRowStripFragment(Value fragment,
                                                       const ConvLoweringState& state,
                                                       PatternRewriter& rewriter,
                                                       Location loc) {
-  auto paddedType = RankedTensorType::get({1, state.numChannelsIn, 1, state.xWidth + 2},
-                                          state.xType.getElementType(),
-                                          state.xType.getEncoding());
-  return createZeroPaddedTensor(fragment, paddedType, {0, 0, 0, 1}, {0, 0, 0, 1}, rewriter, loc);
+  auto paddedType = RankedTensorType::get(
+    {1, state.numChannelsIn, 1, state.xWidth + state.padWidthBegin + state.padWidthEnd},
+    state.xType.getElementType(),
+    state.xType.getEncoding());
+  return createZeroPaddedTensor(fragment,
+                                paddedType,
+                                {0, 0, 0, state.padWidthBegin},
+                                {0, 0, 0, state.padWidthEnd},
+                                rewriter,
+                                loc);
 }
 
 static Value createRowStripWindowSourceRowTable(const ConvLoweringState& state, PatternRewriter& rewriter) {
@@ -2561,7 +2598,8 @@ static Value createRowStripWindowSourceRowTable(const ConvLoweringState& state,
   values.reserve(tableType.getNumElements());
   for (int64_t outputRow = 0; outputRow < state.outHeight; ++outputRow) {
     for (int64_t kernelRow = 0; kernelRow < state.wHeight; ++kernelRow) {
-      int64_t sourceRow = outputRow + kernelRow - state.padHeightBegin;
+      int64_t sourceRow =
+        outputRow * state.strideHeight + kernelRow * state.dilationHeight - state.padHeightBegin;
       sourceRow = std::clamp(sourceRow, int64_t {0}, state.xHeight - 1);
       values.push_back(rewriter.getIndexAttr(sourceRow));
     }
@@ -2595,6 +2633,26 @@ static Value extractProjectedRowStripWindowRow(Value rowStripStorage,
   return extractRowStripFragment(rowStripStorage, state.xType, sourceRow, rewriter, loc);
 }
 
+static Value extractDenseConvWindowRow(Value denseInput,
+                                       Value sourceRowTable,
+                                       const ConvLoweringState& state,
+                                       Value outputHeight,
+                                       Value kernelRow,
+                                       PatternRewriter& rewriter,
+                                       Location loc) {
+  Value tableIndex = createRowStripWindowTableIndex(outputHeight, kernelRow, state, rewriter, loc);
+  Value sourceRow = tensor::ExtractOp::create(rewriter, loc, sourceRowTable, ValueRange {tableIndex}).getResult();
+  auto fragmentType = getRowStripFragmentType(state.xType);
+  SmallVector<OpFoldResult> offsets {
+    rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), sourceRow, rewriter.getIndexAttr(0)};
+  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
+                                   rewriter.getIndexAttr(state.numChannelsIn),
+                                   rewriter.getIndexAttr(1),
+                                   rewriter.getIndexAttr(state.xWidth)};
+  return tensor::ExtractSliceOp::create(
+    rewriter, loc, fragmentType, denseInput, offsets, sizes, getUnitStrides(rewriter, 4));
+}
+
 static FailureOr<Value> createRowStripWindowMaskTable(const ConvLoweringState& state, PatternRewriter& rewriter) {
   auto elementType = state.xType.getElementType();
   auto floatType = dyn_cast<FloatType>(elementType);
@@ -2611,7 +2669,8 @@ static FailureOr<Value> createRowStripWindowMaskTable(const ConvLoweringState& s
   values.reserve(tableType.getNumElements());
   for (int64_t outputRow = 0; outputRow < state.outHeight; ++outputRow) {
     for (int64_t kernelRow = 0; kernelRow < state.wHeight; ++kernelRow) {
-      int64_t sourceRow = outputRow + kernelRow - state.padHeightBegin;
+      int64_t sourceRow =
+        outputRow * state.strideHeight + kernelRow * state.dilationHeight - state.padHeightBegin;
       Attribute value = (sourceRow < 0 || sourceRow >= state.xHeight) ? zero : one;
       for (int64_t channel = 0; channel < state.numChannelsIn; ++channel)
         for (int64_t width = 0; width < state.xWidth; ++width)
@@ -2645,15 +2704,20 @@ static Value extractProjectedRowStripWindowMask(Value maskTable,
                                         getUnitStrides(rewriter, 4));
 }
 
-static FailureOr<Value> createNchwRowStripConvWindow(Value rowStripStorage,
-                                                     const ConvLoweringState& state,
-                                                     Value outputHeight,
-                                                     PatternRewriter& rewriter,
-                                                     Location loc) {
+static FailureOr<Value> createConvInputWindow(Value input,
+                                              const ConvLoweringState& state,
+                                              Value outputHeight,
+                                              PatternRewriter& rewriter,
+                                              Location loc) {
   auto fragmentType = getRowStripFragmentType(state.xType);
-  auto paddedWindowType = RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.xWidth + 2},
-                                                state.xType.getElementType(),
-                                                state.xType.getEncoding());
+  auto inputType = dyn_cast<RankedTensorType>(input.getType());
+  const bool denseInput = inputType == state.xType;
+  if (!denseInput && inputType != getRowStripStorageType(state.xType))
+    return failure();
+  auto paddedWindowType = RankedTensorType::get(
+    {1, state.numChannelsIn, state.wHeight, state.xWidth + state.padWidthBegin + state.padWidthEnd},
+    state.xType.getElementType(),
+    state.xType.getEncoding());
   Value sourceRowTable = createRowStripWindowSourceRowTable(state, rewriter);
   FailureOr<Value> maskTable = createRowStripWindowMaskTable(state, rewriter);
   if (failed(maskTable))
@@ -2664,8 +2728,10 @@ static FailureOr<Value> createNchwRowStripConvWindow(Value rowStripStorage,
   Value window = initWindow;
   for (int64_t kernelRowIndex = 0; kernelRowIndex < state.wHeight; ++kernelRowIndex) {
     Value kernelRow = getOrCreateIndexConstant(rewriter, anchorOp, kernelRowIndex);
-    Value sourceRow =
-      extractProjectedRowStripWindowRow(rowStripStorage, sourceRowTable, state, outputHeight, kernelRow, rewriter, loc);
+    Value sourceRow = denseInput
+                        ? extractDenseConvWindowRow(input, sourceRowTable, state, outputHeight, kernelRow, rewriter, loc)
+                        : extractProjectedRowStripWindowRow(
+                            input, sourceRowTable, state, outputHeight, kernelRow, rewriter, loc);
     Value mask = extractProjectedRowStripWindowMask(*maskTable, state, outputHeight, kernelRow, rewriter, loc);
     Value semanticRow = spatial::SpatVMulOp::create(rewriter, loc, fragmentType, sourceRow, mask).getResult();
     Value paddedRow = createHorizontallyPaddedRowStripFragment(semanticRow, state, rewriter, loc);
@@ -2680,7 +2746,9 @@ static FailureOr<Value> createNchwRowStripConvWindow(Value rowStripStorage,
                                            SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
                                                                       rewriter.getIndexAttr(state.numChannelsIn),
                                                                       rewriter.getIndexAttr(1),
-                                                                      rewriter.getIndexAttr(state.xWidth + 2)},
+                                                                      rewriter.getIndexAttr(
+                                                                        state.xWidth + state.padWidthBegin
+                                                                        + state.padWidthEnd)},
                                            getUnitStrides(rewriter, 4));
   }
   return window;
@@ -2698,12 +2766,13 @@ static FailureOr<Value> createNchwRowStripConvPatchRow(Value paddedWindow,
                                          state.xType.getEncoding());
   auto rowType = RankedTensorType::get({1, patchSize}, state.xType.getElementType(), state.xType.getEncoding());
   Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
+  Value inputWidthOffset = affineMulConst(rewriter, loc, outputWidth, state.strideWidth, anchorOp);
   Value patch = createConvInputPatch(paddedWindow,
                                      patchType,
                                      c0,
                                      c0,
                                      c0,
-                                     outputWidth,
+                                     inputWidthOffset,
                                      state.dilationHeight,
                                      state.dilationWidth,
                                      rewriter,
@@ -2713,6 +2782,57 @@ static FailureOr<Value> createNchwRowStripConvPatchRow(Value paddedWindow,
     .getResult();
 }
 
+static FailureOr<Value> createPaddedConvOutputTile(Value paddedPatchRow,
+                                                   Value tileWeights,
+                                                   int64_t numKSlices,
+                                                   int64_t xbarDim,
+                                                   PatternRewriter& rewriter,
+                                                   Location loc) {
+  auto elementType = cast<RankedTensorType>(paddedPatchRow.getType()).getElementType();
+  auto paddedRowType = RankedTensorType::get({1, xbarDim}, elementType);
+  auto weightElementType = cast<RankedTensorType>(tileWeights.getType()).getElementType();
+  auto paddedWeightTileType = RankedTensorType::get({xbarDim, xbarDim}, weightElementType);
+
+  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+  Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
+  Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
+  Value cNumKSlices = getOrCreateIndexConstant(rewriter, anchorOp, numKSlices);
+  Value cXbar = getOrCreateIndexConstant(rewriter, anchorOp, xbarDim);
+  auto createPiece = [&](Value kSlice, Location pieceLoc) -> Value {
+    Value kOffset = arith::MulIOp::create(rewriter, pieceLoc, kSlice, cXbar);
+    SmallVector<OpFoldResult> aOffsets {rewriter.getIndexAttr(0), kOffset};
+    SmallVector<OpFoldResult> aSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)};
+    Value aTile = extractStaticSliceOrIdentity(
+      rewriter, pieceLoc, paddedPatchRow, paddedRowType, aOffsets, aSizes, getUnitStrides(rewriter, 2));
+    SmallVector<OpFoldResult> bOffsets {kOffset, rewriter.getIndexAttr(0)};
+    SmallVector<OpFoldResult> bSizes {rewriter.getIndexAttr(xbarDim), rewriter.getIndexAttr(xbarDim)};
+    Value bTile = extractStaticSliceOrIdentity(
+      rewriter, pieceLoc, tileWeights, paddedWeightTileType, bOffsets, bSizes, getUnitStrides(rewriter, 2));
+    return spatial::SpatVMMOp::create(rewriter, pieceLoc, paddedRowType, bTile, aTile).getResult();
+  };
+
+  Value tileResult = createPiece(c0, loc);
+  if (numKSlices == 1)
+    return tileResult;
+
+  auto kLoop = buildNormalizedScfFor(
+    rewriter,
+    loc,
+    c1,
+    cNumKSlices,
+    c1,
+    ValueRange {tileResult},
+    [&](OpBuilder&, Location reduceLoc, Value kSlice, ValueRange reduceIterArgs, SmallVectorImpl<Value>& reduceYielded) {
+      Value piece = createPiece(kSlice, reduceLoc);
+      reduceYielded.push_back(
+        spatial::SpatVAddOp::create(rewriter, reduceLoc, paddedRowType, reduceIterArgs.front(), piece).getResult());
+      return success();
+    });
+  if (failed(kLoop))
+    return failure();
+  return kLoop->results.front();
+}
+
 static FailureOr<Value> createPaddedConvOutputRow(Value patchRow,
                                                   const ConvLoweringState& state,
                                                   Value paddedWeights,
@@ -2727,67 +2847,241 @@ static FailureOr<Value> createPaddedConvOutputRow(Value patchRow,
   auto rowType = RankedTensorType::get({1, state.numChannelsOut}, elementType);
   auto paddedRowType = RankedTensorType::get({1, xbarDim}, elementType);
   auto paddedPatchRowType = RankedTensorType::get({1, paddedK}, elementType);
-  auto paddedWeightTileType = RankedTensorType::get({xbarDim, xbarDim}, state.wType.getElementType());
+  auto tileWeightsType = RankedTensorType::get({paddedK, xbarDim}, state.wType.getElementType());
+  const int64_t outputTileCount = ceilIntegerDivide(state.numChannelsOut, xbarDim);
 
   Value paddedPatchRow = patchRow;
   if (patchSize != paddedK)
     paddedPatchRow = createZeroPaddedTensor(
       paddedPatchRow, paddedPatchRowType, {0, 0}, {0, paddedK - patchSize}, rewriter, loc);
 
-  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
-  Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
-  Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
-  Value cNumKSlices = getOrCreateIndexConstant(rewriter, anchorOp, numKSlices);
-  Value cXbar = getOrCreateIndexConstant(rewriter, anchorOp, xbarDim);
-  auto createPiece = [&](Value kSlice, Location pieceLoc) -> Value {
-    Value kOffset = arith::MulIOp::create(rewriter, pieceLoc, kSlice, cXbar);
-    SmallVector<OpFoldResult> aOffsets {rewriter.getIndexAttr(0), kOffset};
-    SmallVector<OpFoldResult> aSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)};
-    Value aTile = extractStaticSliceOrIdentity(
-      rewriter, pieceLoc, paddedPatchRow, paddedRowType, aOffsets, aSizes, getUnitStrides(rewriter, 2));
-    SmallVector<OpFoldResult> bOffsets {kOffset, rewriter.getIndexAttr(0)};
-    SmallVector<OpFoldResult> bSizes {rewriter.getIndexAttr(xbarDim), rewriter.getIndexAttr(xbarDim)};
-    Value bTile = extractStaticSliceOrIdentity(
-      rewriter, pieceLoc, paddedWeights, paddedWeightTileType, bOffsets, bSizes, getUnitStrides(rewriter, 2));
-    return spatial::SpatVMMOp::create(rewriter, pieceLoc, paddedRowType, bTile, aTile).getResult();
+  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));
   };
 
-  Value rowResult = createPiece(c0, loc);
-  if (numKSlices > 1) {
-    auto kLoop = buildNormalizedScfFor(
-      rewriter,
-      loc,
-      c1,
-      cNumKSlices,
-      c1,
-      ValueRange {rowResult},
-      [&](OpBuilder&, Location reduceLoc, Value kSlice, ValueRange reduceIterArgs, SmallVectorImpl<Value>& reduceYielded) {
-        Value piece = createPiece(kSlice, reduceLoc);
-        reduceYielded.push_back(
-          spatial::SpatVAddOp::create(rewriter, reduceLoc, paddedRowType, reduceIterArgs.front(), piece).getResult());
-        return success();
-      });
-    if (failed(kLoop))
+  if (outputTileCount == 1) {
+    FailureOr<Value> rowResult = createPaddedConvOutputTile(
+      paddedPatchRow, getTileWeights(0), numKSlices, xbarDim, rewriter, loc);
+    if (failed(rowResult))
       return failure();
-    rowResult = kLoop->results.front();
+    if (paddedBias)
+      rowResult = spatial::SpatVAddOp::create(rewriter, loc, paddedRowType, *rowResult, paddedBias).getResult();
+    if (state.numChannelsOut == xbarDim)
+      return *rowResult;
+
+    SmallVector<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
+    SmallVector<OpFoldResult> outputSizes {
+      rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut)};
+    return tensor::ExtractSliceOp::create(
+             rewriter, loc, rowType, *rowResult, outputOffsets, outputSizes, getUnitStrides(rewriter, 2))
+      .getResult();
   }
 
-  if (paddedBias)
-    rowResult = spatial::SpatVAddOp::create(rewriter, loc, paddedRowType, rowResult, paddedBias).getResult();
-  if (state.numChannelsOut == xbarDim)
-    return rowResult;
+  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 = createPaddedConvOutputTile(
+      paddedPatchRow, getTileWeights(outputTile), numKSlices, 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));
+  }
 
   SmallVector<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
   SmallVector<OpFoldResult> outputSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut)};
   return tensor::ExtractSliceOp::create(
-           rewriter, loc, rowType, rowResult, outputOffsets, outputSizes, getUnitStrides(rewriter, 2))
+           rewriter, loc, rowType, paddedOutput, outputOffsets, outputSizes, getUnitStrides(rewriter, 2))
     .getResult();
 }
 
+static bool rowStripOutputFitsOneCore(const ConvGeometry& geometry) {
+  const int64_t inputTileCount = ceilIntegerDivide(geometry.k, geometry.xbarSize);
+  const int64_t outputTileCount = ceilIntegerDivide(geometry.c, geometry.xbarSize);
+  return inputTileCount * outputTileCount <= static_cast<int64_t>(crossbarCountInCore.getValue());
+}
+
+static bool rowStripOutputTileFitsOneCore(const ConvGeometry& geometry) {
+  return ceilIntegerDivide(geometry.k, geometry.xbarSize)
+         <= static_cast<int64_t>(crossbarCountInCore.getValue());
+}
+
+static FailureOr<Value> createOutputChannelTiledRowStripConvOutput(const ConvLoweringState& state,
+                                                                   Value paddedWeights,
+                                                                   int64_t paddedK,
+                                                                   int64_t numKSlices,
+                                                                   int64_t xbarDim,
+                                                                   PatternRewriter& rewriter,
+                                                                   Location loc) {
+  const int64_t outputTileCount = ceilIntegerDivide(state.numChannelsOut, xbarDim);
+  const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
+  auto elementType = state.outType.getElementType();
+  auto paddedPatchRowType = RankedTensorType::get({1, paddedK}, elementType);
+  auto tileWeightsType = RankedTensorType::get({paddedK, xbarDim}, state.wType.getElementType());
+  SmallVector<Value> outputTiles;
+  outputTiles.reserve(outputTileCount);
+
+  for (int64_t outputTile = 0; outputTile < outputTileCount; ++outputTile) {
+    const int64_t channelOffset = outputTile * xbarDim;
+    const int64_t tileChannels = std::min(xbarDim, state.numChannelsOut - channelOffset);
+    auto tileRowType = RankedTensorType::get({1, tileChannels}, elementType);
+    auto tilePixelType = RankedTensorType::get({1, tileChannels, 1, 1}, elementType);
+    auto tileFragmentType = RankedTensorType::get({1, tileChannels, 1, state.outWidth}, elementType);
+    auto tileStorageType = spatial::getGraphBatchPhysicalResultType(state.outHeight, tileFragmentType);
+    SmallVector<OpFoldResult> weightOffsets {
+      rewriter.getIndexAttr(outputTile), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
+    SmallVector<OpFoldResult> weightSizes {
+      rewriter.getIndexAttr(1), rewriter.getIndexAttr(paddedK), rewriter.getIndexAttr(xbarDim)};
+    Value tileWeights = extractStaticSliceOrIdentity(
+      rewriter, loc, paddedWeights, tileWeightsType, weightOffsets, weightSizes, getUnitStrides(rewriter, 3));
+
+    auto tileBatch = createSpatComputeBatch(
+      rewriter,
+      loc,
+      TypeRange {tileStorageType},
+      state.outHeight,
+      ValueRange {tileWeights},
+      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, tileFragmentType.getShape(), elementType);
+        auto widthLoop = buildNormalizedScfFor(
+          rewriter,
+          loc,
+          c0,
+          cOutWidth,
+          c1,
+          ValueRange {fragmentInit},
+          [&](OpBuilder&,
+              Location widthLoc,
+              Value widthIndex,
+              ValueRange widthIterArgs,
+              SmallVectorImpl<Value>& widthYielded) {
+            FailureOr<Value> patchRow =
+              createNchwRowStripConvPatchRow(*inputWindow, state, widthIndex, rewriter, widthLoc);
+            if (failed(patchRow))
+              return failure();
+            Value paddedPatchRow = *patchRow;
+            if (patchSize != paddedK)
+              paddedPatchRow = createZeroPaddedTensor(
+                paddedPatchRow, paddedPatchRowType, {0, 0}, {0, paddedK - patchSize}, rewriter, widthLoc);
+            FailureOr<Value> paddedOutputRow = createPaddedConvOutputTile(
+              paddedPatchRow, args.weights.front(), numKSlices, xbarDim, rewriter, widthLoc);
+            if (failed(paddedOutputRow))
+              return failure();
+            Value outputRow = *paddedOutputRow;
+            if (tileChannels != xbarDim) {
+              SmallVector<OpFoldResult> rowOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
+              SmallVector<OpFoldResult> rowSizes {
+                rewriter.getIndexAttr(1), rewriter.getIndexAttr(tileChannels)};
+              outputRow = tensor::ExtractSliceOp::create(rewriter,
+                                                         widthLoc,
+                                                         tileRowType,
+                                                         outputRow,
+                                                         rowOffsets,
+                                                         rowSizes,
+                                                         getUnitStrides(rewriter, 2));
+            }
+            Value outputPixel = tensor::ExpandShapeOp::create(
+              rewriter, widthLoc, tilePixelType, outputRow, SmallVector<ReassociationIndices> {{0}, {1, 2, 3}});
+            SmallVector<OpFoldResult> rowOffsets {
+              rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex};
+            SmallVector<OpFoldResult> rowSizes {rewriter.getIndexAttr(1),
+                                               rewriter.getIndexAttr(tileChannels),
+                                               rewriter.getIndexAttr(1),
+                                               rewriter.getIndexAttr(1)};
+            Value nextFragment = tensor::InsertSliceOp::create(rewriter,
+                                                                widthLoc,
+                                                                outputPixel,
+                                                                widthIterArgs.front(),
+                                                                rowOffsets,
+                                                                rowSizes,
+                                                                getUnitStrides(rewriter, 4));
+            widthYielded.push_back(nextFragment);
+            return success();
+          });
+        if (failed(widthLoop))
+          return failure();
+        publishGraphBatchPhysicalFragment(
+          rewriter, loc, widthLoop->results.front(), args.outputs.front(), args.lane);
+        return success();
+      });
+    if (failed(tileBatch))
+      return failure();
+    outputTiles.push_back(tileBatch->getResult(0));
+  }
+
+  auto fragmentType = getRowStripFragmentType(state.outType);
+  auto outputStorageType = getRowStripStorageType(state.outType);
+  auto assemblyBatch = createSpatComputeBatch(rewriter,
+                                              loc,
+                                              TypeRange {outputStorageType},
+                                              state.outHeight,
+                                              {},
+                                              ValueRange(outputTiles),
+                                              [&](detail::SpatComputeBatchBodyArgs args) {
+                                                Value fragment = tensor::EmptyOp::create(
+                                                  rewriter, loc, fragmentType.getShape(), elementType);
+                                                for (int64_t outputTile = 0; outputTile < outputTileCount; ++outputTile) {
+                                                  const int64_t channelOffset = outputTile * xbarDim;
+                                                  const int64_t tileChannels =
+                                                    std::min(xbarDim, state.numChannelsOut - channelOffset);
+                                                  auto tileFragmentType = RankedTensorType::get(
+                                                    {1, tileChannels, 1, state.outWidth}, elementType);
+                                                  FailureOr<Value> tileFragment = extractGraphBatchPhysicalFragment(
+                                                    rewriter, loc, args.inputs[outputTile], args.lane, tileFragmentType);
+                                                  if (failed(tileFragment))
+                                                    return failure();
+                                                  SmallVector<OpFoldResult> offsets {rewriter.getIndexAttr(0),
+                                                                                   rewriter.getIndexAttr(channelOffset),
+                                                                                   rewriter.getIndexAttr(0),
+                                                                                   rewriter.getIndexAttr(0)};
+                                                  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
+                                                                                 rewriter.getIndexAttr(tileChannels),
+                                                                                 rewriter.getIndexAttr(1),
+                                                                                 rewriter.getIndexAttr(state.outWidth)};
+                                                  fragment = tensor::InsertSliceOp::create(rewriter,
+                                                                                          loc,
+                                                                                          *tileFragment,
+                                                                                          fragment,
+                                                                                          offsets,
+                                                                                          sizes,
+                                                                                          getUnitStrides(rewriter, 4));
+                                                }
+                                                insertRowStripFragment(
+                                                  fragment, args.outputs.front(), state.outType, args.lane, rewriter, loc);
+                                                return success();
+                                              });
+  if (failed(assemblyBatch))
+    return failure();
+  Value output = assemblyBatch->getResult(0);
+  if (state.hasBias)
+    return applyRowStripBiasAdd(output, state.outType, state.b, rewriter, loc);
+  return output;
+}
+
 static FailureOr<Value>
 createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRewriter& rewriter, Location loc) {
   ConvGeometry geometry = buildConvGeometry(state);
-  if (state.group != 1 || state.batchSize != 1 || geometry.c > geometry.xbarSize)
+  if (state.group != 1 || state.batchSize != 1 || !rowStripOutputTileFitsOneCore(geometry))
     return failure();
 
   auto weightDenseAttr = getHostConstDenseElementsAttr(state.w);
@@ -2803,12 +3097,17 @@ createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRe
   auto elementType = state.outType.getElementType();
   auto fragmentType = getRowStripFragmentType(state.outType);
   auto outputPixelType = RankedTensorType::get({1, state.numChannelsOut, 1, 1}, elementType);
-  auto patchType = RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.wWidth}, state.xType.getElementType());
-  auto patchRowType = RankedTensorType::get({1, patchSize}, state.xType.getElementType());
   auto outputStorageType = getRowStripStorageType(state.outType);
 
-  PreparedConvInput preparedInput = standard::prepareInputForIm2Col(state, rewriter, loc);
-  Value paddedWeights = standard::createPaddedInputKTiledWeightConstant(weightDenseAttr, state, paddedK, xbarDim, rewriter);
+  Value paddedWeights = state.numChannelsOut <= xbarDim
+                          ? standard::createPaddedInputKTiledWeightConstant(
+                              weightDenseAttr, state, paddedK, xbarDim, rewriter)
+                          : standard::createPaddedOutputChannelTiledWeightConstant(
+                              weightDenseAttr, state, paddedK, xbarDim, rewriter);
+  if (!rowStripOutputFitsOneCore(geometry))
+    return createOutputChannelTiledRowStripConvOutput(
+      state, paddedWeights, paddedK, numKSlices, xbarDim, rewriter, loc);
+
   FailureOr<Value> paddedBias = failure();
   if (state.hasBias)
     paddedBias = createPaddedBiasRowConstant(state, xbarDim, rewriter);
@@ -2821,13 +3120,16 @@ createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRe
     TypeRange {outputStorageType},
     state.outHeight,
     ValueRange {paddedWeights},
-    state.hasBias ? ValueRange {preparedInput.value, *paddedBias} : ValueRange {preparedInput.value},
+    state.hasBias ? ValueRange {state.x, *paddedBias} : 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);
-      Value inputHeightOffset = affineMulConst(rewriter, loc, args.lane, state.strideHeight, anchorOp);
+      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);
       auto widthLoop = buildNormalizedScfFor(
         rewriter,
@@ -2837,20 +3139,11 @@ createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRe
         c1,
         ValueRange {fragmentInit},
         [&](OpBuilder&, Location widthLoc, Value widthIndex, ValueRange widthIterArgs, SmallVectorImpl<Value>& widthYielded) {
-          Value inputWidthOffset = affineMulConst(rewriter, widthLoc, widthIndex, state.strideWidth, anchorOp);
-          Value patch = createConvInputPatch(args.inputs.front(),
-                                             patchType,
-                                             c0,
-                                             c0,
-                                             inputHeightOffset,
-                                             inputWidthOffset,
-                                             state.dilationHeight,
-                                             state.dilationWidth,
-                                             rewriter,
-                                             widthLoc);
-          Value patchRow = tensor::CollapseShapeOp::create(
-            rewriter, widthLoc, patchRowType, patch, SmallVector<ReassociationIndices> {{0}, {1, 2, 3}});
-          FailureOr<Value> outputRow = createPaddedConvOutputRow(patchRow,
+          FailureOr<Value> patchRow =
+            createNchwRowStripConvPatchRow(*inputWindow, state, widthIndex, rewriter, widthLoc);
+          if (failed(patchRow))
+            return failure();
+          FailureOr<Value> outputRow = createPaddedConvOutputRow(*patchRow,
                                                                  state,
                                                                  args.weights.front(),
                                                                  state.hasBias ? args.inputs[1] : Value(),
@@ -2931,7 +3224,7 @@ static FailureOr<Value> createConvOutputFromNchwRowStripFragments(Value rowStrip
       Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
       Value cOutWidth = getOrCreateIndexConstant(rewriter, anchorOp, state.outWidth);
       auto fragmentType = getRowStripFragmentType(state.outType);
-      FailureOr<Value> inputWindow = createNchwRowStripConvWindow(args.inputs.front(), state, args.lane, rewriter, loc);
+      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);
@@ -3818,7 +4111,7 @@ LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp) {
   analysis.barrierKind = DistributedConvBarrierKind::UnsupportedConsumer;
   analysis.barrierDetail = "selected row-strip layout";
   ConvGeometry geometry = buildConvGeometry(*state);
-  if (geometry.c > geometry.xbarSize)
+  if (!rowStripOutputTileFitsOneCore(geometry))
     return failure();
   ConvLoweringDecision decision = chooseConvLoweringStrategy(geometry, *requestedStrategy, analysis);
   if (decision.strategy == PimConvLoweringDepthwise && !depthwise::canUseStructuredRewrite(*state)

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

@@ -16,6 +16,8 @@
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -24,7 +26,7 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace {
 
-static Value materializeTileTensor(ConversionPatternRewriter& rewriter, Location loc, Value tile) {
+static Value materializeTileTensor(PatternRewriter& rewriter, Location loc, Value tile) {
   auto tileType = cast<RankedTensorType>(tile.getType());
   Value empty = tensor::EmptyOp::create(rewriter, loc, tileType.getShape(), tileType.getElementType());
   return insertStaticSlice(rewriter, loc, tile, empty, getZeroOffsets(rewriter, tileType.getRank()));
@@ -228,6 +230,23 @@ struct PoolToSpatialComputeBase : public OpConversionPattern<PoolOp> {
       }
     }
 
+    if constexpr (std::is_same_v<PoolOp, ONNXMaxPoolSingleOutOp>) {
+      if (batchSize == 1) {
+        auto plan = spatial::SpatMaxPool2DPlanOp::create(
+          rewriter,
+          loc,
+          outType,
+          x,
+          rewriter.getDenseI64ArrayAttr({kernelHeight, kernelWidth}),
+          rewriter.getDenseI64ArrayAttr({padTop, padLeft, padBottom, padRight}),
+          rewriter.getDenseI64ArrayAttr({strideHeight, strideWidth}),
+          rewriter.getDenseI64ArrayAttr({dilationHeight, dilationWidth}),
+          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;
     const int64_t outputPatchCount = batchSize * outputHeight * outputWidth;
@@ -396,6 +415,220 @@ struct PoolToSpatialCompute<ONNXAveragePoolOp>
 
 } // namespace
 
+LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp 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))
+    return failure();
+  if (llvm::any_of(planOp.getKernelShape(), [](int64_t value) { return value <= 0; })
+      || llvm::any_of(planOp.getStrides(), [](int64_t value) { return value <= 0; })
+      || llvm::any_of(planOp.getDilations(), [](int64_t value) { return value <= 0; }))
+    return failure();
+  return success();
+}
+
+static Value createClampedPoolIndexTable(PatternRewriter& rewriter,
+                                         Operation* anchorOp,
+                                         int64_t outputSize,
+                                         int64_t kernelSize,
+                                         int64_t stride,
+                                         int64_t dilation,
+                                         int64_t padBegin,
+                                         int64_t inputSize) {
+  auto tableType = RankedTensorType::get({outputSize * kernelSize}, rewriter.getIndexType());
+  SmallVector<Attribute> values;
+  values.reserve(tableType.getNumElements());
+  for (int64_t output = 0; output < outputSize; ++output)
+    for (int64_t kernel = 0; kernel < kernelSize; ++kernel)
+      values.push_back(rewriter.getIndexAttr(
+        std::clamp(output * stride + kernel * dilation - padBegin, int64_t {0}, inputSize - 1)));
+  return getOrCreateConstant(rewriter, anchorOp, DenseElementsAttr::get(tableType, values), tableType);
+}
+
+static Value extractPoolIndex(PatternRewriter& rewriter,
+                              Location loc,
+                              Operation* anchorOp,
+                              Value table,
+                              Value outputIndex,
+                              int64_t kernelIndex,
+                              int64_t kernelSize) {
+  Value tableIndex = arith::MulIOp::create(
+    rewriter, loc, outputIndex, getOrCreateIndexConstant(rewriter, anchorOp, kernelSize));
+  if (kernelIndex != 0)
+    tableIndex = arith::AddIOp::create(
+      rewriter, loc, tableIndex, getOrCreateIndexConstant(rewriter, anchorOp, kernelIndex));
+  return tensor::ExtractOp::create(rewriter, loc, table, tableIndex);
+}
+
+FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
+                                            std::optional<Value> rowStripInput,
+                                            PatternRewriter& rewriter) {
+  if (failed(canLowerMaxPoolPlanToRowStrip(planOp)))
+    return failure();
+
+  Location loc = planOp.getLoc();
+  auto inputType = cast<RankedTensorType>(planOp.getInput().getType());
+  auto outputType = cast<RankedTensorType>(planOp.getOutput().getType());
+  const int64_t channels = inputType.getDimSize(1);
+  const int64_t inputHeight = inputType.getDimSize(2);
+  const int64_t inputWidth = inputType.getDimSize(3);
+  const int64_t outputHeight = outputType.getDimSize(2);
+  const int64_t outputWidth = outputType.getDimSize(3);
+  const int64_t kernelHeight = planOp.getKernelShape()[0];
+  const int64_t kernelWidth = planOp.getKernelShape()[1];
+  Value input = rowStripInput.value_or(planOp.getInput());
+  auto actualInputType = dyn_cast<RankedTensorType>(input.getType());
+  const bool physicalInput = actualInputType == getRowStripStorageType(inputType);
+  if (!physicalInput && actualInputType != inputType)
+    return failure();
+
+  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+  Value rowTable = createClampedPoolIndexTable(rewriter,
+                                               anchorOp,
+                                               outputHeight,
+                                               kernelHeight,
+                                               planOp.getStrides()[0],
+                                               planOp.getDilations()[0],
+                                               planOp.getPads()[0],
+                                               inputHeight);
+  Value columnTable = createClampedPoolIndexTable(rewriter,
+                                                  anchorOp,
+                                                  outputWidth,
+                                                  kernelWidth,
+                                                  planOp.getStrides()[1],
+                                                  planOp.getDilations()[1],
+                                                  planOp.getPads()[1],
+                                                  inputWidth);
+  auto inputFragmentType = getRowStripFragmentType(inputType);
+  auto outputFragmentType = getRowStripFragmentType(outputType);
+  auto outputStorageType = getRowStripStorageType(outputType);
+  auto tileType = RankedTensorType::get({1, channels, 1, 1}, outputType.getElementType());
+  auto batch = createSpatComputeBatch(
+    rewriter,
+    loc,
+    TypeRange {outputStorageType},
+    outputHeight,
+    {},
+    ValueRange {input},
+    [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
+      SmallVector<Value> inputRows;
+      inputRows.reserve(kernelHeight);
+      for (int64_t kernelRow = 0; kernelRow < kernelHeight; ++kernelRow) {
+        Value sourceRow =
+          extractPoolIndex(rewriter, loc, anchorOp, rowTable, args.lane, kernelRow, kernelHeight);
+        if (physicalInput) {
+          inputRows.push_back(
+            extractRowStripFragment(args.inputs.front(), inputType, sourceRow, rewriter, loc));
+        }
+        else {
+          SmallVector<OpFoldResult> offsets {
+            rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), sourceRow, rewriter.getIndexAttr(0)};
+          inputRows.push_back(tensor::ExtractSliceOp::create(rewriter,
+                                                            loc,
+                                                            inputFragmentType,
+                                                            args.inputs.front(),
+                                                            offsets,
+                                                            SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
+                                                                                       rewriter.getIndexAttr(channels),
+                                                                                       rewriter.getIndexAttr(1),
+                                                                                       rewriter.getIndexAttr(inputWidth)},
+                                                            getUnitStrides(rewriter, 4)));
+        }
+      }
+
+      auto windowType = RankedTensorType::get(
+        {1, channels, kernelHeight, inputWidth}, inputType.getElementType(), inputType.getEncoding());
+      Value window = tensor::EmptyOp::create(
+        rewriter, loc, windowType.getShape(), windowType.getElementType());
+      for (int64_t kernelRow = 0; kernelRow < kernelHeight; ++kernelRow) {
+        SmallVector<OpFoldResult> offsets {rewriter.getIndexAttr(0),
+                                           rewriter.getIndexAttr(0),
+                                           rewriter.getIndexAttr(kernelRow),
+                                           rewriter.getIndexAttr(0)};
+        window = tensor::InsertSliceOp::create(rewriter,
+                                               loc,
+                                               inputRows[kernelRow],
+                                               window,
+                                               offsets,
+                                               SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
+                                                                          rewriter.getIndexAttr(channels),
+                                                                          rewriter.getIndexAttr(1),
+                                                                          rewriter.getIndexAttr(inputWidth)},
+                                               getUnitStrides(rewriter, 4));
+      }
+
+      Value outputInit = tensor::EmptyOp::create(
+        rewriter, loc, outputFragmentType.getShape(), outputFragmentType.getElementType());
+      Operation* bodyAnchor = rewriter.getInsertionBlock()->getParentOp();
+      Value c0 = getOrCreateIndexConstant(rewriter, bodyAnchor, 0);
+      Value c1 = getOrCreateIndexConstant(rewriter, bodyAnchor, 1);
+      Value cOutputWidth = getOrCreateIndexConstant(rewriter, bodyAnchor, outputWidth);
+      auto outputLoop = buildNormalizedScfFor(
+        rewriter,
+        loc,
+        c0,
+        cOutputWidth,
+        c1,
+        ValueRange {outputInit},
+        [&](OpBuilder&, Location nestedLoc, Value outputColumn, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
+          Value reduced;
+          for (int64_t kernelRow = 0; kernelRow < kernelHeight; ++kernelRow) {
+            for (int64_t kernelColumn = 0; kernelColumn < kernelWidth; ++kernelColumn) {
+              Value sourceColumn = extractPoolIndex(rewriter,
+                                                    nestedLoc,
+                                                    bodyAnchor,
+                                                    columnTable,
+                                                    outputColumn,
+                                                    kernelColumn,
+                                                    kernelWidth);
+              SmallVector<OpFoldResult> offsets {
+                rewriter.getIndexAttr(0),
+                rewriter.getIndexAttr(0),
+                rewriter.getIndexAttr(kernelRow),
+                sourceColumn};
+              Value point = tensor::ExtractSliceOp::create(rewriter,
+                                                           nestedLoc,
+                                                           tileType,
+                                                           window,
+                                                           offsets,
+                                                           SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
+                                                                                      rewriter.getIndexAttr(channels),
+                                                                                      rewriter.getIndexAttr(1),
+                                                                                      rewriter.getIndexAttr(1)},
+                                                           getUnitStrides(rewriter, 4));
+              reduced = reduced ? spatial::SpatVMaxOp::create(rewriter, nestedLoc, tileType, reduced, point).getResult()
+                                : materializeTileTensor(rewriter, nestedLoc, point);
+            }
+          }
+          SmallVector<OpFoldResult> outputOffsets {
+            rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), outputColumn};
+          Value updated = tensor::InsertSliceOp::create(rewriter,
+                                                        nestedLoc,
+                                                        reduced,
+                                                        iterArgs.front(),
+                                                        outputOffsets,
+                                                        SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
+                                                                                   rewriter.getIndexAttr(channels),
+                                                                                   rewriter.getIndexAttr(1),
+                                                                                   rewriter.getIndexAttr(1)},
+                                                        getUnitStrides(rewriter, 4));
+          yielded.push_back(updated);
+          return success();
+        });
+      if (failed(outputLoop))
+        return failure();
+      insertRowStripFragment(
+        outputLoop->results.front(), args.outputs.front(), outputType, args.lane, rewriter, loc);
+      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/PlanLowering.hpp

@@ -18,4 +18,11 @@ lowerSelectedConv2DPlan(spatial::SpatConv2DPlanOp planOp,
 mlir::LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp);
 mlir::LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp);
 
+mlir::LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp);
+
+mlir::FailureOr<mlir::Value>
+lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
+                           std::optional<mlir::Value> rowStripInput,
+                           mlir::PatternRewriter& rewriter);
+
 } // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp

@@ -38,6 +38,8 @@ static bool usesSelectedRowStrip(Operation* user, llvm::DenseMap<Value, Selected
     return getSelectedLayout(layouts, biasAddPlan.getResult()) == SelectedLayout::NchwRowStrip;
   if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(user))
     return getSelectedLayout(layouts, convPlan.getResult()) == SelectedLayout::NchwRowStrip;
+  if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(user))
+    return getSelectedLayout(layouts, maxPoolPlan.getResult()) == SelectedLayout::NchwRowStrip;
   return false;
 }
 
@@ -60,6 +62,8 @@ static bool canConsumeRowStripAsUser(Operation* user) {
   }
   if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(user))
     return succeeded(canConsumeAndProduceRowStrip(convPlan));
+  if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(user))
+    return succeeded(canLowerMaxPoolPlanToRowStrip(maxPoolPlan));
   return false;
 }
 
@@ -70,7 +74,6 @@ static bool hasRowStripConsumer(Value value) {
   return false;
 }
 
-
 static bool canSelectConvRowStrip(spatial::SpatConv2DPlanOp convPlan,
                                   llvm::DenseMap<Value, SelectedLayout>& layouts) {
   SelectedLayout inputLayout = getSelectedLayout(layouts, convPlan.getInput());
@@ -83,9 +86,6 @@ static SelectedLayout chooseConvLayout(spatial::SpatConv2DPlanOp convPlan,
                                        llvm::DenseMap<Value, SelectedLayout>& layouts) {
   if (!canSelectConvRowStrip(convPlan, layouts))
     return SelectedLayout::DenseNchw;
-  if (getSelectedLayout(layouts, convPlan.getInput()) != SelectedLayout::NchwRowStrip
-      && !hasRowStripConsumer(convPlan.getResult()))
-    return SelectedLayout::DenseNchw;
   if (!allUsersCanHandleRowStrip(convPlan.getResult(), layouts))
     return SelectedLayout::DenseNchw;
   return SelectedLayout::NchwRowStrip;
@@ -116,6 +116,11 @@ static SelectedLayout chooseBiasAddLayout(spatial::SpatBiasAddPlanOp biasAddPlan
   return SelectedLayout::NchwRowStrip;
 }
 
+static SelectedLayout chooseMaxPoolLayout(spatial::SpatMaxPool2DPlanOp maxPoolPlan) {
+  return succeeded(canLowerMaxPoolPlanToRowStrip(maxPoolPlan)) ? SelectedLayout::NchwRowStrip
+                                                               : SelectedLayout::DenseNchw;
+}
+
 static spatial::SpatBlueprintOp insertRowStripBlueprint(IRRewriter& rewriter, Value value) {
   auto outputType = cast<RankedTensorType>(value.getType());
   auto [offsets, sizes] = buildRowStripMetadata(outputType);
@@ -208,6 +213,14 @@ struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass,
           }
           continue;
         }
+        if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(&op)) {
+          SelectedLayout selected = chooseMaxPoolLayout(maxPoolPlan);
+          if (layouts[maxPoolPlan.getResult()] != selected) {
+            layouts[maxPoolPlan.getResult()] = selected;
+            changed = true;
+          }
+          continue;
+        }
       }
     }
 
@@ -219,6 +232,8 @@ struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass,
         producedValue = biasAddPlan.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
         continue;
 

src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp

@@ -29,6 +29,11 @@ static bool isUsedOnlyAsExplicitHostOperand(Value value) {
   });
 }
 
+static bool isUsedOnlyByExtractSlices(Value value) {
+  return !value.use_empty()
+         && llvm::all_of(value.getUsers(), [](Operation* user) { return isa<tensor::ExtractSliceOp>(user); });
+}
+
 static FailureOr<unsigned> getDirectReturnOperandIndex(OpResult result) {
   if (!result.hasOneUse())
     return failure();
@@ -51,12 +56,14 @@ collectFragmentAssemblyCopiesFromBlueprint(spatial::SpatBlueprintOp blueprint,
     return blueprint.emitOpError("fragment assembly lowering requires static ranked tensor results");
 
   std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
+  std::optional<ArrayRef<int64_t>> sourceSlotsAttr = blueprint.getFragmentSourceSlots();
   std::optional<ArrayRef<int64_t>> fragmentStridesAttr = blueprint.getFragmentStrides();
-  if (!operandIndicesAttr || !fragmentStridesAttr)
+  if (!operandIndicesAttr || !sourceSlotsAttr || !fragmentStridesAttr)
     return blueprint.emitOpError(
-      "fragment assembly lowering requires explicit operand indices and unit strides");
+      "fragment assembly lowering requires explicit operand indices, source slots, and unit strides");
 
   ArrayRef<int64_t> operandIndices = *operandIndicesAttr;
+  ArrayRef<int64_t> sourceSlots = *sourceSlotsAttr;
   std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
   if (!sourceOffsetsAttr)
     return blueprint.emitOpError("fragment assembly lowering requires explicit source offsets");
@@ -110,7 +117,11 @@ collectFragmentAssemblyCopiesFromBlueprint(spatial::SpatBlueprintOp blueprint,
             copy.sourceType = sourceType;
             copy.hostTargetIndex = hostTargetIndex;
             copy.lane = lane;
-            copy.sourceByteOffset = (sourceOffsets[fragmentIndex] + relativeSourceOffset) * static_cast<int64_t>(elementSize);
+            copy.sourceByteOffset =
+              (getFragmentAssemblySourceElementOffset(
+                 sourceType, sourceSlots[fragmentIndex], sourceOffsets[fragmentIndex])
+               + relativeSourceOffset)
+              * static_cast<int64_t>(elementSize);
             copy.hostByteOffset = hostElementOffset * static_cast<int64_t>(elementSize);
             copy.byteSize = chunkElements * static_cast<int64_t>(elementSize);
             copies.push_back(copy);
@@ -178,8 +189,8 @@ collectTopLevelFragmentAssemblyCopies(OpResult result, RankedTensorType packedRe
       if (operandIndices[fragmentIndex] != static_cast<int64_t>(use.getOperandNumber()))
         continue;
 
-      int64_t sourceElementOffset =
-        sourceSlots[fragmentIndex] * payloadElementCount + sourceOffsets[fragmentIndex];
+      int64_t sourceElementOffset = getFragmentAssemblySourceElementOffset(
+        packedResultType, sourceSlots[fragmentIndex], sourceOffsets[fragmentIndex]);
       int64_t lane = sourceElementOffset / payloadElementCount;
       if (lane < 0 || lane >= static_cast<int64_t>(laneCount))
         return failure();
@@ -357,6 +368,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatSchedul
     rewriter.createBlock(&coreBatchOp.getBody(), coreBatchOp.getBody().end(), TypeRange(blockArgTypes), blockArgLocs);
 
   IRMapping mapper;
+  SmallPtrSet<Value, 4> hostResidentTensors;
   rewriter.setInsertionPointToStart(newBlock);
   auto oldLaneArg = computeBatchOp.getLaneArgument();
   if (!oldLaneArg)
@@ -523,8 +535,10 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatSchedul
       if (isa_and_present<memref::GetGlobalOp>(toTensorOp.getBuffer().getDefiningOp())) {
         Operation* cloned = rewriter.clone(op, mapper);
         auto clonedTensor = cloned->getResult(0);
-        if (isUsedOnlyAsExplicitHostOperand(toTensorOp.getResult())) {
+        if (isUsedOnlyAsExplicitHostOperand(toTensorOp.getResult())
+            || isUsedOnlyByExtractSlices(toTensorOp.getResult())) {
           mapper.map(toTensorOp.getResult(), clonedTensor);
+          hostResidentTensors.insert(toTensorOp.getResult());
           continue;
         }
         auto clonedType = cast<ShapedType>(clonedTensor.getType());
@@ -542,6 +556,28 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatSchedul
       }
     }
 
+    if (auto extractSlice = dyn_cast<tensor::ExtractSliceOp>(op);
+        extractSlice && hostResidentTensors.contains(extractSlice.getSource())) {
+      Operation* cloned = rewriter.clone(op, mapper);
+      Value hostSlice = cloned->getResult(0);
+      auto outputBuffer = createEmptyTensorFromShaped(rewriter, loc, cast<ShapedType>(hostSlice.getType()));
+      Value zeroOffset = getOrCreateIndexConstant(rewriter, coreBatchOp.getOperation(), 0);
+      auto sizeAttr = getTensorSizeInBytesAttr(rewriter, coreBatchOp.getOperation(), hostSlice);
+      if (failed(sizeAttr))
+        return failure();
+      auto copied = pim::PimMemCopyHostToDevOp::create(rewriter,
+                                                       loc,
+                                                       outputBuffer.getType(),
+                                                       zeroOffset,
+                                                       zeroOffset,
+                                                       outputBuffer,
+                                                       hostSlice,
+                                                       *sizeAttr)
+                      .getOutput();
+      mapper.map(extractSlice.getResult(), copied);
+      continue;
+    }
+
     for (auto [operandIndex, operand] : llvm::enumerate(op.getOperands())) {
       if (!isa<TensorType>(operand.getType()) || mapper.contains(operand))
         continue;

src/PIM/Conversion/SpatialToPim/Common.cpp

@@ -215,6 +215,14 @@ forEachContiguousDestinationChunk(ArrayRef<int64_t> destShape,
   return visit(visit, 0);
 }
 
+int64_t getFragmentAssemblySourceElementOffset(RankedTensorType sourceType,
+                                               int64_t sourceSlot,
+                                               int64_t sourceOffset) {
+  assert(sourceType.getRank() > 0 && sourceType.hasStaticShape()
+         && "fragment assembly source must have a static leading slot dimension");
+  return sourceSlot * (sourceType.getNumElements() / sourceType.getDimSize(0)) + sourceOffset;
+}
+
 static mlir::Value
 createSteppedOffset(OpBuilder& builder, Location loc, mlir::Value start, mlir::Value index,
                     int64_t stepBytes, Operation *constantAnchor) {

src/PIM/Conversion/SpatialToPim/Common.hpp

@@ -65,6 +65,10 @@ forEachContiguousDestinationChunk(llvm::ArrayRef<int64_t> destShape,
                                   llvm::function_ref<mlir::LogicalResult(llvm::ArrayRef<int64_t>, int64_t, int64_t)>
                                     callback);
 
+int64_t getFragmentAssemblySourceElementOffset(mlir::RankedTensorType sourceType,
+                                               int64_t sourceSlot,
+                                               int64_t sourceOffset);
+
 struct FragmentAssemblyCopy {
   mlir::Value source;
   mlir::RankedTensorType sourceType;

src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp

@@ -44,13 +44,15 @@ static FailureOr<Value> lowerFragmentAssemblyBlueprint(IRRewriter& rewriter,
 
   std::optional<StringRef> modeAttr = blueprint.getMode();
   std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
+  std::optional<ArrayRef<int64_t>> sourceSlotsAttr = blueprint.getFragmentSourceSlots();
   std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
   std::optional<ArrayRef<int64_t>> fragmentStridesAttr = blueprint.getFragmentStrides();
-  if (!modeAttr || *modeAttr != "fragment_assembly" || !operandIndicesAttr || !sourceOffsetsAttr
-      || !fragmentStridesAttr)
+  if (!modeAttr || *modeAttr != "fragment_assembly" || !operandIndicesAttr || !sourceSlotsAttr
+      || !sourceOffsetsAttr || !fragmentStridesAttr)
     return blueprint.emitOpError("fragment assembly lowering requires explicit fragment metadata");
 
   ArrayRef<int64_t> operandIndices = *operandIndicesAttr;
+  ArrayRef<int64_t> sourceSlots = *sourceSlotsAttr;
   ArrayRef<int64_t> sourceOffsets = *sourceOffsetsAttr;
   ArrayRef<int64_t> flatOffsets = blueprint.getFragmentOffsets();
   ArrayRef<int64_t> flatSizes = blueprint.getFragmentSizes();
@@ -102,7 +104,10 @@ static FailureOr<Value> lowerFragmentAssemblyBlueprint(IRRewriter& rewriter,
             copy.source = source;
             copy.sourceType = sourceType;
             copy.sourceByteOffset =
-              (sourceOffsets[fragmentIndex] + relativeSourceOffset) * static_cast<int64_t>(elementSize);
+              (getFragmentAssemblySourceElementOffset(
+                 sourceType, sourceSlots[fragmentIndex], sourceOffsets[fragmentIndex])
+               + relativeSourceOffset)
+              * static_cast<int64_t>(elementSize);
             copy.hostByteOffset = hostElementOffset * static_cast<int64_t>(elementSize);
             copy.byteSize = chunkElements * static_cast<int64_t>(elementSize);
             copies.push_back(copy);

src/PIM/Conversion/SpatialToPim/ReturnPathNormalization.cpp

@@ -608,11 +608,12 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
     for (auto [blueprint, operandNumber] : *fragmentAssemblyUses) {
       rewriter.setInsertionPointAfterValue(storedValue);
       std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
+      std::optional<ArrayRef<int64_t>> sourceSlotsAttr = blueprint.getFragmentSourceSlots();
       std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
       std::optional<ArrayRef<int64_t>> stridesAttr = blueprint.getFragmentStrides();
-      if (!operandIndicesAttr || !sourceOffsetsAttr || !stridesAttr) {
+      if (!operandIndicesAttr || !sourceSlotsAttr || !sourceOffsetsAttr || !stridesAttr) {
         blueprint.emitOpError(
-          "fragment assembly lowering requires explicit operand, source-offset, and stride metadata");
+          "fragment assembly lowering requires explicit operand, source-slot, source-offset, and stride metadata");
         return ReturnPathLoweringResult::Failure;
       }
 
@@ -626,6 +627,7 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
       }
 
       ArrayRef<int64_t> operandIndices = *operandIndicesAttr;
+      ArrayRef<int64_t> sourceSlots = *sourceSlotsAttr;
       ArrayRef<int64_t> sourceOffsets = *sourceOffsetsAttr;
       ArrayRef<int64_t> flatOffsets = blueprint.getFragmentOffsets();
       ArrayRef<int64_t> flatSizes = blueprint.getFragmentSizes();
@@ -668,7 +670,9 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
                                    elementSize,
                                    producerOp,
                                    "fragment assembly host offset");
-            auto sourceOffset = getCheckedByteOffset(sourceOffsets[fragmentIndex] + relativeSourceOffset,
+            int64_t sourceElementOffset = getFragmentAssemblySourceElementOffset(
+              sourceType, sourceSlots[fragmentIndex], sourceOffsets[fragmentIndex]);
+            auto sourceOffset = getCheckedByteOffset(sourceElementOffset + relativeSourceOffset,
                                                      elementSize,
                                                      producerOp,
                                                      "fragment assembly source offset");

src/PIM/Conversion/SpatialToPim/SpatialToPim.td

@@ -12,7 +12,7 @@ include "src/Accelerators/PIM/Dialect/Pim/Pim.td"
 def spatToPimVMM : Pat<
   (SpatVMMOp:$srcOpRes $weight, $vector),
   (PimVMMOp $weight, $vector,
-    (NativeCodeCall<"onnx_mlir::getBestOutputTensorFromOperandsOrAllocate($_builder, $0.getDefiningOp())"> $srcOpRes))
+    (NativeCodeCall<"tensor::EmptyOp::create($_builder, $_loc, cast<ShapedType>($0.getType()).getShape(), cast<ShapedType>($0.getType()).getElementType())"> $srcOpRes))
 >;
 
 def spatToPimVVDMul : Pat<

src/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.cpp

@@ -105,6 +105,18 @@ static FailureOr<SmallVector<int64_t>> getStaticMemRefStrides(MemRefType type) {
   return strides;
 }
 
+static bool haveEquivalentStrides(MemRefType type, ArrayRef<int64_t> lhs,
+                                  ArrayRef<int64_t> rhs) {
+  if (lhs.size() != rhs.size()
+      || lhs.size() != static_cast<size_t>(type.getRank()))
+    return false;
+  for (auto [dimension, strides] : llvm::enumerate(llvm::zip(lhs, rhs)))
+    if (type.getDimSize(dimension) != 1
+        && std::get<0>(strides) != std::get<1>(strides))
+      return false;
+  return true;
+}
+
 static FailureOr<SmallVector<int64_t>> getProvenMemRefStrides(Value value) {
   llvm::SmallPtrSet<Value, 8> visiting;
   std::function<FailureOr<SmallVector<int64_t>>(Value)> prove =
@@ -229,9 +241,19 @@ static FailureOr<SmallVector<int64_t>> getProvenMemRefStrides(Value value) {
         return strides;
       }
       auto result = dyn_cast<OpResult>(current);
+      SmallVector<Value, 4> alternatives;
       SmallVector<Region *> regions;
       if (result) {
-        if (auto selection = dyn_cast<scf::IndexSwitchOp>(result.getOwner()))
+        if (auto loop = dyn_cast<scf::ForOp>(result.getOwner())) {
+          auto yield = dyn_cast<scf::YieldOp>(loop.getBody()->getTerminator());
+          if (!yield || result.getResultNumber() >= loop.getInitArgs().size()
+              || result.getResultNumber() >= yield.getNumOperands()) {
+            visiting.erase(current);
+            return failure();
+          }
+          alternatives.push_back(loop.getInitArgs()[result.getResultNumber()]);
+          alternatives.push_back(yield.getOperand(result.getResultNumber()));
+        } else if (auto selection = dyn_cast<scf::IndexSwitchOp>(result.getOwner()))
           for (Region &region : selection->getRegions())
             regions.push_back(&region);
         else if (auto selection = dyn_cast<scf::IfOp>(result.getOwner())) {
@@ -239,19 +261,23 @@ static FailureOr<SmallVector<int64_t>> getProvenMemRefStrides(Value value) {
           regions.push_back(&selection.getElseRegion());
         }
       }
-      if (regions.empty()) {
-        visiting.erase(current);
-        return failure();
-      }
-      std::optional<SmallVector<int64_t>> common;
       for (Region *region : regions) {
         auto yield = dyn_cast<scf::YieldOp>(region->front().getTerminator());
         if (!yield || result.getResultNumber() >= yield.getNumOperands()) {
           visiting.erase(current);
           return failure();
         }
-        auto strides = prove(yield.getOperand(result.getResultNumber()));
-        if (failed(strides) || (common && *common != *strides)) {
+        alternatives.push_back(yield.getOperand(result.getResultNumber()));
+      }
+      if (alternatives.empty()) {
+        visiting.erase(current);
+        return failure();
+      }
+      std::optional<SmallVector<int64_t>> common;
+      for (Value alternative : alternatives) {
+        auto strides = prove(alternative);
+        if (failed(strides)
+            || (common && !haveEquivalentStrides(type, *common, *strides))) {
           visiting.erase(current);
           return failure();
         }
@@ -314,7 +340,7 @@ static FailureOr<int64_t> getContiguousSuffixRank(Value value, ArrayRef<int64_t>
   int64_t expectedStride = 1;
   int64_t contiguousSuffixRank = 0;
   for (int64_t dim = type.getRank() - 1; dim >= 0; --dim) {
-    if ((*strides)[dim] != expectedStride)
+    if (copyShape[dim] != 1 && (*strides)[dim] != expectedStride)
       break;
     ++contiguousSuffixRank;
     auto nextStride = checkedPositiveMul(expectedStride, copyShape[dim]);
@@ -405,7 +431,7 @@ analyzeCopyRewrite(Value target, Value source, Value targetOffset, Value sourceO
   auto targetBytes = getShapedByteSize(targetType);
   auto sourceBytes = getShapedByteSize(sourceType);
   if (targetType.getElementType() == sourceType.getElementType() && succeeded(targetBytes) && succeeded(sourceBytes)
-      && *targetBytes == size && *sourceBytes == size) {
+      && size <= *targetBytes && size <= *sourceBytes) {
     auto targetSuffixRank = getContiguousSuffixRank(target, targetType.getShape());
     auto sourceSuffixRank = getContiguousSuffixRank(source, sourceType.getShape());
     if (succeeded(targetSuffixRank) && succeeded(sourceSuffixRank)

src/PIM/Dialect/Pim/Transforms/Bufferization/PimBufferizationPass.cpp

@@ -296,11 +296,7 @@ void PimBufferizationPass::runOnOperation() {
       });
   });
   moduleOp.walk([&](pim::PimCoreBatchOp coreBatchOp) {
-    llvm::SmallVector<unsigned, 2> lanes;
-    lanes.push_back(0);
-    if (coreBatchOp.getLaneCount() > 1)
-      lanes.push_back(static_cast<unsigned>(coreBatchOp.getLaneCount() - 1));
-    for (unsigned lane : lanes) {
+    for (unsigned lane = 0; lane < coreBatchOp.getLaneCount(); ++lane) {
       StaticValueKnowledge knowledge = seedCoreBatchKnowledge(coreBatchOp, lane);
       (void) walkPimCoreBlockStructurally(
         coreBatchOp.getBody().front(), knowledge, [&](Operation& op, const StaticValueKnowledge& opKnowledge) {

src/PIM/Dialect/Spatial/Spatial.td

@@ -288,6 +288,25 @@ def SpatReluPlanOp : SpatOp<"relu_plan", []> {
   let hasVerifier = 1;
 }
 
+def SpatMaxPool2DPlanOp : SpatOp<"max_pool2d_plan", []> {
+  let summary = "Layout-aware 2D NCHW MaxPool planning op";
+
+  let arguments = (ins
+    SpatTensor:$input,
+    DenseI64ArrayAttr:$kernelShape,
+    DenseI64ArrayAttr:$pads,
+    DenseI64ArrayAttr:$strides,
+    DenseI64ArrayAttr:$dilations,
+    StrAttr:$logicalLayout
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let hasVerifier = 1;
+}
+
 def SpatBiasAddPlanOp : SpatOp<"bias_add_plan", []> {
   let summary = "Layout-aware Conv-style bias add planning op";
 

src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp

@@ -459,6 +459,26 @@ LogicalResult SpatReluPlanOp::verify() {
   return success();
 }
 
+LogicalResult SpatMaxPool2DPlanOp::verify() {
+  if (failed(verifyPlanTensorTypes(getOperation(), getInput(), getOutput(), "spat.max_pool2d_plan")))
+    return failure();
+  auto inputType = dyn_cast<RankedTensorType>(getInput().getType());
+  auto outputType = dyn_cast<RankedTensorType>(getOutput().getType());
+  if (!inputType.hasStaticShape() || !outputType.hasStaticShape() || inputType.getRank() != 4
+      || outputType.getRank() != 4)
+    return emitError("requires static rank-4 input and output tensors");
+  if (getLogicalLayout() != "nchw")
+    return emitError("requires logical layout \"nchw\"");
+  if (getKernelShape().size() != 2 || getStrides().size() != 2 || getDilations().size() != 2)
+    return emitError("requires two kernel, stride, and dilation values");
+  if (getPads().size() != 4)
+    return emitError("requires four pad values");
+  if (inputType.getDimSize(0) != outputType.getDimSize(0)
+      || inputType.getDimSize(1) != outputType.getDimSize(1))
+    return emitError("requires matching input/output batch and channel dimensions");
+  return success();
+}
+
 LogicalResult SpatBiasAddPlanOp::verify() {
   if (failed(verifyPlanTensorTypes(getOperation(), getInput(), getOutput(), "spat.bias_add_plan")))
     return failure();

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp

@@ -285,6 +285,8 @@ static std::optional<DeferredAssemblySourceTransform> getSourceTransform(
 
 static FailureOr<std::optional<DeferredInsertAssemblyTemplate>>
 analyzeInsertAssembly(const DeferredProgramTemplate &program) {
+  if (program.specializationCount != 1)
+    return std::optional<DeferredInsertAssemblyTemplate>();
   auto finalInsert = program.yieldedValue.getDefiningOp<tensor::InsertSliceOp>();
   if (!finalInsert || program.leaves.empty())
     return std::optional<DeferredInsertAssemblyTemplate>();

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp

@@ -424,10 +424,12 @@ retargetBlueprint(DeferredTransferPlan& plan, SpatBlueprintOp blueprint, GraphBa
         return blueprint.emitOpError(
           "phase 2 Blueprint source has no scheduled publication"), failure();
       source = (*producer)->published;
-      slot = (*producer)->publishedSlotStart + graphLane - (*producer)->laneStart;
-      if (slot < (*producer)->publishedSlotStart
-          || slot >= (*producer)->publishedSlotStart
-                     + (*producer)->publishedSlotCount)
+      int64_t publicationSlotStart = (*producer)->scheduled->isBatch()
+                                       ? (*producer)->publishedSlotStart
+                                       : 0;
+      slot = publicationSlotStart + graphLane - (*producer)->laneStart;
+      if (slot < publicationSlotStart
+          || slot >= publicationSlotStart + (*producer)->publishedSlotCount)
         return blueprint.emitOpError(
           "phase 2 Blueprint slot is outside its scheduled publication window"), failure();
     }

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp

@@ -780,7 +780,7 @@ ComputeGraph buildComputeGraph(Operation* entryOp) {
         if (auto batch = dyn_cast<SpatComputeBatch>(&op)) {
           if (isUsedAsWeightOnly(batch.getOperation()))
             continue;
-          size_t chunkCount = getBatchChunkTargetCount(batch.getLaneCount());
+          size_t chunkCount = getBatchChunkTargetCount(batch);
           for (size_t chunkIndex = 0; chunkIndex < chunkCount; ++chunkIndex) {
             ComputeInstance instance = getBatchChunkForIndex(batch, chunkIndex);
             size_t index = graph.nodes.size();

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp

@@ -1,6 +1,7 @@
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 
+#include <algorithm>
 #include <limits>
 #include <optional>
 
@@ -18,14 +19,8 @@ size_t getSchedulingCpuBudget() {
   return std::numeric_limits<size_t>::max();
 }
 
-size_t getBatchChunkTargetCount(int32_t laneCount) {
+static BatchChunkRange getBatchChunkRange(int32_t laneCount, size_t chunkCount, size_t chunkIndex) {
   assert(laneCount > 0 && "laneCount must be positive");
-  return std::min(static_cast<size_t>(laneCount), getSchedulingCpuBudget());
-}
-
-BatchChunkRange getBatchChunkRange(int32_t laneCount, size_t chunkIndex) {
-  assert(laneCount > 0 && "laneCount must be positive");
-  size_t chunkCount = getBatchChunkTargetCount(laneCount);
   assert(chunkIndex < chunkCount && "chunkIndex out of range");
 
   size_t laneCountSize = static_cast<size_t>(laneCount);
@@ -38,11 +33,22 @@ BatchChunkRange getBatchChunkRange(int32_t laneCount, size_t chunkIndex) {
   return {static_cast<uint32_t>(start), static_cast<uint32_t>(count)};
 }
 
-size_t getBatchChunkIndexForLane(int32_t laneCount, uint32_t lane) {
+size_t getBatchChunkTargetCount(SpatComputeBatch batch) {
+  int32_t laneCount = batch.getLaneCount();
+  assert(laneCount > 0 && "laneCount must be positive");
+  return std::min(static_cast<size_t>(laneCount), getSchedulingCpuBudget());
+}
+
+BatchChunkRange getBatchChunkRange(SpatComputeBatch batch, size_t chunkIndex) {
+  return getBatchChunkRange(batch.getLaneCount(), getBatchChunkTargetCount(batch), chunkIndex);
+}
+
+size_t getBatchChunkIndexForLane(SpatComputeBatch batch, uint32_t lane) {
+  int32_t laneCount = batch.getLaneCount();
   assert(laneCount > 0 && "laneCount must be positive");
   assert(lane < static_cast<uint32_t>(laneCount) && "lane out of range");
 
-  size_t chunkCount = getBatchChunkTargetCount(laneCount);
+  size_t chunkCount = getBatchChunkTargetCount(batch);
   size_t laneCountSize = static_cast<size_t>(laneCount);
   size_t baseChunkSize = laneCountSize / chunkCount;
   size_t remainder = laneCountSize % chunkCount;
@@ -56,12 +62,12 @@ size_t getBatchChunkIndexForLane(int32_t laneCount, uint32_t lane) {
 }
 
 ComputeInstance getBatchChunkForIndex(SpatComputeBatch batch, size_t chunkIndex) {
-  BatchChunkRange chunk = getBatchChunkRange(batch.getLaneCount(), chunkIndex);
+  BatchChunkRange chunk = getBatchChunkRange(batch, chunkIndex);
   return {batch.getOperation(), chunk.laneStart, chunk.laneCount};
 }
 
 ComputeInstance getBatchChunkForLane(SpatComputeBatch batch, uint32_t lane) {
-  return getBatchChunkForIndex(batch, getBatchChunkIndexForLane(batch.getLaneCount(), lane));
+  return getBatchChunkForIndex(batch, getBatchChunkIndexForLane(batch, lane));
 }
 
 llvm::SmallVector<ComputeInstance, 4>
@@ -74,8 +80,8 @@ getBatchChunksForRange(SpatComputeBatch batch, uint32_t laneStart, uint32_t lane
   assert(laneEnd >= laneStart && "lane range overflow");
   assert(laneEnd <= static_cast<uint32_t>(batch.getLaneCount()) && "lane range out of bounds");
 
-  size_t firstChunk = getBatchChunkIndexForLane(batch.getLaneCount(), laneStart);
-  size_t lastChunk = getBatchChunkIndexForLane(batch.getLaneCount(), laneEnd - 1);
+  size_t firstChunk = getBatchChunkIndexForLane(batch, laneStart);
+  size_t lastChunk = getBatchChunkIndexForLane(batch, laneEnd - 1);
   chunks.reserve(lastChunk - firstChunk + 1);
   for (size_t chunkIndex = firstChunk; chunkIndex <= lastChunk; ++chunkIndex)
     chunks.push_back(getBatchChunkForIndex(batch, chunkIndex));

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.hpp

@@ -27,9 +27,9 @@ struct BatchChunkRange {
 };
 
 size_t getSchedulingCpuBudget();
-size_t getBatchChunkTargetCount(int32_t laneCount);
-BatchChunkRange getBatchChunkRange(int32_t laneCount, size_t chunkIndex);
-size_t getBatchChunkIndexForLane(int32_t laneCount, uint32_t lane);
+size_t getBatchChunkTargetCount(SpatComputeBatch batch);
+BatchChunkRange getBatchChunkRange(SpatComputeBatch batch, size_t chunkIndex);
+size_t getBatchChunkIndexForLane(SpatComputeBatch batch, uint32_t lane);
 ComputeInstance getBatchChunkForIndex(SpatComputeBatch batch, size_t chunkIndex);
 ComputeInstance getBatchChunkForLane(SpatComputeBatch batch, uint32_t lane);
 llvm::SmallVector<ComputeInstance, 4>

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp

@@ -240,17 +240,17 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
     Time bestOeft = std::numeric_limits<Time>::max();
     unsigned int bestOverlapCount = 0;
     size_t bestCenterDistance = std::numeric_limits<size_t>::max();
+    size_t smallestCrossbarUnion = std::numeric_limits<size_t>::max();
     bool crossbarRejected = false;
 
     for (size_t processor = 0; processor < processorCount; ++processor) {
       unsigned int overlapCount = countCrossbarOverlap(processorCrossbars[processor], graph.nodes[task].crossbarUsage);
-      if (!graph.nodes[task].crossbarUsage.empty()
-          && getCrossbarUnionSize(processorCrossbars[processor], graph.nodes[task].crossbarUsage)
-               > options.crossbarCapacity) {
+      size_t crossbarUnion = getCrossbarUnionSize(processorCrossbars[processor], graph.nodes[task].crossbarUsage);
+      smallestCrossbarUnion = std::min(smallestCrossbarUnion, crossbarUnion);
+      if (!graph.nodes[task].crossbarUsage.empty() && crossbarUnion > options.crossbarCapacity) {
         crossbarRejected = true;
         continue;
       }
-
       Time dataReady = 0;
       for (const auto& [pred, comm] : graph.predecessors[task]) {
         const ScheduledTask& predSchedule = schedules[pred];
@@ -314,9 +314,15 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
 
     if (bestProcessor == std::numeric_limits<size_t>::max()) {
       if (crossbarRejected) {
+        const ComputeInstance& instance = graph.nodes[task].instance;
         std::string message =
-          llvm::formatv("PEFT scheduler: no valid processor for task {0}; crossbar capacity {1} is exhausted",
+          llvm::formatv("PEFT scheduler: no valid processor for task {0} (lanes {1}..{2}, {3} distinct weights); "
+                        "smallest processor union is {4}, exceeding crossbar capacity {5}",
                         graph.nodes[task].originalOrder,
+                        instance.laneStart,
+                        instance.laneStart + instance.laneCount,
+                        graph.nodes[task].crossbarUsage.size(),
+                        smallestCrossbarUnion,
                         options.crossbarCapacity)
             .str();
         llvm::report_fatal_error(llvm::StringRef(message));