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

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)