llvm-project/mlir/test/lib/Dialect/Affine/TestAffineLoopParametricTiling.cpp
River Riddle 3655069234 [mlir] Move the Builtin FuncOp to the Func dialect
This commit moves FuncOp out of the builtin dialect, and into the Func
dialect. This move has been planned in some capacity from the moment
we made FuncOp an operation (years ago). This commit handles the
functional aspects of the move, but various aspects are left untouched
to ease migration: func::FuncOp is re-exported into mlir to reduce
the actual API churn, the assembly format still accepts the unqualified
`func`. These temporary measures will remain for a little while to
simplify migration before being removed.

Differential Revision: https://reviews.llvm.org/D121266
2022-03-16 17:07:03 -07:00

94 lines
3.4 KiB
C++

//= TestAffineLoopParametricTiling.cpp -- Parametric Affine loop tiling pass =//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file implements a test pass to test parametric tiling of perfectly
// nested affine for loops.
//
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/LoopUtils.h"
#include "mlir/Dialect/Affine/Passes.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
using namespace mlir;
#define DEBUG_TYPE "test-affine-parametric-tile"
namespace {
struct TestAffineLoopParametricTiling
: public PassWrapper<TestAffineLoopParametricTiling,
OperationPass<FuncOp>> {
StringRef getArgument() const final { return "test-affine-parametric-tile"; }
StringRef getDescription() const final {
return "Tile affine loops using SSA values as tile sizes";
}
void runOnOperation() override;
};
} // namespace
/// Checks if the function enclosing the loop nest has any arguments passed to
/// it, which can be used as tiling parameters. Assumes that atleast 'n'
/// arguments are passed, where 'n' is the number of loops in the loop nest.
static void checkIfTilingParametersExist(ArrayRef<AffineForOp> band) {
assert(!band.empty() && "no loops in input band");
AffineForOp topLoop = band[0];
if (FuncOp funcOp = dyn_cast<FuncOp>(topLoop->getParentOp()))
assert(funcOp.getNumArguments() >= band.size() && "Too few tile sizes");
}
/// Captures tiling parameters, which are expected to be passed as arguments
/// to the function enclosing the loop nest. Also checks if the required
/// parameters are of index type. This approach is temporary for testing
/// purposes.
static void getTilingParameters(ArrayRef<AffineForOp> band,
SmallVectorImpl<Value> &tilingParameters) {
AffineForOp topLoop = band[0];
Region *funcOpRegion = topLoop->getParentRegion();
unsigned nestDepth = band.size();
for (BlockArgument blockArgument :
funcOpRegion->getArguments().take_front(nestDepth)) {
if (blockArgument.getArgNumber() < nestDepth) {
assert(blockArgument.getType().isIndex() &&
"expected tiling parameters to be of index type.");
tilingParameters.push_back(blockArgument);
}
}
}
void TestAffineLoopParametricTiling::runOnOperation() {
// Bands of loops to tile.
std::vector<SmallVector<AffineForOp, 6>> bands;
getTileableBands(getOperation(), &bands);
// Tile each band.
for (MutableArrayRef<AffineForOp> band : bands) {
// Capture the tiling parameters from the arguments to the function
// enclosing this loop nest.
SmallVector<AffineForOp, 6> tiledNest;
SmallVector<Value, 6> tilingParameters;
// Check if tiling parameters are present.
checkIfTilingParametersExist(band);
// Get function arguments as tiling parameters.
getTilingParameters(band, tilingParameters);
(void)tilePerfectlyNestedParametric(band, tilingParameters, &tiledNest);
}
}
namespace mlir {
namespace test {
void registerTestAffineLoopParametricTilingPass() {
PassRegistration<TestAffineLoopParametricTiling>();
}
} // namespace test
} // namespace mlir