Unexpected invariant now it's clear (batched in the first tensor rank)

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

@@ -346,6 +346,52 @@ inline void createParallelInsertSliceIntoBatchOutput(mlir::PatternRewriter& rewr
   mlir::tensor::ParallelInsertSliceOp::create(rewriter, loc, source, dest, offsets, sizes, strides);
 }
 
+inline void publishGraphBatchPhysicalFragment(mlir::PatternRewriter& rewriter,
+                                              mlir::Location loc,
+                                              mlir::Value fragment,
+                                              mlir::Value output,
+                                              mlir::Value physicalSlot) {
+  auto fragmentType = mlir::cast<mlir::RankedTensorType>(fragment.getType());
+  mlir::SmallVector<mlir::OpFoldResult> offsets {physicalSlot};
+  mlir::SmallVector<mlir::OpFoldResult> sizes {rewriter.getIndexAttr(1)};
+  mlir::SmallVector<mlir::OpFoldResult> strides {rewriter.getIndexAttr(1)};
+  for (int64_t dim : fragmentType.getShape()) {
+    offsets.push_back(rewriter.getIndexAttr(0));
+    sizes.push_back(rewriter.getIndexAttr(dim));
+    strides.push_back(rewriter.getIndexAttr(1));
+  }
+  createParallelInsertSliceIntoBatchOutput(rewriter, loc, fragment, output, offsets, sizes, strides);
+}
+
+inline mlir::FailureOr<mlir::Value>
+extractGraphBatchPhysicalFragment(mlir::PatternRewriter& rewriter,
+                                  mlir::Location loc,
+                                  mlir::Value physicalBatch,
+                                  mlir::OpFoldResult slot,
+                                  mlir::RankedTensorType fragmentType) {
+  if (fragmentType.getRank() == 0)
+    return mlir::failure();
+  auto physicalType = mlir::dyn_cast<mlir::RankedTensorType>(physicalBatch.getType());
+  if (!physicalType || physicalType.getRank() != fragmentType.getRank() + 1)
+    return mlir::failure();
+  mlir::SmallVector<int64_t> selectedShape {1};
+  llvm::append_range(selectedShape, fragmentType.getShape());
+  auto selectedType = mlir::RankedTensorType::get(selectedShape, fragmentType.getElementType(), fragmentType.getEncoding());
+  mlir::SmallVector<mlir::OpFoldResult> offsets {slot};
+  mlir::SmallVector<mlir::OpFoldResult> sizes {rewriter.getIndexAttr(1)};
+  mlir::SmallVector<mlir::OpFoldResult> strides {rewriter.getIndexAttr(1)};
+  for (int64_t dim : fragmentType.getShape()) {
+    offsets.push_back(rewriter.getIndexAttr(0));
+    sizes.push_back(rewriter.getIndexAttr(dim));
+    strides.push_back(rewriter.getIndexAttr(1));
+  }
+  mlir::Value selected = mlir::tensor::ExtractSliceOp::create(rewriter, loc, selectedType, physicalBatch, offsets, sizes, strides);
+  mlir::SmallVector<mlir::ReassociationIndices> reassociation {{0, 1}};
+  for (int64_t dim = 2; dim <= fragmentType.getRank(); ++dim)
+    reassociation.push_back({dim});
+  return mlir::tensor::CollapseShapeOp::create(rewriter, loc, fragmentType, selected, reassociation).getResult();
+}
+
 template <typename BodyFn>
 mlir::Value materializeOrComputeUnary(mlir::Value input,
                                       mlir::RankedTensorType resultType,

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

@@ -19,9 +19,7 @@ RankedTensorType getRowStripFragmentType(RankedTensorType logicalType) {
 }
 
 RankedTensorType getRowStripStorageType(RankedTensorType logicalType) {
-  return RankedTensorType::get({logicalType.getDimSize(2), logicalType.getDimSize(1), 1, logicalType.getDimSize(3)},
-                               logicalType.getElementType(),
-                               logicalType.getEncoding());
+  return spatial::getGraphBatchPhysicalResultType(logicalType.getDimSize(2), getRowStripFragmentType(logicalType));
 }
 
 std::pair<SmallVector<int64_t>, SmallVector<int64_t>> buildRowStripMetadata(RankedTensorType type) {
@@ -39,29 +37,12 @@ std::pair<SmallVector<int64_t>, SmallVector<int64_t>> buildRowStripMetadata(Rank
   return {offsets, sizes};
 }
 
-SmallVector<OpFoldResult> buildRowStripFragmentOffsets(PatternRewriter& rewriter, OpFoldResult row) {
-  return {row, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
-}
-
-SmallVector<OpFoldResult> buildRowStripFragmentSizes(PatternRewriter& rewriter, RankedTensorType logicalType) {
-  return {rewriter.getIndexAttr(1),
-          rewriter.getIndexAttr(logicalType.getDimSize(1)),
-          rewriter.getIndexAttr(1),
-          rewriter.getIndexAttr(logicalType.getDimSize(3))};
-}
-
 Value extractRowStripFragment(Value storage,
                               RankedTensorType logicalType,
                               OpFoldResult row,
                               PatternRewriter& rewriter,
                               Location loc) {
-  return tensor::ExtractSliceOp::create(rewriter,
-                                        loc,
-                                        getRowStripFragmentType(logicalType),
-                                        storage,
-                                        buildRowStripFragmentOffsets(rewriter, row),
-                                        buildRowStripFragmentSizes(rewriter, logicalType),
-                                        getUnitStrides(rewriter, 4));
+  return *extractGraphBatchPhysicalFragment(rewriter, loc, storage, row, getRowStripFragmentType(logicalType));
 }
 
 void insertRowStripFragment(Value fragment,
@@ -70,13 +51,11 @@ void insertRowStripFragment(Value fragment,
                             OpFoldResult row,
                             PatternRewriter& rewriter,
                             Location loc) {
-  createParallelInsertSliceIntoBatchOutput(rewriter,
-                                           loc,
-                                           fragment,
-                                           output,
-                                           buildRowStripFragmentOffsets(rewriter, row),
-                                           buildRowStripFragmentSizes(rewriter, logicalType),
-                                           getUnitStrides(rewriter, 4));
+  assert(fragment.getType() == getRowStripFragmentType(logicalType));
+  assert(output.getType() == getRowStripStorageType(logicalType));
+  auto slot = dyn_cast<Value>(row);
+  assert(slot && "row-strip graph publication requires a dynamic physical slot");
+  publishGraphBatchPhysicalFragment(rewriter, loc, fragment, output, slot);
 }
 
 FailureOr<Value> createPerChannelConstantFragment(DenseElementsAttr denseAttr,
@@ -145,30 +124,23 @@ FailureOr<Value> createRowStripStorageFromRows(Value rows,
 }
 
 FailureOr<Value>
-materializeRowStripStorageToDense(Value storage, RankedTensorType logicalType, PatternRewriter& rewriter, Location loc) {
+createRowStripAssemblyBlueprint(Value storage, RankedTensorType logicalType, PatternRewriter& rewriter, Location loc) {
   auto storageType = dyn_cast<RankedTensorType>(storage.getType());
   if (!storageType || storageType != getRowStripStorageType(logicalType))
     return failure();
 
-  auto batchOp = createSpatComputeBatch(
-    rewriter, loc, TypeRange {logicalType}, logicalType.getDimSize(2), {}, ValueRange {storage},
-    [&](detail::SpatComputeBatchBodyArgs args) {
-      Value fragment = extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc);
-      createParallelInsertSliceIntoBatchOutput(rewriter,
-                                               loc,
-                                               fragment,
-                                               args.outputs.front(),
-                                               SmallVector<OpFoldResult> {rewriter.getIndexAttr(0),
-                                                                          rewriter.getIndexAttr(0),
-                                                                          args.lane,
-                                                                          rewriter.getIndexAttr(0)},
-                                               buildRowStripFragmentSizes(rewriter, logicalType),
-                                               getUnitStrides(rewriter, 4));
-      return success();
-    });
-  if (failed(batchOp))
-    return failure();
-  return batchOp->getResult(0);
+  auto [offsets, sizes] = buildRowStripMetadata(logicalType);
+  int64_t height = logicalType.getDimSize(2);
+  SmallVector<int64_t> operandIndices(height, 0), sourceSlots, sourceOffsets(height, 0), strides(height * 4, 1);
+  for (int64_t row = 0; row < height; ++row)
+    sourceSlots.push_back(row);
+  return spatial::SpatBlueprintOp::create(rewriter, loc, logicalType, storage, ValueRange {},
+    rewriter.getStringAttr("nchw"), rewriter.getStringAttr("nchw_row_strip"),
+    rewriter.getDenseI64ArrayAttr(offsets), rewriter.getDenseI64ArrayAttr(sizes),
+    rewriter.getStringAttr("nchw_row_strip_fragments"), rewriter.getStringAttr("fragment_assembly"),
+    rewriter.getDenseI64ArrayAttr(operandIndices), rewriter.getDenseI64ArrayAttr(sourceSlots),
+    rewriter.getDenseI64ArrayAttr(sourceOffsets), rewriter.getDenseI64ArrayAttr(strides),
+    rewriter.getStringAttr("disjoint"), rewriter.getStringAttr("complete")).getOutput();
 }
 
 FailureOr<Value>

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

@@ -50,7 +50,7 @@ mlir::FailureOr<mlir::Value> createRowStripStorageFromRows(mlir::Value rows,
                                                            mlir::PatternRewriter& rewriter,
                                                            mlir::Location loc);
 
-mlir::FailureOr<mlir::Value> materializeRowStripStorageToDense(mlir::Value storage,
+mlir::FailureOr<mlir::Value> createRowStripAssemblyBlueprint(mlir::Value storage,
                                                                mlir::RankedTensorType logicalType,
                                                                mlir::PatternRewriter& rewriter,
                                                                mlir::Location loc);

src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp

@@ -75,7 +75,7 @@ materializeRowStripToDense(const RowStripPhysicalValue& rowStripValue, Location
   auto [expectedOffsets, expectedSizes] = buildRowStripMetadata(rowStripValue.logicalType);
   if (!llvm::equal(rowStripValue.fragmentOffsets, expectedOffsets) || !llvm::equal(rowStripValue.fragmentSizes, expectedSizes))
     return failure();
-  return materializeRowStripStorageToDense(rowStripValue.storage, rowStripValue.logicalType, rewriter, loc);
+  return createRowStripAssemblyBlueprint(rowStripValue.storage, rowStripValue.logicalType, rewriter, loc);
 }
 
 struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, OperationPass<ModuleOp>> {
@@ -277,6 +277,8 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
         continue;
       }
       if (auto blueprintOp = dyn_cast<spatial::SpatBlueprintOp>(&op)) {
+        if (std::optional<StringRef> mode = blueprintOp.getMode(); mode && *mode == "fragment_assembly")
+          continue;
         if (blueprintOp.getPhysicalLayout() == kDenseLayout) {
           rewriter.replaceOp(blueprintOp, blueprintOp.getInput());
           continue;
@@ -366,11 +368,15 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
     moduleOp.walk([&](Operation* op) {
       if (isa<ONNXEntryPointOp>(op))
         return;
-      if (isa<spatial::SpatConv2DPlanOp,
-              spatial::SpatBiasAddPlanOp,
-              spatial::SpatReluPlanOp,
-              spatial::SpatBlueprintOp,
-              spatial::SpatMaterializeLayoutOp>(op)
+      if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
+        if (std::optional<StringRef> mode = blueprint.getMode(); mode && *mode == "fragment_assembly")
+          return;
+        op->emitOpError("planning blueprint must not remain after LowerSpatialPlans");
+        hasIllegalOps = true;
+      } else if (isa<spatial::SpatConv2DPlanOp,
+                     spatial::SpatBiasAddPlanOp,
+                     spatial::SpatReluPlanOp,
+                     spatial::SpatMaterializeLayoutOp>(op)
           || op->getDialect()->getNamespace() == "onnx") {
         op->emitOpError("operation must not remain after LowerSpatialPlans");
         hasIllegalOps = true;

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

@@ -1448,12 +1448,10 @@ static FailureOr<Value> reconstructDepthwiseGemmRows(Value pieces,
             Value laneIndex = createOrFoldAffineApply(
               rewriter, tileLoc, (d0 * tiling.totalPatches) + d1, ValueRange {channelTileIndex, patchIndex}, anchorOp);
             auto rowTileType = RankedTensorType::get({1, tiling.tileOutputChannels}, piecesType.getElementType());
-            SmallVector<OpFoldResult> pieceOffsets {laneIndex, rewriter.getIndexAttr(0)};
-            SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(1),
-                                                  rewriter.getIndexAttr(tiling.tileOutputChannels)};
-            Value rowTile = tensor::ExtractSliceOp::create(
-              rewriter, tileLoc, rowTileType, piecesArg, pieceOffsets, pieceSizes, getUnitStrides(rewriter, 2));
-            Value rowNext = insertOutputTile(rowTile, rowAcc, channelTileIndex, tiling, rewriter, tileLoc);
+            FailureOr<Value> rowTile = extractGraphBatchPhysicalFragment(rewriter, tileLoc, piecesArg, laneIndex, rowTileType);
+            if (failed(rowTile))
+              return failure();
+            Value rowNext = insertOutputTile(*rowTile, rowAcc, channelTileIndex, tiling, rewriter, tileLoc);
             tileYielded.push_back(rowNext);
             return success();
           });
@@ -1556,8 +1554,9 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter&
 
   auto gemmOutType =
     RankedTensorType::get({tiling->totalPatches, state.outType.getDimSize(1)}, state.outType.getElementType());
-  auto piecesType = RankedTensorType::get({tiling->totalPatches * tiling->numChannelTiles, tiling->tileOutputChannels},
-                                          state.outType.getElementType());
+  auto rowTileType = RankedTensorType::get({1, tiling->tileOutputChannels}, state.outType.getElementType());
+  auto piecesType = spatial::getGraphBatchPhysicalResultType(
+    tiling->totalPatches * tiling->numChannelTiles, rowTileType);
   auto paddedInputType = cast<RankedTensorType>(paddedInput.getType());
   auto inputTileType =
     RankedTensorType::get({1, tiling->channelsPerTile, state.wType.getDimSize(2), state.wType.getDimSize(3)},
@@ -1626,7 +1625,6 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter&
                                             *tiling,
                                             rewriter,
                                             loc);
-      auto rowTileType = RankedTensorType::get({1, tiling->tileOutputChannels}, state.outType.getElementType());
       Value rowTile = spatial::SpatVMMOp::create(rewriter, loc, rowTileType, weightTile, inputTile).getResult();
       if (args.inputs.size() > 1) {
         Value biasArg = pickInputByRank(/*rank=*/2);
@@ -1637,11 +1635,7 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter&
         Value biasTile = tiling->numChannelTiles == 1 ? biasArg : createBiasTile(biasArg, channelTileIndex, *tiling, rewriter, loc);
         rowTile = spatial::SpatVAddOp::create(rewriter, loc, rowTileType, rowTile, biasTile).getResult();
       }
-      SmallVector<OpFoldResult> outputOffsets {args.lane, rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> outputSizes {rewriter.getIndexAttr(1),
-                                             rewriter.getIndexAttr(tiling->tileOutputChannels)};
-      createParallelInsertSliceIntoBatchOutput(
-        rewriter, loc, rowTile, args.outputs.front(), outputOffsets, outputSizes, getUnitStrides(rewriter, 2));
+      publishGraphBatchPhysicalFragment(rewriter, loc, rowTile, args.outputs.front(), args.lane);
       return success();
     });
   if (failed(batchOp))
@@ -1650,14 +1644,10 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter&
   auto nhwcType = RankedTensorType::get(
     {state.xType.getDimSize(0), state.outType.getDimSize(2), state.outType.getDimSize(3), state.outType.getDimSize(1)},
     state.outType.getElementType());
-  Value collectedRows = batchOp->getResult(0);
-  if (tiling->numChannelTiles != 1) {
-    auto reconstructedRows =
-      reconstructDepthwiseGemmRows(batchOp->getResult(0), piecesType, gemmOutType, *tiling, rewriter, loc);
-    if (failed(reconstructedRows))
-      return failure();
-    collectedRows = *reconstructedRows;
-  }
+  auto reconstructedRows = reconstructDepthwiseGemmRows(batchOp->getResult(0), piecesType, gemmOutType, *tiling, rewriter, loc);
+  if (failed(reconstructedRows))
+    return failure();
+  Value collectedRows = *reconstructedRows;
 
   return createCollectedConvOutput(ValueRange {collectedRows},
                                    state.outType,

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

@@ -251,10 +251,7 @@ static FailureOr<spatial::SpatComputeBatch> createVmmBatch(Value a,
         rewriter, loc, args.weights.front(), bTileType, bOffsets, bSizes, unitStrides);
       Value piece = spatial::SpatVMMOp::create(rewriter, loc, pieceType, bTile, aTile).getResult();
 
-      SmallVector<OpFoldResult> pieceOffsets {args.lane, rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())};
-      createParallelInsertSliceIntoBatchOutput(
-        rewriter, loc, piece, args.outputs.front(), pieceOffsets, pieceSizes, unitStrides);
+      publishGraphBatchPhysicalFragment(rewriter, loc, piece, args.outputs.front(), args.lane);
     });
   if (failed(batchOp))
     return failure();
@@ -401,11 +398,7 @@ static FailureOr<spatial::SpatComputeBatch> createVvdmulBatch(Value a,
       Value bVector = extractDynamicGemmBColumn(args.inputs[1], column, vectorType, rewriter, loc);
       Value scalar = spatial::SpatVVDMulOp::create(rewriter, loc, scalarType, aVector, bVector).getResult();
 
-      SmallVector<OpFoldResult> outputOffsets {args.lane, rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-      SmallVector<OpFoldResult> unitStrides = getUnitStrides(rewriter, 2);
-      createParallelInsertSliceIntoBatchOutput(
-        rewriter, loc, scalar, args.outputs.front(), outputOffsets, scalarSizes, unitStrides);
+      publishGraphBatchPhysicalFragment(rewriter, loc, scalar, args.outputs.front(), args.lane);
     });
   if (failed(batchOp))
     return failure();
@@ -447,15 +440,14 @@ static FailureOr<spatial::SpatCompute> createDynamicGemmOutputCompute(Value scal
         Value row = createDynamicGemmBatchRow(lane, numOutCols, rewriter, nestedLoc);
         Value column =
           onnx_mlir::affineModConst(rewriter, nestedLoc, lane, numOutCols, rewriter.getInsertionBlock()->getParentOp());
-        SmallVector<OpFoldResult> scalarOffsets {lane, rewriter.getIndexAttr(0)};
         SmallVector<OpFoldResult> scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
         SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-        Value scalar = tensor::ExtractSliceOp::create(
-                         rewriter, nestedLoc, scalarType, pieces, scalarOffsets, scalarSizes, unitStrides)
-                         .getResult();
+        FailureOr<Value> scalar = extractGraphBatchPhysicalFragment(rewriter, nestedLoc, pieces, lane, scalarType);
+        if (failed(scalar))
+          return failure();
         if (alpha != 1.0f) {
           Value alphaTensor = createScalarTensorConstant(scalarType, alpha, rewriter, nestedLoc);
-          scalar = spatial::SpatVMulOp::create(rewriter, nestedLoc, scalarType, scalar, alphaTensor).getResult();
+          *scalar = spatial::SpatVMulOp::create(rewriter, nestedLoc, scalarType, *scalar, alphaTensor).getResult();
         }
         if (biasArg) {
           Value biasScalar =
@@ -465,11 +457,11 @@ static FailureOr<spatial::SpatCompute> createDynamicGemmOutputCompute(Value scal
             biasScalar =
               spatial::SpatVMulOp::create(rewriter, nestedLoc, scalarType, biasScalar, betaTensor).getResult();
           }
-          scalar = spatial::SpatVAddOp::create(rewriter, nestedLoc, scalarType, scalar, biasScalar).getResult();
+          *scalar = spatial::SpatVAddOp::create(rewriter, nestedLoc, scalarType, *scalar, biasScalar).getResult();
         }
         SmallVector<OpFoldResult> outputOffsets {row, column};
         Value outputNext =
-          tensor::InsertSliceOp::create(rewriter, nestedLoc, scalar, outputAcc, outputOffsets, scalarSizes, unitStrides)
+          tensor::InsertSliceOp::create(rewriter, nestedLoc, *scalar, outputAcc, outputOffsets, scalarSizes, unitStrides)
             .getResult();
         yielded.push_back(outputNext);
         return success();
@@ -505,14 +497,13 @@ static Value extractReductionPiece(Value partialPiecesArg,
                                    int64_t numOutRows,
                                    ConversionPatternRewriter& rewriter,
                                    Location loc) {
-  SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-  SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(numOutRows),
-                                        rewriter.getIndexAttr(crossbarSize.getValue())};
+  SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
+  SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())};
   SmallVector<OpFoldResult> pieceOffsets {
-    createPartialGroupOffset(hSlice, kSlice, numKSlices, numOutRows, rewriter, loc), rewriter.getIndexAttr(0)};
-  return tensor::ExtractSliceOp::create(
-           rewriter, loc, pieceType, partialPiecesArg, pieceOffsets, pieceSizes, unitStrides)
-    .getResult();
+    createPartialGroupOffset(hSlice, kSlice, numKSlices, numOutRows, rewriter, loc), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
+  auto selectedType = RankedTensorType::get({numOutRows, 1, static_cast<int64_t>(crossbarSize.getValue())}, pieceType.getElementType());
+  Value selected = tensor::ExtractSliceOp::create(rewriter, loc, selectedType, partialPiecesArg, pieceOffsets, pieceSizes, unitStrides);
+  return tensor::CollapseShapeOp::create(rewriter, loc, pieceType, selected, SmallVector<ReassociationIndices> {{0, 1}, {2}});
 }
 
 static Value reducePartialPiecesForHSlice(Value partialPiecesArg,
@@ -730,7 +721,7 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
       return failure();
     }
 
-    auto scalarPiecesType = RankedTensorType::get({laneCount64, 1}, outType.getElementType());
+    auto scalarPiecesType = spatial::getGraphBatchPhysicalResultType(laneCount64, RankedTensorType::get({1, 1}, outType.getElementType()));
     auto batchOp = createVvdmulBatch(a, b, aType, bType, scalarPiecesType, outType, rewriter, loc);
     if (failed(batchOp))
       return failure();
@@ -802,8 +793,8 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
     return failure();
   }
 
-  auto partialPiecesType =
-    RankedTensorType::get({laneCount64, static_cast<int64_t>(crossbarSize.getValue())}, outType.getElementType());
+  auto partialPiecesType = spatial::getGraphBatchPhysicalResultType(
+    laneCount64, RankedTensorType::get({1, static_cast<int64_t>(crossbarSize.getValue())}, outType.getElementType()));
   auto batchOp =
     createVmmBatch(a, b, aType, paddedBType, partialPiecesType, numOutRows, numKSlices, numOutHSlices, rewriter, loc);
   if (failed(batchOp))

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

@@ -398,10 +398,7 @@ static FailureOr<spatial::SpatComputeBatch> createBatchedVmmBatch(Value a,
         args.weights.front(), bBatchShape, outputBatchShape, batch, kOffset, hOffset, bTileType, rewriter, loc);
       Value piece = spatial::SpatVMMOp::create(rewriter, loc, pieceType, bTile, aTile).getResult();
 
-      SmallVector<OpFoldResult> pieceOffsets {args.lane, rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())};
-      createParallelInsertSliceIntoBatchOutput(
-        rewriter, loc, piece, args.outputs.front(), pieceOffsets, pieceSizes, getUnitStrides(rewriter, 2));
+      publishGraphBatchPhysicalFragment(rewriter, loc, piece, args.outputs.front(), args.lane);
     });
   if (failed(batchOp))
     return failure();
@@ -506,10 +503,7 @@ static FailureOr<spatial::SpatComputeBatch> createBatchedVvdmulBatch(Value a,
       Value bVector = extractDynamicBatchedBColumn(
         args.inputs[1], bBatchShape, outputBatchShape, batch, column, vectorType, rewriter, loc);
       Value scalar = spatial::SpatVVDMulOp::create(rewriter, loc, scalarType, aVector, bVector).getResult();
-      SmallVector<OpFoldResult> outputOffsets {args.lane, rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-      createParallelInsertSliceIntoBatchOutput(
-        rewriter, loc, scalar, args.outputs.front(), outputOffsets, scalarSizes, getUnitStrides(rewriter, 2));
+      publishGraphBatchPhysicalFragment(rewriter, loc, scalar, args.outputs.front(), args.lane);
     });
   if (failed(batchOp))
     return failure();
@@ -548,14 +542,13 @@ static FailureOr<Value> createBatchedDynamicOutputCompute(Value scalarPieces,
           Value batchLane = affineModConst(rewriter, nestedLoc, lane, numOutRows * numOutCols, anchorOp);
           Value row = affineFloorDivConst(rewriter, nestedLoc, batchLane, numOutCols, anchorOp);
           Value column = affineModConst(rewriter, nestedLoc, batchLane, numOutCols, anchorOp);
-          SmallVector<OpFoldResult> scalarOffsets {lane, rewriter.getIndexAttr(0)};
-          SmallVector<OpFoldResult> scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-          Value scalar = tensor::ExtractSliceOp::create(
-            rewriter, nestedLoc, scalarType, pieces, scalarOffsets, scalarSizes, getUnitStrides(rewriter, 2));
+          FailureOr<Value> scalar = extractGraphBatchPhysicalFragment(rewriter, nestedLoc, pieces, lane, scalarType);
+          if (failed(scalar))
+            return failure();
           Value expanded = tensor::ExpandShapeOp::create(rewriter,
                                                          nestedLoc,
                                                          outputScalarType,
-                                                         scalar,
+                                                         *scalar,
                                                          SmallVector<ReassociationIndices> {
                                                            {0},
                                                            {1, 2}
@@ -596,10 +589,11 @@ static Value extractBatchedReductionPiece(Value partialPiecesArg,
   Value kOffset = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), kSlice * numOutRows);
   Value batchAndHSlice = arith::AddIOp::create(rewriter, loc, batchOffset, hOffset);
   Value pieceOffset = arith::AddIOp::create(rewriter, loc, batchAndHSlice, kOffset);
-  SmallVector<OpFoldResult> offsets {pieceOffset, rewriter.getIndexAttr(0)};
-  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(crossbarSize.getValue())};
-  return tensor::ExtractSliceOp::create(
-    rewriter, loc, pieceType, partialPiecesArg, offsets, sizes, getUnitStrides(rewriter, 2));
+  SmallVector<OpFoldResult> offsets {pieceOffset, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
+  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())};
+  auto selectedType = RankedTensorType::get({numOutRows, 1, static_cast<int64_t>(crossbarSize.getValue())}, pieceType.getElementType());
+  Value selected = tensor::ExtractSliceOp::create(rewriter, loc, selectedType, partialPiecesArg, offsets, sizes, getUnitStrides(rewriter, 3));
+  return tensor::CollapseShapeOp::create(rewriter, loc, pieceType, selected, SmallVector<ReassociationIndices> {{0, 1}, {2}});
 }
 
 static Value reduceBatchedPartialPiecesForHSlice(Value partialPiecesArg,
@@ -1019,8 +1013,8 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
       if (succeeded(paddedRhs)) {
         Value paddedLhs = createPaddedInputCompute(plan.lhs, paddedLhsType, rewriter, loc);
         const int64_t laneCount = plan.batch * plan.m * numKSlices * numOutHSlices;
-        auto partialPiecesType = RankedTensorType::get({laneCount, static_cast<int64_t>(crossbarSize.getValue())},
-                                                       shapeInfo->outType.getElementType());
+        auto partialPiecesType = spatial::getGraphBatchPhysicalResultType(
+          laneCount, RankedTensorType::get({1, static_cast<int64_t>(crossbarSize.getValue())}, shapeInfo->outType.getElementType()));
         auto batchOp = createBatchedVmmBatch(paddedLhs,
                                              *paddedRhs,
                                              paddedLhsType,
@@ -1061,7 +1055,8 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
       }
     }
     const int64_t laneCount = plan.batch * plan.m * plan.n;
-    auto scalarPiecesType = RankedTensorType::get({laneCount, 1}, shapeInfo->outType.getElementType());
+    auto scalarPiecesType = spatial::getGraphBatchPhysicalResultType(
+      laneCount, RankedTensorType::get({1, 1}, shapeInfo->outType.getElementType()));
     auto batchOp = createBatchedVvdmulBatch(plan.lhs,
                                             plan.lhsBatchShape,
                                             plan.rhs,

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

@@ -139,9 +139,7 @@ static RankedTensorType getReducedSliceType(RankedTensorType inputType, ArrayRef
 }
 
 static RankedTensorType getLanePackedKeepdimsType(int64_t laneCount, RankedTensorType leafType) {
-  SmallVector<int64_t> shape(leafType.getShape().begin(), leafType.getShape().end());
-  shape.front() = laneCount;
-  return RankedTensorType::get(shape, leafType.getElementType(), leafType.getEncoding());
+  return spatial::getGraphBatchPhysicalResultType(laneCount, leafType);
 }
 
 static SmallVector<int64_t> getKeptAxes(ArrayRef<bool> reducedAxes) {
@@ -191,12 +189,9 @@ static FailureOr<Value> buildReduceMeanKeepdimsBatch(Value input,
 
   SmallVector<OpFoldResult> sliceOffsets;
   SmallVector<OpFoldResult> sliceSizes;
-  SmallVector<OpFoldResult> insertOffsets;
-  SmallVector<OpFoldResult> insertSizes(inputType.getRank(), rewriter.getIndexAttr(1));
   SmallVector<OpFoldResult> unitStrides = getUnitStrides(rewriter, inputType.getRank());
   sliceOffsets.reserve(inputType.getRank());
   sliceSizes.reserve(inputType.getRank());
-  insertOffsets.reserve(inputType.getRank());
 
   auto batchOp =
     createSpatComputeBatch(rewriter,
@@ -209,7 +204,6 @@ static FailureOr<Value> buildReduceMeanKeepdimsBatch(Value input,
                              size_t keptAxisIndex = 0;
                              sliceOffsets.clear();
                              sliceSizes.clear();
-                             insertOffsets.clear();
                              for (auto [axis, isReduced] : llvm::enumerate(reducedAxes)) {
                                if (isReduced) {
                                  sliceOffsets.push_back(rewriter.getIndexAttr(0));
@@ -224,14 +218,10 @@ static FailureOr<Value> buildReduceMeanKeepdimsBatch(Value input,
                                sliceSizes.push_back(rewriter.getIndexAttr(1));
                              }
 
-                             insertOffsets.push_back(args.lane);
-                             insertOffsets.append(inputType.getRank() - 1, rewriter.getIndexAttr(0));
-
                              Value slice = tensor::ExtractSliceOp::create(
                                rewriter, loc, sliceType, args.inputs.front(), sliceOffsets, sliceSizes, unitStrides);
                              Value reduced = spatial::SpatVAvgOp::create(rewriter, loc, leafType, slice).getResult();
-                             createParallelInsertSliceIntoBatchOutput(
-                               rewriter, loc, reduced, args.outputs.front(), insertOffsets, insertSizes, unitStrides);
+                             publishGraphBatchPhysicalFragment(rewriter, loc, reduced, args.outputs.front(), args.lane);
                            });
   if (failed(batchOp))
     return failure();
@@ -276,10 +266,11 @@ static Value buildKeepdimsFromLanePackedBatch(Value batchValue,
   if (!pendingLeadingReducedAxes.empty())
     expandCompactToKeepdims.back().append(pendingLeadingReducedAxes.begin(), pendingLeadingReducedAxes.end());
 
+  if (batchType.getNumElements() != batchType.getDimSize(0))
+    return {};
   auto reshapeCompute =
     createSpatCompute<1>(rewriter, loc, TypeRange {keepdimsType}, {}, ValueRange {batchValue}, [&](Value input) {
-      auto flatType =
-        RankedTensorType::get({batchType.getDimSize(0)}, batchType.getElementType(), batchType.getEncoding());
+      auto flatType = RankedTensorType::get({batchType.getNumElements()}, batchType.getElementType(), batchType.getEncoding());
       Value flat = tensor::CollapseShapeOp::create(rewriter, loc, flatType, input, collapseToFlat);
       Value compact = flat;
       if (compactKeptType != flatType)

src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp

@@ -134,6 +134,7 @@ static spatial::SpatBlueprintOp insertRowStripBlueprint(IRRewriter& rewriter, Va
                                               nullptr,
                                               nullptr,
                                               nullptr,
+                                              nullptr,
                                               nullptr);
 }