Fix bug in visitDivExpr, visitModExpr and visitMulExpr
#145290
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
@llvm/pr-subscribers-mlir-affine @llvm/pr-subscribers-mlir Author: Arnab Dutta (arnab-polymage) ChangesWhenever the result of a div or mod affine expression is a constant expression, place the value in the constant index of the flattened expression instead of adding it as a local expression. Full diff: https://github.com/llvm/llvm-project/pull/145290.diff 1 Files Affected:
diff --git a/mlir/lib/IR/AffineExpr.cpp b/mlir/lib/IR/AffineExpr.cpp
index cc81f9d19aca7..feedef46c66b8 100644
--- a/mlir/lib/IR/AffineExpr.cpp
+++ b/mlir/lib/IR/AffineExpr.cpp
@@ -1177,10 +1177,10 @@ static AffineExpr getSemiAffineExprFromFlatForm(ArrayRef<int64_t> flatExprs,
if (flatExprs[numDims + numSymbols + it.index()] == 0)
continue;
AffineExpr expr = it.value();
- auto binaryExpr = dyn_cast<AffineBinaryOpExpr>(expr);
- if (!binaryExpr)
- continue;
-
+ assert(isa<AffineBinaryOpExpr>(expr) &&
+ "local expression cannot be a dimension, symbol or a constant -- it "
+ "should be a binary op expression");
+ auto binaryExpr = cast<AffineBinaryOpExpr>(expr);
AffineExpr lhs = binaryExpr.getLHS();
AffineExpr rhs = binaryExpr.getRHS();
if (!((isa<AffineDimExpr>(lhs) || isa<AffineSymbolExpr>(lhs)) &&
@@ -1348,6 +1348,11 @@ LogicalResult SimpleAffineExprFlattener::visitModExpr(AffineBinaryOpExpr expr) {
AffineExpr divisorExpr = getAffineExprFromFlatForm(rhs, numDims, numSymbols,
localExprs, context);
AffineExpr modExpr = dividendExpr % divisorExpr;
+ if (auto constModExpr = dyn_cast<AffineConstantExpr>(modExpr)) {
+ std::fill(lhs.begin(), lhs.end(), 0);
+ lhs[getConstantIndex()] = constModExpr.getValue();
+ return success();
+ }
return addLocalVariableSemiAffine(modLhs, rhs, modExpr, lhs, lhs.size());
}
@@ -1482,6 +1487,11 @@ LogicalResult SimpleAffineExprFlattener::visitDivExpr(AffineBinaryOpExpr expr,
AffineExpr b = getAffineExprFromFlatForm(rhs, numDims, numSymbols,
localExprs, context);
AffineExpr divExpr = isCeil ? a.ceilDiv(b) : a.floorDiv(b);
+ if (auto constDivExpr = dyn_cast<AffineConstantExpr>(divExpr)) {
+ std::fill(lhs.begin(), lhs.end(), 0);
+ lhs[getConstantIndex()] = constDivExpr.getValue();
+ return success();
+ }
return addLocalVariableSemiAffine(divLhs, rhs, divExpr, lhs, lhs.size());
}
|
|
@llvm/pr-subscribers-mlir-core Author: Arnab Dutta (arnab-polymage) ChangesWhenever the result of a div or mod affine expression is a constant expression, place the value in the constant index of the flattened expression instead of adding it as a local expression. Full diff: https://github.com/llvm/llvm-project/pull/145290.diff 1 Files Affected:
diff --git a/mlir/lib/IR/AffineExpr.cpp b/mlir/lib/IR/AffineExpr.cpp
index cc81f9d19aca7..feedef46c66b8 100644
--- a/mlir/lib/IR/AffineExpr.cpp
+++ b/mlir/lib/IR/AffineExpr.cpp
@@ -1177,10 +1177,10 @@ static AffineExpr getSemiAffineExprFromFlatForm(ArrayRef<int64_t> flatExprs,
if (flatExprs[numDims + numSymbols + it.index()] == 0)
continue;
AffineExpr expr = it.value();
- auto binaryExpr = dyn_cast<AffineBinaryOpExpr>(expr);
- if (!binaryExpr)
- continue;
-
+ assert(isa<AffineBinaryOpExpr>(expr) &&
+ "local expression cannot be a dimension, symbol or a constant -- it "
+ "should be a binary op expression");
+ auto binaryExpr = cast<AffineBinaryOpExpr>(expr);
AffineExpr lhs = binaryExpr.getLHS();
AffineExpr rhs = binaryExpr.getRHS();
if (!((isa<AffineDimExpr>(lhs) || isa<AffineSymbolExpr>(lhs)) &&
@@ -1348,6 +1348,11 @@ LogicalResult SimpleAffineExprFlattener::visitModExpr(AffineBinaryOpExpr expr) {
AffineExpr divisorExpr = getAffineExprFromFlatForm(rhs, numDims, numSymbols,
localExprs, context);
AffineExpr modExpr = dividendExpr % divisorExpr;
+ if (auto constModExpr = dyn_cast<AffineConstantExpr>(modExpr)) {
+ std::fill(lhs.begin(), lhs.end(), 0);
+ lhs[getConstantIndex()] = constModExpr.getValue();
+ return success();
+ }
return addLocalVariableSemiAffine(modLhs, rhs, modExpr, lhs, lhs.size());
}
@@ -1482,6 +1487,11 @@ LogicalResult SimpleAffineExprFlattener::visitDivExpr(AffineBinaryOpExpr expr,
AffineExpr b = getAffineExprFromFlatForm(rhs, numDims, numSymbols,
localExprs, context);
AffineExpr divExpr = isCeil ? a.ceilDiv(b) : a.floorDiv(b);
+ if (auto constDivExpr = dyn_cast<AffineConstantExpr>(divExpr)) {
+ std::fill(lhs.begin(), lhs.end(), 0);
+ lhs[getConstantIndex()] = constDivExpr.getValue();
+ return success();
+ }
return addLocalVariableSemiAffine(divLhs, rhs, divExpr, lhs, lhs.size());
}
|
arnab-polymage
force-pushed
the
ornib/mod_div_bug
branch
from
f51da54 to
7c184d5
Compare
bondhugula
reviewed
Jun 23, 2025
bondhugula
requested changes
Jun 23, 2025
CoTinker
reviewed
Jun 23, 2025
arnab-polymage
force-pushed
the
ornib/mod_div_bug
branch
from
7c184d5 to
75ab97d
Compare
bondhugula
reviewed
Jun 26, 2025
Which constraint? |
arnab-polymage
force-pushed
the
ornib/mod_div_bug
branch
from
75ab97d to
17ce1f0
Compare
arnab-polymage
changed the title
visitDivExpr and visitModExprvisitDivExpr, visitModExpr and visitMulExpr
CoTinker
requested changes
Jun 27, 2025
arnab-polymage
force-pushed
the
ornib/mod_div_bug
branch
from
17ce1f0 to
5fc2915
Compare
arnab-polymage
force-pushed
the
ornib/mod_div_bug
branch
from
5fc2915 to
7d262f5
Compare
arnab-polymage
force-pushed
the
ornib/mod_div_bug
branch
from
7d262f5 to
c2493e6
Compare
lipracer
reviewed
Jun 27, 2025
arnab-polymage
force-pushed
the
ornib/mod_div_bug
branch
from
c2493e6 to
c140442
Compare
bondhugula
reviewed
Jun 27, 2025
bondhugula
reviewed
Jun 27, 2025
Whenever the result of a div or mod affine expression is a constant expression, place the value in the constant index of the flattened expression instead of adding it as a local expression.
arnab-polymage
force-pushed
the
ornib/mod_div_bug
branch
from
c140442 to
f7715ad
Compare
bondhugula
reviewed
Jun 27, 2025
| @@ -418,6 +418,14 @@ class SimpleAffineExprFlattener | |||
| AffineExpr localExpr); | |||
|
|
|||
| private: | |||
| /// Flatten binary expression `expr` and add it to `result`. If `expr` is a | |||
| /// dimension, symbol or constant, we add it to appropriate index in `result`. | |||
| /// Otherwise we add it in the local variable section. `lhs` and `rhs` are the | |||
There was a problem hiding this comment.
It's unclear here what operands of expr are! An AffineExpr doesn't have any operands. Rephrase.
bondhugula
reviewed
Jun 27, 2025
| @@ -1274,6 +1273,27 @@ SimpleAffineExprFlattener::SimpleAffineExprFlattener(unsigned numDims, | |||
| operandExprStack.reserve(8); | |||
| } | |||
|
|
|||
| LogicalResult SimpleAffineExprFlattener::addExprToFlattenedList( | |||
There was a problem hiding this comment.
I think this is completely misnamed. Is this specific to semi-affine exprs? In fact, addLocalIdSemiAffine in any of its overrides doesn't even use lhs and rhs. So all of these arguments aren't making sense to me.
LogicalResult SimpleAffineExprFlattener::addLocalIdSemiAffine(
ArrayRef<int64_t> lhs, ArrayRef<int64_t> rhs, AffineExpr localExpr) {
for (SmallVector<int64_t, 8> &subExpr : operandExprStack)
subExpr.insert(subExpr.begin() + getLocalVarStartIndex() + numLocals, 0);
localExprs.push_back(localExpr);
++numLocals;
// lhs and rhs are not used here; an override of this method uses them.
return success();
}
bondhugula
reviewed
Jun 27, 2025
| // CHECK-DAG: #[[$MAP:.*]] = affine_map<()[s0] -> (13 mod s0)> | ||
| // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0) -> (d0 * 2)> | ||
| // CHECK-LABEL: semiaffine_simplification_local_expr_folded_into_non_binary_expr | ||
| func.func @semiaffine_simplification_local_expr_folded_into_non_binary_expr(%arg0: memref<?x?xf32>) -> (index, index) { |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Whenever the result of a div, mod or mul affine expression is a constant expression, place the value in the constant index of the flattened expression instead of adding it as a local expression.