[InsertPipelinePragma] move loop unrolling and canonicalization to this pass; inner loops of the pipelined loop are automatically fully unrolled now

This commit is contained in:
Hanchen Ye 2020-12-14 12:58:15 -06:00
parent 31ce83be83
commit c24b418c4a
2 changed files with 18 additions and 22 deletions

View File

@ -7,7 +7,6 @@
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/Passes.h"
#include "mlir/IR/Builders.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/LoopUtils.h"
using namespace std;
@ -110,26 +109,6 @@ void ArrayPartition::runOnOperation() {
auto func = getOperation();
auto builder = OpBuilder(func);
// If the current loop is annotated as pipeline, all intter loops are
// automatically unrolled.
for (auto forOp : func.getOps<mlir::AffineForOp>()) {
auto outermost = getPipelineLoop(forOp);
outermost.walk([&](mlir::AffineForOp loop) {
if (loop != outermost)
loopUnrollFull(loop);
});
}
// Canonicalize the analyzed IR.
OwningRewritePatternList patterns;
auto *context = &getContext();
for (auto *op : context->getRegisteredOperations())
op->getCanonicalizationPatterns(patterns, context);
Operation *op = getOperation();
applyPatternsAndFoldGreedily(op->getRegions(), std::move(patterns));
// Apply array partition.
for (auto forOp : func.getOps<mlir::AffineForOp>()) {
auto outermost = getPipelineLoop(forOp);

View File

@ -6,6 +6,7 @@
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/Passes.h"
#include "mlir/IR/Builders.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/LoopUtils.h"
using namespace std;
@ -23,7 +24,7 @@ void InsertPipelinePragma::runOnOperation() {
auto func = getOperation();
auto builder = OpBuilder(func);
// Walk through all functions and loops.
// Walk through loops in the function.
for (auto forOp : func.getOps<mlir::AffineForOp>()) {
SmallVector<mlir::AffineForOp, 4> nestedLoops;
forOp.walk([&](mlir::AffineForOp loop) { nestedLoops.push_back(loop); });
@ -33,7 +34,23 @@ void InsertPipelinePragma::runOnOperation() {
targetLoop = *std::next(nestedLoops.begin(), insertLevel);
targetLoop.setAttr("pipeline", builder.getBoolAttr(true));
// All intter loops of the pipelined loop are automatically unrolled.
targetLoop.walk([&](mlir::AffineForOp loop) {
if (loop != targetLoop)
loopUnrollFull(loop);
});
}
// Canonicalize the IR after loop unrolling.
OwningRewritePatternList patterns;
auto *context = &getContext();
for (auto *op : context->getRegisteredOperations())
op->getCanonicalizationPatterns(patterns, context);
applyPatternsAndFoldGreedily(func.getOperation()->getRegions(),
std::move(patterns));
}
std::unique_ptr<mlir::Pass> scalehls::createInsertPipelinePragmaPass() {