72 lines · cpp
1//===- TestLoopParametricTiling.cpp --- Parametric 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 pass to parametrically tile nests of standard loops.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Dialect/SCF/IR/SCF.h"14#include "mlir/Dialect/SCF/Utils/Utils.h"15#include "mlir/IR/Builders.h"16#include "mlir/Pass/Pass.h"17 18using namespace mlir;19 20namespace {21 22// Extracts fixed-range loops for top-level loop nests with ranges defined in23// the pass constructor. Assumes loops are permutable.24class SimpleParametricLoopTilingPass25 : public PassWrapper<SimpleParametricLoopTilingPass, OperationPass<>> {26public:27 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SimpleParametricLoopTilingPass)28 29 StringRef getArgument() const final {30 return "test-extract-fixed-outer-loops";31 }32 StringRef getDescription() const final {33 return "test application of parametric tiling to the outer loops so that "34 "the ranges of outer loops become static";35 }36 SimpleParametricLoopTilingPass() = default;37 SimpleParametricLoopTilingPass(const SimpleParametricLoopTilingPass &) {}38 explicit SimpleParametricLoopTilingPass(ArrayRef<int64_t> outerLoopSizes) {39 sizes = outerLoopSizes;40 }41 42 void runOnOperation() override {43 if (sizes.empty()) {44 emitError(45 UnknownLoc::get(&getContext()),46 "missing `test-outer-loop-sizes` pass-option for outer loop sizes");47 signalPassFailure();48 return;49 }50 getOperation()->walk([this](scf::ForOp op) {51 // Ignore nested loops.52 if (op->getParentRegion()->getParentOfType<scf::ForOp>())53 return;54 extractFixedOuterLoops(op, sizes);55 });56 }57 58 ListOption<int64_t> sizes{59 *this, "test-outer-loop-sizes",60 llvm::cl::desc(61 "fixed number of iterations that the outer loops should have")};62};63} // namespace64 65namespace mlir {66namespace test {67void registerSimpleParametricTilingPass() {68 PassRegistration<SimpleParametricLoopTilingPass>();69}70} // namespace test71} // namespace mlir72