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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-tosa

Author: Victor Perez (victor-eds)

<details>
<summary>Changes</summary>

`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.

---
Full diff: https://github.com/llvm/llvm-project/pull/212073.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp (+4-1) 
- (modified) mlir/test/Dialect/Tosa/canonicalize.mlir (+25) 


``````````diff
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index 19f35757bdd83..8937875df7fc5 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -1796,7 +1796,7 @@ OpFoldResult IntDivOp::fold(FoldAdaptor adaptor) {
   }
 
   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 @@ std::optional<APInt> mulInt(APInt lhs, APInt rhs, int32_t shift,
 
 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>();
diff --git a/mlir/test/Dialect/Tosa/canonicalize.mlir b/mlir/test/Dialect/Tosa/canonicalize.mlir
index 585589ac1405a..7c9dd260ca497 100644
--- a/mlir/test/Dialect/Tosa/canonicalize.mlir
+++ b/mlir/test/Dialect/Tosa/canonicalize.mlir
@@ -1441,6 +1441,31 @@ func.func @no_fold_mul_result_exceeds_i32() -> tensor<i32> {
 
 // -----
 
+// 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>

``````````

</details>


https://github.com/llvm/llvm-project/pull/212073