MLIR  22.0.0git
LoopUnroll.cpp
Go to the documentation of this file.
1 //===- LoopUnroll.cpp - Code to perform loop unrolling --------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file implements loop unrolling.
10 //
11 //===----------------------------------------------------------------------===//
12 
14 
18 #include "llvm/Support/CommandLine.h"
19 #include <optional>
20 
21 namespace mlir {
22 namespace affine {
23 #define GEN_PASS_DEF_AFFINELOOPUNROLL
24 #include "mlir/Dialect/Affine/Passes.h.inc"
25 } // namespace affine
26 } // namespace mlir
27 
28 #define DEBUG_TYPE "affine-loop-unroll"
29 
30 using namespace mlir;
31 using namespace mlir::affine;
32 
33 namespace {
34 
35 // TODO: this is really a test pass and should be moved out of dialect
36 // transforms.
37 
38 /// Loop unrolling pass. Unrolls all innermost loops unless full unrolling and a
39 /// full unroll threshold was specified, in which case, fully unrolls all loops
40 /// with trip count less than the specified threshold. The latter is for testing
41 /// purposes, especially for testing outer loop unrolling.
42 struct LoopUnroll : public affine::impl::AffineLoopUnrollBase<LoopUnroll> {
43  // Callback to obtain unroll factors; if this has a callable target, takes
44  // precedence over command-line argument or passed argument.
45  const std::function<unsigned(AffineForOp)> getUnrollFactor;
46 
47  LoopUnroll() : getUnrollFactor(nullptr) {}
48  LoopUnroll(const LoopUnroll &other)
49 
50  = default;
51  explicit LoopUnroll(
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) {
56  if (unrollFactor)
57  this->unrollFactor = *unrollFactor;
58  this->unrollUpToFactor = unrollUpToFactor;
59  this->unrollFull = unrollFull;
60  }
61 
62  void runOnOperation() override;
63 
64  /// Unroll this for op. Returns failure if nothing was done.
65  LogicalResult runOnAffineForOp(AffineForOp forOp);
66 };
67 } // namespace
68 
69 /// Returns true if no other affine.for ops are nested within `op`.
70 static bool isInnermostAffineForOp(AffineForOp op) {
71  return !op.getBody()
72  ->walk([&](AffineForOp nestedForOp) {
73  return WalkResult::interrupt();
74  })
75  .wasInterrupted();
76 }
77 
78 /// Gathers loops that have no affine.for's nested within.
79 static void gatherInnermostLoops(FunctionOpInterface f,
81  f.walk([&](AffineForOp forOp) {
82  if (isInnermostAffineForOp(forOp))
83  loops.push_back(forOp);
84  });
85 }
86 
87 void LoopUnroll::runOnOperation() {
88  FunctionOpInterface func = getOperation();
89  if (func.isExternal())
90  return;
91 
92  if (unrollFull && unrollFullThreshold.hasValue()) {
93  // Store short loops as we walk.
95 
96  // Gathers all loops with trip count <= minTripCount. Do a post order walk
97  // so that loops are gathered from innermost to outermost (or else
98  // unrolling an outer one may delete gathered inner ones).
99  getOperation().walk([&](AffineForOp forOp) {
100  std::optional<uint64_t> tripCount = getConstantTripCount(forOp);
101  if (tripCount && *tripCount <= unrollFullThreshold)
102  loops.push_back(forOp);
103  });
104  for (auto forOp : loops)
105  (void)loopUnrollFull(forOp);
106  return;
107  }
108 
109  // If the call back is provided, we will recurse until no loops are found.
111  for (unsigned i = 0; i < numRepetitions || getUnrollFactor; i++) {
112  loops.clear();
113  gatherInnermostLoops(func, loops);
114  if (loops.empty())
115  break;
116  bool unrolled = false;
117  for (auto forOp : loops)
118  unrolled |= succeeded(runOnAffineForOp(forOp));
119  if (!unrolled)
120  // Break out if nothing was unrolled.
121  break;
122  }
123 }
124 
125 /// Unrolls a 'affine.for' op. Returns success if the loop was unrolled,
126 /// failure otherwise. The default unroll factor is 4.
127 LogicalResult LoopUnroll::runOnAffineForOp(AffineForOp forOp) {
128  // Use the function callback if one was provided.
129  if (getUnrollFactor)
130  return loopUnrollByFactor(forOp, getUnrollFactor(forOp),
131  /*annotateFn=*/nullptr, cleanUpUnroll);
132  // Unroll completely if full loop unroll was specified.
133  if (unrollFull)
134  return loopUnrollFull(forOp);
135  // Otherwise, unroll by the given unroll factor.
136  if (unrollUpToFactor)
137  return loopUnrollUpToFactor(forOp, unrollFactor);
138  return loopUnrollByFactor(forOp, unrollFactor, /*annotateFn=*/nullptr,
139  cleanUpUnroll);
140 }
141 
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);
149 }
static void gatherInnermostLoops(FunctionOpInterface f, SmallVectorImpl< AffineForOp > &loops)
Gathers loops that have no affine.for's nested within.
Definition: LoopUnroll.cpp:79
static bool isInnermostAffineForOp(AffineForOp op)
Returns true if no other affine.for ops are nested within op.
Definition: LoopUnroll.cpp:70
static WalkResult interrupt()
Definition: WalkResult.h:46
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.
Definition: LoopUtils.cpp:871
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.
Definition: LoopUtils.cpp:995
LogicalResult loopUnrollUpToFactor(AffineForOp forOp, uint64_t unrollFactor)
Unrolls this loop by the specified unroll factor or its trip count, whichever is lower.
Definition: LoopUtils.cpp:886
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.
Definition: LoopUnroll.cpp:143
Include the generated interface declarations.