[Mlir-commits] [mlir] [mlir][arith] Fold signed quotient and remainder recomposition (PR #236005)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Oct 28 23:52:05 PDT 2026
https://github.com/meinie0826 updated https://github.com/llvm/llvm-project/pull/236005
>From f44da65a0d9dc6a71d211cc9ac67efa71a174385 Mon Sep 17 00:00:00 2001
From: meinie <929449726 at qq.com>
Date: Tue, 27 Oct 2026 13:30:07 +0800
Subject: [PATCH 1/2] [mlir][arith] Fold signed quotient and remainder
recomposition
---
mlir/lib/Dialect/Arith/IR/ArithOps.cpp | 23 ++
mlir/test/Dialect/Arith/canonicalize.mlir | 264 ++++++++++++++++++++++
2 files changed, 287 insertions(+)
diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index 19e9caca6567a6..7a90e146d019ca 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -466,6 +466,23 @@ Value mlir::arith::getZeroConstant(OpBuilder &builder, Location loc,
// AddIOp
//===----------------------------------------------------------------------===//
+// Recompose a dividend from its signed quotient and remainder. Both
+// operations must use the same dividend and divisor.
+static Value recomposeSignedDividend(Value product, Value remainder) {
+ auto mul = product.getDefiningOp<arith::MulIOp>();
+ auto rem = remainder.getDefiningOp<arith::RemSIOp>();
+ if (!mul || !rem)
+ return {};
+
+ for (unsigned i = 0; i != 2; ++i) {
+ auto div = mul->getOperand(i).getDefiningOp<arith::DivSIOp>();
+ if (div && div.getLhs() == rem.getLhs() && div.getRhs() == rem.getRhs() &&
+ mul->getOperand(1 - i) == rem.getRhs())
+ return rem.getLhs();
+ }
+ return {};
+}
+
OpFoldResult arith::AddIOp::fold(FoldAdaptor adaptor) {
// addi(x, 0) -> x
if (matchPattern(adaptor.getRhs(), m_Zero()))
@@ -481,6 +498,12 @@ OpFoldResult arith::AddIOp::fold(FoldAdaptor adaptor) {
if (getLhs() == sub.getRhs())
return sub.getLhs();
+ // addi(muli(divsi(x, d), d), remsi(x, d)) -> x
+ if (Value dividend = recomposeSignedDividend(getLhs(), getRhs()))
+ return dividend;
+ if (Value dividend = recomposeSignedDividend(getRhs(), getLhs()))
+ return dividend;
+
return constFoldBinaryOp<IntegerAttr>(
adaptor.getOperands(),
[](APInt a, const APInt &b) { return std::move(a) + b; });
diff --git a/mlir/test/Dialect/Arith/canonicalize.mlir b/mlir/test/Dialect/Arith/canonicalize.mlir
index 8f9f9df5f9fe3a..42d924948d425d 100644
--- a/mlir/test/Dialect/Arith/canonicalize.mlir
+++ b/mlir/test/Dialect/Arith/canonicalize.mlir
@@ -5579,3 +5579,267 @@ func.func @subi_subi_lhs_rhs_lhs_dynamic(%a : tensor<?xi32>, %b : tensor<?xi32>)
%1 = arith.subi %0, %a : tensor<?xi32>
return %1 : tensor<?xi32>
}
+
+// -----
+
+// CHECK-LABEL: @recompose_signed_dividend_0_0
+// CHECK-NEXT: return %arg0
+func.func @recompose_signed_dividend_0_0(%x: i32, %d: i32) -> i32 {
+ %q = arith.divsi %x, %d : i32
+ %m = arith.muli %q, %d : i32
+ %r = arith.remsi %x, %d : i32
+ %sum = arith.addi %m, %r : i32
+ return %sum : i32
+}
+
+// -----
+
+// CHECK-LABEL: @recompose_signed_dividend_0_1
+// CHECK-NEXT: return %arg0
+func.func @recompose_signed_dividend_0_1(%x: i32, %d: i32) -> i32 {
+ %q = arith.divsi %x, %d : i32
+ %m = arith.muli %d, %q : i32
+ %r = arith.remsi %x, %d : i32
+ %sum = arith.addi %m, %r : i32
+ return %sum : i32
+}
+
+// -----
+
+// CHECK-LABEL: @recompose_signed_dividend_1_0
+// CHECK-NEXT: return %arg0
+func.func @recompose_signed_dividend_1_0(%x: i32, %d: i32) -> i32 {
+ %q = arith.divsi %x, %d : i32
+ %m = arith.muli %q, %d : i32
+ %r = arith.remsi %x, %d : i32
+ %sum = arith.addi %r, %m : i32
+ return %sum : i32
+}
+
+// -----
+
+// CHECK-LABEL: @recompose_signed_dividend_1_1
+// CHECK-NEXT: return %arg0
+func.func @recompose_signed_dividend_1_1(%x: i32, %d: i32) -> i32 {
+ %q = arith.divsi %x, %d : i32
+ %m = arith.muli %d, %q : i32
+ %r = arith.remsi %x, %d : i32
+ %sum = arith.addi %r, %m : i32
+ return %sum : i32
+}
+
+// -----
+
+// CHECK-LABEL: @recompose_signed_dividend_index
+// CHECK-NEXT: return %arg0
+func.func @recompose_signed_dividend_index(%x: index, %d: index) -> index {
+ %q = arith.divsi %x, %d : index
+ %m = arith.muli %q, %d : index
+ %r = arith.remsi %x, %d : index
+ %sum = arith.addi %m, %r : index
+ return %sum : index
+}
+
+// -----
+
+// CHECK-LABEL: @recompose_signed_dividend_i1
+// CHECK-NEXT: return %arg0
+func.func @recompose_signed_dividend_i1(%x: i1, %d: i1) -> i1 {
+ %q = arith.divsi %x, %d : i1
+ %m = arith.muli %q, %d : i1
+ %r = arith.remsi %x, %d : i1
+ %sum = arith.addi %m, %r : i1
+ return %sum : i1
+}
+
+// -----
+
+// CHECK-LABEL: @recompose_signed_dividend_i8
+// CHECK-NEXT: return %arg0
+func.func @recompose_signed_dividend_i8(%x: i8, %d: i8) -> i8 {
+ %q = arith.divsi %x, %d : i8
+ %m = arith.muli %q, %d : i8
+ %r = arith.remsi %x, %d : i8
+ %sum = arith.addi %m, %r : i8
+ return %sum : i8
+}
+
+// -----
+
+// CHECK-LABEL: @recompose_signed_dividend_i64
+// CHECK-NEXT: return %arg0
+func.func @recompose_signed_dividend_i64(%x: i64, %d: i64) -> i64 {
+ %q = arith.divsi %x, %d : i64
+ %m = arith.muli %q, %d : i64
+ %r = arith.remsi %x, %d : i64
+ %sum = arith.addi %m, %r : i64
+ return %sum : i64
+}
+
+// -----
+
+// CHECK-LABEL: @recompose_signed_dividend_vector
+// CHECK-NEXT: return %arg0
+func.func @recompose_signed_dividend_vector(%x: vector<4xi32>, %d: vector<4xi32>) -> vector<4xi32> {
+ %q = arith.divsi %x, %d : vector<4xi32>
+ %m = arith.muli %q, %d : vector<4xi32>
+ %r = arith.remsi %x, %d : vector<4xi32>
+ %sum = arith.addi %m, %r : vector<4xi32>
+ return %sum : vector<4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @recompose_signed_dividend_scalable_vector
+// CHECK-NEXT: return %arg0
+func.func @recompose_signed_dividend_scalable_vector(%x: vector<[4]xi32>, %d: vector<[4]xi32>) -> vector<[4]xi32> {
+ %q = arith.divsi %x, %d : vector<[4]xi32>
+ %m = arith.muli %q, %d : vector<[4]xi32>
+ %r = arith.remsi %x, %d : vector<[4]xi32>
+ %sum = arith.addi %m, %r : vector<[4]xi32>
+ return %sum : vector<[4]xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @recompose_signed_dividend_tensor
+// CHECK-NEXT: return %arg0
+func.func @recompose_signed_dividend_tensor(%x: tensor<2x3xi32>, %d: tensor<2x3xi32>) -> tensor<2x3xi32> {
+ %q = arith.divsi %x, %d : tensor<2x3xi32>
+ %m = arith.muli %q, %d : tensor<2x3xi32>
+ %r = arith.remsi %x, %d : tensor<2x3xi32>
+ %sum = arith.addi %m, %r : tensor<2x3xi32>
+ return %sum : tensor<2x3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @recompose_signed_dividend_dynamic_tensor
+// CHECK-NEXT: return %arg0
+func.func @recompose_signed_dividend_dynamic_tensor(%x: tensor<?xi32>, %d: tensor<?xi32>) -> tensor<?xi32> {
+ %q = arith.divsi %x, %d : tensor<?xi32>
+ %m = arith.muli %q, %d : tensor<?xi32>
+ %r = arith.remsi %x, %d : tensor<?xi32>
+ %sum = arith.addi %m, %r : tensor<?xi32>
+ return %sum : tensor<?xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @recompose_signed_dividend_exact
+// CHECK-NEXT: return %arg0
+func.func @recompose_signed_dividend_exact(%x: i32, %d: i32) -> i32 {
+ %q = arith.divsi %x, %d exact : i32
+ %m = arith.muli %q, %d : i32
+ %r = arith.remsi %x, %d : i32
+ %sum = arith.addi %m, %r : i32
+ return %sum : i32
+}
+
+// -----
+
+// CHECK-LABEL: @recompose_signed_dividend_nsw
+// CHECK-NEXT: return %arg0
+func.func @recompose_signed_dividend_nsw(%x: i32, %d: i32) -> i32 {
+ %q = arith.divsi %x, %d : i32
+ %m = arith.muli %q, %d overflow<nsw> : i32
+ %r = arith.remsi %x, %d : i32
+ %sum = arith.addi %m, %r overflow<nsw> : i32
+ return %sum : i32
+}
+
+// -----
+
+// CHECK-LABEL: @recompose_signed_dividend_nuw
+// CHECK-NEXT: return %arg0
+func.func @recompose_signed_dividend_nuw(%x: i32, %d: i32) -> i32 {
+ %q = arith.divsi %x, %d : i32
+ %m = arith.muli %q, %d overflow<nuw> : i32
+ %r = arith.remsi %x, %d : i32
+ %sum = arith.addi %m, %r overflow<nuw> : i32
+ return %sum : i32
+}
+
+// -----
+
+// CHECK-LABEL: @recompose_signed_dividend_different_dividend
+// CHECK: arith.divsi
+// CHECK: arith.remsi
+// CHECK: arith.addi
+func.func @recompose_signed_dividend_different_dividend(%x: i32, %d: i32, %other: i32) -> i32 {
+ %q = arith.divsi %x, %d : i32
+ %m = arith.muli %q, %d : i32
+ %r = arith.remsi %other, %d : i32
+ %sum = arith.addi %m, %r : i32
+ return %sum : i32
+}
+
+// -----
+
+// CHECK-LABEL: @recompose_signed_dividend_different_divisor
+// CHECK: arith.divsi
+// CHECK: arith.remsi
+// CHECK: arith.addi
+func.func @recompose_signed_dividend_different_divisor(%x: i32, %d: i32, %other: i32) -> i32 {
+ %q = arith.divsi %x, %d : i32
+ %m = arith.muli %q, %d : i32
+ %r = arith.remsi %x, %other : i32
+ %sum = arith.addi %m, %r : i32
+ return %sum : i32
+}
+
+// -----
+
+// CHECK-LABEL: @recompose_signed_dividend_different_multiplier
+// CHECK: arith.divsi
+// CHECK: arith.remsi
+// CHECK: arith.addi
+func.func @recompose_signed_dividend_different_multiplier(%x: i32, %d: i32, %other: i32) -> i32 {
+ %q = arith.divsi %x, %d : i32
+ %m = arith.muli %q, %other : i32
+ %r = arith.remsi %x, %d : i32
+ %sum = arith.addi %m, %r : i32
+ return %sum : i32
+}
+
+// -----
+
+// CHECK-LABEL: @recompose_signed_dividend_unsigned_division
+// CHECK: arith.divui
+// CHECK: arith.remsi
+// CHECK: arith.addi
+func.func @recompose_signed_dividend_unsigned_division(%x: i32, %d: i32) -> i32 {
+ %q = arith.divui %x, %d : i32
+ %m = arith.muli %q, %d : i32
+ %r = arith.remsi %x, %d : i32
+ %sum = arith.addi %m, %r : i32
+ return %sum : i32
+}
+
+// -----
+
+// CHECK-LABEL: @recompose_signed_dividend_unsigned_remainder
+// CHECK: arith.divsi
+// CHECK: arith.remui
+// CHECK: arith.addi
+func.func @recompose_signed_dividend_unsigned_remainder(%x: i32, %d: i32) -> i32 {
+ %q = arith.divsi %x, %d : i32
+ %m = arith.muli %q, %d : i32
+ %r = arith.remui %x, %d : i32
+ %sum = arith.addi %m, %r : i32
+ return %sum : i32
+}
+
+// -----
+
+// CHECK-LABEL: @recompose_signed_dividend_floor_division
+// CHECK: arith.floordivsi
+// CHECK: arith.remsi
+// CHECK: arith.addi
+func.func @recompose_signed_dividend_floor_division(%x: i32, %d: i32) -> i32 {
+ %q = arith.floordivsi %x, %d : i32
+ %m = arith.muli %q, %d : i32
+ %r = arith.remsi %x, %d : i32
+ %sum = arith.addi %m, %r : i32
+ return %sum : i32
+}
>From d66036a06a75401a733c2bbdd229c367e9ce16cd Mon Sep 17 00:00:00 2001
From: meinie0826 <929449726 at qq.com>
Date: Thu, 29 Oct 2026 14:51:52 +0800
Subject: [PATCH 2/2] [mlir][arith] Consolidate divrem recomposition tests
---
mlir/test/Dialect/Arith/canonicalize.mlir | 93 +++++------------------
1 file changed, 21 insertions(+), 72 deletions(-)
diff --git a/mlir/test/Dialect/Arith/canonicalize.mlir b/mlir/test/Dialect/Arith/canonicalize.mlir
index 42d924948d425d..58e859be95be12 100644
--- a/mlir/test/Dialect/Arith/canonicalize.mlir
+++ b/mlir/test/Dialect/Arith/canonicalize.mlir
@@ -5582,50 +5582,18 @@ func.func @subi_subi_lhs_rhs_lhs_dynamic(%a : tensor<?xi32>, %b : tensor<?xi32>)
// -----
-// CHECK-LABEL: @recompose_signed_dividend_0_0
-// CHECK-NEXT: return %arg0
-func.func @recompose_signed_dividend_0_0(%x: i32, %d: i32) -> i32 {
+// CHECK-LABEL: @recompose_signed_dividend_operand_orders
+// CHECK-NEXT: return %arg0, %arg0, %arg0, %arg0
+func.func @recompose_signed_dividend_operand_orders(%x: i32, %d: i32) -> (i32, i32, i32, i32) {
%q = arith.divsi %x, %d : i32
- %m = arith.muli %q, %d : i32
%r = arith.remsi %x, %d : i32
- %sum = arith.addi %m, %r : i32
- return %sum : i32
-}
-
-// -----
-
-// CHECK-LABEL: @recompose_signed_dividend_0_1
-// CHECK-NEXT: return %arg0
-func.func @recompose_signed_dividend_0_1(%x: i32, %d: i32) -> i32 {
- %q = arith.divsi %x, %d : i32
- %m = arith.muli %d, %q : i32
- %r = arith.remsi %x, %d : i32
- %sum = arith.addi %m, %r : i32
- return %sum : i32
-}
-
-// -----
-
-// CHECK-LABEL: @recompose_signed_dividend_1_0
-// CHECK-NEXT: return %arg0
-func.func @recompose_signed_dividend_1_0(%x: i32, %d: i32) -> i32 {
- %q = arith.divsi %x, %d : i32
- %m = arith.muli %q, %d : i32
- %r = arith.remsi %x, %d : i32
- %sum = arith.addi %r, %m : i32
- return %sum : i32
-}
-
-// -----
-
-// CHECK-LABEL: @recompose_signed_dividend_1_1
-// CHECK-NEXT: return %arg0
-func.func @recompose_signed_dividend_1_1(%x: i32, %d: i32) -> i32 {
- %q = arith.divsi %x, %d : i32
- %m = arith.muli %d, %q : i32
- %r = arith.remsi %x, %d : i32
- %sum = arith.addi %r, %m : i32
- return %sum : i32
+ %qd = arith.muli %q, %d : i32
+ %dq = arith.muli %d, %q : i32
+ %sum0 = arith.addi %qd, %r : i32
+ %sum1 = arith.addi %dq, %r : i32
+ %sum2 = arith.addi %r, %qd : i32
+ %sum3 = arith.addi %r, %dq : i32
+ return %sum0, %sum1, %sum2, %sum3 : i32, i32, i32, i32
}
// -----
@@ -5726,38 +5694,19 @@ func.func @recompose_signed_dividend_dynamic_tensor(%x: tensor<?xi32>, %d: tenso
// -----
-// CHECK-LABEL: @recompose_signed_dividend_exact
-// CHECK-NEXT: return %arg0
-func.func @recompose_signed_dividend_exact(%x: i32, %d: i32) -> i32 {
- %q = arith.divsi %x, %d exact : i32
- %m = arith.muli %q, %d : i32
- %r = arith.remsi %x, %d : i32
- %sum = arith.addi %m, %r : i32
- return %sum : i32
-}
-
-// -----
-
-// CHECK-LABEL: @recompose_signed_dividend_nsw
-// CHECK-NEXT: return %arg0
-func.func @recompose_signed_dividend_nsw(%x: i32, %d: i32) -> i32 {
- %q = arith.divsi %x, %d : i32
- %m = arith.muli %q, %d overflow<nsw> : i32
- %r = arith.remsi %x, %d : i32
- %sum = arith.addi %m, %r overflow<nsw> : i32
- return %sum : i32
-}
-
-// -----
-
-// CHECK-LABEL: @recompose_signed_dividend_nuw
-// CHECK-NEXT: return %arg0
-func.func @recompose_signed_dividend_nuw(%x: i32, %d: i32) -> i32 {
+// CHECK-LABEL: @recompose_signed_dividend_flags
+// CHECK-NEXT: return %arg0, %arg0, %arg0
+func.func @recompose_signed_dividend_flags(%x: i32, %d: i32) -> (i32, i32, i32) {
%q = arith.divsi %x, %d : i32
- %m = arith.muli %q, %d overflow<nuw> : i32
+ %q_exact = arith.divsi %x, %d exact : i32
%r = arith.remsi %x, %d : i32
- %sum = arith.addi %m, %r overflow<nuw> : i32
- return %sum : i32
+ %m_exact = arith.muli %q_exact, %d : i32
+ %sum_exact = arith.addi %m_exact, %r : i32
+ %m_nsw = arith.muli %q, %d overflow<nsw> : i32
+ %sum_nsw = arith.addi %m_nsw, %r overflow<nsw> : i32
+ %m_nuw = arith.muli %q, %d overflow<nuw> : i32
+ %sum_nuw = arith.addi %m_nuw, %r overflow<nuw> : i32
+ return %sum_exact, %sum_nsw, %sum_nuw : i32, i32, i32
}
// -----
More information about the Mlir-commits
mailing list