Implement operands for the lower and upper bounds of the for statement.
This revamps implementation of the loop bounds in the ForStmt, using general representation that supports operands. The frequent case of constant bounds is supported
via special access methods.
This also includes:
- Operand iterators for the Statement class.
- OpPointer::is() method to query the class of the Operation.
- Support for the bound shorthand notation parsing and printing.
- Validity checks for the bound operands used as dim ids and symbols
I didn't mean this CL to be so large. It just happened this way, as one thing led to another.
PiperOrigin-RevId: 210204858
diff --git a/lib/Transforms/LoopUnroll.cpp b/lib/Transforms/LoopUnroll.cpp
index ffb2947..6a5c667 100644
--- a/lib/Transforms/LoopUnroll.cpp
+++ b/lib/Transforms/LoopUnroll.cpp
@@ -19,6 +19,7 @@
//
//===----------------------------------------------------------------------===//
+#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/CFGFunction.h"
@@ -128,8 +129,10 @@
ShortLoopGatherer(unsigned minTripCount) : minTripCount(minTripCount) {}
void visitForStmt(ForStmt *forStmt) {
- auto lb = forStmt->getLowerBound()->getValue();
- auto ub = forStmt->getUpperBound()->getValue();
+ if (!forStmt->hasConstantBounds())
+ return;
+ auto lb = forStmt->getConstantLowerBound();
+ auto ub = forStmt->getConstantUpperBound();
auto step = forStmt->getStep();
if ((ub - lb) / step + 1 <= minTripCount)
@@ -174,18 +177,19 @@
return loopUnrollByFactor(forStmt, 4);
}
-// Unrolls this loop completely.
+// Unrolls this loop completely. Fails assertion if loop bounds are
+// non-constant.
bool LoopUnroll::loopUnrollFull(ForStmt *forStmt) {
- auto lb = forStmt->getLowerBound()->getValue();
- auto ub = forStmt->getUpperBound()->getValue();
+ auto lb = forStmt->getConstantLowerBound();
+ auto ub = forStmt->getConstantUpperBound();
auto step = forStmt->getStep();
- // Builder to add constants need for the unrolled iterator.
+ // Builder to add constants needed for the unrolled iterator.
auto *mlFunc = forStmt->findFunction();
MLFuncBuilder funcTopBuilder(&mlFunc->front());
// Builder to insert the unrolled bodies. We insert right after the
- /// ForStmt we're unrolling.
+ // ForStmt we're unrolling.
MLFuncBuilder builder(forStmt->getBlock(), ++StmtBlock::iterator(forStmt));
// Unroll the contents of 'forStmt'.
@@ -218,8 +222,11 @@
if (unrollFactor == 1 || forStmt->getStatements().empty())
return false;
- auto lb = forStmt->getLowerBound()->getValue();
- auto ub = forStmt->getUpperBound()->getValue();
+ if (!forStmt->hasConstantBounds())
+ return false;
+
+ auto lb = forStmt->getConstantLowerBound();
+ auto ub = forStmt->getConstantUpperBound();
auto step = forStmt->getStep();
int64_t tripCount = (int64_t)ceilf((ub - lb + 1) / (float)step);
@@ -234,16 +241,16 @@
DenseMap<const MLValue *, MLValue *> operandMap;
MLFuncBuilder builder(forStmt->getBlock(), ++StmtBlock::iterator(forStmt));
auto *cleanupForStmt = cast<ForStmt>(builder.clone(*forStmt, operandMap));
- cleanupForStmt->setLowerBound(builder.getConstantExpr(
- lb + (tripCount - tripCount % unrollFactor) * step));
+ cleanupForStmt->setConstantLowerBound(
+ lb + (tripCount - tripCount % unrollFactor) * step);
}
// Builder to insert unrolled bodies right after the last statement in the
// body of 'forStmt'.
MLFuncBuilder builder(forStmt, StmtBlock::iterator(forStmt->end()));
forStmt->setStep(step * unrollFactor);
- forStmt->setUpperBound(builder.getConstantExpr(
- lb + (tripCount - tripCount % unrollFactor - 1) * step));
+ forStmt->setConstantUpperBound(
+ lb + (tripCount - tripCount % unrollFactor - 1) * step);
// Keep a pointer to the last statement in the original block so that we know
// what to clone (since we are doing this in-place).