forked from google/heir
/
FullLoopUnroll.cpp
33 lines (26 loc) · 1017 Bytes
/
FullLoopUnroll.cpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
#include "include/Transforms/FullLoopUnroll/FullLoopUnroll.h"
#include "mlir/include/mlir/Dialect/Affine/IR/AffineOps.h" // from @llvm-project
#include "mlir/include/mlir/Dialect/Affine/LoopUtils.h" // from @llvm-project
#include "mlir/include/mlir/Transforms/Passes.h" // from @llvm-project
namespace mlir {
namespace heir {
#define GEN_PASS_DEF_FULLLOOPUNROLL
#include "include/Transforms/FullLoopUnroll/FullLoopUnroll.h.inc"
struct FullLoopUnroll : impl::FullLoopUnrollBase<FullLoopUnroll> {
using FullLoopUnrollBase::FullLoopUnrollBase;
void runOnOperation() override {
auto walkResult =
getOperation()->walk<WalkOrder::PostOrder>([&](affine::AffineForOp op) {
auto result = mlir::affine::loopUnrollFull(op);
if (failed(result)) {
return WalkResult::interrupt();
}
return WalkResult::advance();
});
if (walkResult.wasInterrupted()) {
signalPassFailure();
}
}
};
} // namespace heir
} // namespace mlir