[Mlir-commits] [mlir] [mlir][arith] Fix muli divisibility inference under wraparound (PR #237342)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-arith

Author: edisongz

<details>
<summary>Changes</summary>

Wrapping only preserves the power-of-two factor of the operands'
divisor product. Keep the full unsigned/signed divisor only under
`nuw`/`nsw`, otherwise reduce to the largest power-of-two divisor
(`v & -v`). Compute the product with APInt to avoid uint64_t overflow.

Fixes [#<!-- -->233210](https://github.com/llvm/llvm-project/issues/233210).

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Arith/IR/InferIntDivisibilityOpInterfaceImpl.cpp (+20-2) 
- (modified) mlir/test/Analysis/DataFlow/integer-divisibility.mlir (+60) 


``````````diff
diff --git a/mlir/lib/Dialect/Arith/IR/InferIntDivisibilityOpInterfaceImpl.cpp b/mlir/lib/Dialect/Arith/IR/InferIntDivisibilityOpInterfaceImpl.cpp
index b23ee108ca4a3d..535981bd2cf421 100644
--- a/mlir/lib/Dialect/Arith/IR/InferIntDivisibilityOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Arith/IR/InferIntDivisibilityOpInterfaceImpl.cpp
@@ -42,6 +42,16 @@ inferBinaryGCDResultDivisibility(OpTy op, ArrayRef<IntegerDivisibility> argDivs,
   setResultDivs(op.getResult(), lhsDiv.getUnion(rhsDiv));
 }
 
+// Overflow-free product of two divisors; falls back to the largest
+// power-of-two factor when the product exceeds uint64_t.
+static uint64_t mulDivisibilities(uint64_t a, uint64_t b) {
+  APInt product = APInt(128, a) * APInt(128, b);
+  if (product.getActiveBits() <= 64)
+    return product.getZExtValue();
+  unsigned trailingZeros = product.countr_zero();
+  return uint64_t(1) << (trailingZeros >= 63 ? 63 : trailingZeros);
+}
+
 void ConstantOp::inferResultDivisibility(ArrayRef<IntegerDivisibility> argDivs,
                                          SetIntDivisibilityFn setResultDivs) {
   auto constAttr = dyn_cast_if_present<IntegerAttr>(getValue());
@@ -88,8 +98,16 @@ void MulIOp::inferResultDivisibility(ArrayRef<IntegerDivisibility> argDivs,
   auto lhsDivisibility = getDivisibilityOfOperand(getLhs(), argDivs[0]);
   auto rhsDivisibility = getDivisibilityOfOperand(getRhs(), argDivs[1]);
 
-  uint64_t mulUDiv = lhsDivisibility.udiv() * rhsDivisibility.udiv();
-  uint64_t mulSDiv = lhsDivisibility.sdiv() * rhsDivisibility.sdiv();
+  uint64_t fullUDiv =
+      mulDivisibilities(lhsDivisibility.udiv(), rhsDivisibility.udiv());
+  uint64_t fullSDiv =
+      mulDivisibilities(lhsDivisibility.sdiv(), rhsDivisibility.sdiv());
+
+  // `v & -v` is the largest power-of-two divisor (all that wrapping keeps);
+  // nuw/nsw keep the full unsigned/signed divisor. Maps the 1/0 sentinels to
+  // themselves.
+  uint64_t mulUDiv = hasNoUnsignedWrap() ? fullUDiv : (fullUDiv & -fullUDiv);
+  uint64_t mulSDiv = hasNoSignedWrap() ? fullSDiv : (fullSDiv & -fullSDiv);
 
   setResultDivs(getResult(), ConstantIntDivisibility(mulUDiv, mulSDiv));
 }
diff --git a/mlir/test/Analysis/DataFlow/integer-divisibility.mlir b/mlir/test/Analysis/DataFlow/integer-divisibility.mlir
index 5db5c679eac829..89a512495b491e 100644
--- a/mlir/test/Analysis/DataFlow/integer-divisibility.mlir
+++ b/mlir/test/Analysis/DataFlow/integer-divisibility.mlir
@@ -150,3 +150,63 @@ func.func @affine_apply_mul_plus_const(%arg0 : index) -> index {
   %1 = "test.int_divisibility"(%0) : (index) -> index
   return %1 : index
 }
+
+// -----
+
+// CHECK-LABEL: @muli_wrapping
+func.func @muli_wrapping(%arg0 : i64) -> i64 {
+  %c6 = arith.constant 6 : i64
+  %0 = arith.muli %arg0, %c6 : i64
+  // 6 has power-of-two factor 2 only, so wrapping preserves just that.
+  // CHECK: divisibility = "udiv = 2, sdiv = 2"
+  %1 = "test.int_divisibility"(%0) : (i64) -> i64
+  return %1 : i64
+}
+
+// -----
+
+// CHECK-LABEL: @muli_nuw
+func.func @muli_nuw(%arg0 : i64) -> i64 {
+  %c6 = arith.constant 6 : i64
+  %0 = arith.muli %arg0, %c6 overflow<nuw> : i64
+  // nuw preserves the full unsigned divisor 6; signed still wraps to 2.
+  // CHECK: divisibility = "udiv = 6, sdiv = 2"
+  %1 = "test.int_divisibility"(%0) : (i64) -> i64
+  return %1 : i64
+}
+
+// -----
+
+// CHECK-LABEL: @muli_nsw
+func.func @muli_nsw(%arg0 : i64) -> i64 {
+  %c6 = arith.constant 6 : i64
+  %0 = arith.muli %arg0, %c6 overflow<nsw> : i64
+  // nsw preserves the full signed divisor 6; unsigned still wraps to 2.
+  // CHECK: divisibility = "udiv = 2, sdiv = 6"
+  %1 = "test.int_divisibility"(%0) : (i64) -> i64
+  return %1 : i64
+}
+
+// -----
+
+// CHECK-LABEL: @muli_nuw_nsw
+func.func @muli_nuw_nsw(%arg0 : i64) -> i64 {
+  %c6 = arith.constant 6 : i64
+  %0 = arith.muli %arg0, %c6 overflow<nsw, nuw> : i64
+  // With both no-wrap flags the full divisor 6 survives in both interpretations.
+  // CHECK: divisibility = "udiv = 6, sdiv = 6"
+  %1 = "test.int_divisibility"(%0) : (i64) -> i64
+  return %1 : i64
+}
+
+// -----
+
+// CHECK-LABEL: @muli_power_of_two
+func.func @muli_power_of_two(%arg0 : i64) -> i64 {
+  %c4 = arith.constant 4 : i64
+  %0 = arith.muli %arg0, %c4 : i64
+  // 4 is already a power of two, so wrapping preserves it unchanged.
+  // CHECK: divisibility = "udiv = 4, sdiv = 4"
+  %1 = "test.int_divisibility"(%0) : (i64) -> i64
+  return %1 : i64
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/237342