1//===- TestLoopPermutation.cpp - Test affine loop permutation -------------===//
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-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// This file implements a pass to test the affine for op permutation utility.
10//
11//===----------------------------------------------------------------------===//
12
13#include "mlir/Dialect/Affine/Analysis/Utils.h"
14#include "mlir/Dialect/Affine/IR/AffineOps.h"
15#include "mlir/Dialect/Affine/LoopUtils.h"
16#include "mlir/Pass/Pass.h"
17
18#define PASS_NAME "test-loop-permutation"
19
20using namespace mlir;
21using namespace mlir::affine;
22
23namespace {
24
25/// This pass applies the permutation on the first maximal perfect nest.
26struct TestLoopPermutation
27 : public PassWrapper<TestLoopPermutation, OperationPass<>> {
28 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLoopPermutation)
29
30 StringRef getArgument() const final { return PASS_NAME; }
31 StringRef getDescription() const final {
32 return "Tests affine loop permutation utility";
33 }
34 TestLoopPermutation() = default;
35 TestLoopPermutation(const TestLoopPermutation &pass) : PassWrapper(pass){};
36
37 void runOnOperation() override;
38
39private:
40 /// Permutation specifying loop i is mapped to permList[i] in
41 /// transformed nest (with i going from outermost to innermost).
42 ListOption<unsigned> permList{*this, "permutation-map",
43 llvm::cl::desc("Specify the loop permutation"),
44 llvm::cl::OneOrMore};
45};
46
47} // namespace
48
49void TestLoopPermutation::runOnOperation() {
50
51 SmallVector<unsigned, 4> permMap(permList.begin(), permList.end());
52
53 SmallVector<AffineForOp, 2> forOps;
54 getOperation()->walk([&](AffineForOp forOp) { forOps.push_back(forOp); });
55
56 for (auto forOp : forOps) {
57 SmallVector<AffineForOp, 6> nest;
58 // Get the maximal perfect nest.
59 getPerfectlyNestedLoops(nest, forOp);
60 // Permute if the nest's size is consistent with the specified
61 // permutation.
62 if (nest.size() >= 2 && nest.size() == permMap.size()) {
63 permuteLoops(nest, permMap);
64 }
65 }
66}
67
68namespace mlir {
69void registerTestLoopPermutationPass() {
70 PassRegistration<TestLoopPermutation>();
71}
72} // namespace mlir
73

source code of mlir/test/lib/Dialect/Affine/TestLoopPermutation.cpp