[arith][mlir] Fixed a bug in CeilDiv with neg values

[bviyer]

When either the divisor or dividend, not both, is negative the compiler gives a positive value, while it should be a negative one. This patch should fix this.

Fixes #89382

[llvmbot]

@llvm/pr-subscribers-mlir-arith

@llvm/pr-subscribers-mlir

Author: Balaji V. Iyer. (bviyer)

Changes

When either the divisor or dividend the compiler
gives a positive value. This patch should fix
this issue.

Fixes #89382

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/Arith/IR/ArithOps.cpp (+9-10)
• (modified) mlir/test/Transforms/constant-fold.mlir (+11)

diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index 9f64a07f31e3af..34df4e686ff67d 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -652,16 +652,15 @@ OpFoldResult arith::CeilDivSIOp::fold(FoldAdaptor adaptor) {
           APInt posB = zero.ssub_ov(b, overflowOrDiv0);
           return signedCeilNonnegInputs(posA, posB, overflowOrDiv0);
         }
-        if (!aGtZero && bGtZero) {
-          // A is negative, b is positive, return - ( -a / b).
-          APInt posA = zero.ssub_ov(a, overflowOrDiv0);
-          APInt div = posA.sdiv_ov(b, overflowOrDiv0);
-          return zero.ssub_ov(div, overflowOrDiv0);
-        }
-        // A is positive, b is negative, return - (a / -b).
-        APInt posB = zero.ssub_ov(b, overflowOrDiv0);
-        APInt div = a.sdiv_ov(posB, overflowOrDiv0);
-        return zero.ssub_ov(div, overflowOrDiv0);
+        // If either divisor or dividend is negative, then take their absolute
+        // value and then do a normal signedCeil Division, but add 1 to bring
+        // the quotient down. In essense, Ceil Division with one of the values
+        // negative works like a floorDivision with negated quotient.
+        APInt posA = aGtZero ? a : zero.ssub_ov(a, overflowOrDiv0);
+        APInt posB = bGtZero ? b : zero.ssub_ov(b, overflowOrDiv0);
+        APInt div = signedCeilNonnegInputs(posA, posB, overflowOrDiv0);
+        APInt res = div.ssub_ov(APInt::getOneBitSet(bits, 0), overflowOrDiv0);
+        return zero.ssub_ov(res, overflowOrDiv0);
       });
 
   return overflowOrDiv0 ? Attribute() : result;
diff --git a/mlir/test/Transforms/constant-fold.mlir b/mlir/test/Transforms/constant-fold.mlir
index 253163f2af9110..dd5f475e946e75 100644
--- a/mlir/test/Transforms/constant-fold.mlir
+++ b/mlir/test/Transforms/constant-fold.mlir
@@ -478,6 +478,17 @@ func.func @simple_arith.ceildivsi() -> (i32, i32, i32, i32, i32) {
 
 // -----
 
+// CHECK-LABEL: simple_arith.ceildivsi_i8
+func.func @simple_arith.ceildivsi_i8() -> (i8) {
+  %0 = arith.constant 7 : i8
+  %1 = arith.constant -128 : i8
+  // CHECK-NEXT: arith.constant -18
+  %2 = arith.ceildivsi %1, %0 : i8
+  return %2 : i8
+}
+
+// -----
+
 // CHECK-LABEL: func @simple_arith.ceildivui
 func.func @simple_arith.ceildivui() -> (i32, i32, i32, i32, i32) {
   // CHECK-DAG: [[C0:%.+]] = arith.constant 0

[banach-space]

This example works fine without your change (-120 instead of -128):

func.func @simple_arith.ceildivsi_i8() -> (i8) {
  %0 = arith.constant 7 : i8
  %1 = arith.constant -120 : i8
  // CHECK-NEXT: arith.constant -18
  %2 = arith.ceildivsi %1, %0 : i8
  return %2 : i8
}

So, is the issue stemming from that fact that 128 is out of range for 8 bit signed integers? I'm basically trying to understand why the current implementation is incorrect in your case.

[bviyer]

This example works fine without your change (-120 instead of -128):

func.func @simple_arith.ceildivsi_i8() -> (i8) {
  %0 = arith.constant 7 : i8
  %1 = arith.constant -120 : i8
  // CHECK-NEXT: arith.constant -18
  %2 = arith.ceildivsi %1, %0 : i8
  return %2 : i8
}

So, is the issue stemming from that fact that 128 is out of range for 8 bit signed integers? I'm basically trying to understand why the current implementation is incorrect in your case.

yes, as far as I understand it.

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[banach-space]

This PR is missing some analysis explaining what the actual issue is. Without that, the proposed update seems rather arbitrary - it's hard for me to see why the proposed update fixes the issue if we don't know what "the issue" is.

When either the divisor or dividend, not both, is negative the compiler gives a positive value, while it should be a negative one.

As I hinted in my message above, the current logic does work for negative values. It breaks when using MININT (e.g. -128forint8`). The solution should be built around that observation.

I've spent some time analysing this and found that overflow flags are effectively ignored. I felt like the easiest way to demonstrate/document/explain this was to post a patch:

• [mlir][arith] Fix overflow bug in arith::CeilDivSIOp::fold #90947

Now, I agree that the folder could still work for the case that you are testing:

func.func @simple_arith.ceildivsi_i8() -> (i8) {
  %0 = arith.constant 7 : i8
  %1 = arith.constant -128 : i8
  %2 = arith.ceildivsi %1, %0 : i8
  return %2 : i8
}

However, first I'd make sure that all overflow flags are checked (and that no folding happens if there's an overflow). Following that, we can iterate and extend the folder to support your case. WDYT?