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
20 using namespace mlir;
21 using namespace mlir::affine;
22
23 namespace {
24
25 /// This pass applies the permutation on the first maximal perfect nest.
26 struct TestLoopPermutation
27 : public PassWrapper<TestLoopPermutation, OperationPass<>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID__anonca63a3fa0111::TestLoopPermutation28 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLoopPermutation)
29
30 StringRef getArgument() const final { return PASS_NAME; }
getDescription__anonca63a3fa0111::TestLoopPermutation31 StringRef getDescription() const final {
32 return "Tests affine loop permutation utility";
33 }
34 TestLoopPermutation() = default;
TestLoopPermutation__anonca63a3fa0111::TestLoopPermutation35 TestLoopPermutation(const TestLoopPermutation &pass) : PassWrapper(pass){};
36
37 void runOnOperation() override;
38
39 private:
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
runOnOperation()49 void 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
68 namespace mlir {
registerTestLoopPermutationPass()69 void registerTestLoopPermutationPass() {
70 PassRegistration<TestLoopPermutation>();
71 }
72 } // namespace mlir
73