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

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

Author: Hardcode84

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>(

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