[mlir][arith] Fold (a * b) / b -> a

[Hardcode84]

If overflow flags allow it.

Alive2 check: https://alive2.llvm.org/ce/z/5XWjWE

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-arith

Author: Ivan Butygin (Hardcode84)

Changes

If overflow flags allow it.

Alive2 check: https://alive2.llvm.org/ce/z/5XWjWE

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/Arith/IR/ArithOps.cpp (+21)
• (modified) mlir/test/Dialect/Arith/canonicalize.mlir (+20)

diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index d8b314a3fa43c0..e486bb678ce33e 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -580,11 +580,29 @@ void arith::MulUIExtendedOp::getCanonicalizationPatterns(
 // DivUIOp
 //===----------------------------------------------------------------------===//
 
+static Value foldDivMul(Value lhs, Value rhs,
+                        arith::IntegerOverflowFlags ovfFlags) {
+  auto mul = lhs.getDefiningOp<mlir::arith::MulIOp>();
+  if (!mul || !bitEnumContainsAll(mul.getOverflowFlags(), ovfFlags))
+    return {};
+
+  if (mul.getLhs() == rhs)
+    return mul.getRhs();
+
+  if (mul.getRhs() == rhs)
+    return mul.getLhs();
+
+  return {};
+}
+
 OpFoldResult arith::DivUIOp::fold(FoldAdaptor adaptor) {
   // divui (x, 1) -> x.
   if (matchPattern(adaptor.getRhs(), m_One()))
     return getLhs();
 
+  if (Value val = foldDivMul(getLhs(), getRhs(), IntegerOverflowFlags::nuw))
+    return val;
+
   // Don't fold if it would require a division by zero.
   bool div0 = false;
   auto result = constFoldBinaryOp<IntegerAttr>(adaptor.getOperands(),
@@ -621,6 +639,9 @@ OpFoldResult arith::DivSIOp::fold(FoldAdaptor adaptor) {
   if (matchPattern(adaptor.getRhs(), m_One()))
     return getLhs();
 
+  if (Value val = foldDivMul(getLhs(), getRhs(), IntegerOverflowFlags::nsw))
+    return val;
+
   // Don't fold if it would overflow or if it requires a division by zero.
   bool overflowOrDiv0 = false;
   auto result = constFoldBinaryOp<IntegerAttr>(
diff --git a/mlir/test/Dialect/Arith/canonicalize.mlir b/mlir/test/Dialect/Arith/canonicalize.mlir
index 6a186a0c6ceca0..1b80086c13eced 100644
--- a/mlir/test/Dialect/Arith/canonicalize.mlir
+++ b/mlir/test/Dialect/Arith/canonicalize.mlir
@@ -2060,6 +2060,26 @@ func.func @test_divf1(%arg0 : f32, %arg1 : f32) -> (f32) {
 
 // -----
 
+// CHECK-LABEL: @test_divui_mul
+//  CHECK-SAME:   (%[[ARG:.*]]: index, %{{.*}}: index)
+//       CHECK:   return %[[ARG]]
+func.func @test_divui_mul(%arg0: index, %arg1: index) -> index {
+  %0 = arith.muli %arg0, %arg1 overflow<nuw>  : index
+  %1 = arith.divui %0, %arg1 : index
+  return %1 : index
+}
+
+// CHECK-LABEL: @test_divsi_mul
+//  CHECK-SAME:   (%[[ARG:.*]]: index, %{{.*}}: index)
+//       CHECK:   return %[[ARG]]
+func.func @test_divsi_mul(%arg0: index, %arg1: index) -> index {
+  %0 = arith.muli %arg1, %arg0 overflow<nsw>  : index
+  %1 = arith.divsi %0, %arg1 : index
+  return %1 : index
+}
+
+// -----
+
 // CHECK-LABEL: @test_cmpf(
 func.func @test_cmpf(%arg0 : f32) -> (i1, i1, i1, i1) {
 //   CHECK-DAG:   %[[T:.*]] = arith.constant true