[Mlir-commits] [mlir] [mlir][tosa] Allow rank-0 vector operands in tosa.apply_scale (PR #199924)

[llvmlistbot at llvm.org]

https://github.com/ziereis created https://github.com/llvm/llvm-project/pull/199924

Allow rank-0 vector operands in tosa.apply_scale.

Assisted by: Claude Code 4.7

>From 621b4e65fc2e8cb9eb8c50289ee5de9a873fc15d Mon Sep 17 00:00:00 2001
From: default <ziereis at roofline.ai>
Date: Wed, 27 May 2026 10:31:15 +0000
Subject: [PATCH] tosa_allow_0d_vectors

---
 mlir/include/mlir/Dialect/Tosa/IR/TosaTypesBase.td  |  2 +-
 mlir/test/Conversion/TosaToArith/tosa-to-arith.mlir | 10 ++++++++++
 mlir/test/Dialect/Tosa/ops.mlir                     |  7 +++++++
 3 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaTypesBase.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaTypesBase.td
index 5b5189c84f4a3..6bc4216c4b78d 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaTypesBase.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaTypesBase.td
@@ -244,7 +244,7 @@ def Tosa_MXFPScaleTensorAtLeast1D : AnyTypeOf<[
 
 class Tosa_TypeLike<list<Type> types, string description = ""> : TypeConstraint<Or<[
      AnyTypeOf<types>.predicate,
-     VectorOfNonZeroRankOf<types>.predicate,
+     VectorOfAnyRankOf<types>.predicate,
      TosaTensorOf<types>.predicate]>,
      description>;
 
diff --git a/mlir/test/Conversion/TosaToArith/tosa-to-arith.mlir b/mlir/test/Conversion/TosaToArith/tosa-to-arith.mlir
index f293138b625da..48d8d52ddc39a 100644
--- a/mlir/test/Conversion/TosaToArith/tosa-to-arith.mlir
+++ b/mlir/test/Conversion/TosaToArith/tosa-to-arith.mlir
@@ -83,6 +83,16 @@ func.func @apply_scale_test_vector(%arg0 : vector<4xi32>, %arg1 : vector<4xi32>,
 
 // -----
 
+// CHECK-LABEL: @apply_scale_test_vector_rank_0
+// SCALE: tosa.apply_scale
+func.func @apply_scale_test_vector_rank_0(%arg0 : vector<i32>, %arg1 : vector<i32>, %arg2 : vector<i8>) -> (vector<i32>) {
+  // CHECK-NOT: "tosa.apply_scale"
+  %res = tosa.apply_scale %arg0, %arg1, %arg2 {rounding_mode = DOUBLE_ROUND} : (vector<i32>, vector<i32>, vector<i8>) -> vector<i32>
+  return %res : vector<i32>
+}
+
+// -----
+
 // CHECK-LABEL: @apply_scale_test_i48
 // SCALE: tosa.apply_scale
 func.func @apply_scale_test_i48(%arg0 : i48, %arg1 : i32, %arg2 : i8) -> (i32) {
diff --git a/mlir/test/Dialect/Tosa/ops.mlir b/mlir/test/Dialect/Tosa/ops.mlir
index 5c368b3da4ff5..d98736a37afa4 100644
--- a/mlir/test/Dialect/Tosa/ops.mlir
+++ b/mlir/test/Dialect/Tosa/ops.mlir
@@ -1777,3 +1777,10 @@ func.func @test_assert_equal_shape() {
   tosa.assert_equal_shape %0, %1 {allow_broadcast = true} : (!tosa.shape<2>, !tosa.shape<2>) -> ()
   return
 }
+
+// -----
+// CHECK-LABEL: test_apply_scale_rank_0_vector
+func.func @test_apply_scale_rank_0_vector(%arg0: vector<i32>, %arg1: vector<i32>, %arg2: vector<i8>) -> vector<i32> {
+  %0 = tosa.apply_scale %arg0, %arg1, %arg2 {rounding_mode = SINGLE_ROUND} : (vector<i32>, vector<i32>, vector<i8>) -> vector<i32>
+  return %0 : vector<i32>
+}