1919#include " mlir/IR/AffineExpr.h"
2020#include " mlir/IR/AffineMap.h"
2121#include " mlir/IR/Builders.h"
22- #include " mlir/IR/BuiltinOps.h"
2322#include " llvm/ADT/DenseMap.h"
2423#include " llvm/Support/CommandLine.h"
2524#include " llvm/Support/Debug.h"
@@ -92,44 +91,40 @@ static void gatherInnermostLoops(FunctionOpInterface f,
9291}
9392
9493void LoopUnroll::runOnOperation () {
95- mlir::ModuleOp module = getOperation ();
96- SmallVector<FunctionOpInterface> funcOps;
97- module.walk ([&](FunctionOpInterface func) { funcOps.push_back (func); });
98- for (auto func : funcOps) {
99- if (func.isExternal ())
100- return ;
101-
102- if (unrollFull && unrollFullThreshold.hasValue ()) {
103- // Store short loops as we walk.
104- SmallVector<AffineForOp, 4 > loops;
105-
106- // Gathers all loops with trip count <= minTripCount. Do a post order walk
107- // so that loops are gathered from innermost to outermost (or else
108- // unrolling an outer one may delete gathered inner ones).
109- getOperation ().walk ([&](AffineForOp forOp) {
110- std::optional<uint64_t > tripCount = getConstantTripCount (forOp);
111- if (tripCount && *tripCount <= unrollFullThreshold)
112- loops.push_back (forOp);
113- });
114- for (auto forOp : loops)
115- (void )loopUnrollFull (forOp);
116- return ;
117- }
118-
119- // If the call back is provided, we will recurse until no loops are found.
94+ FunctionOpInterface func = getOperation ();
95+ if (func.isExternal ())
96+ return ;
97+
98+ if (unrollFull && unrollFullThreshold.hasValue ()) {
99+ // Store short loops as we walk.
120100 SmallVector<AffineForOp, 4 > loops;
121- for (unsigned i = 0 ; i < numRepetitions || getUnrollFactor; i++) {
122- loops.clear ();
123- gatherInnermostLoops (func, loops);
124- if (loops.empty ())
125- break ;
126- bool unrolled = false ;
127- for (auto forOp : loops)
128- unrolled |= succeeded (runOnAffineForOp (forOp));
129- if (!unrolled)
130- // Break out if nothing was unrolled.
131- break ;
132- }
101+
102+ // Gathers all loops with trip count <= minTripCount. Do a post order walk
103+ // so that loops are gathered from innermost to outermost (or else
104+ // unrolling an outer one may delete gathered inner ones).
105+ getOperation ().walk ([&](AffineForOp forOp) {
106+ std::optional<uint64_t > tripCount = getConstantTripCount (forOp);
107+ if (tripCount && *tripCount <= unrollFullThreshold)
108+ loops.push_back (forOp);
109+ });
110+ for (auto forOp : loops)
111+ (void )loopUnrollFull (forOp);
112+ return ;
113+ }
114+
115+ // If the call back is provided, we will recurse until no loops are found.
116+ SmallVector<AffineForOp, 4 > loops;
117+ for (unsigned i = 0 ; i < numRepetitions || getUnrollFactor; i++) {
118+ loops.clear ();
119+ gatherInnermostLoops (func, loops);
120+ if (loops.empty ())
121+ break ;
122+ bool unrolled = false ;
123+ for (auto forOp : loops)
124+ unrolled |= succeeded (runOnAffineForOp (forOp));
125+ if (!unrolled)
126+ // Break out if nothing was unrolled.
127+ break ;
133128 }
134129}
135130
@@ -150,7 +145,7 @@ LogicalResult LoopUnroll::runOnAffineForOp(AffineForOp forOp) {
150145 cleanUpUnroll);
151146}
152147
153- std::unique_ptr<OperationPass<ModuleOp> > mlir::affine::createLoopUnrollPass (
148+ std::unique_ptr<Pass > mlir::affine::createLoopUnrollPass (
154149 int unrollFactor, bool unrollUpToFactor, bool unrollFull,
155150 const std::function<unsigned (AffineForOp)> &getUnrollFactor) {
156151 return std::make_unique<LoopUnroll>(
0 commit comments