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();
     });