[Mlir-commits] [mlir] 07a97fd - [mlir][tosa] Fix tosa-to-linalg assertion on rescale with an output wider than 32 bits (#233508)

[llvmlistbot at llvm.org]

Author: Stanislav Bardyuk
Date: 2026-10-20T08:20:52+01:00
New Revision: 07a97fd6b9130d77aa0faee58f1d96ac721567f8

URL: https://github.com/llvm/llvm-project/commit/07a97fd6b9130d77aa0faee58f1d96ac721567f8
DIFF: https://github.com/llvm/llvm-project/commit/07a97fd6b9130d77aa0faee58f1d96ac721567f8.diff

LOG: [mlir][tosa] Fix tosa-to-linalg assertion on rescale with an output wider than 32 bits (#233508)

The rescale lowering does its math in i32 and then saturates to the
output type, so it only works for outputs up to 32 bits. That assumption
was only checked with an assert inside the linalg.generic body builder.
The op verifier accepts any integer output type, so a tosa.rescale with
an i64 (or i48) result passes verification and then crashes the pass:

```
TosaToLinalg.cpp:1952: Assertion `outBitWidth <= 32 && "Unexpected output zeropoint bitwidth"' failed.
```

Check the output element type up front, next to the other early checks
in RescaleConverter, and decline to match when it is not an integer type
of at most 32 bits. The pass then reports "failed to legalize operation
'tosa.rescale'" like it already does for the other unsupported
configurations. Inputs wider than 32 bits keep lowering as before.

Fixes #224020

Added: 
    

Modified: 
    mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
    mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
index 25b4b360723b2..867ebbd6ce5ef 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
@@ -1825,6 +1825,12 @@ class RescaleConverter : public OpRewritePattern<tosa::RescaleOp> {
     if (!isa<IntegerType>(inputTy.getElementType()))
       return rewriter.notifyMatchFailure(op, "only support integer type");
 
+    // The generic body computes in i32 and saturates to the output type.
+    auto outputElemTy = dyn_cast<IntegerType>(outputTy.getElementType());
+    if (!outputElemTy || outputElemTy.getWidth() > 32)
+      return rewriter.notifyMatchFailure(
+          op, "only support integer output types up to 32 bits");
+
     SmallVector<Value> dynDims;
     for (int i = 0; i < outputTy.getRank(); i++) {
       if (outputTy.isDynamicDim(i)) {

diff  --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir
index 57116f03ad34f..3a0b6b87eb1df 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir
@@ -44,6 +44,32 @@ func.func @rescale_unsupported_type(%arg0: tensor<13x21x3x!quant.uniform<u8:f32,
 
 // -----
 
+// CHECK-LABEL: @rescale_unsupported_output_width
+func.func @rescale_unsupported_output_width(%arg0: tensor<2xi32>) -> tensor<2xi64> {
+  %multiplier = "tosa.const"() <{values = dense<1073741824> : tensor<1xi32> }> : () -> tensor<1xi32>
+  %shift = "tosa.const"() <{values = dense<30> : tensor<1xi8> }> : () -> tensor<1xi8>
+  %input_zp = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
+  %output_zp = "tosa.const"() <{values = dense<0> : tensor<1xi64>}> : () -> tensor<1xi64>
+  // expected-error at +1 {{failed to legalize operation 'tosa.rescale'}}
+  %0 = tosa.rescale %arg0, %multiplier, %shift, %input_zp, %output_zp scale32(true) rounding_mode<SINGLE_ROUND> per_channel(false) input_unsigned(false) output_unsigned(false) : (tensor<2xi32>, tensor<1xi32>, tensor<1xi8>, tensor<1xi32>, tensor<1xi64>) -> tensor<2xi64>
+  return %0 : tensor<2xi64>
+}
+
+// -----
+
+// CHECK-LABEL: @rescale_unsupported_i48_output
+func.func @rescale_unsupported_i48_output(%arg0: tensor<2xi8>) -> tensor<2xi48> {
+  %multiplier = "tosa.const"() <{values = dense<1073741824> : tensor<1xi32> }> : () -> tensor<1xi32>
+  %shift = "tosa.const"() <{values = dense<30> : tensor<1xi8> }> : () -> tensor<1xi8>
+  %input_zp = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+  %output_zp = "tosa.const"() <{values = dense<0> : tensor<1xi48>}> : () -> tensor<1xi48>
+  // expected-error at +1 {{failed to legalize operation 'tosa.rescale'}}
+  %0 = tosa.rescale %arg0, %multiplier, %shift, %input_zp, %output_zp scale32(true) rounding_mode<SINGLE_ROUND> per_channel(false) input_unsigned(false) output_unsigned(false) : (tensor<2xi8>, tensor<1xi32>, tensor<1xi8>, tensor<1xi8>, tensor<1xi48>) -> tensor<2xi48>
+  return %0 : tensor<2xi48>
+}
+
+// -----
+
 func.func @test_add_2d_
diff erent_ranks(%arg0: tensor<3x4xf32>, %arg1: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> {
   // expected-error at +1 {{'tosa.add' op operands don't have matching ranks}}
   %0 = "tosa.add"(%arg0, %arg1) : (tensor<3x4xf32>, tensor<2x3x4xf32>) -> tensor<2x3x4xf32>