Skip to content

[mlir][AffineExpr] Order arguments in the commutative affine exprs #146895

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jul 4, 2025
Merged

[mlir][AffineExpr] Order arguments in the commutative affine exprs #146895

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0fd934f
[mlir][AffineExpr] Order arguments in the commutative affine exprs
Hardcode84 Jul 3, 2025
67adece
"drop constants case"
Hardcode84 Jul 3, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 36 additions & 2 deletions mlir/lib/IR/AffineExpr.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -784,16 +784,48 @@ static AffineExpr simplifyAdd(AffineExpr lhs, AffineExpr rhs) {
return nullptr;
}

static std::pair<AffineExpr, AffineExpr>
orderCommutativeArgs(AffineExpr expr1, AffineExpr expr2) {
auto sym1 = dyn_cast<AffineSymbolExpr>(expr1);
auto sym2 = dyn_cast<AffineSymbolExpr>(expr2);
// Try to order by symbol/dim position first
if (sym1 && sym2)
return sym1.getPosition() < sym2.getPosition() ? std::pair{expr1, expr2}
: std::pair{expr2, expr1};

auto dim1 = dyn_cast<AffineDimExpr>(expr1);
auto dim2 = dyn_cast<AffineDimExpr>(expr2);
if (dim1 && dim2)
return dim1.getPosition() < dim2.getPosition() ? std::pair{expr1, expr2}
: std::pair{expr2, expr1};

// Put dims before symbols
if (dim1 && sym2)
return {dim1, sym2};

if (sym1 && dim2)
return {dim2, sym1};

// Move constants to the right
if (isa<AffineConstantExpr>(expr1) && !isa<AffineConstantExpr>(expr2))
return {expr2, expr1};
Copy link
Member

@ftynse ftynse Jul 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you double-check composition with special cases. IIRC, in multiplication, we wanted to put the constant on the LHS, and there was something special about -1 as well.

Copy link
Contributor Author

@Hardcode84 Hardcode84 Jul 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks like simplifyMul/simplifyAdd already handles it, so I can just drop this case.


// Otherwise, keep original order
return {expr1, expr2};
}

AffineExpr AffineExpr::operator+(int64_t v) const {
return *this + getAffineConstantExpr(v, getContext());
}
AffineExpr AffineExpr::operator+(AffineExpr other) const {
if (auto simplified = simplifyAdd(*this, other))
return simplified;

auto [lhs, rhs] = orderCommutativeArgs(*this, other);

StorageUniquer &uniquer = getContext()->getAffineUniquer();
return uniquer.get<AffineBinaryOpExprStorage>(
/*initFn=*/{}, static_cast<unsigned>(AffineExprKind::Add), *this, other);
/*initFn=*/{}, static_cast<unsigned>(AffineExprKind::Add), lhs, rhs);
}

/// Simplify a multiply expression. Return nullptr if it can't be simplified.
Expand Down Expand Up @@ -856,9 +888,11 @@ AffineExpr AffineExpr::operator*(AffineExpr other) const {
if (auto simplified = simplifyMul(*this, other))
return simplified;

auto [lhs, rhs] = orderCommutativeArgs(*this, other);

StorageUniquer &uniquer = getContext()->getAffineUniquer();
return uniquer.get<AffineBinaryOpExprStorage>(
/*initFn=*/{}, static_cast<unsigned>(AffineExprKind::Mul), *this, other);
/*initFn=*/{}, static_cast<unsigned>(AffineExprKind::Mul), lhs, rhs);
}

// Unary minus, delegate to operator*.
Expand Down
4 changes: 2 additions & 2 deletions mlir/test/Dialect/Affine/simplify-structures.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -508,7 +508,7 @@ func.func @test_not_trivially_true_or_false_returning_three_results() -> (index,
// -----

// Test simplification of mod expressions.
// CHECK-DAG: #[[$MOD:.*]] = affine_map<()[s0, s1, s2, s3, s4] -> (s3 + s4 * s1 + (s0 - s1) mod s2)>
// CHECK-DAG: #[[$MOD:.*]] = affine_map<()[s0, s1, s2, s3, s4] -> (s1 * s4 + s3 + (s0 - s1) mod s2)>
// CHECK-DAG: #[[$SIMPLIFIED_MOD_RHS:.*]] = affine_map<()[s0, s1, s2, s3] -> (s3 mod (s2 - s0 * s1))>
// CHECK-DAG: #[[$MODULO_AND_PRODUCT:.*]] = affine_map<()[s0, s1, s2, s3] -> (s0 * s1 + s3 - (-s0 + s3) mod s2)>
// CHECK-LABEL: func @semiaffine_simplification_mod
Expand Down Expand Up @@ -547,7 +547,7 @@ func.func @semiaffine_simplification_floordiv_and_ceildiv(%arg0: index, %arg1: i

// Test simplification of product expressions.
// CHECK-DAG: #[[$PRODUCT:.*]] = affine_map<()[s0, s1, s2, s3, s4] -> (s3 + s4 + (s0 - s1) * s2)>
// CHECK-DAG: #[[$SUM_OF_PRODUCTS:.*]] = affine_map<()[s0, s1, s2, s3, s4] -> (s2 + s2 * s0 + s3 + s3 * s0 + s3 * s1 + s4 + s4 * s1)>
// CHECK-DAG: #[[$SUM_OF_PRODUCTS:.*]] = affine_map<()[s0, s1, s2, s3, s4] -> (s0 * s2 + s0 * s3 + s1 * s3 + s1 * s4 + s2 + s3 + s4)>
// CHECK-LABEL: func @semiaffine_simplification_product
// CHECK-SAME: (%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index, %[[ARG5:.*]]: index)
func.func @semiaffine_simplification_product(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index) -> (index, index) {
Expand Down
2 changes: 1 addition & 1 deletion mlir/test/IR/affine-map.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,7 @@
#map44 = affine_map<(i, j) -> (i - 2*j, j * 6 floordiv 4)>

// Simplifications
// CHECK: #map{{[0-9]*}} = affine_map<(d0, d1, d2)[s0] -> (d0 + d1 + d2 + 1, d2 + d1, (d0 * s0) * 8)>
// CHECK: #map{{[0-9]*}} = affine_map<(d0, d1, d2)[s0] -> (d0 + d1 + d2 + 1, d1 + d2, (d0 * s0) * 8)>
#map45 = affine_map<(i, j, k) [N] -> (1 + i + 3 + j - 3 + k, k + 5 + j - 5, 2*i*4*N)>

// CHECK: #map{{[0-9]*}} = affine_map<(d0, d1, d2) -> (0, d1, d0 * 2, 0)>
Expand Down
23 changes: 23 additions & 0 deletions mlir/unittests/IR/AffineExprTest.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,20 @@ TEST(AffineExprTest, constantFolding) {
ASSERT_EQ(cminfloordivcn1.getKind(), AffineExprKind::FloorDiv);
}

TEST(AffineExprTest, commutative) {
MLIRContext ctx;
OpBuilder b(&ctx);
auto c2 = b.getAffineConstantExpr(1);
auto d0 = b.getAffineDimExpr(0);
auto d1 = b.getAffineDimExpr(1);
auto s0 = b.getAffineSymbolExpr(0);
auto s1 = b.getAffineSymbolExpr(1);

ASSERT_EQ(d0 * d1, d1 * d0);
ASSERT_EQ(s0 + s1, s1 + s0);
ASSERT_EQ(s0 * c2, c2 * s0);
}

TEST(AffineExprTest, divisionSimplification) {
MLIRContext ctx;
OpBuilder b(&ctx);
Expand Down Expand Up @@ -147,3 +161,12 @@ TEST(AffineExprTest, simpleAffineExprFlattenerRegression) {
ASSERT_TRUE(isa<AffineConstantExpr>(result));
ASSERT_EQ(cast<AffineConstantExpr>(result).getValue(), 7);
}

TEST(AffineExprTest, simplifyCommutative) {
MLIRContext ctx;
OpBuilder b(&ctx);
auto s0 = b.getAffineSymbolExpr(0);
auto s1 = b.getAffineSymbolExpr(1);

ASSERT_EQ(s0 * s1 - s1 * s0 + 1, 1);
}
Loading