[MLIR][Affine] Enhance simplifyAdd for AffineExpr mod (#146492)

Currently AffineExpr Add has ability to optimize `"s1 + (s1 // c * -c)"
to "s1 % c"`,
but can not optimize `"(s0 + s1) + (s1 // c * -c)"`. 
This patch provide an opportunity to do this simplification, let it can
be simplified to `"s0 + s1 % c"`.
This commit is contained in:
XiangZhang
2025-07-02 11:08:58 +08:00
committed by GitHub
parent eb07f0d4a9
commit aa1d9a4c31
2 changed files with 15 additions and 12 deletions

View File

@@ -781,6 +781,15 @@ static AffineExpr simplifyAdd(AffineExpr lhs, AffineExpr rhs) {
if (isPositiveRhs && lhs == llrhs && rlrhs == -rrhs) {
return lhs % rlrhs;
}
// Try simplify lhs's last operand with rhs. e.g:
// (s0 * 64 + s1) + (s1 // c * -c) --->
// s0 * 64 + (s1 + s1 // c * -c) -->
// s0 * 64 + s1 % c
if (lBinOpExpr && lBinOpExpr.getKind() == AffineExprKind::Add) {
if (auto simplified = simplifyAdd(lBinOpExpr.getRHS(), rhs))
return lBinOpExpr.getLHS() + simplified;
}
return nullptr;
}