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

[Ian Tayler Lessa]

https://github.com/IanTaylerLessa-arm updated https://github.com/llvm/llvm-project/pull/132786

>From 9391bbbd5364e84909999d97d278258e0904afe7 Mon Sep 17 00:00:00 2001
From: Ian Tayler Lessa <ian.taylerlessa at arm.com>
Date: Fri, 21 Mar 2025 15:42:22 +0000
Subject: [PATCH] [mlir][tosa] Avoid overflow in reduction folders

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 <ian.taylerlessa at arm.com>
---
 mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td |  6 ++---
 mlir/test/Dialect/Tosa/constant-op-fold.mlir | 25 ++++++++++++++++++++
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
index bd01c5f704b6b..bf0823c489b98 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 3f7de9357297e..39bd84f2fe74e 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> {