Fix to add default value to ConstantOfShape (onnx#3174)

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

Author: AlexandreEichenberger

src/Dialect/ONNX/Transforms/Decompose.cpp

@@ -607,6 +607,21 @@ Value normalizeConstantOp(
   return createONNX.constant(denseAttr);
 }
 
+ElementsAttr reshapeElementsAttrToRank0WithDefaultValue(
+    PatternRewriter &rewriter, Value shape, Attribute val) {
+  if (!val) {
+    // Default is 0.0 in float32. It is not created by default in the ONNX
+    // getValue() as the ONNX td does not define a default value. So explicitly
+    // create a dense array of 1 zero value here.
+    Type elementType = rewriter.getF32Type();
+    RankedTensorType tensorType = RankedTensorType::get({1}, elementType);
+    FloatAttr floatAttr = rewriter.getFloatAttr(elementType, 0.0);
+    val = DenseElementsAttr::get(tensorType, floatAttr.getValue());
+  }
+  return OnnxElementsAttrBuilder(shape.getContext())
+      .reshape(cast<ElementsAttr>(val), {});
+}
+
 } // namespace onnx_mlir
 
 namespace {

src/Dialect/ONNX/Transforms/Decompose.td

@@ -68,8 +68,10 @@ def createScalarDenseAttrRank0
 // Create a scalar DenseElementsAttr of tensor<dtype> from an ElementsAttr.
 // The input ElementsAttr must have only one element. Otherwise only the first
 // element is used to create the scalar DenseElementsAttr
-def ReshapeElementsAttrToRank0 : NativeCodeCall<
-    "onnx_mlir::OnnxElementsAttrBuilder($0.getContext()).reshape(cast<ElementsAttr>($0), {})">;
+// When no attribute is provided ($1 is nullptr), then a default value attribute
+// is created (float32 0.0) according to ONNX specs.
+def ReshapeElementsAttrToRank0WithDefaultValue: NativeCodeCall<
+    "onnx_mlir::reshapeElementsAttrToRank0WithDefaultValue($_builder, $0, $1)">;
 
 def ReplaceSequenceAt :  NativeCodeCall<
     "onnx_mlir::replaceSequenceAt($_builder, $_loc, $0)">;
@@ -585,7 +587,7 @@ def ConvTransposeOpPattern2: Pattern<
 
 def ConstantOfShapePattern: Pat<
   (ONNXConstantOfShapeOp:$res $shape, $value),
-  (ONNXExpandOp (ONNXConstantOpFromDenseAttr (ReshapeElementsAttrToRank0 $value)),
+  (ONNXExpandOp (ONNXConstantOpFromDenseAttr (ReshapeElementsAttrToRank0WithDefaultValue $shape, $value)),
                 $shape)
 >;