vgg8 6.88 vs 7.89

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

@@ -2899,20 +2899,60 @@ static FailureOr<Value> createPixelMajorConvPatchRow(Value paddedWindow,
     .getResult();
 }
 
-static FailureOr<Value> createConvOutputTile(Value patchRow,
-                                             Value& partialInputScratch,
-                                             Value tileWeights,
-                                             int64_t patchSize,
-                                             int64_t numKSlices,
-                                             int64_t xbarDim,
-                                             PatternRewriter& rewriter,
-                                             Location loc) {
-  auto elementType = cast<RankedTensorType>(patchRow.getType()).getElementType();
+static FailureOr<SmallVector<Value>> createConvInputTiles(Value paddedWindow,
+                                                          const ConvLoweringState& state,
+                                                          Value outputWidth,
+                                                          Value& partialInputScratch,
+                                                          int64_t patchSize,
+                                                          int64_t numKSlices,
+                                                          int64_t xbarDim,
+                                                          PatternRewriter& rewriter,
+                                                          Location loc) {
+  auto elementType = state.xType.getElementType();
   auto paddedRowType = RankedTensorType::get({1, xbarDim}, elementType);
-  auto weightElementType = cast<RankedTensorType>(tileWeights.getType()).getElementType();
-  auto paddedWeightTileType = RankedTensorType::get({xbarDim, xbarDim}, weightElementType);
+  SmallVector<Value> inputTiles;
+  inputTiles.reserve(numKSlices);
+
+  if (state.numChannelsIn % xbarDim == 0) {
+    Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+    auto inputTileType = RankedTensorType::get(
+      {1, 1, 1, xbarDim}, elementType, state.xType.getEncoding());
+    for (int64_t kSlice = 0; kSlice < numKSlices; ++kSlice) {
+      const int64_t linearOffset = kSlice * xbarDim;
+      const int64_t kernelPixel = linearOffset / state.numChannelsIn;
+      const int64_t kernelRow = kernelPixel / state.wWidth;
+      const int64_t kernelColumn = kernelPixel % state.wWidth;
+      const int64_t channelOffset = linearOffset % state.numChannelsIn;
+      Value inputWidthOffset =
+        affineMulConst(rewriter, loc, outputWidth, state.strideWidth, anchorOp);
+      inputWidthOffset = affineAddConst(
+        rewriter, loc, inputWidthOffset, kernelColumn * state.dilationWidth, anchorOp);
+      Value inputTile = tensor::ExtractSliceOp::create(
+        rewriter,
+        loc,
+        inputTileType,
+        paddedWindow,
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(0),
+                                   rewriter.getIndexAttr(kernelRow),
+                                   inputWidthOffset,
+                                   rewriter.getIndexAttr(channelOffset)},
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
+                                   rewriter.getIndexAttr(1),
+                                   rewriter.getIndexAttr(1),
+                                   rewriter.getIndexAttr(xbarDim)},
+        getUnitStrides(rewriter, 4));
+      inputTiles.push_back(tensor::CollapseShapeOp::create(
+        rewriter, loc, paddedRowType, inputTile, SmallVector<ReassociationIndices> {{0, 1, 2}, {3}})
+                             .getResult());
+    }
+    return inputTiles;
+  }
+
+  FailureOr<Value> patchRow =
+    createPixelMajorConvPatchRow(paddedWindow, state, outputWidth, rewriter, loc);
+  if (failed(patchRow))
+    return failure();
 
-  Value tileResult;
   for (int64_t kSlice = 0; kSlice < numKSlices; ++kSlice) {
     const int64_t kOffset = kSlice * xbarDim;
     const int64_t sliceSize = std::min(xbarDim, patchSize - kOffset);
@@ -2921,7 +2961,7 @@ static FailureOr<Value> createConvOutputTile(Value patchRow,
       inputTile = extractStaticSliceOrIdentity(
         rewriter,
         loc,
-        patchRow,
+        *patchRow,
         paddedRowType,
         SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(kOffset)},
         SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)},
@@ -2934,7 +2974,7 @@ static FailureOr<Value> createConvOutputTile(Value patchRow,
       Value partial = extractStaticSliceOrIdentity(
         rewriter,
         loc,
-        patchRow,
+        *patchRow,
         partialType,
         SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(kOffset)},
         SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(sliceSize)},
@@ -2949,6 +2989,26 @@ static FailureOr<Value> createConvOutputTile(Value patchRow,
         getUnitStrides(rewriter, 2));
       inputTile = partialInputScratch;
     }
+    inputTiles.push_back(inputTile);
+  }
+  return inputTiles;
+}
+
+static FailureOr<Value> createConvOutputTile(ValueRange inputTiles,
+                                             Value tileWeights,
+                                             int64_t outputChannels,
+                                             int64_t xbarDim,
+                                             PatternRewriter& rewriter,
+                                             Location loc) {
+  auto elementType = cast<RankedTensorType>(inputTiles.front().getType()).getElementType();
+  auto paddedRowType = RankedTensorType::get({1, xbarDim}, elementType);
+  auto resultType = RankedTensorType::get({1, outputChannels}, elementType);
+  auto weightElementType = cast<RankedTensorType>(tileWeights.getType()).getElementType();
+  auto paddedWeightTileType = RankedTensorType::get({xbarDim, xbarDim}, weightElementType);
+
+  Value tileResult;
+  for (auto [kSlice, inputTile] : llvm::enumerate(inputTiles)) {
+    const int64_t kOffset = static_cast<int64_t>(kSlice) * xbarDim;
     SmallVector<OpFoldResult> bOffsets {
       rewriter.getIndexAttr(kOffset), rewriter.getIndexAttr(0)};
     SmallVector<OpFoldResult> bSizes {rewriter.getIndexAttr(xbarDim), rewriter.getIndexAttr(xbarDim)};
@@ -2956,26 +3016,32 @@ static FailureOr<Value> createConvOutputTile(Value patchRow,
       rewriter, loc, tileWeights, paddedWeightTileType, bOffsets, bSizes, getUnitStrides(rewriter, 2));
     Value piece = spatial::SpatVMMOp::create(
       rewriter, loc, paddedRowType, bTile, inputTile).getResult();
+    if (outputChannels != xbarDim)
+      piece = tensor::ExtractSliceOp::create(
+        rewriter,
+        loc,
+        resultType,
+        piece,
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(outputChannels)},
+        getUnitStrides(rewriter, 2));
     tileResult = tileResult
                    ? spatial::SpatVAddOp::create(
-                       rewriter, loc, paddedRowType, tileResult, piece).getResult()
+                       rewriter, loc, resultType, tileResult, piece).getResult()
                    : piece;
   }
   return tileResult;
 }
 
-static FailureOr<Value> createConvOutputRow(Value patchRow,
-                                            Value& partialInputScratch,
-                                            int64_t patchSize,
+static FailureOr<Value> createConvOutputRow(ValueRange inputTiles,
                                             int64_t paddedK,
                                             int64_t outputChannels,
                                             Value paddedWeights,
                                             Value bias,
-                                            int64_t numKSlices,
                                             int64_t xbarDim,
                                             PatternRewriter& rewriter,
                                             Location loc) {
-  auto elementType = cast<RankedTensorType>(patchRow.getType()).getElementType();
+  auto elementType = cast<RankedTensorType>(inputTiles.front().getType()).getElementType();
   auto rowType = RankedTensorType::get({1, outputChannels}, elementType);
   auto tileWeightsType =
     RankedTensorType::get({paddedK, xbarDim},
@@ -2995,19 +3061,10 @@ static FailureOr<Value> createConvOutputRow(Value patchRow,
 
   if (outputTileCount == 1) {
     FailureOr<Value> rowResult = createConvOutputTile(
-      patchRow, partialInputScratch, getTileWeights(0), patchSize, numKSlices, xbarDim, rewriter, loc);
+      inputTiles, getTileWeights(0), outputChannels, xbarDim, rewriter, loc);
     if (failed(rowResult))
       return failure();
     Value validRow = *rowResult;
-    if (outputChannels != xbarDim)
-      validRow = tensor::ExtractSliceOp::create(
-        rewriter,
-        loc,
-        rowType,
-        validRow,
-        SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
-        SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(outputChannels)},
-        getUnitStrides(rewriter, 2));
     if (bias)
       validRow = spatial::SpatVAddOp::create(rewriter, loc, rowType, validRow, bias).getResult();
     return validRow;
@@ -3018,7 +3075,7 @@ static FailureOr<Value> createConvOutputRow(Value patchRow,
   Value paddedOutput = tensor::EmptyOp::create(rewriter, loc, paddedOutputType.getShape(), elementType);
   for (int64_t outputTile = 0; outputTile < outputTileCount; ++outputTile) {
     FailureOr<Value> tileResult = createConvOutputTile(
-      patchRow, partialInputScratch, getTileWeights(outputTile), patchSize, numKSlices, xbarDim, rewriter, loc);
+      inputTiles, getTileWeights(outputTile), xbarDim, xbarDim, rewriter, loc);
     if (failed(tileResult))
       return failure();
     SmallVector<OpFoldResult> tileOffsets {
@@ -3107,19 +3164,20 @@ static FailureOr<Value> createOutputChannelTiledRowStripConvOutput(const ConvLow
             Value widthIndex,
             ValueRange widthIterArgs,
             SmallVectorImpl<Value>& widthYielded) {
-          FailureOr<Value> patchRow =
-            createPixelMajorConvPatchRow(*inputWindow, state, widthIndex, rewriter, widthLoc);
-          if (failed(patchRow))
-            return failure();
           Value partialInputScratch = hasPartialInputTile ? widthIterArgs[1] : Value();
-          FailureOr<Value> paddedOutputRow = createConvOutputTile(*patchRow,
-                                                                 partialInputScratch,
-                                                                 tileWeights,
-                                                                 patchSize,
-                                                                 numKSlices,
-                                                                 xbarDim,
-                                                                 rewriter,
-                                                                 widthLoc);
+          FailureOr<SmallVector<Value>> inputTiles = createConvInputTiles(*inputWindow,
+                                                                          state,
+                                                                          widthIndex,
+                                                                          partialInputScratch,
+                                                                          patchSize,
+                                                                          numKSlices,
+                                                                          xbarDim,
+                                                                          rewriter,
+                                                                          widthLoc);
+          if (failed(inputTiles))
+            return failure();
+          FailureOr<Value> paddedOutputRow =
+            createConvOutputTile(*inputTiles, tileWeights, xbarDim, xbarDim, rewriter, widthLoc);
           if (failed(paddedOutputRow))
             return failure();
           if (state.hasBias)
@@ -3221,19 +3279,23 @@ static FailureOr<Value>
         c1,
         widthLoopInit,
         [&](OpBuilder&, Location widthLoc, Value widthIndex, ValueRange widthIterArgs, SmallVectorImpl<Value>& widthYielded) {
-          FailureOr<Value> patchRow =
-            createPixelMajorConvPatchRow(*inputWindow, state, widthIndex, rewriter, widthLoc);
-          if (failed(patchRow))
-            return failure();
           Value partialInputScratch = hasPartialInputTile ? widthIterArgs[1] : Value();
-          FailureOr<Value> outputRow = createConvOutputRow(*patchRow,
-                                                           partialInputScratch,
-                                                           patchSize,
+          FailureOr<SmallVector<Value>> inputTiles = createConvInputTiles(*inputWindow,
+                                                                          state,
+                                                                          widthIndex,
+                                                                          partialInputScratch,
+                                                                          patchSize,
+                                                                          numKSlices,
+                                                                          xbarDim,
+                                                                          rewriter,
+                                                                          widthLoc);
+          if (failed(inputTiles))
+            return failure();
+          FailureOr<Value> outputRow = createConvOutputRow(*inputTiles,
                                                            paddedK,
                                                            state.numChannelsOut,
                                                            args.weights.front(),
                                                            state.hasBias ? args.inputs[1] : Value(),
-                                                           numKSlices,
                                                            xbarDim,
                                                            rewriter,
                                                            widthLoc);
@@ -3334,20 +3396,23 @@ static FailureOr<Value> createConvOutputFromPixelMajorRowStripFragments(Value ro
         c1,
         widthLoopInit,
         [&](OpBuilder&, Location widthLoc, Value widthIndex, ValueRange widthIterArgs, SmallVectorImpl<Value>& widthYielded) {
-          FailureOr<Value> patchRow =
-            createPixelMajorConvPatchRow(*inputWindow, state, widthIndex, rewriter, widthLoc);
-          if (failed(patchRow))
-            return failure();
-
           Value partialInputScratch = hasPartialInputTile ? widthIterArgs[1] : Value();
-          FailureOr<Value> outputRow = createConvOutputRow(*patchRow,
-                                                           partialInputScratch,
-                                                           patchSize,
+          FailureOr<SmallVector<Value>> inputTiles = createConvInputTiles(*inputWindow,
+                                                                          state,
+                                                                          widthIndex,
+                                                                          partialInputScratch,
+                                                                          patchSize,
+                                                                          numKSlices,
+                                                                          xbarDim,
+                                                                          rewriter,
+                                                                          widthLoc);
+          if (failed(inputTiles))
+            return failure();
+          FailureOr<Value> outputRow = createConvOutputRow(*inputTiles,
                                                            paddedK,
                                                            state.numChannelsOut,
                                                            args.weights.front(),
                                                            state.hasBias ? args.inputs[1] : Value(),
-                                                           numKSlices,
                                                            xbarDim,
                                                            rewriter,
                                                            widthLoc);
@@ -3746,7 +3811,8 @@ static FailureOr<ConvLoweringState> analyzeConvLoweringState(ONNXConvOp convOp,
   state.wWidth = state.wType.getDimSize(3);
   state.outHeight = state.outType.getDimSize(2);
   state.outWidth = state.outType.getDimSize(3);
-  state.hasBias = state.b && !isa<ONNXNoneOp>(state.b.getDefiningOp());
+  state.hasBias =
+    state.b && !isa<ONNXNoneOp>(state.b.getDefiningOp()) && !isZeroSplatHostConstant(state.b);
 
   if (state.numChannelsIn % state.group != 0) {
     convOp.emitOpError() << "requires input channels " << state.numChannelsIn << " to be divisible by group "
@@ -3872,7 +3938,7 @@ static FailureOr<ConvLoweringState> analyzeConvLoweringState(spatial::SpatConv2D
   state.wWidth = state.wType.getDimSize(3);
   state.outHeight = state.outType.getDimSize(2);
   state.outWidth = state.outType.getDimSize(3);
-  state.hasBias = static_cast<bool>(planOp.getBias());
+  state.hasBias = planOp.getBias() && !isZeroSplatHostConstant(planOp.getBias());
 
   if (state.numChannelsIn % state.group != 0 || state.numChannelsOut % state.group != 0)
     return planOp.emitOpError("requires input and output channels divisible by group"), failure();