[Mlir-commits] [mlir] [mlir][tosa] Fix reduce-transpose sub-byte handling (PR #208446)

[Iliyan Georgiev]

https://github.com/iliyan-georgiev-arm updated https://github.com/llvm/llvm-project/pull/208446

>From 8da7e598baf10207aa150110c11590eb0bfefe63 Mon Sep 17 00:00:00 2001
From: Iliyan Georgiev <Iliyan.Georgiev at arm.com>
Date: Wed, 8 Jul 2026 14:30:22 +0100
Subject: [PATCH] [mlir][tosa] Fix reduce-transpose sub-byte handling

- Fix element size calculation when processing sub-byte types
- Add tests for i1 and i4

Signed-off-by: Iliyan Georgiev <Iliyan.Georgiev at arm.com>
Change-Id: I06988256384550738becf5691d5880366293d58d
---
 .../Tosa/Transforms/TosaReduceTransposes.cpp   |  3 +--
 .../Dialect/Tosa/tosa-reduce-transposes.mlir   | 18 ++++++++++++++++++
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaReduceTransposes.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaReduceTransposes.cpp
index 0bd415040f22f..ffd48e8ab16ae 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaReduceTransposes.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaReduceTransposes.cpp
@@ -233,8 +233,7 @@ TosaReduceTransposes::transposeDenseAttribute(DenseElementsAttr input,
   // perms: dim 0→2, dim 1→0, dim 2→1, giving source position
   // calculated as 1*inputStrides[2] + 1*inputStrides[0] + 2*inputStrides[1]
   // = 1*1 + 1*12 + 2*4 = 21
-
-  size_t elementSize = oldType.getElementTypeBitWidth() / 8;
+  size_t elementSize = llvm::divideCeil(oldType.getElementTypeBitWidth(), 8);
   int64_t numElements = oldType.getNumElements();
 
   SmallVector<char> outputBuffer(numElements * elementSize);
diff --git a/mlir/test/Dialect/Tosa/tosa-reduce-transposes.mlir b/mlir/test/Dialect/Tosa/tosa-reduce-transposes.mlir
index 0ad3068011f7f..46e4f504af0ec 100644
--- a/mlir/test/Dialect/Tosa/tosa-reduce-transposes.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-reduce-transposes.mlir
@@ -597,3 +597,21 @@ func.func @test_unimplemented_static_diverges_to_one_nullifying_one_non_nullifyi
   %result = tosa.transpose %add {perms = array<i32: 0, 3, 1, 2>}: (tensor<1x3x4x2xi32>) -> tensor<1x2x3x4xi32>
   return %result : tensor<1x2x3x4xi32>
 }
+
+// CHECK-LABEL: @test_transpose_bool
+// CHECK: %{{.*}} = "tosa.const"()
+// CHECK-SAME{LITERAL}: dense<[[true, false], [false, false], [false, true]]>
+func.func @test_transpose_bool() -> tensor<3x2xi1> {
+  %0 = "tosa.const"() <{values = dense<[[true, false, false], [false, false, true]]> : tensor<2x3xi1>}> : () -> tensor<2x3xi1>
+  %1 = tosa.transpose %0 {perms = array<i32: 1, 0>} : (tensor<2x3xi1>) -> tensor<3x2xi1>
+  return %1 : tensor<3x2xi1>
+}
+
+// CHECK-LABEL: @test_transpose_i4
+// CHECK: %{{.*}} = "tosa.const"()
+// CHECK-SAME{LITERAL}: dense<[[1, 4], [2, 5], [3, 6]]>
+func.func @test_transpose_i4() -> tensor<3x2xi4> {
+  %0 = "tosa.const"() <{values = dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi4>}> : () -> tensor<2x3xi4>
+  %1 = tosa.transpose %0 {perms = array<i32: 1, 0>} : (tensor<2x3xi4>) -> tensor<3x2xi4>
+  return %1 : tensor<3x2xi4>
+}
\ No newline at end of file