[Mlir-commits] [mlir] [arith][mlir] Fixed a bug in CeilDiv with neg values (PR #90855)

[Balaji V. Iyer.]

https://github.com/bviyer created https://github.com/llvm/llvm-project/pull/90855

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

Fixes #89382

>From ff3923654be490b7601acb1bfef11657a50e0f61 Mon Sep 17 00:00:00 2001
From: "Balaji V. Iyer" <bviyer at gmail.com>
Date: Thu, 2 May 2024 08:43:48 -0500
Subject: [PATCH] [arith][mlir] Fixed a bug in CeilDiv with neg values

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

Fixes #89382
---
 mlir/lib/Dialect/Arith/IR/ArithOps.cpp  | 19 +++++++++----------
 mlir/test/Transforms/constant-fold.mlir | 11 +++++++++++
 2 files changed, 20 insertions(+), 10 deletions(-)

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