diff --git a/include/Transforms/Passes.h b/include/Transforms/Passes.h index 41da49a..75120f1 100644 --- a/include/Transforms/Passes.h +++ b/include/Transforms/Passes.h @@ -17,6 +17,7 @@ namespace scalehls { std::unique_ptr createPragmaDSEPass(); std::unique_ptr createAffineLoopPerfectionPass(); +std::unique_ptr createPartialAffineLoopTilePass(); void registerTransformsPasses(); diff --git a/include/Transforms/Passes.td b/include/Transforms/Passes.td index 5491d64..e2621aa 100644 --- a/include/Transforms/Passes.td +++ b/include/Transforms/Passes.td @@ -23,13 +23,32 @@ def PragmaDSE : Pass<"pragma-dse", "ModuleOp"> { def AffineLoopPerfection : Pass<"affine-loop-perfection", "ModuleOp"> { let summary = "Try to perfect a nested loop"; let description = [{ - This affine-loop-perfection pass will try to perfect affine loops. + This affine-loop-perfection pass will try to perfect all affine loops. Specifically, this pass will move operations under non-innermost loops into - innermost loop and create an affine if region to ensure the correctness of - the transformation. + innermost loop and create affine if regions to ensure the correctness of the + movement. }]; let constructor = "mlir::scalehls::createAffineLoopPerfectionPass()"; } +def PartialAffineLoopTile : Pass<"partial-affine-loop-tile", "ModuleOp"> { + let summary = "Try to perfect a nested loop"; + let description = [{ + This partial-affine-loop-tile pass will try to tile the nested loops. The + difference with the official affine-loop-tile pass is this pass will only + tile the first "tileLevel" outermost loop levels rather than all loops + levels. + }]; + + let constructor = "mlir::scalehls::createPartialAffineLoopTilePass()"; + + let options = [ + Option<"tileLevel", "tile-level", "unsigned", /*default=*/"1", + "Positive number: the level of loops to be tiles">, + Option<"tileSize", "tile-size", "unsigned", /*default=*/"2", + "Positive number: the size of tiling"> + ]; +} + #endif // SCALEHLS_TRANSFORMS_PASSES_TD diff --git a/lib/Transforms/PartialAffineLoopTile.cpp b/lib/Transforms/PartialAffineLoopTile.cpp new file mode 100644 index 0000000..30b29fc --- /dev/null +++ b/lib/Transforms/PartialAffineLoopTile.cpp @@ -0,0 +1,71 @@ +//===------------------------------------------------------------*- C++ -*-===// +// +//===----------------------------------------------------------------------===// + +#include "Transforms/Passes.h" +#include "mlir/Analysis/LoopAnalysis.h" +#include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/Affine/Passes.h" +#include "mlir/IR/Builders.h" +#include "mlir/Transforms/LoopUtils.h" + +using namespace std; +using namespace mlir; +using namespace scalehls; + +namespace { +struct PartialAffineLoopTile + : public PartialAffineLoopTileBase { + void runOnOperation() override; + void getTileSizes(ArrayRef band, + SmallVectorImpl *tileSizes); +}; +} // namespace + +void PartialAffineLoopTile::runOnOperation() { + // Walk through all functions and loops. + for (auto func : getOperation().getOps()) { + // Bands of loops to tile. + std::vector> bands; + getTileableBands(func, &bands); + + // Tile each band. + for (auto &band : bands) { + // Truncate band and only keep first tileLevel loops. + size_t realTileLevel = band.size(); + if (realTileLevel > tileLevel) { + band.resize(tileLevel); + realTileLevel = tileLevel; + } + + // Set up tile sizes; fill missing tile sizes at the end with default tile + // size or tileSize if one was provided. + SmallVector tileSizes; + tileSizes.assign(band.size(), tileSize); + + SmallVector tiledNest; + if (failed(tilePerfectlyNested(band, tileSizes, &tiledNest))) + return signalPassFailure(); + + // Permute loop order to move the tiled loop to the innermost of the + // perfect nested loop. + SmallVector nestedLoops; + getPerfectlyNestedLoops(nestedLoops, tiledNest.front()); + + SmallVector permMap; + for (size_t i = 0, e = nestedLoops.size(); i < e; ++i) { + if (i < realTileLevel) + permMap.push_back(i); + else if (i < 2 * realTileLevel) + permMap.push_back(e + i - 2 * realTileLevel); + else + permMap.push_back(i - realTileLevel); + } + permuteLoops(nestedLoops, permMap); + } + } +} + +std::unique_ptr scalehls::createPartialAffineLoopTilePass() { + return std::make_unique(); +} diff --git a/test/Transforms/PartialAffineLoopTile/test_for.mlir b/test/Transforms/PartialAffineLoopTile/test_for.mlir new file mode 100644 index 0000000..5ab9474 --- /dev/null +++ b/test/Transforms/PartialAffineLoopTile/test_for.mlir @@ -0,0 +1,6 @@ +// RUN: scalehls-opt -partial-affine-loop-tile %s | FileCheck %s + +// CHECK-LABEL: func @test_for +func @test_for() { + return +}