18 #include "llvm/Support/CommandLine.h"
23 #define GEN_PASS_DEF_AFFINELOOPUNROLL
24 #include "mlir/Dialect/Affine/Passes.h.inc"
28 #define DEBUG_TYPE "affine-loop-unroll"
42 struct LoopUnroll :
public affine::impl::AffineLoopUnrollBase<LoopUnroll> {
45 const std::function<unsigned(AffineForOp)> getUnrollFactor;
47 LoopUnroll() : getUnrollFactor(
nullptr) {}
48 LoopUnroll(
const LoopUnroll &other)
52 std::optional<unsigned> unrollFactor = std::nullopt,
53 bool unrollUpToFactor =
false,
bool unrollFull =
false,
54 const std::function<
unsigned(AffineForOp)> &getUnrollFactor =
nullptr)
55 : getUnrollFactor(getUnrollFactor) {
57 this->unrollFactor = *unrollFactor;
58 this->unrollUpToFactor = unrollUpToFactor;
59 this->unrollFull = unrollFull;
62 void runOnOperation()
override;
65 LogicalResult runOnAffineForOp(AffineForOp forOp);
72 ->walk([&](AffineForOp nestedForOp) {
81 f.walk([&](AffineForOp forOp) {
83 loops.push_back(forOp);
87 void LoopUnroll::runOnOperation() {
88 FunctionOpInterface func = getOperation();
89 if (func.isExternal())
92 if (unrollFull && unrollFullThreshold.hasValue()) {
99 getOperation().walk([&](AffineForOp forOp) {
101 if (tripCount && *tripCount <= unrollFullThreshold)
102 loops.push_back(forOp);
104 for (
auto forOp : loops)
111 for (
unsigned i = 0; i < numRepetitions || getUnrollFactor; i++) {
116 bool unrolled =
false;
117 for (
auto forOp : loops)
118 unrolled |= succeeded(runOnAffineForOp(forOp));
127 LogicalResult LoopUnroll::runOnAffineForOp(AffineForOp forOp) {
131 nullptr, cleanUpUnroll);
136 if (unrollUpToFactor)
142 std::unique_ptr<InterfacePass<FunctionOpInterface>>
144 int unrollFactor,
bool unrollUpToFactor,
bool unrollFull,
145 const std::function<
unsigned(AffineForOp)> &getUnrollFactor) {
146 return std::make_unique<LoopUnroll>(
147 unrollFactor == -1 ? std::nullopt : std::optional<unsigned>(unrollFactor),
148 unrollUpToFactor, unrollFull, getUnrollFactor);
static void gatherInnermostLoops(FunctionOpInterface f, SmallVectorImpl< AffineForOp > &loops)
Gathers loops that have no affine.for's nested within.
static bool isInnermostAffineForOp(AffineForOp op)
Returns true if no other affine.for ops are nested within op.
static WalkResult interrupt()
std::optional< uint64_t > getConstantTripCount(AffineForOp forOp)
Returns the trip count of the loop if it's a constant, std::nullopt otherwise.
LogicalResult loopUnrollFull(AffineForOp forOp)
Unrolls this for operation completely if the trip count is known to be constant.
LogicalResult loopUnrollByFactor(AffineForOp forOp, uint64_t unrollFactor, function_ref< void(unsigned, Operation *, OpBuilder)> annotateFn=nullptr, bool cleanUpUnroll=false)
Unrolls this for operation by the specified unroll factor.
LogicalResult loopUnrollUpToFactor(AffineForOp forOp, uint64_t unrollFactor)
Unrolls this loop by the specified unroll factor or its trip count, whichever is lower.
std::unique_ptr< InterfacePass< FunctionOpInterface > > createLoopUnrollPass(int unrollFactor=-1, bool unrollUpToFactor=false, bool unrollFull=false, const std::function< unsigned(AffineForOp)> &getUnrollFactor=nullptr)
Creates a loop unrolling pass with the provided parameters.
Include the generated interface declarations.