From d32df0f63d27bb40981eed14a9e6aebcad4d884b Mon Sep 17 00:00:00 2001
From: Lei Zhang <antiagainst@gmail.com>
Date: Fri, 7 Oct 2022 19:42:08 +0000
Subject: [PATCH] [mlir][arith] Expose dedicated API for expanding ceil/floor
 division

This allows more precise control over which patterns to pick to
expand arithmetic ops. Previously ceil/floor division epxansion
is only available together with various min/max op expansion.

Reviewed By: ThomasRaoux

Differential Revision: https://reviews.llvm.org/D135479
---
 mlir/include/mlir/Dialect/Arith/Transforms/Passes.h |  7 +++++--
 mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp     | 11 ++++++++---
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Arith/Transforms/Passes.h b/mlir/include/mlir/Dialect/Arith/Transforms/Passes.h
index a987bcd..5e441a7 100644
--- a/mlir/include/mlir/Dialect/Arith/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/Arith/Transforms/Passes.h
@@ -31,10 +31,13 @@ std::unique_ptr<Pass> createConstantBufferizePass(uint64_t alignment = 0);
 void populateWideIntEmulationPatterns(WideIntEmulationConverter &typeConverter,
                                       RewritePatternSet &patterns);
 
-/// Add patterns to expand Arith ops for LLVM lowering.
+/// Add patterns to expand Arith ceil/floor division ops.
+void populateCeilFloorDivExpandOpsPatterns(RewritePatternSet &patterns);
+
+/// Add patterns to expand Arith ops.
 void populateArithExpandOpsPatterns(RewritePatternSet &patterns);
 
-/// Create a pass to legalize Arith ops for LLVM lowering.
+/// Create a pass to legalize Arith ops.
 std::unique_ptr<Pass> createArithExpandOpsPass();
 
 /// Create a pass to replace signed ops with unsigned ones where they are proven
diff --git a/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp b/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
index 5a1793a..1a5911c 100644
--- a/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
@@ -225,12 +225,17 @@ struct ArithExpandOpsPass
 
 } // namespace
 
+void mlir::arith::populateCeilFloorDivExpandOpsPatterns(
+    RewritePatternSet &patterns) {
+  patterns
+      .add<CeilDivSIOpConverter, CeilDivUIOpConverter, FloorDivSIOpConverter>(
+          patterns.getContext());
+}
+
 void mlir::arith::populateArithExpandOpsPatterns(RewritePatternSet &patterns) {
+  populateCeilFloorDivExpandOpsPatterns(patterns);
   // clang-format off
   patterns.add<
-    CeilDivSIOpConverter,
-    CeilDivUIOpConverter,
-    FloorDivSIOpConverter,
     MaxMinFOpConverter<MaxFOp, arith::CmpFPredicate::UGT>,
     MaxMinFOpConverter<MinFOp, arith::CmpFPredicate::ULT>,
     MaxMinIOpConverter<MaxSIOp, arith::CmpIPredicate::sgt>,
-- 
2.7.4