Split convolution into explicit padding and unpaded convolution. (#82)

* Split convolution into explicit padding and unpaded convolution.

* Refactor code. Add test.

src/dialect/onnx/onnx.td

@@ -111,6 +111,7 @@ def ONNXGemmNoBiasOp: ONNX_Op<"GemmNoBias",
 
 def ONNXConvNoBiasOp:ONNX_Op<"ConvNoBias",
     [NoSideEffect, DeclareOpInterfaceMethods<ShapeInferenceOpInterface>]> {
+  let hasCanonicalizer = 1;
   let summary = "ONNX Conv operation with no Bias operand.";
   let description = [{
     "The convolution operator consumes an input tensor and a filter, and"

src/pass/onnx_combine.td

@@ -39,14 +39,14 @@ def GemmTransB : NativeCodeCall<"$_builder.getI64IntegerAttr(0)">;
 // onnx.add(onnx.matmul(%X, %Y), %Z) = onnx.Gemm(%X, %Y, %Z)
 def MulAddToGemmOptPattern : Pat<(ONNXAddOp (ONNXMatMulOp:$res $m1, $m2), $m3),
                                  (ONNXGemmOp $m1, $m2, $m3, (GemmAlpha), (GemmBeta), (GemmTransA), (GemmTransB)),
-				 [(HasOneUse $res), (HasRankOf<2> $m1), (HasRankOf<2> $m2)]>;
+                                 [(HasOneUse $res), (HasRankOf<2> $m1), (HasRankOf<2> $m2)]>;
 
 // ONNX_Op (onnx.Identity (%X)) = ONNX_Op (%X)
 def IdentityEliminationPattern : Pat<(ONNXIdentityOp $arg),
-				     (replaceWithValue $arg)>;
+                                     (replaceWithValue $arg)>;
 
 def ConstantPadPattern : Pat<(ONNXPadConstantValueOp $m1, (ONNXConstantOp:$res $v1, $v2), $m2, $m3),
                              (ONNXPadConstantValuePadOp $m1, $v2, $m2, $m3),
-			     [(HasOneUse $res)]>;
+                             [(HasOneUse $res)]>;
 
 #endif // ONNX_COMBINE

src/pass/onnx_rewrite.cpp

@@ -263,6 +263,103 @@ struct ReduceSumSquareOpPattern : public RewritePattern {
     return matchSuccess();
   };
 };
+
+//===----------------------------------------------------------------------===//
+// Rewrite:
+// %0 = onnx.ConvNoBiasOp(%D : tensor<DShape>, %K)
+//     {pads = [b0, b1, ... bK, e0, e1, ..., eK]} ->
+//         tensor<OutShape>
+//
+// as:
+// %0 = onnx.PadConstantValuePasOp(%D)
+//     {pads = [0, 0, b0, b1, ... bK, 0, 0, e0, e1, ..., eK]} ->
+//     tensor<DPaddedShape>
+// %1 = onnx.ConvNoBias(%0 : tensor<DPaddedShape>, %K) {pads = [0, ..., 0]} ->
+//     tensor<OutShape>
+//===----------------------------------------------------------------------===//
+struct SplitConvOpPattern : public RewritePattern {
+  SplitConvOpPattern(MLIRContext *context)
+      : RewritePattern(ONNXConvNoBiasOp::getOperationName(),
+                       {ONNXPadConstantValuePadOp::getOperationName(),
+                        ONNXConvNoBiasOp::getOperationName()},
+                       1, context) {}
+
+  PatternMatchResult matchAndRewrite(Operation *op,
+                                     PatternRewriter &rewriter) const override {
+    auto loc = op->getLoc();
+
+    // If convolution does not use padding then no rewrite is required.
+    ONNXConvNoBiasOp convOp = llvm::dyn_cast<ONNXConvNoBiasOp>(op);
+    auto padsAttribute = convOp.padsAttr();
+    if (!padsAttribute)
+      return matchFailure();
+
+    // If auto_pad is VALID then no padding happens and no rewrite isrequired.
+    auto autoPad = convOp.auto_pad();
+    if (autoPad == "VALID")
+      return matchFailure();
+
+    auto data = op->getOperands()[0];
+    auto inputShape = data.getType().cast<TensorType>().getShape();
+
+    // Dimensionality of the input:
+    //              inputRank
+    //      |----------------------|
+    // D : (N x C x D1 x D2 x ... DK)
+    //              |______________|
+    //                  inputDims
+    //
+    int64_t inputRank = inputShape.size();
+    int64_t inputDims = inputRank - 2;
+
+    // If all pads values are equal to zero then no rewrite is required.
+    bool allZeros = true;
+    for (auto padsValue : padsAttribute.getValue()) {
+      if (padsValue.cast<IntegerAttr>().getInt() > 0) {
+        allZeros = false;
+        break;
+      }
+    }
+
+    if (allZeros)
+      return matchFailure();
+
+    // Create padding vector for the explicit padding op attribute.
+    SmallVector<int64_t, 4> pads(2 * inputRank, 0);
+    SmallVector<int64_t, 4> outPaddedShape(inputRank, 0);
+    outPaddedShape[0] = inputShape[0];
+    outPaddedShape[1] = inputShape[1];
+    for (int i = 0; i < inputDims; ++i) {
+      int64_t beginPad =
+          padsAttribute.getValue()[i].cast<IntegerAttr>().getInt();
+      int64_t endPad =
+          padsAttribute.getValue()[inputDims + i].cast<IntegerAttr>().getInt();
+      pads[i + 2] = beginPad;
+      pads[inputRank + i + 2] = endPad;
+      outPaddedShape[i + 2] += beginPad + inputShape[i + 2] + endPad;
+    }
+
+    // Create padding operation.
+    auto inputElemType = data.getType().cast<TensorType>().getElementType();
+    ONNXPadConstantValuePadOp paddingOp =
+        rewriter.create<ONNXPadConstantValuePadOp>(
+            loc, RankedTensorType::get(outPaddedShape, inputElemType), data,
+            rewriter.getI64ArrayAttr(pads), FloatAttr::get(inputElemType, 0),
+            StringAttr::get("constant", loc->getContext()));
+
+    SmallVector<int64_t, 4> newConvPads(2 * inputDims, 0);
+    auto tensorType = (*op->result_type_begin()).cast<TensorType>();
+    ONNXConvNoBiasOp newConvOp = rewriter.create<ONNXConvNoBiasOp>(
+            loc, tensorType, paddingOp.getResult(), convOp.getOperands()[1],
+            convOp.auto_padAttr(), convOp.dilationsAttr(),
+            convOp.groupAttr(), convOp.kernel_shapeAttr(),
+            rewriter.getI64ArrayAttr(newConvPads),
+            convOp.stridesAttr());
+
+    rewriter.replaceOp(op, newConvOp.getResult());
+    return matchSuccess();
+  };
+};
 } // end anonymous namespace
 
 /// on the ONNXReduceL1Op.
@@ -293,3 +390,9 @@ void ONNXReduceSumSquareOp::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {
   results.insert<ReduceSumSquareOpPattern>(context);
 }
+
+/// on the ONNXReduceSumSquareOp.
+void ONNXConvNoBiasOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &results, MLIRContext *context) {
+  results.insert<SplitConvOpPattern>(context);
+}

test/mlir/onnx/onnx_canonicalization.mlir

@@ -92,3 +92,12 @@ func @test_constant_pad(%arg0 : tensor<?x?xf32>) -> tensor<*xf32> {
   %2 = "onnx.PadConstantValue"(%arg0, %0) {constant_value=0. : f32, mode = "constant"} : (tensor<?x?xf32>, tensor<?xi64>)-> tensor<*xf32>
   "std.return"(%2) : (tensor<*xf32>) -> ()
 }
+
+// CHECK-LABEL: @test_conv_split(%{{.*}}: tensor<1x9x32x64xf32>, %{{.*}}: tensor<5x9x6x7xf32>) -> tensor<*xf32> {
+func @test_conv_split(%arg0 : tensor<1x9x32x64xf32>, %arg1 : tensor<5x9x6x7xf32>) -> tensor<*xf32> {
+  %0 = "onnx.ConvNoBias"(%arg0, %arg1) {auto_pad = "NOTSET", group = 1 : i64, pads = [2, 3, 4, 5]} : (tensor<1x9x32x64xf32>, tensor<5x9x6x7xf32>) -> tensor<*xf32>
+  "std.return"(%0) : (tensor<*xf32>) -> ()
+  // CHECK-NEXT: %0 = "onnx.PadConstatValuePad"(%arg0) {constant_value = 0.000000e+00 : f32, mode = "constant", pads = [0, 0, 2, 3, 0, 0, 4, 5]} : (tensor<1x9x32x64xf32>) -> tensor<1x9x38x72xf32>
+  // CHECK-NEXT: %1 = "onnx.ConvNoBias"(%0, %arg1) {auto_pad = "NOTSET", group = 1 : i64, pads = [0, 0, 0, 0]} : (tensor<1x9x38x72xf32>, tensor<5x9x6x7xf32>) -> tensor<*xf32>
+  // CHECK-NEXT: return %1 : tensor<*xf32>
+}