[mlir][tosa] Remove NegateOp to SubOp and 48-bit promotion in TosaToLinalg (#170622)

The patch motivated by Tosa Conformance test negate_32x45x49_i16_full failure.

TosaToLinalg pass has an optimization to transfer Tosa Negate to Sub if the zero points are zeros. However, when the input value is minimum negative number, the transformation will cause the underflow. By removing the transformation, if zp = 0 it would do the promotion to avoid the underflow.

Promotion types could be from int32 to int48. TOSA negate specification does not mention support for int48. Should we consider removing the promotion to int48 to stay aligned with the TOSA spec?

Author: ShivaChen

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

@@ -200,13 +200,6 @@ static Value createLinalgBodyCalculationForElementwiseOp(
       return arith::NegFOp::create(rewriter, loc, resultTypes, args[0]);
 
     if (isa<IntegerType>(elementTy)) {
-      if (hasInZp && hasOutZp && !inZp && !outZp) {
-        auto constant = arith::ConstantOp::create(
-            rewriter, loc, IntegerAttr::get(elementTy, 0));
-        return arith::SubIOp::create(rewriter, loc, resultTypes, constant,
-                                     args[0]);
-      }
-
       Value zpAddValue;
       Type intermediateType;
       // Compute the maximum value that can occur in the intermediate buffer.
@@ -221,14 +214,11 @@ static Value createLinalgBodyCalculationForElementwiseOp(
             std::abs(zpAdd) + 1;
 
         // Convert that maximum value into the maximum bitwidth needed to
-        // represent it. We assume 48-bit numbers may be supported further in
-        // the pipeline.
+        // represent it.
         if (maxValue <= APInt::getSignedMaxValue(16).getSExtValue()) {
           intermediateBitWidth = 16;
         } else if (maxValue <= APInt::getSignedMaxValue(32).getSExtValue()) {
           intermediateBitWidth = 32;
-        } else if (maxValue <= APInt::getSignedMaxValue(48).getSExtValue()) {
-          intermediateBitWidth = 48;
         }
 
         intermediateType = rewriter.getIntegerType(intermediateBitWidth);

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir

@@ -666,7 +666,13 @@ func.func @test_simple_i32(%arg0: tensor<1xi32>, %unsigned: tensor<1xui32>, %uns
   // CHECK: linalg.generic
   // CHECK: ^bb0(%[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32):
   // CHECK: [[ZERO:%.+]] = arith.constant 0
-  // CHECK: arith.subi [[ZERO]], %[[ARG1]]
+  // CHECK: [[EXT:%.+]] = arith.extsi %[[IN:.*]] : i32 to i64
+  // CHECK: [[SUB:%.+]] = arith.subi [[ZERO]], [[EXT]]
+  // CHECK: [[MIN:%.+]] = arith.constant -2147483648
+  // CHECK: [[MAX:%.+]] = arith.constant 2147483647
+  // CHECK: [[LBOUND:%.+]] = arith.maxsi [[MIN]], [[SUB]]
+  // CHECK: [[UBOUND:%.+]] = arith.minsi [[MAX]], [[LBOUND]]
+  // CHECK: [[TRUNC:%.+]] = arith.trunci [[UBOUND]]
   %in_zp = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
   %out_zp = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
   %5 = tosa.negate %arg0, %in_zp, %out_zp : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
@@ -889,8 +895,13 @@ func.func @test_negate_quantized(%arg0: tensor<1xi8>) -> () {
   // CHECK: linalg.generic
   // CHECK: ^bb0(%[[BBARG0:.+]]: i8,
   // CHECK: [[ZERO:%.+]] = arith.constant 0
-  // CHECK: [[SUB:%.+]] = arith.subi [[ZERO]],
-  // CHECK: linalg.yield [[SUB]]
+  // CHECK: [[EXT:%.+]] = arith.extsi %[[IN:.*]] : i8 to i16
+  // CHECK: [[SUB:%.+]] = arith.subi [[ZERO]], [[EXT]]
+  // CHECK: [[MIN:%.+]] = arith.constant -128
+  // CHECK: [[MAX:%.+]] = arith.constant 127
+  // CHECK: [[LBOUND:%.+]] = arith.maxsi [[MIN]], [[SUB]]
+  // CHECK: [[UBOUND:%.+]] = arith.minsi [[MAX]], [[LBOUND]]
+  // CHECK: [[TRUNC:%.+]] = arith.trunci [[UBOUND]]
   %in_zp4 = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
   %out_zp4 = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
   %4 = tosa.negate %arg0, %in_zp4, %out_zp4 : (tensor<1xi8>, tensor<1xi8>, tensor<1xi8>) -> tensor<1xi8>
@@ -2610,7 +2621,13 @@ func.func @test_0d_input(%arg0: tensor<i32>) -> () {
   // CHECK: linalg.generic
   // CHECK: ^bb0(%[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32):
   // CHECK: [[ZERO:%.+]] = arith.constant 0
-  // CHECK: arith.subi [[ZERO]], %[[ARG1]]
+  // CHECK: [[EXT:%.+]] = arith.extsi %[[IN:.*]] : i32 to i64
+  // CHECK: [[SUB:%.+]] = arith.subi [[ZERO]], [[EXT]]
+  // CHECK: [[MIN:%.+]] = arith.constant -2147483648
+  // CHECK: [[MAX:%.+]] = arith.constant 2147483647
+  // CHECK: [[LBOUND:%.+]] = arith.maxsi [[MIN]], [[SUB]]
+  // CHECK: [[UBOUND:%.+]] = arith.minsi [[MAX]], [[LBOUND]]
+  // CHECK: [[TRUNC:%.+]] = arith.trunci [[UBOUND]]
   %in_zp = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
   %out_zp = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
   %5 = tosa.negate %arg0, %in_zp, %out_zp : (tensor<i32>, tensor<1xi32>, tensor<1xi32>) -> tensor<i32>