[mlir][tosa] Avoid overflow in reduction folders (llvm#132786)

Avoid operations that can overflow in constant folders for
`tosa.reduce_max` and `tosa.reduce_min`

Includes tests to avoid regressions

Signed-off-by: Ian Tayler Lessa <[email protected]>

Author: IanTaylerLessa-arm

mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td

@@ -1728,8 +1728,7 @@ def Tosa_ReduceMaxOp : Tosa_InferTensorTypeOp<"reduce_max"> {
 
     /// Return the max of the two integer operands
     static inline APInt calcOneElement(APInt leftOperand, APInt rightOperand) {
-      const llvm::APInt subtractRes = leftOperand - rightOperand;
-      return (!subtractRes.isNegative()) ? leftOperand : rightOperand;
+      return (leftOperand.sge(rightOperand)) ? leftOperand : rightOperand;
     }
   }];
 }
@@ -1769,8 +1768,7 @@ def Tosa_ReduceMinOp : Tosa_InferTensorTypeOp<"reduce_min"> {
 
     /// Return the min of the two integer operands
     static inline APInt calcOneElement(APInt leftOperand, APInt rightOperand) {
-      const llvm::APInt subtractRes = leftOperand - rightOperand;
-      return (!subtractRes.isNegative()) ? rightOperand : leftOperand;
+      return (leftOperand.sle(rightOperand)) ? leftOperand : rightOperand;
     }
   }];
 }

mlir/test/Dialect/Tosa/constant-op-fold.mlir

@@ -883,6 +883,18 @@ func.func @reduce_max_constant() -> tensor<1x1x1xi32> {
   return %0 : tensor<1x1x1xi32>
 }
 
+// -----
+
+func.func @reduce_max_constant_no_overflow() -> tensor<1xi8> {
+  // CHECK-LABEL:   func.func @reduce_max_constant_no_overflow() -> tensor<1xi8> {
+  // CHECK:           %[[VAL_0:.*]] = "tosa.const"() <{values = dense<120> : tensor<1xi8>}> : () -> tensor<1xi8>
+  // CHECK:           return %[[VAL_0]] : tensor<1xi8>
+  // CHECK:         }
+  %const = "tosa.const"() <{values = dense<[-127, 120, -126]> : tensor<3xi8>}> : () -> tensor<3xi8>
+  %0 = tosa.reduce_max %const {axis = 0 : i32} : (tensor<3xi8>) -> tensor<1xi8>
+  return %0 : tensor<1xi8>
+}
+
 // -----
 
   func.func @reduce_min_constant() -> tensor<1x3xi32> {
@@ -968,6 +980,19 @@ func.func @reduce_min_constant() -> tensor<1x1x1xi32> {
   return %0 : tensor<1x1x1xi32>
 }
 
+// -----
+
+func.func @reduce_min_constant_no_overflow() -> tensor<1xi8> {
+  // CHECK-LABEL:   func.func @reduce_min_constant_no_overflow() -> tensor<1xi8> {
+  // CHECK:           %[[VAL_0:.*]] = "tosa.const"() <{values = dense<-127> : tensor<1xi8>}> : () -> tensor<1xi8>
+  // CHECK:           return %[[VAL_0]] : tensor<1xi8>
+  // CHECK:         }
+  %const = "tosa.const"() <{values = dense<[-127, 120, -126]> : tensor<3xi8>}> : () -> tensor<3xi8>
+  %0 = tosa.reduce_min %const {axis = 0 : i32} : (tensor<3xi8>) -> tensor<1xi8>
+  return %0 : tensor<1xi8>
+}
+
+
 // -----
 
 func.func @reduce_any_constant() -> tensor<1x3xi1> {