[mlir][arith] fix wrong floordivsi fold (#83079)

Fixs #83079

Author: lipracer

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);
       });
 

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