[Mlir-commits] [mlir] [mlir][tosa] Fix i64 negate lowering with constant zero points (PR #237808)

[Sayan Saha]

https://github.com/sahas3 created https://github.com/llvm/llvm-project/pull/237808

When both zero points of tosa.negate are constant, the TosaToLinalg lowering picks the narrowest intermediate type that can hold maxInt(inputBitWidth) + |inZp + outZp| + 1. For i64 inputs this sum overflows int64_t, which is undefined behavior and in practice wraps to a negative value. The bound then compares below INT16_MAX, so the lowering selects an i16 intermediate and emits an invalid `arith.extsi` from i64 to i16.

The fix is to compute the zero point sum and the bound with 128-bit APInt so that neither can overflow. i64 inputs now keep the 64-bit intermediate, as in the non-constant zero point path, and the intermediate type chosen for narrower inputs is unchanged.

Assisted By: Claude

>From d8b75e863c155209882ca64d69c03f3a0b76c6a4 Mon Sep 17 00:00:00 2001
From: Sayan Saha <sayans at mathworks.com>
Date: Sun, 1 Nov 2026 15:32:06 -0500
Subject: [PATCH] [mlir][tosa] Fix i64 negate lowering with constant zero
 points

When both zero points of tosa.negate are constant, the TosaToLinalg
lowering picks the narrowest intermediate type that can hold
maxInt(inputBitWidth) + |inZp + outZp| + 1. For i64 inputs this sum
overflows int64_t, which is undefined behavior and in practice wraps
to a negative value. The bound then compares below INT16_MAX, so the
lowering selects an i16 intermediate and emits an invalid
`arith.extsi` from i64 to i16.

The fix is to compute the zero point sum and the bound with 128-bit
APInt so that neither can overflow. i64 inputs now keep the 64-bit
intermediate, as in the non-constant zero point path, and the
intermediate type chosen for narrower inputs is unchanged.

Assisted By: Claude
---
 .../Conversion/TosaToLinalg/TosaToLinalg.cpp  | 18 +++++----
 .../TosaToLinalg/tosa-to-linalg.mlir          | 38 +++++++++++++++++++
 2 files changed, 49 insertions(+), 7 deletions(-)

diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
index 1886226b1e8ff0..f0a4dc9f806e37 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
@@ -227,22 +227,26 @@ static Value createLinalgBodyCalculationForElementwiseOp(
 
       if (hasInZp && hasOutZp) {
         // Compute the maximum value that can occur in the intermediate buffer.
-        const int64_t zpAdd = inZp + outZp;
-        const int64_t maxValue =
-            APInt::getSignedMaxValue(inputBitWidth).getSExtValue() +
-            std::abs(zpAdd) + 1;
+        // Use 128 bits so that the computation cannot overflow for 64-bit
+        // inputs and zero points.
+        const APInt zpAdd = APInt(128, inZp, /*isSigned=*/true) +
+                            APInt(128, outZp, /*isSigned=*/true);
+        const APInt maxValue =
+            APInt::getSignedMaxValue(inputBitWidth).sext(128) + zpAdd.abs() + 1;
 
         // Convert that maximum value into the maximum bitwidth needed to
         // represent it.
-        if (maxValue <= APInt::getSignedMaxValue(16).getSExtValue()) {
+        if (maxValue.sle(APInt::getSignedMaxValue(16).getSExtValue())) {
           intermediateBitWidth = 16;
-        } else if (maxValue <= APInt::getSignedMaxValue(32).getSExtValue()) {
+        } else if (maxValue.sle(APInt::getSignedMaxValue(32).getSExtValue())) {
           intermediateBitWidth = 32;
         }
 
         intermediateType = rewriter.getIntegerType(intermediateBitWidth);
         zpAddValue = arith::ConstantOp::create(
-            rewriter, loc, rewriter.getIntegerAttr(intermediateType, zpAdd));
+            rewriter, loc,
+            rewriter.getIntegerAttr(intermediateType,
+                                    zpAdd.trunc(intermediateBitWidth)));
       } else {
         intermediateType = rewriter.getIntegerType(intermediateBitWidth);
         Value arg1 = args[1];
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
index d1c39ea86658ad..78adf54f572dce 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
@@ -3033,6 +3033,44 @@ func.func @negate_i64_no_noop_cast(%arg0: tensor<4xi64>, %in_zp: tensor<1xi64>,
 
 // -----
 
+// CHECK-LABEL: @negate_i64_const_zp
+// CHECK: linalg.generic
+// CHECK: ^bb0(%[[IN:.*]]: i64, %{{.*}}: i64):
+// CHECK: %[[ZP_ADD:.*]] = arith.constant 0 : i64
+// CHECK: %[[SUB:.*]] = arith.subi %[[ZP_ADD]], %[[IN]] : i64
+// CHECK: %[[MIN:.*]] = arith.constant -9223372036854775808 : i64
+// CHECK: %[[MAX:.*]] = arith.constant 9223372036854775807 : i64
+// CHECK: %[[LOWER:.*]] = arith.maxsi %[[MIN]], %[[SUB]] : i64
+// CHECK: %[[CLAMP:.*]] = arith.minsi %[[MAX]], %[[LOWER]] : i64
+// CHECK: linalg.yield %[[CLAMP]] : i64
+func.func @negate_i64_const_zp(%arg0: tensor<1x1xi64>) -> tensor<1x1xi64> {
+  %zp = "tosa.const"() <{values = dense<0> : tensor<1xi64>}> : () -> tensor<1xi64>
+  %neg = tosa.negate %arg0, %zp, %zp : (tensor<1x1xi64>, tensor<1xi64>, tensor<1xi64>) -> tensor<1x1xi64>
+  return %neg : tensor<1x1xi64>
+}
+
+// -----
+
+// CHECK-LABEL: @negate_i16_const_zp
+// CHECK: linalg.generic
+// CHECK: ^bb0(%[[IN:.*]]: i16, %{{.*}}: i16):
+// CHECK: %[[ZP_ADD:.*]] = arith.constant 0 : i32
+// CHECK: %[[EXT:.*]] = arith.extsi %[[IN]] : i16 to i32
+// CHECK: %[[SUB:.*]] = arith.subi %[[ZP_ADD]], %[[EXT]] : i32
+// CHECK: %[[MIN:.*]] = arith.constant -32768 : i32
+// CHECK: %[[MAX:.*]] = arith.constant 32767 : i32
+// CHECK: %[[LOWER:.*]] = arith.maxsi %[[MIN]], %[[SUB]] : i32
+// CHECK: %[[CLAMP:.*]] = arith.minsi %[[MAX]], %[[LOWER]] : i32
+// CHECK: %[[TRUNC:.*]] = arith.trunci %[[CLAMP]] : i32 to i16
+// CHECK: linalg.yield %[[TRUNC]] : i16
+func.func @negate_i16_const_zp(%arg0: tensor<1x1xi16>) -> tensor<1x1xi16> {
+  %zp = "tosa.const"() <{values = dense<0> : tensor<1xi16>}> : () -> tensor<1xi16>
+  %neg = tosa.negate %arg0, %zp, %zp : (tensor<1x1xi16>, tensor<1xi16>, tensor<1xi16>) -> tensor<1x1xi16>
+  return %neg : tensor<1x1xi16>
+}
+
+// -----
+
 // A non-constant mul shift or negate zero point is a single-element tensor: it
 // does not take part in broadcasting and every iteration reads its element 0.