[Mlir-commits] [mlir] [mlir][arith] Fold `(a * b) / b` (PR #121534)

[Ivan Butygin]

https://github.com/Hardcode84 created https://github.com/llvm/llvm-project/pull/121534

If overflow flags allow it.

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

>From dca3903165834ebd375fbdad38906e79ceae2e59 Mon Sep 17 00:00:00 2001
From: Ivan Butygin <ivan.butygin at gmail.com>
Date: Fri, 3 Jan 2025 02:31:26 +0100
Subject: [PATCH] [mlir][arith] Fold `(a * b) / b`

Alive2 check: https://alive2.llvm.org/ce/z/5XWjWE
---
 mlir/lib/Dialect/Arith/IR/ArithOps.cpp    | 21 +++++++++++++++++++++
 mlir/test/Dialect/Arith/canonicalize.mlir | 20 ++++++++++++++++++++
 2 files changed, 41 insertions(+)

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