fix: no fold for mul for zero const if output shape is dynamic

Author: ehsan-toosi

mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp

@@ -718,7 +718,7 @@ OpFoldResult MulOp::fold(FoldAdaptor adaptor) {
   auto lhsTy = llvm::dyn_cast<RankedTensorType>(lhs.getType());
   auto rhsTy = llvm::dyn_cast<RankedTensorType>(rhs.getType());
   auto resultTy = llvm::dyn_cast<RankedTensorType>(getType());
-  if (!lhsTy || !rhsTy || !resultTy)
+  if (!lhsTy || !rhsTy || !resultTy || !resultTy.hasStaticShape())
     return {};
 
   auto resultETy = resultTy.getElementType();

mlir/test/Dialect/Tosa/canonicalize.mlir

@@ -301,6 +301,16 @@ func.func @mul_zero_broadcast(%arg0: tensor<2x3xf32>) -> (tensor<2x3xf32>, tenso
   return %1, %2 : tensor<2x3xf32>, tensor<2x3xf32>
 }
 
+// CHECK-LABEL: @mul_zero_broadcast_dynamic_result
+func.func @mul_zero_broadcast_dynamic_result(%arg0: tensor<?x3xf32>) -> (tensor<?x3xf32>, tensor<?x3xf32>) {
+  // CHECK: tosa.mul
+  // CHECK: tosa.mul
+  %zeros = "tosa.const"() {value = dense<0.0> : tensor<1x1xf32>} : () -> tensor<1x1xf32>
+  %1 = tosa.mul %arg0, %zeros {shift = 0 : i8} : (tensor<?x3xf32>, tensor<1x1xf32>) -> tensor<?x3xf32>
+  %2 = tosa.mul %zeros, %arg0 {shift = 0 : i8} : (tensor<1x1xf32>, tensor<?x3xf32>) -> tensor<?x3xf32>
+  return %1, %2 : tensor<?x3xf32>, tensor<?x3xf32>
+}
+
 // CHECK-LABEL: @select_same_value
 func.func @select_same_value(%arg0: tensor<2x3xi1>, %arg1: tensor<2x3xi32>) -> tensor<2x3xi32> {
   %0 = tosa.select %arg0, %arg1, %arg1 : (tensor<2x3xi1>, tensor<2x3xi32>, tensor<2x3xi32>) -> tensor<2x3xi32>