[Mlir-commits] [mlir] [mlir][tosa] Avoid overflow in reduction folders (PR #132786)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-tosa

Author: None (IanTaylerLessa-arm)

<details>
<summary>Changes</summary>

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

Includes tests to avoid regressions

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


2 Files Affected:

- (modified) mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td (+2-4) 
- (modified) mlir/test/Dialect/Tosa/constant-op-fold.mlir (+25) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
index 14e15173de7bc..1ae6463179970 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
@@ -1731,8 +1731,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;
     }
   }];
 }
@@ -1772,8 +1771,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;
     }
   }];
 }
diff --git a/mlir/test/Dialect/Tosa/constant-op-fold.mlir b/mlir/test/Dialect/Tosa/constant-op-fold.mlir
index 8ac1e177ae4d4..b8b8e8d69fb16 100644
--- a/mlir/test/Dialect/Tosa/constant-op-fold.mlir
+++ b/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> {

``````````

</details>


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