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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: ziereis

<details>
<summary>Changes</summary>

I was facing a bug that can be reproduced this way:

```mlir
 // RUN:  mlir-opt --transform-interpreter tosa_apply_scale_rank0_repro.mlir
  #map        = affine_map<(d0) -> (d0)>
  #map_scalar = affine_map<(d0) -> ()>

  func.func @<!-- -->repro(%input: tensor<64xi32>, %scalar_t: tensor<i32>,
                   %out_init: tensor<64xi8>) -> tensor<64xi8> {
    %c31_i8     = arith.constant 31 : i8
    %cScale_i32 = arith.constant -1010580540 : i32

    %tile_out = linalg.generic
      { indexing_maps = [#map, #map_scalar, #map],
        iterator_types = ["parallel"] }
      ins(%input, %scalar_t : tensor<64xi32>, tensor<i32>)
      outs(%out_init : tensor<64xi8>) {
      ^bb0(%in: i32, %in_4: i32, %out: i8):
        %scaled = tosa.apply_scale %in_4, %cScale_i32, %c31_i8
                  { rounding_mode = SINGLE_ROUND } : (i32, i32, i8) -> i32
        %as_i8  = arith.trunci %scaled : i32 to i8
        linalg.yield %as_i8 : i8
    } -> tensor<64xi8>

    return %tile_out : tensor<64xi8>
  }

  module attributes {transform.with_named_sequence} {
    transform.named_sequence @<!-- -->__transform_main(%arg0: !transform.any_op {transform.readonly}) {
      %func = transform.structured.match ops{["func.func"]} in %arg0
              : (!transform.any_op) -> !transform.any_op
      %op   = transform.structured.match ops{["linalg.generic"]} in %arg0
              : (!transform.any_op) -> !transform.any_op
      transform.structured.vectorize %op vector_sizes [64] : !transform.any_op
      transform.apply_patterns to %func {
        transform.apply_patterns.vector.transfer_permutation_patterns
        transform.apply_patterns.vector.sink_ops
      } : !transform.any_op
      transform.yield
    }
  }
```

Which lead to the following error:

```mlir
  error: 'tosa.apply_scale' op operand #<!-- -->0 must be signless-integer-like, but got 'vector<i32>'
  note: see current operation: %5 = "tosa.apply_scale"(%4, %1, %0) ... : (vector<i32>, vector<i32>, vector<i8>) -> vector<i32>
```

After this patch it gives the following IR:

```mlir
  func.func @<!-- -->repro(%arg0: tensor<64xi32>, %arg1: tensor<i32>, %arg2: tensor<64xi8>) -> tensor<64xi8> {
    %cst = arith.constant dense<31> : vector<i8>
    %cst_0 = arith.constant dense<-1010580540> : vector<i32>
    %c0 = arith.constant 0 : index
    %0 = ub.poison : i32
    %1 = vector.transfer_read %arg1[], %0 : tensor<i32>, vector<i32>
    %2 = tosa.apply_scale %1, %cst_0, %cst {rounding_mode = SINGLE_ROUND} : (vector<i32>, vector<i32>, vector<i8>) -> vector<i32>
    %3 = arith.trunci %2 : vector<i32> to vector<i8>
    %4 = vector.broadcast %3 : vector<i8> to vector<64xi8>
    %5 = vector.transfer_write %4, %arg2[%c0] {in_bounds = [true]} : vector<64xi8>, tensor<64xi8>
    return %5 : tensor<64xi8>
  }
  ```
    
Claude suggested the best fix for this is to ease the "no 0 rank vectors" restriction on the tosa.apply_scale op so this is what this PR includes. I honestly can't 100% judge if this is the right way to fix this but i couldn't find anything that would speak against it.

Assisted by: Claude Code 4.7

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


3 Files Affected:

- (modified) mlir/include/mlir/Dialect/Tosa/IR/TosaTypesBase.td (+1-1) 
- (modified) mlir/test/Conversion/TosaToArith/tosa-to-arith.mlir (+10) 
- (modified) mlir/test/Dialect/Tosa/ops.mlir (+7) 


``````````diff
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>
+}

``````````

</details>


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