[mlir][tosa] Fix mul/intdiv folds crashing on dynamic-shaped results (#212073)

`tosa.mul` and `tosa.intdiv` fold constant splat operands by building a
`DenseElementsAttr` of the result type, which asserts when that type has
a dynamic shape. When the operands are constants but the result type is
dynamic, folding aborted. Guard `mulBinaryFolder` and the `IntDivOp`
splat path on a static result shape, matching the other fold paths in
these ops.

Signed-off-by: Víctor Pérez Carrasco <victor.pc.upm@gmail.com>

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

diff --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index 19f3575..8937875 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -1796,7 +1796,7 @@
   }
 
   if (rhsAttr && lhsAttr && rhsAttr.isSplat() && lhsAttr.isSplat() &&
-      llvm::isa<IntegerType>(resultETy)) {
+      llvm::isa<IntegerType>(resultETy) && resultTy.hasStaticShape()) {
     APInt l = lhsAttr.getSplatValue<APInt>();
     APInt r = rhsAttr.getSplatValue<APInt>();
     if (!r.isZero()) {
@@ -1841,6 +1841,9 @@
 
 DenseElementsAttr mulBinaryFolder(DenseElementsAttr lhs, DenseElementsAttr rhs,
                                   RankedTensorType ty, int32_t shift) {
+  // A constant result can only be built for a statically-shaped type.
+  if (!ty.hasStaticShape())
+    return {};
   if (rhs && lhs && rhs.isSplat() && lhs.isSplat()) {
     if (llvm::isa<IntegerType>(ty.getElementType())) {
       APInt l = lhs.getSplatValue<APInt>();

mlir/test/Dialect/Tosa/canonicalize.mlir

diff --git a/mlir/test/Dialect/Tosa/canonicalize.mlir b/mlir/test/Dialect/Tosa/canonicalize.mlir
index 585589a..7c9dd26 100644
--- a/mlir/test/Dialect/Tosa/canonicalize.mlir
+++ b/mlir/test/Dialect/Tosa/canonicalize.mlir
@@ -1441,6 +1441,31 @@
 
 // -----
 
+// A dynamically-shaped result cannot be a constant; folding must bail instead
+// of building a DenseElementsAttr of a non-static shape (which would assert).
+// CHECK-LABEL: @no_fold_mul_dynamic_result
+// CHECK: tosa.mul
+func.func @no_fold_mul_dynamic_result() -> tensor<?xf32> {
+    %0 = "tosa.const"() <{values = dense<2.0> : tensor<4xf32>}> : () -> tensor<4xf32>
+    %1 = "tosa.const"() <{values = dense<3.0> : tensor<4xf32>}> : () -> tensor<4xf32>
+    %2 = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+    %3 = tosa.mul %0, %1, %2 : (tensor<4xf32>, tensor<4xf32>, tensor<1xi8>) -> tensor<?xf32>
+    return %3 : tensor<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @no_fold_intdiv_dynamic_result
+// CHECK: tosa.intdiv
+func.func @no_fold_intdiv_dynamic_result() -> tensor<?xi32> {
+    %0 = "tosa.const"() <{values = dense<6> : tensor<4xi32>}> : () -> tensor<4xi32>
+    %1 = "tosa.const"() <{values = dense<2> : tensor<4xi32>}> : () -> tensor<4xi32>
+    %2 = tosa.intdiv %0, %1 : (tensor<4xi32>, tensor<4xi32>) -> tensor<?xi32>
+    return %2 : tensor<?xi32>
+}
+
+// -----
+
 // CHECK-LABEL: @test_fold_i1_to_i32_cast
 // CHECK: %[[OUT:.*]] = "tosa.const"() <{values = dense<1> : tensor<i32>}> : () -> tensor<i32>
 // CHECK: return %[[OUT]] : tensor<i32>