NFC. Move out and expose affine expression simplification utility out of AffineOps lib #69813
Conversation
… of AffineOps lib Move out trivial affine expression simplification out of AffineOps library. Expose it from libIR. Users of such methods shouldn't have to rely on the AffineOps dialect. For eg., with this change, the method can be used now from lib/Analysis/ (FlatLinearConstraints) as well as AffineOps dialect canonicalization. This way those one won't need to depend on AffineOps for some simplification of affine expressions.
|
@llvm/pr-subscribers-mlir-core @llvm/pr-subscribers-mlir-affine Author: Uday Bondhugula (bondhugula) ChangesMove out trivial affine expression simplification out of AffineOps library. Expose it from libIR. Users of such methods shouldn't have to rely on the AffineOps dialect. For eg., with this change, the method can be used now from lib/Analysis/ (FlatLinearConstraints) as well as AffineOps dialect canonicalization. This way those one won't need to depend on AffineOps for some simplification of affine expressions. Full diff: https://github.com/llvm/llvm-project/pull/69813.diff 3 Files Affected:
diff --git a/mlir/include/mlir/IR/AffineExpr.h b/mlir/include/mlir/IR/AffineExpr.h
index 8ced8770591ee8c..69e02c94ef2708d 100644
--- a/mlir/include/mlir/IR/AffineExpr.h
+++ b/mlir/include/mlir/IR/AffineExpr.h
@@ -353,6 +353,20 @@ void bindSymbolsList(MLIRContext *ctx, MutableArrayRef<AffineExprTy> exprs) {
e = getAffineSymbolExpr(idx++, ctx);
}
+/// Get a lower or upper (depending on `isUpper`) bound for `expr` while using
+/// the constant lower and upper bounds for its inputs provided in
+/// `constLowerBounds` and `constUpperBounds`. Return std::nullopt if such a
+/// bound can't be computed. This method only handles simple sum of product
+/// expressions (w.r.t constant coefficients) so as to not depend on anything
+/// heavyweight in `Analysis`. Expressions of the form: c0*d0 + c1*d1 + c2*s0 +
+/// ... + c_n are handled. Expressions involving floordiv, ceildiv, mod or
+/// semi-affine ones will lead a none being returned.
+std::optional<int64_t>
+getBoundForAffineExpr(AffineExpr expr, unsigned numDims, unsigned numSymbols,
+ ArrayRef<std::optional<int64_t>> constLowerBounds,
+ ArrayRef<std::optional<int64_t>> constUpperBounds,
+ bool isUpper);
+
} // namespace mlir
namespace llvm {
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index f2b3171c1ab837b..fe869171ffbd2a6 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -700,93 +700,6 @@ static std::optional<int64_t> getUpperBound(Value iv) {
return forOp.getConstantUpperBound() - 1;
}
-/// Get a lower or upper (depending on `isUpper`) bound for `expr` while using
-/// the constant lower and upper bounds for its inputs provided in
-/// `constLowerBounds` and `constUpperBounds`. Return std::nullopt if such a
-/// bound can't be computed. This method only handles simple sum of product
-/// expressions (w.r.t constant coefficients) so as to not depend on anything
-/// heavyweight in `Analysis`. Expressions of the form: c0*d0 + c1*d1 + c2*s0 +
-/// ... + c_n are handled. Expressions involving floordiv, ceildiv, mod or
-/// semi-affine ones will lead std::nullopt being returned.
-static std::optional<int64_t>
-getBoundForExpr(AffineExpr expr, unsigned numDims, unsigned numSymbols,
- ArrayRef<std::optional<int64_t>> constLowerBounds,
- ArrayRef<std::optional<int64_t>> constUpperBounds,
- bool isUpper) {
- // Handle divs and mods.
- if (auto binOpExpr = expr.dyn_cast<AffineBinaryOpExpr>()) {
- // If the LHS of a floor or ceil is bounded and the RHS is a constant, we
- // can compute an upper bound.
- if (binOpExpr.getKind() == AffineExprKind::FloorDiv) {
- auto rhsConst = binOpExpr.getRHS().dyn_cast<AffineConstantExpr>();
- if (!rhsConst || rhsConst.getValue() < 1)
- return std::nullopt;
- auto bound = getBoundForExpr(binOpExpr.getLHS(), numDims, numSymbols,
- constLowerBounds, constUpperBounds, isUpper);
- if (!bound)
- return std::nullopt;
- return mlir::floorDiv(*bound, rhsConst.getValue());
- }
- if (binOpExpr.getKind() == AffineExprKind::CeilDiv) {
- auto rhsConst = binOpExpr.getRHS().dyn_cast<AffineConstantExpr>();
- if (rhsConst && rhsConst.getValue() >= 1) {
- auto bound =
- getBoundForExpr(binOpExpr.getLHS(), numDims, numSymbols,
- constLowerBounds, constUpperBounds, isUpper);
- if (!bound)
- return std::nullopt;
- return mlir::ceilDiv(*bound, rhsConst.getValue());
- }
- return std::nullopt;
- }
- if (binOpExpr.getKind() == AffineExprKind::Mod) {
- // lhs mod c is always <= c - 1 and non-negative. In addition, if `lhs` is
- // bounded such that lb <= lhs <= ub and lb floordiv c == ub floordiv c
- // (same "interval"), then lb mod c <= lhs mod c <= ub mod c.
- auto rhsConst = binOpExpr.getRHS().dyn_cast<AffineConstantExpr>();
- if (rhsConst && rhsConst.getValue() >= 1) {
- int64_t rhsConstVal = rhsConst.getValue();
- auto lb = getBoundForExpr(binOpExpr.getLHS(), numDims, numSymbols,
- constLowerBounds, constUpperBounds,
- /*isUpper=*/false);
- auto ub = getBoundForExpr(binOpExpr.getLHS(), numDims, numSymbols,
- constLowerBounds, constUpperBounds, isUpper);
- if (ub && lb &&
- floorDiv(*lb, rhsConstVal) == floorDiv(*ub, rhsConstVal))
- return isUpper ? mod(*ub, rhsConstVal) : mod(*lb, rhsConstVal);
- return isUpper ? rhsConstVal - 1 : 0;
- }
- }
- }
- // Flatten the expression.
- SimpleAffineExprFlattener flattener(numDims, numSymbols);
- flattener.walkPostOrder(expr);
- ArrayRef<int64_t> flattenedExpr = flattener.operandExprStack.back();
- // TODO: Handle local variables. We can get hold of flattener.localExprs and
- // get bound on the local expr recursively.
- if (flattener.numLocals > 0)
- return std::nullopt;
- int64_t bound = 0;
- // Substitute the constant lower or upper bound for the dimensional or
- // symbolic input depending on `isUpper` to determine the bound.
- for (unsigned i = 0, e = numDims + numSymbols; i < e; ++i) {
- if (flattenedExpr[i] > 0) {
- auto &constBound = isUpper ? constUpperBounds[i] : constLowerBounds[i];
- if (!constBound)
- return std::nullopt;
- bound += *constBound * flattenedExpr[i];
- } else if (flattenedExpr[i] < 0) {
- auto &constBound = isUpper ? constLowerBounds[i] : constUpperBounds[i];
- if (!constBound)
- return std::nullopt;
- bound += *constBound * flattenedExpr[i];
- }
- }
- // Constant term.
- bound += flattenedExpr.back();
- return bound;
-}
-
/// Determine a constant upper bound for `expr` if one exists while exploiting
/// values in `operands`. Note that the upper bound is an inclusive one. `expr`
/// is guaranteed to be less than or equal to it.
@@ -805,9 +718,9 @@ static std::optional<int64_t> getUpperBound(AffineExpr expr, unsigned numDims,
if (auto constExpr = expr.dyn_cast<AffineConstantExpr>())
return constExpr.getValue();
- return getBoundForExpr(expr, numDims, numSymbols, constLowerBounds,
- constUpperBounds,
- /*isUpper=*/true);
+ return getBoundForAffineExpr(expr, numDims, numSymbols, constLowerBounds,
+ constUpperBounds,
+ /*isUpper=*/true);
}
/// Determine a constant lower bound for `expr` if one exists while exploiting
@@ -829,9 +742,9 @@ static std::optional<int64_t> getLowerBound(AffineExpr expr, unsigned numDims,
if (auto constExpr = expr.dyn_cast<AffineConstantExpr>()) {
lowerBound = constExpr.getValue();
} else {
- lowerBound = getBoundForExpr(expr, numDims, numSymbols, constLowerBounds,
- constUpperBounds,
- /*isUpper=*/false);
+ lowerBound = getBoundForAffineExpr(expr, numDims, numSymbols,
+ constLowerBounds, constUpperBounds,
+ /*isUpper=*/false);
}
return lowerBound;
}
@@ -970,14 +883,14 @@ static void simplifyMinOrMaxExprWithOperands(AffineMap &map,
lowerBounds.push_back(constExpr.getValue());
upperBounds.push_back(constExpr.getValue());
} else {
- lowerBounds.push_back(getBoundForExpr(e, map.getNumDims(),
- map.getNumSymbols(),
- constLowerBounds, constUpperBounds,
- /*isUpper=*/false));
- upperBounds.push_back(getBoundForExpr(e, map.getNumDims(),
- map.getNumSymbols(),
- constLowerBounds, constUpperBounds,
- /*isUpper=*/true));
+ lowerBounds.push_back(
+ getBoundForAffineExpr(e, map.getNumDims(), map.getNumSymbols(),
+ constLowerBounds, constUpperBounds,
+ /*isUpper=*/false));
+ upperBounds.push_back(
+ getBoundForAffineExpr(e, map.getNumDims(), map.getNumSymbols(),
+ constLowerBounds, constUpperBounds,
+ /*isUpper=*/true));
}
}
diff --git a/mlir/lib/IR/AffineExpr.cpp b/mlir/lib/IR/AffineExpr.cpp
index 7eccbca4e6e7a1a..4b7ec89a842bd65 100644
--- a/mlir/lib/IR/AffineExpr.cpp
+++ b/mlir/lib/IR/AffineExpr.cpp
@@ -1438,3 +1438,83 @@ AffineExpr mlir::simplifyAffineExpr(AffineExpr expr, unsigned numDims,
assert(flattener.operandExprStack.empty());
return simplifiedExpr;
}
+
+std::optional<int64_t> mlir::getBoundForAffineExpr(
+ AffineExpr expr, unsigned numDims, unsigned numSymbols,
+ ArrayRef<std::optional<int64_t>> constLowerBounds,
+ ArrayRef<std::optional<int64_t>> constUpperBounds, bool isUpper) {
+ // Handle divs and mods.
+ if (auto binOpExpr = expr.dyn_cast<AffineBinaryOpExpr>()) {
+ // If the LHS of a floor or ceil is bounded and the RHS is a constant, we
+ // can compute an upper bound.
+ if (binOpExpr.getKind() == AffineExprKind::FloorDiv) {
+ auto rhsConst = binOpExpr.getRHS().dyn_cast<AffineConstantExpr>();
+ if (!rhsConst || rhsConst.getValue() < 1)
+ return std::nullopt;
+ auto bound =
+ getBoundForAffineExpr(binOpExpr.getLHS(), numDims, numSymbols,
+ constLowerBounds, constUpperBounds, isUpper);
+ if (!bound)
+ return std::nullopt;
+ return mlir::floorDiv(*bound, rhsConst.getValue());
+ }
+ if (binOpExpr.getKind() == AffineExprKind::CeilDiv) {
+ auto rhsConst = binOpExpr.getRHS().dyn_cast<AffineConstantExpr>();
+ if (rhsConst && rhsConst.getValue() >= 1) {
+ auto bound =
+ getBoundForAffineExpr(binOpExpr.getLHS(), numDims, numSymbols,
+ constLowerBounds, constUpperBounds, isUpper);
+ if (!bound)
+ return std::nullopt;
+ return mlir::ceilDiv(*bound, rhsConst.getValue());
+ }
+ return std::nullopt;
+ }
+ if (binOpExpr.getKind() == AffineExprKind::Mod) {
+ // lhs mod c is always <= c - 1 and non-negative. In addition, if `lhs` is
+ // bounded such that lb <= lhs <= ub and lb floordiv c == ub floordiv c
+ // (same "interval"), then lb mod c <= lhs mod c <= ub mod c.
+ auto rhsConst = binOpExpr.getRHS().dyn_cast<AffineConstantExpr>();
+ if (rhsConst && rhsConst.getValue() >= 1) {
+ int64_t rhsConstVal = rhsConst.getValue();
+ auto lb = getBoundForAffineExpr(binOpExpr.getLHS(), numDims, numSymbols,
+ constLowerBounds, constUpperBounds,
+ /*isUpper=*/false);
+ auto ub =
+ getBoundForAffineExpr(binOpExpr.getLHS(), numDims, numSymbols,
+ constLowerBounds, constUpperBounds, isUpper);
+ if (ub && lb &&
+ floorDiv(*lb, rhsConstVal) == floorDiv(*ub, rhsConstVal))
+ return isUpper ? mod(*ub, rhsConstVal) : mod(*lb, rhsConstVal);
+ return isUpper ? rhsConstVal - 1 : 0;
+ }
+ }
+ }
+ // Flatten the expression.
+ SimpleAffineExprFlattener flattener(numDims, numSymbols);
+ flattener.walkPostOrder(expr);
+ ArrayRef<int64_t> flattenedExpr = flattener.operandExprStack.back();
+ // TODO: Handle local variables. We can get hold of flattener.localExprs and
+ // get bound on the local expr recursively.
+ if (flattener.numLocals > 0)
+ return std::nullopt;
+ int64_t bound = 0;
+ // Substitute the constant lower or upper bound for the dimensional or
+ // symbolic input depending on `isUpper` to determine the bound.
+ for (unsigned i = 0, e = numDims + numSymbols; i < e; ++i) {
+ if (flattenedExpr[i] > 0) {
+ auto &constBound = isUpper ? constUpperBounds[i] : constLowerBounds[i];
+ if (!constBound)
+ return std::nullopt;
+ bound += *constBound * flattenedExpr[i];
+ } else if (flattenedExpr[i] < 0) {
+ auto &constBound = isUpper ? constLowerBounds[i] : constUpperBounds[i];
+ if (!constBound)
+ return std::nullopt;
+ bound += *constBound * flattenedExpr[i];
+ }
+ }
+ // Constant term.
+ bound += flattenedExpr.back();
+ return bound;
+}
|
|
I'm not sure I understand why something like FlatLinearConstraints would call out to this IR method instead of Presburger methods to do the same. Im not very supportive of doing such analysis on AffineMap. In my view, anything that requires a flattening of the AffineMap should be done be done on the flattened form, i.e. to get the bounds, first flatten it and then use Presburger methods to get these bounds (can be restricted to fast checks only). I do not have strong opinions on this, and would like to hear what others think. |
joker-eph
dismissed
their stale review
More discussion needed following comment from @Groverkss
@Groverkss Note that this method, which is being moved ( Anyone wanting to do any basic simplification on an affine expression (by exploiting lower or upper bounds on it) will need this method - |
I double-checked the code - for one of the paths, it does use the flattener, but it also works through the AffineExpr tree. From a layering standpoint, lib/IR/AffineExpr.cpp itself contains that flattener and also a method that depends on it (mlir::simplifyAffineExpr). So from a layering and consistency standpoint, adding |
Thank you for the example on getSliceBounds. I understand what you mean better now. I do see that the method works on native form, but I also looked at the implementation, and it essentially flattens the AffineMap and then does analysis on it. I know it handles some divs/mods on the tree form, but that can be handled with no overhead in the flattened form. What I'm asking is that instead of having such methods that flatten and do analysis, the user should do explicit conversion to the flattened form (concert to Presburger) and then do analysis there. I do not want to end up in case where we duplicate a lot of simplification functionality directly in AffineMap. This is just a discussion that I felt should be had at some point, and I'm okay with landing this and coming and removing these methods later if we come to a consensus. |
If other methods already do this, then let's land this and have this discussion later on discord/discourse. Thank you for the patch and the detailed replies. |
… of AffineOps lib (llvm#69813) Move out trivial affine expression simplification out of AffineOps library. Expose it from libIR. Users of such methods shouldn't have to rely on the AffineOps dialect. For eg., with this change, the method can be used now from lib/Analysis/ (FlatLinearConstraints) as well as AffineOps dialect canonicalization. This way those one won't need to depend on AffineOps for some simplification of affine expressions.
Move out trivial affine expression simplification out of AffineOps library. Expose it from libIR. Users of such methods shouldn't have to rely on the AffineOps dialect. For eg., with this change, the method can be used now from lib/Analysis/ (FlatLinearConstraints) as well as AffineOps dialect canonicalization.
This way those one won't need to depend on AffineOps for some simplification of affine expressions.