[Mlir-commits] [mlir] [MLIR][Tosa] Fix crash in TransposeOp::fold for unranked result type (PR #188996)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

The identity-transpose folder returned getInput1() directly when the permutation is the identity (0, 1, ..., N-1). However, the input and result types may differ — for example the result may be unranked (tensor<*xi32>) while the input is ranked (tensor<3x2xi32>). Returning a value of the wrong type triggers the "incorrect fold result type" assertion.

Fix: bail out of folding when the input type does not match the result type. The transpose_is_reshape canonicalization pattern handles the case where the permutation is an identity but types differ.

Fixes #<!-- -->187974

Assisted-by: Claude Code

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp (+6) 
- (modified) mlir/test/Dialect/Tosa/canonicalize.mlir (+13) 


``````````diff
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index 70e77839f2486..5742af4681b58 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -1969,6 +1969,12 @@ OpFoldResult TransposeOp::fold(FoldAdaptor adaptor) {
   if (!llvm::equal(llvm::seq<int32_t>(0, perms.size()), perms))
     return {};
 
+  // Only fold when the input type matches the result type. They may differ
+  // when the result type is unranked (e.g. tensor<*xi32>) while the input is
+  // ranked.
+  if (getInput1().getType() != resultTy)
+    return {};
+
   return getInput1();
 }
 
diff --git a/mlir/test/Dialect/Tosa/canonicalize.mlir b/mlir/test/Dialect/Tosa/canonicalize.mlir
index bf53f06be3e07..4b9992be8b193 100644
--- a/mlir/test/Dialect/Tosa/canonicalize.mlir
+++ b/mlir/test/Dialect/Tosa/canonicalize.mlir
@@ -850,6 +850,19 @@ func.func @transpose_no_op(%arg0: tensor<3x4x5x6xf32>) -> tensor<3x4x5x6xf32> {
 
 // -----
 
+// Do not fold identity transpose to the input when result type differs from
+// input type (e.g. unranked result type); the transpose_is_reshape pattern
+// handles it instead. (https://github.com/llvm/llvm-project/issues/187974)
+// CHECK-LABEL: @transpose_no_op_unranked_result
+func.func @transpose_no_op_unranked_result(%arg0: tensor<3x2xi32>) -> tensor<*xi32> {
+  // CHECK: tosa.reshape
+  // CHECK-NOT: return %arg0
+  %0 = tosa.transpose %arg0 {perms = array<i32: 0, 1>} : (tensor<3x2xi32>) -> tensor<*xi32>
+  return %0 : tensor<*xi32>
+}
+
+// -----
+
 // CHECK-LABEL: @transpose_is_reshape
 func.func @transpose_is_reshape(%arg0: tensor<1x4x5x1xf32>) -> tensor<1x4x1x5xf32> {
   // CHECK: %[[CONST0:.+]] = tosa.const_shape {values = dense<[1, 4, 1, 5]> : tensor<4xindex>} : () -> !tosa.shape<4>

``````````

</details>


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