[Mlir-commits] [mlir] [mlir][arith] Fold signed quotient and remainder recomposition (PR #236005)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Oct 26 22:30:27 PDT 2026
https://github.com/meinie0826 created https://github.com/llvm/llvm-project/pull/236005
## Summary
Fold signed quotient/remainder recomposition in `arith.addi`:
```mlir
%q = arith.divsi %x, %d : index
%p = arith.muli %q, %d : index
%r = arith.remsi %x, %d : index
%result = arith.addi %p, %r : index
```
The result is `%x`. This removes redundant reconstruction of a linearized index. The fold matches either operand order of the multiplication and addition and requires identical SSA dividends and divisors. It applies elementwise to integer, index, vector, and tensor types.
For defined signed division, `x = (x / d) * d + x % d`, including negative operands. Neither the signed product nor the sum overflows. Division by zero and minimum signed value divided by -1 already have undefined behavior. Refining poison produced by `exact` or integer overflow flags to the dividend is valid; no operation is moved or introduced.
Regression tests cover operand permutations, supported types, flags, and mismatched operands or division/remainder kinds. The Arith lit suite passes locally (30 tests). Exhaustive arithmetic checks also cover all 1,396,044 defined operand pairs at widths 1 through 10. No performance measurement is claimed.
AI tools assisted with implementation, test generation, and this description.
>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] [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
+}
More information about the Mlir-commits
mailing list