Skip to content

Commit 0861464

Browse files
linuxlonelyeaglelinuxlonelyeagle
committed
use InterfacePass implement it.
1 parent 824b650 commit 0861464

File tree

5 files changed

+44
-50
lines changed

5 files changed

+44
-50
lines changed

‎mlir/include/mlir/Dialect/Affine/Passes.h‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@
1919

2020
namespace mlir {
2121

22-
class ModuleOp;
2322
namespace func {
2423
class FuncOp;
2524
} // namespace func
@@ -94,7 +93,7 @@ std::unique_ptr<OperationPass<func::FuncOp>> createLoopTilingPass();
9493
/// factors supplied through other means. If -1 is passed as the unrollFactor
9594
/// and no callback is provided, anything passed from the command-line (if at
9695
/// all) or the default unroll factor is used (LoopUnroll:kDefaultUnrollFactor).
97-
std::unique_ptr<OperationPass<mlir::ModuleOp>> createLoopUnrollPass(
96+
std::unique_ptr<Pass> createLoopUnrollPass(
9897
int unrollFactor = -1, bool unrollUpToFactor = false,
9998
bool unrollFull = false,
10099
const std::function<unsigned(AffineForOp)> &getUnrollFactor = nullptr);

‎mlir/include/mlir/Dialect/Affine/Passes.td‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -199,7 +199,7 @@ def AffineLoopTiling : Pass<"affine-loop-tile", "func::FuncOp"> {
199199
];
200200
}
201201

202-
def AffineLoopUnroll : Pass<"affine-loop-unroll", "ModuleOp"> {
202+
def AffineLoopUnroll : InterfacePass<"affine-loop-unroll", "FunctionOpInterface"> {
203203
let summary = "Unroll affine loops";
204204
let constructor = "mlir::affine::createLoopUnrollPass()";
205205
let options = [

‎mlir/lib/Dialect/Affine/Transforms/LoopUnroll.cpp‎

Lines changed: 34 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@
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

9493
void 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>(

‎mlir/test/Dialect/Affine/unroll.mlir‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
1-
// RUN: mlir-opt -allow-unregistered-dialect %s -affine-loop-unroll="unroll-full" | FileCheck %s --check-prefix UNROLL-FULL
2-
// RUN: mlir-opt -allow-unregistered-dialect %s -affine-loop-unroll="unroll-full unroll-full-threshold=2" | FileCheck %s --check-prefix SHORT
3-
// RUN: mlir-opt -allow-unregistered-dialect %s -affine-loop-unroll="unroll-factor=4" | FileCheck %s --check-prefix UNROLL-BY-4
4-
// RUN: mlir-opt -allow-unregistered-dialect %s -affine-loop-unroll="unroll-factor=1" | FileCheck %s --check-prefix UNROLL-BY-1
5-
// RUN: mlir-opt -allow-unregistered-dialect %s -affine-loop-unroll="unroll-factor=5 cleanup-unroll=true" | FileCheck %s --check-prefix UNROLL-CLEANUP-LOOP
1+
// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll{unroll-full=true}),gpu.module(gpu.func(affine-loop-unroll{unroll-full=true})))" | FileCheck %s --check-prefix UNROLL-FULL
2+
// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll{unroll-full=true unroll-full-threshold=2}),gpu.module(gpu.func(affine-loop-unroll{unroll-full=true unroll-full-threshold=2})))" | FileCheck %s --check-prefix SHORT
3+
// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll{unroll-factor=4}),gpu.module(gpu.func(affine-loop-unroll{unroll-factor=4})))" | FileCheck %s --check-prefix UNROLL-BY-4
4+
// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll{unroll-factor=1}),gpu.module(gpu.func(affine-loop-unroll{unroll-factor=1})))" | FileCheck %s --check-prefix UNROLL-BY-1
5+
// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll{unroll-factor=5 cleanup-unroll=true}),gpu.module(gpu.func(affine-loop-unroll{unroll-factor=5 cleanup-unroll=true})))" | FileCheck %s --check-prefix UNROLL-CLEANUP-LOOP
66

77
// UNROLL-FULL-DAG: [[$MAP0:#map[0-9]*]] = affine_map<(d0) -> (d0 + 1)>
88
// UNROLL-FULL-DAG: [[$MAP1:#map[0-9]*]] = affine_map<(d0) -> (d0 + 2)>

‎mlir/test/Dialect/SCF/loop-unroll.mlir‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,9 @@
33
// RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2 loop-depth=0' | FileCheck %s --check-prefix UNROLL-OUTER-BY-2
44
// RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2 loop-depth=1' | FileCheck %s --check-prefix UNROLL-INNER-BY-2
55
// RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2 annotate=true' | FileCheck %s --check-prefix UNROLL-BY-2-ANNOTATE
6-
// RUN: mlir-opt %s --affine-loop-unroll='unroll-factor=6 unroll-up-to-factor=true' | FileCheck %s --check-prefix UNROLL-UP-TO
7-
// RUN: mlir-opt %s --affine-loop-unroll='unroll-factor=5 cleanup-unroll=true' | FileCheck %s --check-prefix CLEANUP-UNROLL-BY-5
8-
// RUN: mlir-opt %s --affine-loop-unroll --split-input-file | FileCheck %s
6+
// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll{unroll-factor=6 unroll-up-to-factor=true}))" | FileCheck %s --check-prefix UNROLL-UP-TO
7+
// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll{unroll-factor=5 cleanup-unroll=true}))" | FileCheck %s --check-prefix CLEANUP-UNROLL-BY-5
8+
// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll))" --split-input-file | FileCheck %s
99

1010
func.func @dynamic_loop_unroll(%arg0 : index, %arg1 : index, %arg2 : index,
1111
%arg3: memref<?xf32>) {

0 commit comments

Comments
 (0)