Conv add const where the constant is a scalar (#3145)

Signed-off-by: Alexandre Eichenberger <[email protected]>

Author: AlexandreEichenberger

src/Dialect/ONNX/ONNXOps/Canonicalize.td

@@ -62,19 +62,19 @@ def subtractOrNeg: NativeCodeCall<
 def getRankOf :
 	NativeCodeCall<"mlir::cast<ShapedType>($0.getType()).getRank()">;
 
-// Create an ArrayAttr of IntergerAttr(s) of [$0].
+// Create an ArrayAttr of IntegerAttr(s) of [$0].
 def createDenseElementsAttrOf : NativeCodeCall<
   "onnx_mlir::createDenseElementsAttrOfNToM($_builder, $0, $0)">;
 
-// Create an ArrayAttr of IntergerAttr(s) of values in [1, N-1].
+// Create an ArrayAttr of IntegerAttr(s) of values in [1, N-1].
 def createDenseElementsAttrOfOneToRankOf : NativeCodeCall<
   "onnx_mlir::createDenseElementsAttrOfNToM($_builder, 1, mlir::cast<ShapedType>($0.getType()).getRank() - 1)">;
 
-// Create an ArrayAttr of IntergerAttr(s) of values in [1, N-2].
+// Create an ArrayAttr of IntegerAttr(s) of values in [1, N-2].
 def createDenseElementsAttrOfOneToRankOfExclusive : NativeCodeCall<
   "onnx_mlir::createDenseElementsAttrOfNToM($_builder, 1, mlir::cast<ShapedType>($0.getType()).getRank() - 2)">;
 
-// Create an ArrayAttr of IntergerAttr(s) of values in [2, rank - 1].
+// Create an ArrayAttr of IntegerAttr(s) of values in [2, rank - 1].
 def createArrayAttrOfTwoToRankOf : NativeCodeCall<
   "onnx_mlir::createArrayAttrOfNToM($_builder, 2, mlir::cast<ShapedType>($0.getType()).getRank() - 1)">;
 
@@ -167,7 +167,7 @@ def HaveSameElementType : Constraint<
 def HaveSameElementTypeBitWidth: Constraint<
     CPred<"(mlir::dyn_cast<ShapedType>($0.getType()).getElementTypeBitWidth() == "
           "mlir::dyn_cast<ShapedType>($1.getType()).getElementTypeBitWidth())">,
-    "has same element type bitwidth">;
+    "has same element type bit-width">;
 
 def  ElementTypeIsNotUnsigned: Constraint<
     CPred<"!mlir::dyn_cast<ShapedType>($_self.getType()).getElementType().isUnsignedInteger()">,
@@ -334,8 +334,10 @@ def FuseAddConvNullBiasPattern: Pat<
   [(HasShapeAndRank:$res),
    (HasNoneType $b),
    (AttributeIsNotNull:$denseAttr),
+   (RankXMinusRankYIs<1> $res, $y),
+   (HasRankGT<0> $y),
    (AllDimsFromAxisToEndAre<1, 1>:$y),
-   (RankXMinusRankYIs<1> $res, $y)]
+   ]
 >;
 
 def FuseAddConvPattern: Pat<
@@ -356,8 +358,9 @@ def FuseAddConvPattern: Pat<
   [(HasShapeAndRank:$res),
    (NotNoneType $b),
    (AttributeIsNotNull:$denseAttr),
-   (AllDimsFromAxisToEndAre<1, 1>:$y),
-   (RankXMinusRankYIs<1> $res, $y)]
+   (RankXMinusRankYIs<1> $res, $y),
+   (HasRankGT<0> $y),
+   (AllDimsFromAxisToEndAre<1, 1>:$y)]
 >;
 
 //===----------------------------------------------------------------------===//
@@ -403,10 +406,11 @@ def FuseMulConvNullBiasPattern: Pat<
    (HasRankGT<1> $w),                  // rank of $w must be at least 2.
    (RankXMinusRankYIs<1> $w, $y),      // rank($y) must be equal to rank($w)-1.
    (HaveSameDim<0> $w, $y),            // the first dimension of $w and $y must be equal.
+   (HasRankGT<0> $y),                  // constant cannot be a scalar.
    (AllDimsFromAxisToEndAre<1, 1>:$y)] // all dimensions of $y must be 1 except for the first one.
 >;
 
-// TODO add pattern for non-null bias with contraints:
+// TODO add pattern for non-null bias with constraints:
 // - bias must be have rank equal to 1 and
 // - bias element data type must be the same as mul constant
 // - bias dimension (0) must be equal to mul constant dim(0)
@@ -904,7 +908,7 @@ def RewriteBatchNormInferenceModeConvPattern1: Pat<
 
 // Special case of BatchNorm whose input shape is [N]. In this case, 'scale',
 // 'bias', 'mean', and 'var' will have shape of [1], according to ONNXBatchNorm
-// decription: https://github.com/onnx/onnx/blob/main/docs/Operators.md#inputs-12.
+// description: https://github.com/onnx/onnx/blob/main/docs/Operators.md#inputs-12.
 // Thus, we need not unsqueeze intermediate results.
 def RewriteBatchNormInferenceModeConvPattern2: Pat<
   (ONNXBatchNormalizationInferenceModeOp:$res
@@ -1089,7 +1093,7 @@ def ShapeTransformComposePattern : Pat<
 
 // In this pattern, the condition in onnx.Where is always false, so we can replace
 // onnx.Where by its "false" value.
-// Condition in this pattern is a comparision between dimension sizes and negative values.
+// Condition in this pattern is a comparison between dimension sizes and negative values.
 // Since dimension sizes are always positive, the condition is evaluated to false.
 
 // This pattern was found in xlm-roberta-base-language-detection model in HuggingFace.

test/mlir/onnx/onnx_canonicalization.mlir

@@ -887,6 +887,30 @@ func.func @test_fuse_add_conv_bias_unranked(%arg0 : tensor<*xf32>, %arg1 : tenso
 
 // -----
 
+// A bug was discovered when the constant being added was a scalar. This test
+// ensures that the compiler does not crash is such cases. Note that the fusion
+// does not occur, as we would need to first expand the constant to the right shape.
+
+func.func @test_fuse_add_conv_with_scalar_const(%arg0 : tensor<1x1x28x28xf32>, %arg1 : tensor<8x1x5x5xf32>) -> tensor<1x8x28x28xf32> {
+    %cst = "onnx.NoValue"() {value} : () -> none
+    %0 = "onnx.Conv"(%arg0, %arg1, %cst) {auto_pad = "SAME_UPPER", dilations = [1, 1], group = 1 : si64, kernel_shape = [5, 5], onnx_node_name = "Convolution28", strides = [1, 1]} : (tensor<1x1x28x28xf32>, tensor<8x1x5x5xf32>, none) -> tensor<1x8x28x28xf32>
+    %1 = "onnx.Constant"() {value = dense<2.0> : tensor<f32>} : () -> tensor<f32>
+    %2 = "onnx.Add"(%0, %1) : (tensor<1x8x28x28xf32>, tensor<f32>) -> tensor<1x8x28x28xf32>
+    onnx.Return %2 : tensor<1x8x28x28xf32>
+
+// mlir2FileCheck.py
+// CHECK-LABEL:  func.func @test_fuse_add_conv_with_scalar_const
+// CHECK-SAME:   ([[PARAM_0_:%.+]]: tensor<1x1x28x28xf32>, [[PARAM_1_:%.+]]: tensor<8x1x5x5xf32>) -> tensor<1x8x28x28xf32> {
+// CHECK-DAG:       [[VAR_0_:%.+]] = onnx.Constant dense<2.000000e+00> : tensor<f32>
+// CHECK-DAG:       [[VAR_1_:%.+]] = "onnx.NoValue"() {value} : () -> none
+// CHECK:           [[VAR_2_:%.+]] = "onnx.Conv"([[PARAM_0_]], [[PARAM_1_]], [[VAR_1_]]) {auto_pad = "SAME_UPPER", dilations = [1, 1], group = 1 : si64, kernel_shape = [5, 5], onnx_node_name = "Convolution28", strides = [1, 1]} : (tensor<1x1x28x28xf32>, tensor<8x1x5x5xf32>, none) -> tensor<1x8x28x28xf32>
+// CHECK:           [[VAR_3_:%.+]] = "onnx.Add"([[VAR_2_]], [[VAR_0_]]) : (tensor<1x8x28x28xf32>, tensor<f32>) -> tensor<1x8x28x28xf32>
+// CHECK:           onnx.Return [[VAR_3_]] : tensor<1x8x28x28xf32>
+// CHECK:         }
+}
+
+// -----
+
 func.func @test_fuse_mul_conv(%arg0: tensor<1x1x28x28xf32>) -> tensor<*xf32> {
     %0 = onnx.Constant dense<[[[[0.0234164055, 0.0228030644], [2.442580e-02, 0.0237577036]]], [[[-0.0410864502, 0.0488203131], [0.164448678, -0.0200194642]]], [[[-4.34581793E-9, 0.025325032], [0.0373019315, 0.165243402]]], [[[-0.0198689923, 0.131284416], [0.0572107285, 2.33985098E-8]]], [[[0.0187684372, -0.148515195], [0.0154875498, 0.019133633]]], [[[0.0176953916, -0.0154658081], [0.0233727545, -0.274110436]]], [[[-0.021181887, 0.0936150252], [0.135688141, -0.0202601217]]], [[[-0.0201558527, 0.0192655921], [0.227748245, -0.196346223]]]]> : tensor<8x1x2x2xf32>
     %1 = "onnx.NoValue"() {value} : () -> none