Extend getConstantTripCount to deal with a larger subset of loop bounds; make loop
unroll/unroll-and-jam more powerful; add additional affine expr builder methods
- use previously added analysis/simplification to infer multiple of unroll
factor trip counts, making loop unroll/unroll-and-jam more general.
- for loop unroll, support bounds that are single result affine map's with the
same set of operands. For unknown loop bounds, loop unroll will now work as
long as trip count can be determined to be a multiple of unroll factor.
- extend getConstantTripCount to deal with single result affine map's with the
same operands. move it to mlir/Analysis/LoopAnalysis.cpp
- add additional builder utility methods for affine expr arithmetic
(difference, mod/floordiv/ceildiv w.r.t postitive constant). simplify code to
use the utility methods.
- move affine analysis routines to AffineAnalysis.cpp/.h from
AffineStructures.cpp/.h.
- Rename LoopUnrollJam to LoopUnrollAndJam to match class name.
- add an additional simplification for simplifyFloorDiv, simplifyCeilDiv
- Rename AffineMap::getNumOperands() getNumInputs: an affine map by itself does
not have operands. Operands are passed to it through affine_apply, from loop
bounds/if condition's, etc., operands are stored in the latter.
This should be sufficiently powerful for now as far as unroll/unroll-and-jam go for TPU
code generation, and can move to other analyses/transformations.
Loop nests like these are now unrolled without any cleanup loop being generated.
for %i = 1 to 100 {
// unroll factor 4: no cleanup loop will be generated.
for %j = (d0) -> (d0) (%i) to (d0) -> (5*d0 + 3) (%i) {
%x = "foo"(%j) : (affineint) -> i32
}
}
for %i = 1 to 100 {
// unroll factor 4: no cleanup loop will be generated.
for %j = (d0) -> (d0) (%i) to (d0) -> (d0 - d mod 4 - 1) (%i) {
%y = "foo"(%j) : (affineint) -> i32
}
}
for %i = 1 to 100 {
for %j = (d0) -> (d0) (%i) to (d0) -> (d0 + 128) (%i) {
%x = "foo"() : () -> i32
}
}
TODO(bondhugula): extend this to LoopUnrollAndJam as well in the next CL (with minor
changes).
PiperOrigin-RevId: 212661212
diff --git a/lib/IR/AffineExpr.cpp b/lib/IR/AffineExpr.cpp
index 6393b3e..f43f288 100644
--- a/lib/IR/AffineExpr.cpp
+++ b/lib/IR/AffineExpr.cpp
@@ -47,6 +47,25 @@
}
}
+AffineExpr *AffineBinaryOpExpr::getSub(AffineExpr *lhs, AffineExpr *rhs,
+ MLIRContext *context) {
+ return getAdd(lhs, getMul(rhs, AffineConstantExpr::get(-1, context), context),
+ context);
+}
+
+/// Get affine expression that is expr incremented by 'inc'.
+AffineExpr *AffineBinaryOpExpr::getAdd(AffineExpr *expr, int64_t rhs,
+ MLIRContext *context) {
+ return get(AffineExpr::Kind::Add, expr, AffineConstantExpr::get(rhs, context),
+ context);
+}
+
+AffineExpr *AffineBinaryOpExpr::getCeilDiv(AffineExpr *lhs, uint64_t rhs,
+ MLIRContext *context) {
+ return get(AffineExpr::Kind::CeilDiv, lhs,
+ AffineConstantExpr::get(rhs, context), context);
+}
+
/// Returns true if this expression is made out of only symbols and
/// constants (no dimensional identifiers).
bool AffineExpr::isSymbolicOrConstant() const {
@@ -101,7 +120,8 @@
}
}
-uint64_t AffineExpr::getKnownGcd() const {
+/// Returns the greatest known integral divisor of this affine expression.
+uint64_t AffineExpr::getLargestKnownDivisor() const {
AffineBinaryOpExpr *binExpr = nullptr;
switch (kind) {
case Kind::SymbolId:
@@ -112,15 +132,17 @@
return std::abs(cast<AffineConstantExpr>(this)->getValue());
case Kind::Mul:
binExpr = cast<AffineBinaryOpExpr>(const_cast<AffineExpr *>(this));
- return binExpr->getLHS()->getKnownGcd() * binExpr->getRHS()->getKnownGcd();
+ return binExpr->getLHS()->getLargestKnownDivisor() *
+ binExpr->getRHS()->getLargestKnownDivisor();
case Kind::Add:
LLVM_FALLTHROUGH;
case Kind::FloorDiv:
case Kind::CeilDiv:
case Kind::Mod:
binExpr = cast<AffineBinaryOpExpr>(const_cast<AffineExpr *>(this));
- return llvm::GreatestCommonDivisor64(binExpr->getLHS()->getKnownGcd(),
- binExpr->getRHS()->getKnownGcd());
+ return llvm::GreatestCommonDivisor64(
+ binExpr->getLHS()->getLargestKnownDivisor(),
+ binExpr->getRHS()->getLargestKnownDivisor());
}
}
@@ -138,17 +160,18 @@
binExpr = cast<AffineBinaryOpExpr>(const_cast<AffineExpr *>(this));
// It's probably not worth optimizing this further (to not traverse the
// whole sub-tree under - it that would require a version of isMultipleOf
- // that on a 'false' return also returns the known GCD).
- return (l = binExpr->getLHS()->getKnownGcd()) % factor == 0 ||
- (u = binExpr->getRHS()->getKnownGcd()) % factor == 0 ||
+ // that on a 'false' return also returns the largest known divisor).
+ return (l = binExpr->getLHS()->getLargestKnownDivisor()) % factor == 0 ||
+ (u = binExpr->getRHS()->getLargestKnownDivisor()) % factor == 0 ||
(l * u) % factor == 0;
case Kind::Add:
case Kind::FloorDiv:
case Kind::CeilDiv:
case Kind::Mod:
binExpr = cast<AffineBinaryOpExpr>(const_cast<AffineExpr *>(this));
- return llvm::GreatestCommonDivisor64(binExpr->getLHS()->getKnownGcd(),
- binExpr->getRHS()->getKnownGcd()) %
+ return llvm::GreatestCommonDivisor64(
+ binExpr->getLHS()->getLargestKnownDivisor(),
+ binExpr->getRHS()->getLargestKnownDivisor()) %
factor ==
0;
}
diff --git a/lib/IR/AffineMap.cpp b/lib/IR/AffineMap.cpp
index 8d56f24..afa35ba 100644
--- a/lib/IR/AffineMap.cpp
+++ b/lib/IR/AffineMap.cpp
@@ -63,7 +63,7 @@
// If only one of them is a symbolic expressions, make it the RHS.
if (isa<AffineConstantExpr>(lhs) ||
(lhs->isSymbolicOrConstant() && !rhs->isSymbolicOrConstant())) {
- return AffineBinaryOpExpr::get(Kind::Add, rhs, lhs, context);
+ return AffineBinaryOpExpr::getAdd(rhs, lhs, context);
}
// At this point, if there was a constant, it would be on the right.
@@ -77,8 +77,8 @@
auto *lBin = dyn_cast<AffineBinaryOpExpr>(lhs);
if (lBin && rhsConst && lBin->getKind() == Kind::Add) {
if (auto *lrhs = dyn_cast<AffineConstantExpr>(lBin->getRHS()))
- return AffineBinaryOpExpr::get(
- Kind::Add, lBin->getLHS(),
+ return AffineBinaryOpExpr::getAdd(
+ lBin->getLHS(),
AffineConstantExpr::get(lrhs->getValue() + rhsConst->getValue(),
context),
context);
@@ -88,10 +88,9 @@
// + d1 into (d0 + d1) + 2.
if (lBin && lBin->getKind() == Kind::Add) {
if (auto *lrhs = dyn_cast<AffineConstantExpr>(lBin->getRHS())) {
- return AffineBinaryOpExpr::get(
- Kind::Add,
- AffineBinaryOpExpr::get(Kind::Add, lBin->getLHS(), rhs, context),
- lrhs, context);
+ return AffineBinaryOpExpr::getAdd(
+ AffineBinaryOpExpr::getAdd(lBin->getLHS(), rhs, context), lrhs,
+ context);
}
}
@@ -115,7 +114,7 @@
// constant. (Note that a constant is trivially symbolic).
if (!rhs->isSymbolicOrConstant() || isa<AffineConstantExpr>(lhs)) {
// At least one of them has to be symbolic.
- return AffineBinaryOpExpr::get(Kind::Mul, rhs, lhs, context);
+ return AffineBinaryOpExpr::getMul(rhs, lhs, context);
}
// At this point, if there was a constant, it would be on the right.
@@ -133,8 +132,8 @@
auto *lBin = dyn_cast<AffineBinaryOpExpr>(lhs);
if (lBin && rhsConst && lBin->getKind() == Kind::Mul) {
if (auto *lrhs = dyn_cast<AffineConstantExpr>(lBin->getRHS()))
- return AffineBinaryOpExpr::get(
- Kind::Mul, lBin->getLHS(),
+ return AffineBinaryOpExpr::getMul(
+ lBin->getLHS(),
AffineConstantExpr::get(lrhs->getValue() * rhsConst->getValue(),
context),
context);
@@ -144,10 +143,9 @@
// * 2) * d1 into (d0 * d1) * 2.
if (lBin && lBin->getKind() == Kind::Mul) {
if (auto *lrhs = dyn_cast<AffineConstantExpr>(lBin->getRHS())) {
- return AffineBinaryOpExpr::get(
- Kind::Mul,
- AffineBinaryOpExpr::get(Kind::Mul, lBin->getLHS(), rhs, context),
- lrhs, context);
+ return AffineBinaryOpExpr::getMul(
+ AffineBinaryOpExpr::getMul(lBin->getLHS(), rhs, context), lrhs,
+ context);
}
}
@@ -167,13 +165,16 @@
// Fold floordiv of a multiply with a constant that is a multiple of the
// divisor. Eg: (i * 128) floordiv 64 = i * 2.
if (rhsConst) {
+ if (rhsConst->getValue() == 1)
+ return lhs;
+
auto *lBin = dyn_cast<AffineBinaryOpExpr>(lhs);
if (lBin && lBin->getKind() == Kind::Mul) {
if (auto *lrhs = dyn_cast<AffineConstantExpr>(lBin->getRHS())) {
// rhsConst is known to be positive if a constant.
if (lrhs->getValue() % rhsConst->getValue() == 0)
- return AffineBinaryOpExpr::get(
- Kind::Mul, lBin->getLHS(),
+ return AffineBinaryOpExpr::getMul(
+ lBin->getLHS(),
AffineConstantExpr::get(lrhs->getValue() / rhsConst->getValue(),
context),
context);
@@ -199,13 +200,16 @@
// Fold ceildiv of a multiply with a constant that is a multiple of the
// divisor. Eg: (i * 128) ceildiv 64 = i * 2.
if (rhsConst) {
+ if (rhsConst->getValue() == 1)
+ return lhs;
+
auto *lBin = dyn_cast<AffineBinaryOpExpr>(lhs);
if (lBin && lBin->getKind() == Kind::Mul) {
if (auto *lrhs = dyn_cast<AffineConstantExpr>(lBin->getRHS())) {
// rhsConst is known to be positive if a constant.
if (lrhs->getValue() % rhsConst->getValue() == 0)
- return AffineBinaryOpExpr::get(
- Kind::Mul, lBin->getLHS(),
+ return AffineBinaryOpExpr::getMul(
+ lBin->getLHS(),
AffineConstantExpr::get(lrhs->getValue() / rhsConst->getValue(),
context),
context);
@@ -230,7 +234,7 @@
// mod 64 is folded to 0, and less trivially, (i*(j*4*(k*32))) mod 128 = 0.
if (rhsConst) {
// rhsConst is known to be positive if a constant.
- if (lhs->getKnownGcd() % rhsConst->getValue() == 0)
+ if (lhs->getLargestKnownDivisor() % rhsConst->getValue() == 0)
return AffineConstantExpr::get(0, context);
}
diff --git a/lib/IR/AsmPrinter.cpp b/lib/IR/AsmPrinter.cpp
index 28a97ea..83b9493 100644
--- a/lib/IR/AsmPrinter.cpp
+++ b/lib/IR/AsmPrinter.cpp
@@ -108,7 +108,7 @@
// Check if the affine map is single dim id or single symbol identity -
// (i)->(i) or ()[s]->(i)
- return boundMap->getNumOperands() == 1 && boundMap->getNumResults() == 1 &&
+ return boundMap->getNumInputs() == 1 && boundMap->getNumResults() == 1 &&
(isa<AffineDimExpr>(boundMap->getResult(0)) ||
isa<AffineSymbolExpr>(boundMap->getResult(0)));
}
diff --git a/lib/IR/Builders.cpp b/lib/IR/Builders.cpp
index a03de2f..f93508d 100644
--- a/lib/IR/Builders.cpp
+++ b/lib/IR/Builders.cpp
@@ -161,18 +161,46 @@
return AffineBinaryOpExpr::get(AffineExpr::Kind::Mul, lhs, rhs, context);
}
+// Most multiply expressions are pure affine (rhs is a constant).
+AffineExpr *Builder::getMulExpr(AffineExpr *lhs, int64_t rhs) {
+ return AffineBinaryOpExpr::get(AffineExpr::Kind::Mul, lhs,
+ getConstantExpr(rhs), context);
+}
+
+AffineExpr *Builder::getSubExpr(AffineExpr *lhs, AffineExpr *rhs) {
+ return getAddExpr(lhs, getMulExpr(rhs, getConstantExpr(-1)));
+}
+
AffineExpr *Builder::getModExpr(AffineExpr *lhs, AffineExpr *rhs) {
return AffineBinaryOpExpr::get(AffineExpr::Kind::Mod, lhs, rhs, context);
}
+// Most modulo expressions are pure affine.
+AffineExpr *Builder::getModExpr(AffineExpr *lhs, uint64_t rhs) {
+ return AffineBinaryOpExpr::get(AffineExpr::Kind::Mod, lhs,
+ getConstantExpr(rhs), context);
+}
+
AffineExpr *Builder::getFloorDivExpr(AffineExpr *lhs, AffineExpr *rhs) {
return AffineBinaryOpExpr::get(AffineExpr::Kind::FloorDiv, lhs, rhs, context);
}
+// Most floordiv expressions are pure affine.
+AffineExpr *Builder::getFloorDivExpr(AffineExpr *lhs, uint64_t rhs) {
+ return AffineBinaryOpExpr::get(AffineExpr::Kind::FloorDiv, lhs,
+ getConstantExpr(rhs), context);
+}
+
AffineExpr *Builder::getCeilDivExpr(AffineExpr *lhs, AffineExpr *rhs) {
return AffineBinaryOpExpr::get(AffineExpr::Kind::CeilDiv, lhs, rhs, context);
}
+// Most ceildiv expressions are pure affine.
+AffineExpr *Builder::getCeilDivExpr(AffineExpr *lhs, uint64_t rhs) {
+ return AffineBinaryOpExpr::get(AffineExpr::Kind::CeilDiv, lhs,
+ getConstantExpr(rhs), context);
+}
+
IntegerSet *Builder::getIntegerSet(unsigned dimCount, unsigned symbolCount,
ArrayRef<AffineExpr *> constraints,
ArrayRef<bool> isEq) {
diff --git a/lib/IR/Statement.cpp b/lib/IR/Statement.cpp
index c8d3336..14c920c 100644
--- a/lib/IR/Statement.cpp
+++ b/lib/IR/Statement.cpp
@@ -278,9 +278,9 @@
ForStmt *ForStmt::create(Location *location, ArrayRef<MLValue *> lbOperands,
AffineMap *lbMap, ArrayRef<MLValue *> ubOperands,
AffineMap *ubMap, int64_t step, MLIRContext *context) {
- assert(lbOperands.size() == lbMap->getNumOperands() &&
+ assert(lbOperands.size() == lbMap->getNumInputs() &&
"lower bound operand count does not match the affine map");
- assert(ubOperands.size() == ubMap->getNumOperands() &&
+ assert(ubOperands.size() == ubMap->getNumInputs() &&
"upper bound operand count does not match the affine map");
unsigned numOperands = lbOperands.size() + ubOperands.size();
@@ -306,11 +306,11 @@
}
const AffineBound ForStmt::getLowerBound() const {
- return AffineBound(*this, 0, lbMap->getNumOperands(), lbMap);
+ return AffineBound(*this, 0, lbMap->getNumInputs(), lbMap);
}
const AffineBound ForStmt::getUpperBound() const {
- return AffineBound(*this, lbMap->getNumOperands(), getNumOperands(), ubMap);
+ return AffineBound(*this, lbMap->getNumInputs(), getNumOperands(), ubMap);
}
void ForStmt::setLowerBound(ArrayRef<MLValue *> operands, AffineMap *map) {
@@ -327,6 +327,18 @@
this->ubMap = map;
}
+void ForStmt::setLowerBoundMap(AffineMap *map) {
+ assert(lbMap->getNumDims() == map->getNumDims() &&
+ lbMap->getNumSymbols() == map->getNumSymbols());
+ this->lbMap = map;
+}
+
+void ForStmt::setUpperBoundMap(AffineMap *map) {
+ assert(ubMap->getNumDims() == map->getNumDims() &&
+ ubMap->getNumSymbols() == map->getNumSymbols());
+ this->ubMap = map;
+}
+
bool ForStmt::hasConstantLowerBound() const {
return lbMap->isSingleConstant();
}
@@ -343,24 +355,6 @@
return ubMap->getSingleConstantResult();
}
-Optional<uint64_t> ForStmt::getConstantTripCount() const {
- // TODO(bondhugula): handle arbitrary lower/upper bounds.
- if (!hasConstantBounds())
- return None;
- int64_t lb = getConstantLowerBound();
- int64_t ub = getConstantUpperBound();
- int64_t step = getStep();
-
- // 0 iteration loops.
- if ((step >= 1 && lb > ub) || (step <= -1 && lb < ub))
- return 0;
-
- uint64_t tripCount = static_cast<uint64_t>((ub - lb + 1) % step == 0
- ? (ub - lb + 1) / step
- : (ub - lb + 1) / step + 1);
- return tripCount;
-}
-
void ForStmt::setConstantLowerBound(int64_t value) {
MLIRContext *context = getContext();
auto *expr = AffineConstantExpr::get(value, context);
@@ -463,9 +457,9 @@
auto *newFor = ForStmt::create(
getLoc(),
- ArrayRef<MLValue *>(operands).take_front(lbMap->getNumOperands()),
- lbMap, ArrayRef<MLValue *>(operands).take_back(ubMap->getNumOperands()),
- ubMap, forStmt->getStep(), context);
+ ArrayRef<MLValue *>(operands).take_front(lbMap->getNumInputs()), lbMap,
+ ArrayRef<MLValue *>(operands).take_back(ubMap->getNumInputs()), ubMap,
+ forStmt->getStep(), context);
// Remember the induction variable mapping.
operandMap[forStmt] = newFor;