Adding shape inference support for Mish op

Signed-off-by: Jagadeesh V <[email protected]>

Author: jagadeeshvx

src/Dialect/ONNX/ONNXOps/Math/ElementwiseUnary.cpp

@@ -395,6 +395,15 @@ LogicalResult ONNXMeanVarianceNormalizationOp::inferShapes(
   return inferShapeForUnaryOps(this->getOperation());
 }
 
+//===----------------------------------------------------------------------===//
+// MishOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult ONNXMishOp::inferShapes(
+    std::function<void(Region &)> doShapeInference) {
+  return inferShapeForUnaryOps(this->getOperation());
+}
+
 //===----------------------------------------------------------------------===//
 // NegOp
 //===----------------------------------------------------------------------===//

src/Dialect/ONNX/ONNXOps/ShapeHelper.hpp

@@ -382,6 +382,7 @@ using ONNXLogOpShapeHelper = ONNXUnaryOpShapeHelper;
 using ONNXLogSoftmaxOpShapeHelper = ONNXUnaryOpShapeHelper;
 using ONNXLpNormalizationOpShapeHelper = ONNXUnaryOpShapeHelper;
 using ONNXMeanVarianceNormalizationOpShapeHelper = ONNXUnaryOpShapeHelper;
+using ONNXMishOpShapeHelper = ONNXUnaryOpShapeHelper;
 using ONNXNegOpShapeHelper = ONNXUnaryOpShapeHelper;
 using ONNXNotOpShapeHelper = ONNXUnaryOpShapeHelper;
 using ONNXRandomNormalLikeOpShapeHelper = ONNXUnaryOpShapeHelper;

src/Dialect/ONNX/ONNXUnsupportedOps.hpp

@@ -47,7 +47,6 @@ UNSUPPORTED_OPS(ONNXLpPoolOp)
 UNSUPPORTED_OPS(ONNXMaxPoolOp)
 UNSUPPORTED_OPS(ONNXMaxUnpoolOp)
 UNSUPPORTED_OPS(ONNXMelWeightMatrixOp)
-UNSUPPORTED_OPS(ONNXMishOp)
 UNSUPPORTED_OPS(ONNXMomentumOp)
 UNSUPPORTED_OPS(ONNXMultinomialOp)
 UNSUPPORTED_OPS(ONNXNegativeLogLikelihoodLossOp)

test/mlir/conversion/onnx_to_krnl/Math/Elementwise_with_canonicalize.mlir

@@ -1662,9 +1662,9 @@ func.func private @test_shrink(%arg0 : tensor<512xf32>) -> tensor<*xf32> {
 
 // -----
 
-func.func @test_mish(%arg0 : tensor<64x128xf32>) -> tensor<64x128xf32> {
-  %0 = "onnx.Mish"(%arg0) : (tensor<64x128xf32>) -> tensor<64x128xf32>
-  return %0 : tensor<64x128xf32>
+func.func @test_mish(%arg0 : tensor<64x128xf32>) -> tensor<*xf32> {
+  %0 = "onnx.Mish"(%arg0) : (tensor<64x128xf32>) -> tensor<*xf32>
+  "func.return"(%0) : (tensor<*xf32>) -> ()
 
 // CHECK-LABEL:  func.func @test_mish
 // CHECK-SAME:   ([[PARAM_0_:%.+]]: memref<64x128xf32>) -> memref<64x128xf32> {

test/mlir/conversion/onnx_to_krnl/Math/Elementwise_with_canonicalize_O3.mlir

@@ -2578,9 +2578,9 @@ func.func private @test_shrink(%arg0 : tensor<512xf32>) -> tensor<*xf32> {
 
 // -----
 
-func.func @test_mish(%arg0 : tensor<64x128xf32>) -> tensor<64x128xf32> {
-  %0 = "onnx.Mish"(%arg0) : (tensor<64x128xf32>) -> tensor<64x128xf32>
-  return %0 : tensor<64x128xf32>
+func.func @test_mish(%arg0 : tensor<64x128xf32>) -> tensor<*xf32> {
+  %0 = "onnx.Mish"(%arg0) : (tensor<64x128xf32>) -> tensor<*xf32>
+  "func.return"(%0) : (tensor<*xf32>) -> ()
 
 // CHECK-LABEL:  func.func @test_mish
 // CHECK-SAME:   ([[PARAM_0_:%.+]]: memref<64x128xf32>) -> memref<64x128xf32> {