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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-tosa

Author: Sayan Saha (sahas3)

<details>
<summary>Changes</summary>

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

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


2 Files Affected:

- (modified) mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp (+11-7) 
- (modified) mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir (+38) 


``````````diff
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.
 

``````````

</details>


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