for stride 5, if concat of convs leading to odd value dim padding to make it even

Author: chaitanyakamarapu

src/Dialect/ONNX/Transforms/Decompose.cpp

@@ -1246,26 +1246,48 @@ Value decomposeConvT1dIntoPhasedConvs(PatternRewriter &rewriter, Location loc,
         RankedTensorType::get(outputShapeLevel1Concat, elementType);
 
     // Below concats result will have the innermost dim as 2.
-    auto convOfmConcat =
+    Value convOfmConcat =
         rewriter.create<ONNXConcatOp>(loc, level1ConcatOutputType,
             ValueRange{reshapeOutputAddOneDimConv1, reshapeOutputAddOneDimConv2,
                 reshapeOutputAddOneDimConv3, reshapeOutputAddOneDimConv4,
                 reshapeOutputAddOneDimConv5},
             -1);
-    SmallVector<int64_t> outputShapeForResult(paddedConvOutputShapeValue);
-    auto dimValueAtLastIndex =
-        paddedConvOutputShapeValue[paddedConvOutputShapeValue.size() - 1] * 5;
-    outputShapeForResult[outputShapeForResult.size() - 1] = dimValueAtLastIndex;
+    // Making the dim2 of concat even by padding one at the end.
+    bool isPaddedToMakeEven = false;
+    if (outputShapeLevel1Concat[2] % 2 != 0) {
+      SmallVector<int64_t> outputShapePadToEven(outputShapeLevel1Concat);
+      outputShapePadToEven[2] = outputShapePadToEven[2] + 1;
+      auto padToEvenOutputShapedType =
+          RankedTensorType::get(outputShapePadToEven, elementType);
+
+      std::array<int64_t, 8> padValueToEven = {0, 0, 0, 0, 0, 0, 1, 0};
+
+      auto onnxPadsToEvenValueConstant =
+          getONNXConstOpFromVector(rewriter, loc, padValueToEven);
+
+      convOfmConcat = rewriter.create<ONNXPadOp>(loc, padToEvenOutputShapedType,
+          convOfmConcat, onnxPadsToEvenValueConstant, onnxPaddingConstantZero,
+          onnxAxisValueConstantNone, rewriter.getStringAttr("constant"));
+      isPaddedToMakeEven = true;
+    }
+    // This is the shape of the five conv merge. Using [2] as this
+    // is convtranspose 1D.
+    SmallVector<int64_t> reshapeOutputShape(paddedConvOutputShapeValue);
+    reshapeOutputShape[2] =
+        (isPaddedToMakeEven ? (paddedConvOutputShapeValue[2] + 1)
+                            : paddedConvOutputShapeValue[2]) *
+        5;
+    ;
 
     auto onnxConstForLastReshape =
-        getONNXConstOpFromVector(rewriter, loc, outputShapeForResult);
+        getONNXConstOpFromVector(rewriter, loc, reshapeOutputShape);
 
-    auto outputTypeBeforeSlice =
-        RankedTensorType::get(outputShapeForResult, elementType);
+    auto reshapeResultType =
+        RankedTensorType::get(reshapeOutputShape, elementType);
     // Result is reshaped back to match the original convtranspose output
     // dimensions
-    auto outputBeforeSlice = rewriter.create<ONNXReshapeOp>(
-        loc, outputTypeBeforeSlice, convOfmConcat, onnxConstForLastReshape);
+    auto reshapeOutput = rewriter.create<ONNXReshapeOp>(
+        loc, reshapeResultType, convOfmConcat, onnxConstForLastReshape);
 
     SmallVector<int64_t> finalSliceOutputShape(convTransposeOutputShape);
     auto finalSliceOutputType = RankedTensorType::get(
@@ -1276,8 +1298,8 @@ Value decomposeConvT1dIntoPhasedConvs(PatternRewriter &rewriter, Location loc,
         {finalSliceOutputShape[finalSliceOutputShape.size() - 1]});
 
     auto finalSlicedOutput = rewriter.create<ONNXSliceOp>(loc,
-        finalSliceOutputType, outputBeforeSlice, startOnnxConstant,
-        endOnnxConstant, axisOnnxConstant, stepOnnxConstant);
+        finalSliceOutputType, reshapeOutput, startOnnxConstant, endOnnxConstant,
+        axisOnnxConstant, stepOnnxConstant);
 
     return finalSlicedOutput;
   }
@@ -1289,13 +1311,13 @@ Value decomposeConvT1dIntoPhasedConvs(PatternRewriter &rewriter, Location loc,
 // The phased convolutions are then merged to get the final output.
 // The number of phases is determined by the strides of the convtranspose op.
 // The num of phases = stride_x * stride_y.
-// The phased convolutions are weights are created by slicing the weights of the
-// convolution in the specified manner and output of convolutions are stiched
-// together to get the final output. If the case where original weights cannot
-// be sliced into conv weights directly, they are padded to make them compatible
-// with the slicing. and subsequently the extra ofm generated by the padded
-// weights are removed.
-// Below shows the high level view of the decomposition.
+// The phased convolutions are weights are created by slicing the weights of
+// the convolution in the specified manner and output of convolutions are
+// stiched together to get the final output. If the case where original
+// weights cannot be sliced into conv weights directly, they are padded to
+// make them compatible with the slicing. and subsequently the extra ofm
+// generated by the padded weights are removed. Below shows the high level
+// view of the decomposition.
 // clang-format off
 /*
  * +---------------+       +-----------+ 
@@ -1406,8 +1428,8 @@ Value decomposeIntoPhasedConvs(PatternRewriter &rewriter, Location loc,
   }
 
   onnx_mlir::MultiDialectBuilder<onnx_mlir::OnnxBuilder> create(rewriter, loc);
-  // If the convTranspose kernel is 3x3, then the weights needs to be padded to
-  // 4x4
+  // If the convTranspose kernel is 3x3, then the weights needs to be padded
+  // to 4x4
   bool needWeightsPadding = (kernelShape[0] == 3 && stridesShape[0] == 2);
   if (needWeightsPadding) {
     std::array<int64_t, 8> weightsPadValue = {0, 0, 0, 0, 0, 0, 0, 0};
@@ -1613,10 +1635,10 @@ Value decomposeIntoPhasedConvs(PatternRewriter &rewriter, Location loc,
     outputShapeLevel1Concat[outputShapeLevel1Concat.size() - 1] = 2;
     auto level1ConcatOutputType =
         RankedTensorType::get(outputShapeLevel1Concat, elementType);
-    // for the case where convtranspose kernel is [4, 4] and with pads [1, 1, 1,
-    // 1] The phased convs output are to be concatenated in the reverse order.
-    // This is observed by looking at the phased conv outputs with respect to
-    // convtranspose output.
+    // for the case where convtranspose kernel is [4, 4] and with pads [1, 1,
+    // 1, 1] The phased convs output are to be concatenated in the reverse
+    // order. This is observed by looking at the phased conv outputs with
+    // respect to convtranspose output.
     bool reverseConcatOrder = (needWeightsPadding || (kernelShape[0] == 4));
     // Below concats result will have the innermost dim as 2.
     auto firstConcat =