forked from OSchip/llvm-project
Move definitions of lopoUnroll* functions to LoopUtils.cpp.
These functions are declared in Transforms/LoopUtils.h (included to the Transforms/Utils library) but were defined in the loop unrolling pass in Transforms/LoopUnroll.cpp. As a result, targets depending only on TransformUtils library but not on Transforms could get link errors. Move the definitions to Transforms/Utils/LoopUtils.cpp where they should actually live. This does not modify any code. PiperOrigin-RevId: 221508882
This commit is contained in:
Alex Zinenko
jpienaar
parent
a30f7ec74f
commit
cb40633969
|
|
@ -169,123 +169,6 @@ bool LoopUnroll::runOnForStmt(ForStmt *forStmt) {
|
|||
return loopUnrollByFactor(forStmt, 4);
|
||||
}
|
||||
|
||||
/// Unrolls this loop completely.
|
||||
bool mlir::loopUnrollFull(ForStmt *forStmt) {
|
||||
Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(*forStmt);
|
||||
if (mayBeConstantTripCount.hasValue()) {
|
||||
uint64_t tripCount = mayBeConstantTripCount.getValue();
|
||||
if (tripCount == 1) {
|
||||
return promoteIfSingleIteration(forStmt);
|
||||
}
|
||||
return loopUnrollByFactor(forStmt, tripCount);
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
/// Unrolls and jams this loop by the specified factor or by the trip count (if
|
||||
/// constant) whichever is lower.
|
||||
bool mlir::loopUnrollUpToFactor(ForStmt *forStmt, uint64_t unrollFactor) {
|
||||
Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(*forStmt);
|
||||
|
||||
if (mayBeConstantTripCount.hasValue() &&
|
||||
mayBeConstantTripCount.getValue() < unrollFactor)
|
||||
return loopUnrollByFactor(forStmt, mayBeConstantTripCount.getValue());
|
||||
return loopUnrollByFactor(forStmt, unrollFactor);
|
||||
}
|
||||
|
||||
/// Unrolls this loop by the specified factor. Returns true if the loop
|
||||
/// is successfully unrolled.
|
||||
bool mlir::loopUnrollByFactor(ForStmt *forStmt, uint64_t unrollFactor) {
|
||||
assert(unrollFactor >= 1 && "unroll factor should be >= 1");
|
||||
|
||||
if (unrollFactor == 1 || forStmt->getStatements().empty())
|
||||
return false;
|
||||
|
||||
auto lbMap = forStmt->getLowerBoundMap();
|
||||
auto ubMap = forStmt->getUpperBoundMap();
|
||||
|
||||
// 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.
|
||||
if (lbMap.getNumResults() != 1 || ubMap.getNumResults() != 1)
|
||||
return false;
|
||||
|
||||
// Same operand list for lower and upper bound for now.
|
||||
// TODO(bondhugula): handle bounds with different operand lists.
|
||||
if (!forStmt->matchingBoundOperandList())
|
||||
return false;
|
||||
|
||||
Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(*forStmt);
|
||||
|
||||
// If the trip count is lower than the unroll factor, no unrolled body.
|
||||
// TODO(bondhugula): option to specify cleanup loop unrolling.
|
||||
if (mayBeConstantTripCount.hasValue() &&
|
||||
mayBeConstantTripCount.getValue() < unrollFactor)
|
||||
return false;
|
||||
|
||||
// Generate the cleanup loop if trip count isn't a multiple of unrollFactor.
|
||||
if (getLargestDivisorOfTripCount(*forStmt) % unrollFactor != 0) {
|
||||
DenseMap<const MLValue *, MLValue *> operandMap;
|
||||
MLFuncBuilder builder(forStmt->getBlock(), ++StmtBlock::iterator(forStmt));
|
||||
auto *cleanupForStmt = cast<ForStmt>(builder.clone(*forStmt, operandMap));
|
||||
auto clLbMap = getCleanupLoopLowerBound(*forStmt, unrollFactor, &builder);
|
||||
assert(clLbMap &&
|
||||
"cleanup loop lower bound map for single result bound maps can "
|
||||
"always be determined");
|
||||
cleanupForStmt->setLowerBoundMap(clLbMap);
|
||||
// Promote the loop body up if this has turned into a single iteration loop.
|
||||
promoteIfSingleIteration(cleanupForStmt);
|
||||
|
||||
// Adjust upper bound.
|
||||
auto unrolledUbMap =
|
||||
getUnrolledLoopUpperBound(*forStmt, unrollFactor, &builder);
|
||||
assert(unrolledUbMap &&
|
||||
"upper bound map can alwayys be determined for an unrolled loop "
|
||||
"with single result bounds");
|
||||
forStmt->setUpperBoundMap(unrolledUbMap);
|
||||
}
|
||||
|
||||
// Scale the step of loop being unrolled by unroll factor.
|
||||
int64_t step = forStmt->getStep();
|
||||
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()));
|
||||
|
||||
// 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 = std::prev(forStmt->end());
|
||||
|
||||
// Unroll the contents of 'forStmt' (append unrollFactor-1 additional copies).
|
||||
for (unsigned i = 1; i < unrollFactor; i++) {
|
||||
DenseMap<const MLValue *, MLValue *> operandMap;
|
||||
|
||||
// If the induction variable is used, create a remapping to the value for
|
||||
// this unrolled instance.
|
||||
if (!forStmt->use_empty()) {
|
||||
// iv' = iv + 1/2/3...unrollFactor-1;
|
||||
auto d0 = builder.getAffineDimExpr(0);
|
||||
auto bumpMap = builder.getAffineMap(1, 0, {d0 + i * step}, {});
|
||||
auto *ivUnroll =
|
||||
builder.create<AffineApplyOp>(forStmt->getLoc(), bumpMap, forStmt)
|
||||
->getResult(0);
|
||||
operandMap[forStmt] = cast<MLValue>(ivUnroll);
|
||||
}
|
||||
|
||||
// Clone the original body of 'forStmt'.
|
||||
for (auto it = forStmt->begin(); it != std::next(srcBlockEnd); it++) {
|
||||
builder.clone(*it, operandMap);
|
||||
}
|
||||
}
|
||||
|
||||
// Promote the loop body up if this has turned into a single iteration loop.
|
||||
promoteIfSingleIteration(forStmt);
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
FunctionPass *mlir::createLoopUnrollPass(int unrollFactor, int unrollFull) {
|
||||
return new LoopUnroll(unrollFactor == -1 ? None
|
||||
: Optional<unsigned>(unrollFactor),
|
||||
|
|
|
|||
|
|
@ -314,3 +314,120 @@ UtilResult mlir::stmtBodySkew(ForStmt *forStmt, ArrayRef<uint64_t> delays,
|
|||
|
||||
return UtilResult::Success;
|
||||
}
|
||||
|
||||
/// Unrolls this loop completely.
|
||||
bool mlir::loopUnrollFull(ForStmt *forStmt) {
|
||||
Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(*forStmt);
|
||||
if (mayBeConstantTripCount.hasValue()) {
|
||||
uint64_t tripCount = mayBeConstantTripCount.getValue();
|
||||
if (tripCount == 1) {
|
||||
return promoteIfSingleIteration(forStmt);
|
||||
}
|
||||
return loopUnrollByFactor(forStmt, tripCount);
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
/// Unrolls and jams this loop by the specified factor or by the trip count (if
|
||||
/// constant) whichever is lower.
|
||||
bool mlir::loopUnrollUpToFactor(ForStmt *forStmt, uint64_t unrollFactor) {
|
||||
Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(*forStmt);
|
||||
|
||||
if (mayBeConstantTripCount.hasValue() &&
|
||||
mayBeConstantTripCount.getValue() < unrollFactor)
|
||||
return loopUnrollByFactor(forStmt, mayBeConstantTripCount.getValue());
|
||||
return loopUnrollByFactor(forStmt, unrollFactor);
|
||||
}
|
||||
|
||||
/// Unrolls this loop by the specified factor. Returns true if the loop
|
||||
/// is successfully unrolled.
|
||||
bool mlir::loopUnrollByFactor(ForStmt *forStmt, uint64_t unrollFactor) {
|
||||
assert(unrollFactor >= 1 && "unroll factor should be >= 1");
|
||||
|
||||
if (unrollFactor == 1 || forStmt->getStatements().empty())
|
||||
return false;
|
||||
|
||||
auto lbMap = forStmt->getLowerBoundMap();
|
||||
auto ubMap = forStmt->getUpperBoundMap();
|
||||
|
||||
// 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.
|
||||
if (lbMap.getNumResults() != 1 || ubMap.getNumResults() != 1)
|
||||
return false;
|
||||
|
||||
// Same operand list for lower and upper bound for now.
|
||||
// TODO(bondhugula): handle bounds with different operand lists.
|
||||
if (!forStmt->matchingBoundOperandList())
|
||||
return false;
|
||||
|
||||
Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(*forStmt);
|
||||
|
||||
// If the trip count is lower than the unroll factor, no unrolled body.
|
||||
// TODO(bondhugula): option to specify cleanup loop unrolling.
|
||||
if (mayBeConstantTripCount.hasValue() &&
|
||||
mayBeConstantTripCount.getValue() < unrollFactor)
|
||||
return false;
|
||||
|
||||
// Generate the cleanup loop if trip count isn't a multiple of unrollFactor.
|
||||
if (getLargestDivisorOfTripCount(*forStmt) % unrollFactor != 0) {
|
||||
DenseMap<const MLValue *, MLValue *> operandMap;
|
||||
MLFuncBuilder builder(forStmt->getBlock(), ++StmtBlock::iterator(forStmt));
|
||||
auto *cleanupForStmt = cast<ForStmt>(builder.clone(*forStmt, operandMap));
|
||||
auto clLbMap = getCleanupLoopLowerBound(*forStmt, unrollFactor, &builder);
|
||||
assert(clLbMap &&
|
||||
"cleanup loop lower bound map for single result bound maps can "
|
||||
"always be determined");
|
||||
cleanupForStmt->setLowerBoundMap(clLbMap);
|
||||
// Promote the loop body up if this has turned into a single iteration loop.
|
||||
promoteIfSingleIteration(cleanupForStmt);
|
||||
|
||||
// Adjust upper bound.
|
||||
auto unrolledUbMap =
|
||||
getUnrolledLoopUpperBound(*forStmt, unrollFactor, &builder);
|
||||
assert(unrolledUbMap &&
|
||||
"upper bound map can alwayys be determined for an unrolled loop "
|
||||
"with single result bounds");
|
||||
forStmt->setUpperBoundMap(unrolledUbMap);
|
||||
}
|
||||
|
||||
// Scale the step of loop being unrolled by unroll factor.
|
||||
int64_t step = forStmt->getStep();
|
||||
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()));
|
||||
|
||||
// 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 = std::prev(forStmt->end());
|
||||
|
||||
// Unroll the contents of 'forStmt' (append unrollFactor-1 additional copies).
|
||||
for (unsigned i = 1; i < unrollFactor; i++) {
|
||||
DenseMap<const MLValue *, MLValue *> operandMap;
|
||||
|
||||
// If the induction variable is used, create a remapping to the value for
|
||||
// this unrolled instance.
|
||||
if (!forStmt->use_empty()) {
|
||||
// iv' = iv + 1/2/3...unrollFactor-1;
|
||||
auto d0 = builder.getAffineDimExpr(0);
|
||||
auto bumpMap = builder.getAffineMap(1, 0, {d0 + i * step}, {});
|
||||
auto *ivUnroll =
|
||||
builder.create<AffineApplyOp>(forStmt->getLoc(), bumpMap, forStmt)
|
||||
->getResult(0);
|
||||
operandMap[forStmt] = cast<MLValue>(ivUnroll);
|
||||
}
|
||||
|
||||
// Clone the original body of 'forStmt'.
|
||||
for (auto it = forStmt->begin(); it != std::next(srcBlockEnd); it++) {
|
||||
builder.clone(*it, operandMap);
|
||||
}
|
||||
}
|
||||
|
||||
// Promote the loop body up if this has turned into a single iteration loop.
|
||||
promoteIfSingleIteration(forStmt);
|
||||
|
||||
return true;
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in New Issue