93d1108801
Summary: Adds the loop unroll transformation for loop::ForOp. Adds support for promoting the body of single-iteration loop::ForOps into its containing block. Adds check tests for loop::ForOps with dynamic and static lower/upper bounds and step. Care was taken to share code (where possible) with the AffineForOp unroll transformation to ease maintenance and potential future transition to a LoopLike construct on which loop transformations for different loop types can implemented. Reviewers: ftynse, nicolasvasilache Reviewed By: ftynse Subscribers: bondhugula, mgorny, zzheng, mehdi_amini, rriddle, jpienaar, shauheen, antiagainst, nicolasvasilache, arpith-jacob, mgester, lucyrfox, aartbik, liufengdb, Joonsoo, grosul1, frgossen, Kayjukh, llvm-commits Tags: #llvm Differential Revision: https://reviews.llvm.org/D79184
69 lines
2.1 KiB
C++
69 lines
2.1 KiB
C++
//===-------- TestLoopUnrolling.cpp --- loop unrolling test pass ----------===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
//
|
|
// This file implements a pass to unroll loops by a specified unroll factor.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/LoopOps/LoopOps.h"
|
|
#include "mlir/IR/Builders.h"
|
|
#include "mlir/Pass/Pass.h"
|
|
#include "mlir/Transforms/LoopUtils.h"
|
|
#include "mlir/Transforms/Passes.h"
|
|
|
|
using namespace mlir;
|
|
|
|
namespace {
|
|
|
|
static unsigned getNestingDepth(Operation *op) {
|
|
Operation *currOp = op;
|
|
unsigned depth = 0;
|
|
while ((currOp = currOp->getParentOp())) {
|
|
if (isa<loop::ForOp>(currOp))
|
|
depth++;
|
|
}
|
|
return depth;
|
|
}
|
|
|
|
class TestLoopUnrollingPass
|
|
: public PassWrapper<TestLoopUnrollingPass, FunctionPass> {
|
|
public:
|
|
TestLoopUnrollingPass() = default;
|
|
TestLoopUnrollingPass(const TestLoopUnrollingPass &) {}
|
|
explicit TestLoopUnrollingPass(uint64_t unrollFactorParam,
|
|
unsigned loopDepthParam) {
|
|
unrollFactor = unrollFactorParam;
|
|
loopDepth = loopDepthParam;
|
|
}
|
|
|
|
void runOnFunction() override {
|
|
FuncOp func = getFunction();
|
|
SmallVector<loop::ForOp, 4> loops;
|
|
func.walk([&](loop::ForOp forOp) {
|
|
if (getNestingDepth(forOp) == loopDepth)
|
|
loops.push_back(forOp);
|
|
});
|
|
for (auto loop : loops) {
|
|
loopUnrollByFactor(loop, unrollFactor);
|
|
}
|
|
}
|
|
Option<uint64_t> unrollFactor{*this, "unroll-factor",
|
|
llvm::cl::desc("Loop unroll factor."),
|
|
llvm::cl::init(1)};
|
|
Option<unsigned> loopDepth{*this, "loop-depth", llvm::cl::desc("Loop depth."),
|
|
llvm::cl::init(0)};
|
|
};
|
|
} // end namespace
|
|
|
|
namespace mlir {
|
|
void registerTestLoopUnrollingPass() {
|
|
PassRegistration<TestLoopUnrollingPass>(
|
|
"test-loop-unrolling", "Tests loop unrolling transformation");
|
|
}
|
|
} // namespace mlir
|