blazingly faster

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

@@ -20,6 +20,7 @@
 
 #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
 #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
 #include "src/Accelerators/PIM/Common/Support/ReportUtils.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
@@ -1355,20 +1356,11 @@ static Value createWeightTile(Value packedWeights,
   SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
                                    rewriter.getIndexAttr(tiling.tileInputRows),
                                    rewriter.getIndexAttr(tiling.tileOutputChannels)};
-  auto sliceType =
-    RankedTensorType::get({1, tiling.tileInputRows, tiling.tileOutputChannels}, packedWeightType.getElementType());
-  Value slice = tensor::ExtractSliceOp::create(
-    rewriter, loc, sliceType, packedWeights, offsets, sizes, getUnitStrides(rewriter, 3));
   auto collapsedType =
     RankedTensorType::get({tiling.tileInputRows, tiling.tileOutputChannels}, packedWeightType.getElementType());
-  return tensor::CollapseShapeOp::create(rewriter,
-                                         loc,
-                                         collapsedType,
-                                         slice,
-                                         SmallVector<ReassociationIndices> {
-                                           {0, 1},
-                                           {2}
-  });
+  return extractMixedSliceOrIdentity(
+    rewriter, loc, packedWeights, collapsedType,
+    {offsets, sizes, getUnitStrides(rewriter, 3)});
 }
 
 static Value createBiasTile(
@@ -1676,8 +1668,6 @@ struct ConvGemmPlan {
   int64_t effectiveMaxParallelPixels;
   int64_t packedNumRows;
 
-  RankedTensorType im2colType;
-  RankedTensorType im2colRowType;
   RankedTensorType gemmInputRowsType;
   RankedTensorType wFlatType;
   RankedTensorType wTransType;
@@ -1718,52 +1708,6 @@ static PreparedConvInput prepareInputForIm2Col(const ConvLoweringState& state,
   return {paddedInputOp.getResult(0), paddedType};
 }
 
-static Value createPaddedRows(Value rows,
-                              RankedTensorType rowsType,
-                              int64_t paddedRows,
-                              PatternRewriter& rewriter,
-                              Location loc) {
-  if (rowsType.getDimSize(0) == paddedRows)
-    return rows;
-
-  auto paddedType =
-    RankedTensorType::get({paddedRows, rowsType.getDimSize(1)}, rowsType.getElementType(), rowsType.getEncoding());
-  return createZeroPaddedTensor(
-    rows, paddedType, {0, 0}, {paddedRows - rowsType.getDimSize(0), 0}, rewriter, loc);
-}
-
-static Value packRowsForParallelGemm(
-  Value rows, RankedTensorType rowsType, int64_t packFactor, PatternRewriter& rewriter, Location loc) {
-  if (packFactor == 1)
-    return rows;
-
-  const int64_t paddedNumRows = ceilIntegerDivide(rowsType.getDimSize(0), packFactor) * packFactor;
-  const int64_t packedNumRows = paddedNumRows / packFactor;
-  const int64_t rowWidth = rowsType.getDimSize(1);
-  auto groupedType =
-    RankedTensorType::get({packedNumRows, packFactor, rowWidth}, rowsType.getElementType(), rowsType.getEncoding());
-  auto packedType =
-    RankedTensorType::get({packedNumRows, packFactor * rowWidth}, rowsType.getElementType(), rowsType.getEncoding());
-
-  Value padded = createPaddedRows(rows, rowsType, paddedNumRows, rewriter, loc);
-  Value grouped = tensor::ExpandShapeOp::create(rewriter,
-                                                loc,
-                                                groupedType,
-                                                padded,
-                                                SmallVector<ReassociationIndices> {
-                                                  {0, 1},
-                                                  {2}
-  });
-  return tensor::CollapseShapeOp::create(rewriter,
-                                         loc,
-                                         packedType,
-                                         grouped,
-                                         SmallVector<ReassociationIndices> {
-                                           {0},
-                                           {1, 2}
-  });
-}
-
 static Value unpackRowsFromParallelGemm(Value packedRows,
                                         RankedTensorType packedRowsType,
                                         int64_t unpackedRows,
@@ -2166,8 +2110,6 @@ buildConvGemmPlan(const ConvLoweringState& state,
 
   auto elemType = state.xType.getElementType();
   auto outElemType = state.outType.getElementType();
-  plan.im2colType = RankedTensorType::get({plan.chunkNumPatches, plan.patchSize}, elemType);
-  plan.im2colRowType = RankedTensorType::get({1, plan.patchSize}, elemType);
   plan.gemmInputRowsType =
     RankedTensorType::get({plan.packedNumRows, plan.effectiveMaxParallelPixels * plan.patchSize}, elemType);
   plan.wFlatType = RankedTensorType::get({state.numChannelsOut, plan.patchSize}, state.wType.getElementType());
@@ -2185,44 +2127,99 @@ static Value createIm2colRows(const ConvLoweringState& state,
                               const ConvGemmPlan& plan,
                               PatternRewriter& rewriter,
                               Location loc) {
-  constexpr size_t numInputs = 1;
-  auto im2colComputeOp =
-    createSpatCompute<numInputs>(rewriter, loc, TypeRange {plan.gemmInputRowsType}, {}, preparedInput.value, [&](Value xArg) {
-      auto elemType = preparedInput.type.getElementType();
-      // Keep the standard im2col view of convolution, flipped so filters sit in
-      // B / crossbar columns:
-      //   A (im2col):  [numPatches, patchSize]  -- one row per output spatial position
-      //   B (weights): [patchSize, cOut]
-      //   Gemm output: [numPatches, cOut]
-      Value im2colInit = tensor::EmptyOp::create(rewriter, loc, plan.im2colType.getShape(), elemType);
+  if (plan.gemmInputRowsType.getDimSize(1) > crossbarSize.getValue()) {
+    assert(plan.effectiveMaxParallelPixels == 1 && "multi-crossbar im2col rows cannot pack pixels");
+    auto compute = createSpatCompute<1>(
+      rewriter, loc, TypeRange {plan.gemmInputRowsType}, {}, preparedInput.value, [&](Value input) {
+        auto elemType = preparedInput.type.getElementType();
+        Value empty = tensor::EmptyOp::create(rewriter, loc, plan.gemmInputRowsType.getShape(), elemType);
+        Operation *anchor = rewriter.getInsertionBlock()->getParentOp();
+        Value c0 = getOrCreateIndexConstant(rewriter, anchor, 0);
+        Value c1 = getOrCreateIndexConstant(rewriter, anchor, 1);
+        Value upper = getOrCreateIndexConstant(rewriter, anchor, plan.chunkNumPatches);
+        auto patchType = RankedTensorType::get(
+          {1, state.numChannelsIn, state.wHeight, state.wWidth}, elemType);
+        auto rowType = RankedTensorType::get({plan.patchSize}, elemType);
+        auto loop = buildNormalizedScfFor(
+          rewriter, loc, c0, upper, c1, ValueRange {empty},
+          [&](OpBuilder &, Location nestedLoc, Value patchIndex, ValueRange iterArgs,
+              SmallVectorImpl<Value> &yielded) {
+            Value batchIndex = affineAddFloorDivConst(
+              rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchor);
+            Value batchPatchIndex = affineAddModConst(
+              rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchor);
+            Value outHeight = affineFloorDivConst(
+              rewriter, nestedLoc, batchPatchIndex, state.outWidth, anchor);
+            Value outWidth = affineModConst(
+              rewriter, nestedLoc, batchPatchIndex, state.outWidth, anchor);
+            Value patch = createConvInputPatch(
+              input, patchType, batchIndex, c0,
+              affineMulConst(rewriter, nestedLoc, outHeight, state.strideHeight, anchor),
+              affineMulConst(rewriter, nestedLoc, outWidth, state.strideWidth, anchor),
+              state.dilationHeight, state.dilationWidth, rewriter, nestedLoc);
+            Value row = tensor::CollapseShapeOp::create(
+              rewriter, nestedLoc, rowType, patch,
+              SmallVector<ReassociationIndices> {{0, 1, 2, 3}});
+            Value next = tensor::InsertSliceOp::create(
+              rewriter, nestedLoc, row, iterArgs.front(),
+              SmallVector<OpFoldResult> {patchIndex, rewriter.getIndexAttr(0)},
+              SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(plan.patchSize)},
+              getUnitStrides(rewriter, 2));
+            yielded.push_back(next);
+            return success();
+          });
+        if (failed(loop))
+          return failure();
+        spatial::SpatYieldOp::create(rewriter, loc, loop->results.front());
+        return success();
+      });
+    assert(succeeded(compute) && "Conv im2col compute construction must succeed");
+    return compute->getResult(0);
+  }
+
+  auto elemType = preparedInput.type.getElementType();
+  auto packedRowType = RankedTensorType::get(
+    {plan.effectiveMaxParallelPixels * plan.patchSize}, elemType, plan.gemmInputRowsType.getEncoding());
+  auto zeroAttr = DenseElementsAttr::get(packedRowType, rewriter.getZeroAttr(elemType));
+  Value zeroRow = getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), zeroAttr, packedRowType);
+  auto im2colComputeOp = createSpatComputeBatch(
+    rewriter,
+    loc,
+    TypeRange {plan.gemmInputRowsType},
+    plan.packedNumRows,
+    {},
+    ValueRange {preparedInput.value, zeroRow},
+    [&](detail::SpatComputeBatchBodyArgs args) {
       Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
       Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
       Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
+      Value cPack = getOrCreateIndexConstant(rewriter, anchorOp, plan.effectiveMaxParallelPixels);
       Value cNumPatches = getOrCreateIndexConstant(rewriter, anchorOp, plan.chunkNumPatches);
+      Value laneStart = affineMulConst(rewriter, loc, args.lane, plan.effectiveMaxParallelPixels, anchorOp);
+      Value remaining = arith::SubIOp::create(rewriter, loc, cNumPatches, laneStart);
+      Value isPartial = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::ult, remaining, cPack);
+      Value lanePatches = arith::SelectOp::create(rewriter, loc, isPartial, remaining, cPack);
+      auto patchType = RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.wWidth}, elemType);
+      auto patchRowType = RankedTensorType::get({plan.patchSize}, elemType);
 
-      auto im2colLoop = buildNormalizedScfFor(
+      auto rowLoop = buildNormalizedScfFor(
         rewriter,
         loc,
         c0,
-        cNumPatches,
+        lanePatches,
         c1,
-        ValueRange {im2colInit},
-        [&](OpBuilder&, Location nestedLoc, Value patchIndex, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
-          Value im2colAcc = iterArgs.front();
+        ValueRange {args.inputs[1]},
+        [&](OpBuilder&, Location nestedLoc, Value copyIndex, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
+          Value patchIndex = arith::AddIOp::create(rewriter, nestedLoc, laneStart, copyIndex);
           Value batchIndex =
             affineAddFloorDivConst(rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchorOp);
           Value batchPatchIndex =
             affineAddModConst(rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchorOp);
           Value outHeightIndex = affineFloorDivConst(rewriter, nestedLoc, batchPatchIndex, state.outWidth, anchorOp);
           Value outWidthIndex = affineModConst(rewriter, nestedLoc, batchPatchIndex, state.outWidth, anchorOp);
-          Value inputHeightOffset =
-            affineMulConst(rewriter, nestedLoc, outHeightIndex, state.strideHeight, anchorOp);
-          Value inputWidthOffset =
-            affineMulConst(rewriter, nestedLoc, outWidthIndex, state.strideWidth, anchorOp);
-
-          auto patchType =
-            RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.wWidth}, elemType);
-          Value patch = createConvInputPatch(xArg,
+          Value inputHeightOffset = affineMulConst(rewriter, nestedLoc, outHeightIndex, state.strideHeight, anchorOp);
+          Value inputWidthOffset = affineMulConst(rewriter, nestedLoc, outWidthIndex, state.strideWidth, anchorOp);
+          Value patch = createConvInputPatch(args.inputs.front(),
                                              patchType,
                                              batchIndex,
                                              c0,
@@ -2232,34 +2229,27 @@ static Value createIm2colRows(const ConvLoweringState& state,
                                              state.dilationWidth,
                                              rewriter,
                                              nestedLoc);
-          Value row = tensor::CollapseShapeOp::create(rewriter,
-                                                      nestedLoc,
-                                                      plan.im2colRowType,
-                                                      patch,
-                                                      SmallVector<ReassociationIndices> {
-                                                        {0},
-                                                        {1, 2, 3}
+          Value patchRow = tensor::CollapseShapeOp::create(rewriter,
+                                                           nestedLoc,
+                                                           patchRowType,
+                                                           patch,
+                                                           SmallVector<ReassociationIndices> {
+                                                             {0, 1, 2, 3}
           });
-
-          SmallVector<OpFoldResult> rowOffsets {patchIndex, rewriter.getIndexAttr(0)};
-          SmallVector<OpFoldResult> rowSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(plan.patchSize)};
-          Value next = tensor::InsertSliceOp::create(
-            rewriter, nestedLoc, row, im2colAcc, rowOffsets, rowSizes, getUnitStrides(rewriter, 2));
+          Value rowOffset = affineMulConst(rewriter, nestedLoc, copyIndex, plan.patchSize, anchorOp);
+          Value next = tensor::InsertSliceOp::create(rewriter,
+                                                     nestedLoc,
+                                                     patchRow,
+                                                     iterArgs.front(),
+                                                     SmallVector<OpFoldResult> {rowOffset},
+                                                     SmallVector<OpFoldResult> {rewriter.getIndexAttr(plan.patchSize)},
+                                                     getUnitStrides(rewriter, 1));
           yielded.push_back(next);
           return success();
         });
-      if (failed(im2colLoop))
+      if (failed(rowLoop))
         return failure();
-
-      Value gemmInputRows = im2colLoop->results.front();
-      // Pack N old im2col rows into one longer row so one GEMM can cover N
-      // pixels in parallel. The corresponding packed weight matrix contains N
-      // block-diagonal copies of W^T, and the packed output must be unpacked
-      // back to one row per spatial patch.
-      if (plan.effectiveMaxParallelPixels != 1)
-        gemmInputRows = packRowsForParallelGemm(gemmInputRows, plan.im2colType, plan.effectiveMaxParallelPixels, rewriter, loc);
-
-      spatial::SpatYieldOp::create(rewriter, loc, gemmInputRows);
+      publishGraphBatchPhysicalFragment(rewriter, loc, rowLoop->results.front(), args.outputs.front(), args.lane);
       return success();
     });
 

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

@@ -17,6 +17,7 @@
 #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
@@ -501,9 +502,9 @@ static Value extractReductionPiece(Value partialPiecesArg,
   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), 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}});
+  return extractMixedSliceOrIdentity(
+    rewriter, loc, partialPiecesArg, pieceType,
+    {pieceOffsets, pieceSizes, unitStrides});
 }
 
 static Value reducePartialPiecesForHSlice(Value partialPiecesArg,

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

@@ -9,6 +9,7 @@
 
 #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
@@ -299,22 +300,14 @@ static Value extractBatchedATile(Value a,
                                  RankedTensorType aTileType,
                                  PatternRewriter& rewriter,
                                  Location loc) {
-  auto aSliceType = RankedTensorType::get({1, 1, aTileType.getDimSize(1)}, aTileType.getElementType());
   Value sourceBatchIndex =
     mapOutputBatchIndexToSourceBatchIndex(outputBatchIndex, sourceBatchShape, outputBatchShape, rewriter, loc);
   SmallVector<OpFoldResult> offsets {OpFoldResult(sourceBatchIndex), row, kOffset};
   SmallVector<OpFoldResult> sizes {
     rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(aTileType.getDimSize(1))};
-  auto slice =
-    tensor::ExtractSliceOp::create(rewriter, loc, aSliceType, a, offsets, sizes, getUnitStrides(rewriter, 3));
-  return tensor::CollapseShapeOp::create(rewriter,
-                                         loc,
-                                         aTileType,
-                                         slice,
-                                         SmallVector<ReassociationIndices> {
-                                           {0, 1},
-                                           {2}
-  });
+  return extractMixedSliceOrIdentity(
+    rewriter, loc, a, aTileType,
+    {offsets, sizes, getUnitStrides(rewriter, 3)});
 }
 
 static Value extractBatchedBTile(Value b,
@@ -326,24 +319,15 @@ static Value extractBatchedBTile(Value b,
                                  RankedTensorType bTileType,
                                  PatternRewriter& rewriter,
                                  Location loc) {
-  auto bSliceType =
-    RankedTensorType::get({1, bTileType.getDimSize(0), bTileType.getDimSize(1)}, bTileType.getElementType());
   Value sourceBatchIndex =
     mapOutputBatchIndexToSourceBatchIndex(outputBatchIndex, sourceBatchShape, outputBatchShape, rewriter, loc);
   SmallVector<OpFoldResult> offsets {OpFoldResult(sourceBatchIndex), kOffset, hOffset};
   SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
                                    rewriter.getIndexAttr(bTileType.getDimSize(0)),
                                    rewriter.getIndexAttr(bTileType.getDimSize(1))};
-  auto slice =
-    tensor::ExtractSliceOp::create(rewriter, loc, bSliceType, b, offsets, sizes, getUnitStrides(rewriter, 3));
-  return tensor::CollapseShapeOp::create(rewriter,
-                                         loc,
-                                         bTileType,
-                                         slice,
-                                         SmallVector<ReassociationIndices> {
-                                           {0, 1},
-                                           {2}
-  });
+  return extractMixedSliceOrIdentity(
+    rewriter, loc, b, bTileType,
+    {offsets, sizes, getUnitStrides(rewriter, 3)});
 }
 
 static Value getBatchLaneIndex(
@@ -448,22 +432,14 @@ static Value extractDynamicBatchedRowVector(Value matrix,
                                             RankedTensorType vectorType,
                                             PatternRewriter& rewriter,
                                             Location loc) {
-  auto rowSliceType = RankedTensorType::get({1, 1, vectorType.getDimSize(1)}, vectorType.getElementType());
   Value sourceBatchIndex =
     mapOutputBatchIndexToSourceBatchIndex(outputBatchIndex, sourceBatchShape, outputBatchShape, rewriter, loc);
   SmallVector<OpFoldResult> offsets {OpFoldResult(sourceBatchIndex), row, rewriter.getIndexAttr(0)};
   SmallVector<OpFoldResult> sizes {
     rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(vectorType.getDimSize(1))};
-  auto rowSlice =
-    tensor::ExtractSliceOp::create(rewriter, loc, rowSliceType, matrix, offsets, sizes, getUnitStrides(rewriter, 3));
-  return tensor::CollapseShapeOp::create(rewriter,
-                                         loc,
-                                         vectorType,
-                                         rowSlice,
-                                         SmallVector<ReassociationIndices> {
-                                           {0, 1},
-                                           {2}
-  });
+  return extractMixedSliceOrIdentity(
+    rewriter, loc, matrix, vectorType,
+    {offsets, sizes, getUnitStrides(rewriter, 3)});
 }
 
 static FailureOr<spatial::SpatComputeBatch> createBatchedVvdmulBatch(Value a,
@@ -519,7 +495,6 @@ static FailureOr<Value> createBatchedDynamicOutputCompute(Value scalarPieces,
   const int64_t numOutRows = outType.getDimSize(1);
   const int64_t numOutCols = outType.getDimSize(2);
   auto scalarType = RankedTensorType::get({1, 1}, outType.getElementType());
-  auto outputScalarType = RankedTensorType::get({1, 1, 1}, outType.getElementType());
 
   auto computeOp = createSpatCompute<1>(
     rewriter, loc, TypeRange {outType}, {}, ValueRange {scalarPieces}, [&](Value pieces) -> LogicalResult {
@@ -545,20 +520,12 @@ static FailureOr<Value> createBatchedDynamicOutputCompute(Value scalarPieces,
           FailureOr<Value> scalar = extractGraphBatchPhysicalFragment(rewriter, nestedLoc, pieces, lane, scalarType);
           if (failed(scalar))
             return failure();
-          Value expanded = tensor::ExpandShapeOp::create(rewriter,
-                                                         nestedLoc,
-                                                         outputScalarType,
-                                                         *scalar,
-                                                         SmallVector<ReassociationIndices> {
-                                                           {0},
-                                                           {1, 2}
-          });
           SmallVector<OpFoldResult> outputOffsets {batch, row, column};
           SmallVector<OpFoldResult> outputSizes = {
             rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
           Value next =
             tensor::InsertSliceOp::create(
-              rewriter, nestedLoc, expanded, outputAcc, outputOffsets, outputSizes, getUnitStrides(rewriter, 3))
+              rewriter, nestedLoc, *scalar, outputAcc, outputOffsets, outputSizes, getUnitStrides(rewriter, 3))
               .getResult();
           yielded.push_back(next);
           return success();
@@ -591,9 +558,9 @@ static Value extractBatchedReductionPiece(Value partialPiecesArg,
   Value pieceOffset = arith::AddIOp::create(rewriter, loc, batchAndHSlice, kOffset);
   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}});
+  return extractMixedSliceOrIdentity(
+    rewriter, loc, partialPiecesArg, pieceType,
+    {offsets, sizes, getUnitStrides(rewriter, 3)});
 }
 
 static Value reduceBatchedPartialPiecesForHSlice(Value partialPiecesArg,
@@ -640,8 +607,6 @@ static FailureOr<Value> createBatchedReductionCompute(Value partialPieces,
       const int64_t numOutHSlices = ceilIntegerDivide(outType.getDimSize(2), crossbarSize.getValue());
       auto pieceType = RankedTensorType::get({numOutRows, static_cast<int64_t>(crossbarSize.getValue())},
                                              partialPiecesType.getElementType());
-      auto outputSliceType = RankedTensorType::get({1, numOutRows, static_cast<int64_t>(crossbarSize.getValue())},
-                                                   partialPiecesType.getElementType());
 
       Value outputInit =
         tensor::EmptyOp::create(rewriter, loc, paddedOutType.getShape(), paddedOutType.getElementType()).getResult();
@@ -671,14 +636,6 @@ static FailureOr<Value> createBatchedReductionCompute(Value partialPieces,
               Value outputAcc = hIterArgs.front();
               Value reduced = reduceBatchedPartialPiecesForHSlice(
                 partialPiecesArg, batch, hSlice, pieceType, numKSlices, numOutHSlices, numOutRows, rewriter, hLoc);
-              Value expandedReduced = tensor::ExpandShapeOp::create(rewriter,
-                                                                    hLoc,
-                                                                    outputSliceType,
-                                                                    reduced,
-                                                                    SmallVector<ReassociationIndices> {
-                                                                      {0, 1},
-                                                                      {2}
-              });
               Value hOffset = affineMulConst(
                 rewriter, hLoc, hSlice, crossbarSize.getValue(), rewriter.getInsertionBlock()->getParentOp());
               SmallVector<OpFoldResult> outputOffsets {batch, rewriter.getIndexAttr(0), hOffset};
@@ -687,7 +644,7 @@ static FailureOr<Value> createBatchedReductionCompute(Value partialPieces,
                                                      rewriter.getIndexAttr(crossbarSize.getValue())};
               Value next =
                 tensor::InsertSliceOp::create(
-                  rewriter, hLoc, expandedReduced, outputAcc, outputOffsets, outputSizes, getUnitStrides(rewriter, 3))
+                  rewriter, hLoc, reduced, outputAcc, outputOffsets, outputSizes, getUnitStrides(rewriter, 3))
                   .getResult();
               hYielded.push_back(next);
               return success();