From ba4ba42048e2f6f52f3c1eac80bc7b7a42990127 Mon Sep 17 00:00:00 2001
From: Noah Goldstein <goldstein.w.n@gmail.com>
Date: Fri, 25 Oct 2024 15:17:28 -0500
Subject: [PATCH 1/2] [MLIR][Arith] Add tests for improving accurate of
 `inferDivU`; NFC

---
 .../test/Dialect/Arith/int-range-interface.mlir | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/mlir/test/Dialect/Arith/int-range-interface.mlir b/mlir/test/Dialect/Arith/int-range-interface.mlir
index 4b04229e5db52..46b49742e4623 100644
--- a/mlir/test/Dialect/Arith/int-range-interface.mlir
+++ b/mlir/test/Dialect/Arith/int-range-interface.mlir
@@ -190,6 +190,19 @@ func.func @div_zero_undefined(%arg0 : index) -> i1 {
     func.return %2 : i1
 }
 
+// CHECK-LABEL: func @div_refine_min
+// CHECK: %[[ret:.*]] = arith.cmpi uge
+// CHECK: return %[[ret]]
+func.func @div_refine_min(%arg0 : index) -> i1 {
+    %c0 = arith.constant 1 : index
+    %c1 = arith.constant 2 : index
+    %c4 = arith.constant 4 : index
+    %0 = arith.andi %arg0, %c1 : index
+    %1 = arith.divui %c4, %0 : index
+    %2 = arith.cmpi uge, %1, %c0 : index
+    func.return %2 : i1
+}
+
 // CHECK-LABEL: func @ceil_divui
 // CHECK: %[[ret:.*]] = arith.cmpi eq
 // CHECK: return %[[ret]]
@@ -271,13 +284,13 @@ func.func @remui_base(%arg0 : index, %arg1 : index ) -> i1 {
 // CHECK: return %[[true]]
 func.func @remui_base_maybe_zero(%arg0 : index, %arg1 : index ) -> i1 {
     %c4 = arith.constant 4 : index
-    %c5 = arith.constant 5 : index    
+    %c5 = arith.constant 5 : index
 
     %0 = arith.minui %arg1, %c4 : index
     %1 = arith.remui %arg0, %0 : index
     %2 = arith.cmpi ult, %1, %c5 : index
     func.return %2 : i1
-}    
+}
 
 // CHECK-LABEL: func @remsi_base
 // CHECK: %[[ret:.*]] = arith.cmpi sge

From c11ca1468b9f6fa35961ba67bc67b404b49aa97e Mon Sep 17 00:00:00 2001
From: Noah Goldstein <goldstein.w.n@gmail.com>
Date: Fri, 25 Oct 2024 15:17:05 -0500
Subject: [PATCH 2/2] [MLIR][Arith] Improve accuracy of `inferDivU`

1) We can always bound the maximum with the numerator.
    - https://alive2.llvm.org/ce/z/PqHvuT
2) Even if denominator min can be zero, we can still bound the minimum
   result with `lhs.umin u/ rhs.umax`.
---
 mlir/lib/Interfaces/Utils/InferIntRangeCommon.cpp | 10 ++++++++--
 mlir/test/Dialect/Arith/int-range-interface.mlir  |  8 ++++----
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/mlir/lib/Interfaces/Utils/InferIntRangeCommon.cpp b/mlir/lib/Interfaces/Utils/InferIntRangeCommon.cpp
index ec9ed87723e1c..a2acf3e732ada 100644
--- a/mlir/lib/Interfaces/Utils/InferIntRangeCommon.cpp
+++ b/mlir/lib/Interfaces/Utils/InferIntRangeCommon.cpp
@@ -298,8 +298,14 @@ static ConstantIntRanges inferDivURange(const ConstantIntRanges &lhs,
     return minMaxBy(udiv, {lhsMin, lhsMax}, {rhsMin, rhsMax},
                     /*isSigned=*/false);
   }
-  // Otherwise, it's possible we might divide by 0.
-  return ConstantIntRanges::maxRange(rhsMin.getBitWidth());
+
+  APInt umin = APInt::getZero(rhsMin.getBitWidth());
+  if (lhsMin.uge(rhsMax) && !rhsMax.isZero())
+    umin = lhsMin.udiv(rhsMax);
+
+  // X u/ Y u<= X.
+  APInt umax = lhsMax;
+  return ConstantIntRanges::fromUnsigned(umin, umax);
 }
 
 ConstantIntRanges
diff --git a/mlir/test/Dialect/Arith/int-range-interface.mlir b/mlir/test/Dialect/Arith/int-range-interface.mlir
index 46b49742e4623..6d66da2fc1eb3 100644
--- a/mlir/test/Dialect/Arith/int-range-interface.mlir
+++ b/mlir/test/Dialect/Arith/int-range-interface.mlir
@@ -178,8 +178,8 @@ func.func @div_bounds_negative(%arg0 : index) -> i1 {
 }
 
 // CHECK-LABEL: func @div_zero_undefined
-// CHECK: %[[ret:.*]] = arith.cmpi ule
-// CHECK: return %[[ret]]
+// CHECK: %[[true:.*]] = arith.constant true
+// CHECK: return %[[true]]
 func.func @div_zero_undefined(%arg0 : index) -> i1 {
     %c0 = arith.constant 0 : index
     %c1 = arith.constant 1 : index
@@ -191,8 +191,8 @@ func.func @div_zero_undefined(%arg0 : index) -> i1 {
 }
 
 // CHECK-LABEL: func @div_refine_min
-// CHECK: %[[ret:.*]] = arith.cmpi uge
-// CHECK: return %[[ret]]
+// CHECK: %[[true:.*]] = arith.constant true
+// CHECK: return %[[true]]
 func.func @div_refine_min(%arg0 : index) -> i1 {
     %c0 = arith.constant 1 : index
     %c1 = arith.constant 2 : index