Merge pull request #458 from Xilinx/liangta.ConstantOfShape_crash

Fix crash for ConstantOfShape decomposition

Author: ge28boj

src/Dialect/ONNX/Transforms/Decompose.td

@@ -517,10 +517,18 @@ def ReduceSumSquareV13OpPattern2
 //         `ONNXExpandOp(ONNXConstantOp {value}, %shape)
 //===----------------------------------------------------------------------===//
 
-def ConstantOfShapePattern: Pat<
+def ConstantOfShapePattern1: Pat<
   (ONNXConstantOfShapeOp:$res $shape, $value),
   (ONNXExpandOp (ONNXConstantOpFromDenseAttr (ReshapeElementsAttrToRank0 $value)),
-                $shape)
+                $shape),
+  [(AttributeIsNotNull:$value)]
+>;
+
+def ConstantOfShapePattern2: Pat<
+  (ONNXConstantOfShapeOp:$res $shape, $value),
+  (ONNXExpandOp (ONNXConstantOpFromDenseAttr (ReshapeElementsAttrToRank0 (createDenseArrayAttrFromSingleAttr (GetZeroFloatAttr)))),
+                $shape),
+  [(AttributeIsNull:$value)]
 >;
 
 //===----------------------------------------------------------------------===//

test/mlir/onnx/onnx_decompose.mlir

@@ -572,11 +572,11 @@ func.func @test_concatfuse_2(%arg0: tensor<?x20xf32>, %arg1: tensor<?x30xf32>) -
 
 // -----
 
-func.func @test_constantofshape(%arg0: tensor<?xi64>) -> tensor<*xi32> {
+func.func @test_constantofshape_1(%arg0: tensor<?xi64>) -> tensor<*xi32> {
   %0 = onnx.ConstantOfShape(%arg0) {value = dense<1> : tensor<1xi32>} : (tensor<?xi64>) -> tensor<*xi32>
   return %0 : tensor<*xi32>
 
-// CHECK-LABEL:  func.func @test_constantofshape
+// CHECK-LABEL:  func.func @test_constantofshape_1
 // CHECK-SAME:   ([[PARAM_0_:%.+]]: tensor<?xi64>) -> tensor<*xi32> {
 // CHECK:           [[VAR_0_:%.+]] = onnx.Constant dense<1> : tensor<i32>
 // CHECK:           [[VAR_1_:%.+]] = "onnx.Expand"([[VAR_0_]], [[PARAM_0_]]) : (tensor<i32>, tensor<?xi64>) -> tensor<*xi32>
@@ -586,6 +586,20 @@ func.func @test_constantofshape(%arg0: tensor<?xi64>) -> tensor<*xi32> {
 
 // -----
 
+func.func @test_constantofshape_2(%arg0: tensor<?xi64>) -> tensor<*xi32> {
+  %0 = onnx.ConstantOfShape(%arg0) : (tensor<?xi64>) -> tensor<*xi32>
+  return %0 : tensor<*xi32>
+
+// CHECK-LABEL:  func.func @test_constantofshape_2
+// CHECK-SAME:   ([[PARAM_0_:%.+]]: tensor<?xi64>) -> tensor<*xi32> {
+// CHECK:           [[VAR_0_:%.+]] = onnx.Constant dense<0.000000e+00> : tensor<f32>
+// CHECK:           [[VAR_1_:%.+]] = "onnx.Expand"([[VAR_0_]], [[PARAM_0_]]) : (tensor<f32>, tensor<?xi64>) -> tensor<*xi32>
+// CHECK:           return [[VAR_1_]] : tensor<*xi32>
+// CHECK:         }
+}
+
+// -----
+
 func.func @test_hardswish_f32(%arg0: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
   %0 = "onnx.HardSwish"(%arg0) : (tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
   return %0 : tensor<?x?x?xf32>