[Mlir-commits] [mlir] [mlir][arith] fix wrong floordivsi fold (#83079) (PR #83248)

[llvmlistbot at llvm.org]

https://github.com/lipracer created https://github.com/llvm/llvm-project/pull/83248

Fixs https://github.com/llvm/llvm-project/issues/83079

>From 911bc2686ed55c9fa13bbd50e7a5430d4be14c61 Mon Sep 17 00:00:00 2001
From: lipracer <lipracer at gmail.com>
Date: Wed, 28 Feb 2024 18:48:40 +0800
Subject: [PATCH] [mlir][arith] fix wrong floordivsi fold (#83079)

Fixs https://github.com/llvm/llvm-project/issues/83079
---
 mlir/lib/Dialect/Arith/IR/ArithOps.cpp | 12 +++++++++---
 mlir/test/Transforms/canonicalize.mlir |  9 +++++++++
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index 0f71c19c23b654..d370b7d04dea7e 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -709,19 +709,25 @@ OpFoldResult arith::FloorDivSIOp::fold(FoldAdaptor adaptor) {
         }
         if (!aGtZero && !bGtZero) {
           // Both negative, return -a / -b.
-          APInt posA = zero.ssub_ov(a, overflowOrDiv0);
-          APInt posB = zero.ssub_ov(b, overflowOrDiv0);
-          return posA.sdiv_ov(posB, overflowOrDiv0);
+          return a.sdiv_ov(b, overflowOrDiv0);
         }
         if (!aGtZero && bGtZero) {
           // A is negative, b is positive, return - ceil(-a, b).
           APInt posA = zero.ssub_ov(a, overflowOrDiv0);
+          if (overflowOrDiv0)
+            return a;
           APInt ceil = signedCeilNonnegInputs(posA, b, overflowOrDiv0);
+          if (overflowOrDiv0)
+            return a;
           return zero.ssub_ov(ceil, overflowOrDiv0);
         }
         // A is positive, b is negative, return - ceil(a, -b).
         APInt posB = zero.ssub_ov(b, overflowOrDiv0);
+        if (overflowOrDiv0)
+          return a;
         APInt ceil = signedCeilNonnegInputs(a, posB, overflowOrDiv0);
+        if (overflowOrDiv0)
+          return a;
         return zero.ssub_ov(ceil, overflowOrDiv0);
       });
 
diff --git a/mlir/test/Transforms/canonicalize.mlir b/mlir/test/Transforms/canonicalize.mlir
index 2cf86b50d432f6..d2c2c12d323892 100644
--- a/mlir/test/Transforms/canonicalize.mlir
+++ b/mlir/test/Transforms/canonicalize.mlir
@@ -989,6 +989,15 @@ func.func @tensor_arith.floordivsi_by_one(%arg0: tensor<4x5xi32>) -> tensor<4x5x
   return %res : tensor<4x5xi32>
 }
 
+// CHECK-LABEL: func @arith.floordivsi_by_one_overflow
+func.func @arith.floordivsi_by_one_overflow() -> i64 {
+  %neg_one = arith.constant -1 : i64
+  %min_int = arith.constant -9223372036854775808 : i64
+  // CHECK: arith.floordivsi
+  %poision = arith.floordivsi %min_int, %neg_one : i64
+  return %poision : i64
+}
+
 // -----
 
 // CHECK-LABEL: func @arith.ceildivsi_by_one