Remove restrictions on ONNXPow (#100)

* Remove restructions on ONNXPow * Test case for onnx.Pow Add a test case for onnx.Pow that shows that no conversion actually happens * Rename the file
Xilinx · Jun 6, 2024 · 8de2a7c · 8de2a7c
1 parent 7166a41
commit 8de2a7c
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 9 deletions.
diff --git a/src/Dialect/ONNX/ONNXOps/Math/ElementwiseBroadcast.cpp b/src/Dialect/ONNX/ONNXOps/Math/ElementwiseBroadcast.cpp
@@ -343,15 +343,7 @@ LogicalResult ONNXOrOp::inferShapes(
 //===----------------------------------------------------------------------===//
 
 LogicalResult ONNXPowOp::verify() {
-  ShapedType lhsTy = getX().getType().cast<ShapedType>();
-  ShapedType rhsTy = getY().getType().cast<ShapedType>();
-  Type rhsETy = rhsTy.getElementType();
-  Type lhsETy = lhsTy.getElementType();
-  if (rhsETy != lhsETy)
-    return emitOpError("Pow with different input type not implemented yet");
-  if (lhsETy.isa<IntegerType>() || lhsETy.isa<IntegerType>())
-    return emitOpError("Integer power not implemented yet");
-  return success();
+  return verifyShapeForBroadcastingOps(getOperation());
 }
 
 LogicalResult ONNXPowOp::inferShapes(

diff --git a/test/mlir/conversion/onnx_to_tosa/Math/Pow-unsupported.mlir b/test/mlir/conversion/onnx_to_tosa/Math/Pow-unsupported.mlir
@@ -0,0 +1,16 @@
+// RUN: onnx-mlir-opt --shape-inference --convert-onnx-to-tosa %s -split-input-file | FileCheck %s
+
+// onnx.Pow with integer exponent is not supported in TOSA.
+// This test checks that the conversion does not fail but, instead, it keeps the
+// original version of the op.
+
+func.func @test_less_broadcast(%arg0: tensor<5xi32>, %arg1: tensor<5xi32>) -> tensor<*xi32> {
+  %0 = "onnx.Pow"(%arg0, %arg1) : (tensor<5xi32>, tensor<5xi32>) -> tensor<*xi32>
+  onnx.Return %0 : tensor<*xi32>
+
+// CHECK:   test_less_broadcast
+// CHECK-SAME: (%[[ARG_0:.*]]: tensor<5xi32>, %[[ARG_1:.*]]: tensor<5xi32>) -> tensor<5xi32>
+// CHECK-NEXT: %[[VAL_0:.*]] = "onnx.Pow"(%[[ARG_0]], %[[ARG_1]]) : (tensor<5xi32>, tensor<5xi32>) -> tensor<5xi32>
+// CHECK-NEXT: onnx.Return %[[VAL_0]] : tensor<5xi32>
+
+}