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

[Balaji V. Iyer.]

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

>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 1/2] [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

>From 7214f242e0ccfd5fdd1bad0e1f19bc72c8087d80 Mon Sep 17 00:00:00 2001
From: "Balaji V. Iyer" <bviyer at gmail.com>
Date: Thu, 2 May 2024 22:00:00 -0500
Subject: [PATCH 2/2] Added more information to a comment

---
 mlir/lib/Dialect/Arith/IR/ArithOps.cpp | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index 34df4e686ff67d..788d20ca507816 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -656,6 +656,8 @@ OpFoldResult arith::CeilDivSIOp::fold(FoldAdaptor adaptor) {
         // 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.
+        // Mathematically, -1 * (abs(a)-1/abs(b) + 1) + 1, which after factoring
+        // out -1 yields -1 * [abs(a)-1/abs(b) + 1 - 1]. This is implemented below.
         APInt posA = aGtZero ? a : zero.ssub_ov(a, overflowOrDiv0);
         APInt posB = bGtZero ? b : zero.ssub_ov(b, overflowOrDiv0);
         APInt div = signedCeilNonnegInputs(posA, posB, overflowOrDiv0);