Fix lowering of onnx.Mul with dynamic shape

Xilinx · Feb 3, 2025 · 4558b04 · 4558b04
1 parent 3d6f0a2
commit 4558b04
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 1 deletion.
diff --git a/src/Conversion/ONNXToTOSA/Math/Elementwise.cpp b/src/Conversion/ONNXToTOSA/Math/Elementwise.cpp
@@ -108,6 +108,14 @@ LogicalResult checkBasicTosaRequirementsForBinaryOps(
   return success();
 }
 
+namespace {
+template <typename OnnxOp>
+void copyResultType(OnnxOp opToCopyFrom, Value &valueToCopyTo) {
+  assert(opToCopyFrom->getNumResults() == 1);
+  valueToCopyTo.setType(opToCopyFrom->getResult(0).getType());
+}
+} // namespace
+
 // Element-wise unary ops lowering to TOSA dialect.
 //===----------------------------------------------------------------------===//
 template <typename ElementwiseUnaryOpONNX, typename ElementwiseUnaryOpTOSA,
@@ -521,7 +529,7 @@ class ONNXHardSigmoidOpLoweringToTOSA
         rewriter.getF32FloatAttr(0),
         rewriter.getF32FloatAttr(oneOverAlpha.convertToDouble()));
     auto mulOp = tosaBuilder.mul(clampOp, constAlpha);
-
+    copyResultType(op, mulOp);
     rewriter.replaceOp(op, {mulOp});
     return success();
   }

diff --git a/test/mlir/conversion/onnx_to_tosa/Math/Elementwise.mlir b/test/mlir/conversion/onnx_to_tosa/Math/Elementwise.mlir
@@ -203,6 +203,16 @@ func.func @test_mul(%arg0: tensor<13x21x1xf32>, %arg1: tensor<13x21x1xf32>) -> t
 
 // -----
 
+func.func @test_mul_dynamic(%arg0: tensor<?x?x?xf32>, %arg1: tensor<13x?x?xf32>) -> tensor<13x?x?xf32> {
+  %0 = "onnx.Mul"(%arg0, %arg1) : (tensor<?x?x?xf32>, tensor<13x?x?xf32>) -> tensor<13x?x?xf32>
+  "func.return"(%0) : (tensor<13x?x?xf32>) -> ()
+// CHECK-LABEL:  func @test_mul_dynamic
+// CHECK-SAME:   ([[PARAM_0_:%.+]]: tensor<?x?x?xf32>, [[PARAM_1_:%.+]]: tensor<13x?x?xf32>) -> tensor<13x?x?xf32> {
+// CHECK-NEXT:      [[VAR_0_:%.+]] = tosa.mul [[PARAM_0_]], [[PARAM_1_]] {shift = 0 : i8} : (tensor<?x?x?xf32>, tensor<13x?x?xf32>) -> tensor<13x?x?xf32>
+}
+
+// -----
+
 func.func @test_mul_rank_broadcast(%arg0: tensor<13x21x1xf32>, %arg1: tensor<21x1xf32>) -> tensor<13x21x1xf32> {
   %0 = "onnx.Mul"(%arg0, %arg1) : (tensor<13x21x1xf32>, tensor<21x1xf32>) -> tensor<13x21x1xf32>
   "func.return"(%0) : (tensor<13x21x1xf32>) -> ()