110 lines · cpp
1//= TestAffineLoopParametricTiling.cpp -- Parametric Affine loop tiling pass =//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8//9// This file implements a test pass to test parametric tiling of perfectly10// nested affine for loops.11//12//===----------------------------------------------------------------------===//13 14#include "mlir/Dialect/Affine/IR/AffineOps.h"15#include "mlir/Dialect/Affine/LoopUtils.h"16#include "mlir/Dialect/Affine/Passes.h"17#include "mlir/Dialect/Func/IR/FuncOps.h"18 19using namespace mlir;20using namespace mlir::affine;21 22#define DEBUG_TYPE "test-affine-parametric-tile"23 24namespace {25struct TestAffineLoopParametricTiling26 : public PassWrapper<TestAffineLoopParametricTiling,27 OperationPass<func::FuncOp>> {28 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestAffineLoopParametricTiling)29 30 StringRef getArgument() const final { return "test-affine-parametric-tile"; }31 StringRef getDescription() const final {32 return "Tile affine loops using SSA values as tile sizes";33 }34 void runOnOperation() override;35};36} // namespace37 38/// Checks if the function enclosing the loop nest has any arguments passed to39/// it, which can be used as tiling parameters. Assumes that atleast 'n'40/// arguments are passed, where 'n' is the number of loops in the loop nest.41static LogicalResult checkIfTilingParametersExist(ArrayRef<AffineForOp> band) {42 assert(!band.empty() && "no loops in input band");43 AffineForOp topLoop = band[0];44 45 if (func::FuncOp funcOp = dyn_cast<func::FuncOp>(topLoop->getParentOp()))46 if (funcOp.getNumArguments() < band.size())47 return topLoop->emitError(48 "too few tile sizes provided in the argument list of the function "49 "which contains the current band");50 return success();51}52 53/// Captures tiling parameters, which are expected to be passed as arguments54/// to the function enclosing the loop nest. Also checks if the required55/// parameters are of index type. This approach is temporary for testing56/// purposes.57static LogicalResult58getTilingParameters(ArrayRef<AffineForOp> band,59 SmallVectorImpl<Value> &tilingParameters) {60 AffineForOp topLoop = band[0];61 Region *funcOpRegion = topLoop->getParentRegion();62 unsigned nestDepth = band.size();63 64 for (BlockArgument blockArgument :65 funcOpRegion->getArguments().take_front(nestDepth)) {66 if (blockArgument.getArgNumber() < nestDepth) {67 if (!blockArgument.getType().isIndex())68 return topLoop->emitError(69 "expected tiling parameters to be of index type");70 tilingParameters.push_back(blockArgument);71 }72 }73 return success();74}75 76void TestAffineLoopParametricTiling::runOnOperation() {77 // Get maximal perfect nest of 'affine.for' ops at the top-level.78 std::vector<SmallVector<AffineForOp, 6>> bands;79 for (AffineForOp forOp : getOperation().getOps<AffineForOp>()) {80 SmallVector<AffineForOp, 6> band;81 getPerfectlyNestedLoops(band, forOp);82 bands.push_back(band);83 }84 85 // Tile each band.86 for (MutableArrayRef<AffineForOp> band : bands) {87 // Capture the tiling parameters from the arguments to the function88 // enclosing this loop nest.89 SmallVector<AffineForOp, 6> tiledNest;90 SmallVector<Value, 6> tilingParameters;91 // Check if tiling parameters are present.92 if (checkIfTilingParametersExist(band).failed())93 return;94 95 // Get function arguments as tiling parameters.96 if (getTilingParameters(band, tilingParameters).failed())97 return;98 99 (void)tilePerfectlyNestedParametric(band, tilingParameters, &tiledNest);100 }101}102 103namespace mlir {104namespace test {105void registerTestAffineLoopParametricTilingPass() {106 PassRegistration<TestAffineLoopParametricTiling>();107}108} // namespace test109} // namespace mlir110