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/Transforms/LoopUnroll.cpp b/lib/Transforms/LoopUnroll.cpp
index e663ce0..2fd8fe2 100644
--- a/lib/Transforms/LoopUnroll.cpp
+++ b/lib/Transforms/LoopUnroll.cpp
@@ -19,18 +19,15 @@
//
//===----------------------------------------------------------------------===//
+#include "mlir/Transforms/Passes.h"
+
+#include "mlir/Analysis/LoopAnalysis.h"
#include "mlir/IR/AffineExpr.h"
-#include "mlir/IR/Attributes.h"
+#include "mlir/IR/AffineMap.h"
#include "mlir/IR/Builders.h"
-#include "mlir/IR/CFGFunction.h"
-#include "mlir/IR/MLFunction.h"
-#include "mlir/IR/Module.h"
-#include "mlir/IR/OperationSet.h"
#include "mlir/IR/StandardOps.h"
-#include "mlir/IR/Statements.h"
#include "mlir/IR/StmtVisitor.h"
#include "mlir/Transforms/Pass.h"
-#include "mlir/Transforms/Passes.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/Support/CommandLine.h"
@@ -128,7 +125,7 @@
ShortLoopGatherer(unsigned minTripCount) : minTripCount(minTripCount) {}
void visitForStmt(ForStmt *forStmt) {
- Optional<uint64_t> tripCount = forStmt->getConstantTripCount();
+ Optional<uint64_t> tripCount = getConstantTripCount(*forStmt);
if (tripCount.hasValue() && tripCount.getValue() <= minTripCount)
loops.push_back(forStmt);
}
@@ -174,12 +171,43 @@
// Unrolls this loop completely. Fails assertion if loop bounds are
// non-constant.
bool LoopUnroll::loopUnrollFull(ForStmt *forStmt) {
- Optional<uint64_t> tripCount = forStmt->getConstantTripCount();
+ Optional<uint64_t> tripCount = getConstantTripCount(*forStmt);
if (tripCount.hasValue())
return loopUnrollByFactor(forStmt, tripCount.getValue());
return false;
}
+/// Returns the upper bound of an unrolled loop with lower bound 'lb' and with
+/// the specified trip count, stride, and unroll factor.
+static AffineMap *getUnrolledLoopUpperBound(AffineMap *lbMap,
+ uint64_t tripCount,
+ unsigned unrollFactor, int64_t step,
+ MLFuncBuilder *builder) {
+ assert(lbMap->getNumResults() == 1);
+ auto *lbExpr = lbMap->getResult(0);
+ // lbExpr + (count - count % unrollFactor - 1) * step).
+ auto *expr = builder->getAddExpr(
+ lbExpr, builder->getConstantExpr(
+ (tripCount - tripCount % unrollFactor - 1) * step));
+ return builder->getAffineMap(lbMap->getNumDims(), lbMap->getNumSymbols(),
+ {expr}, {});
+}
+
+/// Returns the lower bound of the cleanup loop when unrolling a loop with lower
+/// bound 'lb' and with the specified trip count, stride, and unroll factor.
+static AffineMap *getCleanupLoopLowerBound(AffineMap *lbMap, uint64_t tripCount,
+ unsigned unrollFactor, int64_t step,
+ MLFuncBuilder *builder) {
+ assert(lbMap->getNumResults() == 1);
+ auto *lbExpr = lbMap->getResult(0);
+ // lbExpr + (count - count % unrollFactor) * step);
+ auto *expr = builder->getAddExpr(
+ lbExpr,
+ builder->getConstantExpr((tripCount - tripCount % unrollFactor) * step));
+ return builder->getAffineMap(lbMap->getNumDims(), lbMap->getNumSymbols(),
+ {expr}, {});
+}
+
/// Unrolls this loop by the specified unroll factor.
bool LoopUnroll::loopUnrollByFactor(ForStmt *forStmt, uint64_t unrollFactor) {
assert(unrollFactor >= 1 && "unroll factor shoud be >= 1");
@@ -187,44 +215,96 @@
if (unrollFactor == 1 || forStmt->getStatements().empty())
return false;
- if (!forStmt->hasConstantBounds())
+ Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(*forStmt);
+
+ if (!mayBeConstantTripCount.hasValue() &&
+ getLargestDivisorOfTripCount(*forStmt) % unrollFactor != 0)
return false;
- int64_t lb = forStmt->getConstantLowerBound();
+ const AffineBound &lb = forStmt->getLowerBound();
+ const AffineBound &ub = forStmt->getLowerBound();
+ auto lbMap = lb.getMap();
+ auto ubMap = lb.getMap();
+
+ // Loops with max/min expressions won't be unrolled here (the output can't be
+ // expressed as an MLFunction in the general case). However, the right way to
+ // do such unrolling for an MLFunction would be to specialize the loop for the
+ // 'hotspot' case and unroll that hotspot case.
+ if (lbMap->getNumResults() != 1 || ubMap->getNumResults() != 1)
+ return false;
+
+ // TODO(bondhugula): handle bounds with different sets of operands.
+ // Same operand list for now.
+ if (lbMap->getNumDims() != ubMap->getNumDims() ||
+ lbMap->getNumSymbols() != ubMap->getNumSymbols())
+ return false;
+ unsigned i, e = lb.getNumOperands();
+ for (i = 0; i < e; i++) {
+ if (lb.getStmtOperand(i).get() != ub.getStmtOperand(i).get())
+ break;
+ }
+ if (i != e)
+ return false;
+
int64_t step = forStmt->getStep();
- uint64_t tripCount = forStmt->getConstantTripCount().getValue();
// If the trip count is lower than the unroll factor, no unrolled body.
// TODO(bondhugula): option to specify cleanup loop unrolling.
- if (tripCount < unrollFactor)
- return true;
+ if (mayBeConstantTripCount.hasValue() &&
+ mayBeConstantTripCount.getValue() < unrollFactor)
+ return false;
// Generate the cleanup loop if trip count isn't a multiple of unrollFactor.
- if (tripCount % unrollFactor) {
+ // If the trip count is unknown, we currently unroll only when the unknown
+ // trip count is known to be a multiple of unroll factor - hence, no cleanup
+ // loop will be necessary in those cases.
+ // TODO(bondhugula): handle generation of cleanup loop for unknown trip count
+ // when it's not known to be a multiple of unroll factor (still for single
+ // result / same operands case).
+ if (mayBeConstantTripCount.hasValue() &&
+ mayBeConstantTripCount.getValue() % unrollFactor != 0) {
+ uint64_t tripCount = mayBeConstantTripCount.getValue();
DenseMap<const MLValue *, MLValue *> operandMap;
MLFuncBuilder builder(forStmt->getBlock(), ++StmtBlock::iterator(forStmt));
auto *cleanupForStmt = cast<ForStmt>(builder.clone(*forStmt, operandMap));
- cleanupForStmt->setConstantLowerBound(
- lb + (tripCount - tripCount % unrollFactor) * step);
+ if (forStmt->hasConstantLowerBound()) {
+ cleanupForStmt->setConstantLowerBound(
+ forStmt->getConstantLowerBound() +
+ (tripCount - tripCount % unrollFactor) * step);
+ } else {
+ cleanupForStmt->setLowerBoundMap(
+ getCleanupLoopLowerBound(forStmt->getLowerBoundMap(), tripCount,
+ unrollFactor, step, &builder));
+ }
// Promote the loop body up if this has turned into a single iteration loop.
promoteIfSingleIteration(cleanupForStmt);
+
+ // The upper bound needs to be adjusted.
+ if (forStmt->hasConstantUpperBound()) {
+ forStmt->setConstantUpperBound(
+ forStmt->getConstantLowerBound() +
+ (tripCount - tripCount % unrollFactor - 1) * step);
+ } else {
+ forStmt->setUpperBoundMap(
+ getUnrolledLoopUpperBound(forStmt->getLowerBoundMap(), tripCount,
+ unrollFactor, step, &builder));
+ }
}
+ // Scale the step of loop being unrolled by unroll factor.
+ forStmt->setStep(step * unrollFactor);
+
// 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->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).
- StmtBlock::iterator srcBlockEnd = --forStmt->end();
+ StmtBlock::iterator srcBlockEnd = std::prev(forStmt->end());
- // Unroll the contents of 'forStmt' (unrollFactor-1 additional copies
- // appended).
+ // Unroll the contents of 'forStmt' (append unrollFactor-1 additional copies).
for (unsigned i = 1; i < unrollFactor; i++) {
- DenseMap<const MLValue *, MLValue *> operandMapping;
+ DenseMap<const MLValue *, MLValue *> operandMap;
// If the induction variable is used, create a remapping to the value for
// this unrolled instance.
@@ -236,15 +316,13 @@
auto *ivUnroll =
builder.create<AffineApplyOp>(forStmt->getLoc(), bumpMap, forStmt)
->getResult(0);
- operandMapping[forStmt] = cast<MLValue>(ivUnroll);
+ operandMap[forStmt] = cast<MLValue>(ivUnroll);
}
- // Clone the original body of the loop (this doesn't include the last stmt).
- for (auto it = forStmt->begin(); it != srcBlockEnd; it++) {
- builder.clone(*it, operandMapping);
+ // Clone the original body of 'forStmt'.
+ for (auto it = forStmt->begin(); it != std::next(srcBlockEnd); it++) {
+ builder.clone(*it, operandMap);
}
- // Clone the last statement in the original body.
- builder.clone(*srcBlockEnd, operandMapping);
}
// Promote the loop body up if this has turned into a single iteration loop.
diff --git a/lib/Transforms/LoopUnrollJam.cpp b/lib/Transforms/LoopUnrollAndJam.cpp
similarity index 87%
rename from lib/Transforms/LoopUnrollJam.cpp
rename to lib/Transforms/LoopUnrollAndJam.cpp
index 6fb6134..7fb3309 100644
--- a/lib/Transforms/LoopUnrollJam.cpp
+++ b/lib/Transforms/LoopUnrollAndJam.cpp
@@ -1,5 +1,4 @@
-//===- LoopUnrollAndJam.cpp - Code to perform loop unroll jam
-//----------------===//
+//===- LoopUnrollAndJam.cpp - Code to perform loop unroll and jam ---------===//
//
// Copyright 2019 The MLIR Authors.
//
@@ -16,10 +15,10 @@
// limitations under the License.
// =============================================================================
//
-// This file implements loop unroll jam for MLFunctions. Unroll and jam is a
+// This file implements loop unroll and jam for MLFunctions. Unroll and jam is a
// transformation that improves locality, in particular, register reuse, while
// also improving instruction level parallelism. The example below shows what it
-// does in nearly the general case. Loop unroll jam currently works if the
+// does in nearly the general case. Loop unroll and jam currently works if the
// bounds of the loops inner to the loop being unroll-jammed do not depend on
// the latter.
//
@@ -44,25 +43,27 @@
// stmt's, bodies of those loops will not be jammed.
//
//===----------------------------------------------------------------------===//
+#include "mlir/Transforms/Passes.h"
+
+#include "mlir/Analysis/LoopAnalysis.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/StandardOps.h"
#include "mlir/IR/StmtVisitor.h"
#include "mlir/Transforms/Pass.h"
-#include "mlir/Transforms/Passes.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/Support/CommandLine.h"
using namespace mlir;
-// Loop unroll jam factor.
+// Loop unroll and jam factor.
static llvm::cl::opt<unsigned>
clUnrollJamFactor("unroll-jam-factor", llvm::cl::Hidden,
llvm::cl::desc("Use this unroll jam factor for all loops"
" (default 4)"));
namespace {
-/// Loop unroll jam pass. For test purposes, this just unroll jams the first
+/// Loop unroll jam pass. Currently, this just unroll jams the first
/// outer loop in an MLFunction.
struct LoopUnrollAndJam : public MLFunctionPass {
Optional<unsigned> unrollJamFactor;
@@ -127,10 +128,8 @@
while (it != End && !isa<ForStmt>(it))
++it;
if (it != subBlockStart)
- // Record the last statement (one behind the iterator) while not
- // changing the iterator position.
- subBlocks.push_back({subBlockStart, (--it)++});
- // Process all for Stmts that appear next.
+ subBlocks.push_back({subBlockStart, std::prev(it)});
+ // Process all for stmts that appear next.
while (it != End && isa<ForStmt>(it))
walkForStmt(cast<ForStmt>(it++));
}
@@ -142,12 +141,14 @@
if (unrollJamFactor == 1 || forStmt->getStatements().empty())
return false;
- if (!forStmt->hasConstantBounds())
+ Optional<uint64_t> mayTripCount = getConstantTripCount(*forStmt).getValue();
+
+ if (!mayTripCount.hasValue())
return false;
+ uint64_t tripCount = mayTripCount.getValue();
int64_t lb = forStmt->getConstantLowerBound();
int64_t step = forStmt->getStep();
- uint64_t tripCount = forStmt->getConstantTripCount().getValue();
// If the trip count is lower than the unroll jam factor, no unrolled body.
// TODO(bondhugula): option to specify cleanup loop unrolling.
@@ -164,7 +165,8 @@
if (tripCount % unrollJamFactor) {
DenseMap<const MLValue *, MLValue *> operandMap;
// Insert the cleanup loop right after 'forStmt'.
- MLFuncBuilder builder(forStmt->getBlock(), ++StmtBlock::iterator(forStmt));
+ MLFuncBuilder builder(forStmt->getBlock(),
+ std::next(StmtBlock::iterator(forStmt)));
auto *cleanupForStmt = cast<ForStmt>(builder.clone(*forStmt, operandMap));
cleanupForStmt->setConstantLowerBound(
lb + (tripCount - tripCount % unrollJamFactor) * step);
@@ -200,13 +202,10 @@
->getResult(0);
operandMapping[forStmt] = cast<MLValue>(ivUnroll);
}
- // Clone the sub-block being unroll-jammed (this doesn't include the last
- // stmt because subBlock.second is inclusive).
- for (auto it = subBlock.first; it != subBlock.second; ++it) {
+ // Clone the sub-block being unroll-jammed.
+ for (auto it = subBlock.first; it != std::next(subBlock.second); ++it) {
builder.clone(*it, operandMapping);
}
- // Clone the last statement of the sub-block.
- builder.clone(*subBlock.second, operandMapping);
}
}
diff --git a/lib/Transforms/LoopUtils.cpp b/lib/Transforms/LoopUtils.cpp
index cc2dc40..2eecb6f 100644
--- a/lib/Transforms/LoopUtils.cpp
+++ b/lib/Transforms/LoopUtils.cpp
@@ -19,16 +19,19 @@
//
//===----------------------------------------------------------------------===//
+#include "mlir/Transforms/Passes.h"
+
+#include "mlir/Analysis/LoopAnalysis.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/StandardOps.h"
#include "mlir/IR/Statements.h"
#include "mlir/IR/StmtVisitor.h"
-#include "mlir/Transforms/Passes.h"
/// Promotes the loop body of a forStmt to its containing block if the forStmt
/// was known to have a single iteration. Returns false otherwise.
+// TODO(bondhugula): extend this for arbitrary affine bounds.
bool mlir::promoteIfSingleIteration(ForStmt *forStmt) {
- Optional<uint64_t> tripCount = forStmt->getConstantTripCount();
+ Optional<uint64_t> tripCount = getConstantTripCount(*forStmt);
if (!tripCount.hasValue() || !forStmt->hasConstantLowerBound())
return false;