[Mlir-commits] [mlir] [mlir][tosa] Fix tosa.transpose identity canonicalization type (PR #188020)

[Hocky Yudhiono]

https://github.com/hockyy created https://github.com/llvm/llvm-project/pull/188020

Fixes #187974. The bug is caused by an assertion to check every fold canonicalization must have the same input and folded type. This MR proposed another way to canonicalize identity `tosa.transpose` operations that has unranked tensor type.

>From b008461f13d87a6e0aa802d3fff49f9282359bef Mon Sep 17 00:00:00 2001
From: Hocky Yudhiono <hocky.yudhiono at gmail.com>
Date: Mon, 23 Mar 2026 19:39:35 +0800
Subject: [PATCH] [mlir][tosa] Fix tosa.transpose identity canonicalization
 type

---
 mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp |  3 +++
 mlir/test/Dialect/Tosa/canonicalize.mlir           | 12 ++++++++++++
 2 files changed, 15 insertions(+)

diff --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index b622cbedec1dc..7c990b54e8119 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -1920,6 +1920,9 @@ OpFoldResult TransposeOp::fold(FoldAdaptor adaptor) {
   if (!llvm::equal(llvm::seq<int32_t>(0, perms.size()), perms))
     return {};
 
+  if (resultTy != getInput1().getType())
+    return {};
+
   return getInput1();
 }
 
diff --git a/mlir/test/Dialect/Tosa/canonicalize.mlir b/mlir/test/Dialect/Tosa/canonicalize.mlir
index 52098413f18d9..4f5aaa776e57e 100644
--- a/mlir/test/Dialect/Tosa/canonicalize.mlir
+++ b/mlir/test/Dialect/Tosa/canonicalize.mlir
@@ -1388,3 +1388,15 @@ func.func @test_canonicalize_narrowing_cast_i32_to_i8_to_i16(%arg0: tensor<13x21
   %1 = tosa.cast %0 : (tensor<13x21x3xi8>) -> tensor<13x21x3xi16>
   return %1 : tensor<13x21x3xi16>
 }
+
+// -----
+
+// CHECK-LABEL: test_canonicalize_different_type_identity_transpose(
+// CHECK: tosa.const_shape
+// CHECK: %[[RESHAPE:.+]] = tosa.reshape
+// CHECK-SAME: -> tensor<*xi32>
+// CHECK: return %[[RESHAPE]]
+func.func @test_canonicalize_different_type_identity_transpose(%arg0: tensor<3x2xi32>) -> tensor<*xi32> {
+  %0 = tosa.transpose %arg0 {perms = array<i32: 0, 1>} : (tensor<3x2xi32>) -> tensor<*xi32>
+  return %0 : tensor<*xi32>
+}