82 lines · cpp
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-exception6//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 TestLoopPermutation27 : 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] in41 /// 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 /// Specify whether to check validity of loop permutation.47 Option<bool> checkValidity{48 *this, "check-validity",49 llvm::cl::desc("Check validity of the loop permutation"),50 llvm::cl::init(false)};51};52 53} // namespace54 55void TestLoopPermutation::runOnOperation() {56 57 SmallVector<unsigned, 4> permMap(permList.begin(), permList.end());58 59 SmallVector<AffineForOp, 2> forOps;60 getOperation()->walk([&](AffineForOp forOp) { forOps.push_back(forOp); });61 62 for (auto forOp : forOps) {63 SmallVector<AffineForOp, 6> nest;64 // Get the maximal perfect nest.65 getPerfectlyNestedLoops(nest, forOp);66 // Permute if the nest's size is consistent with the specified67 // permutation.68 if (nest.size() >= 2 && nest.size() == permMap.size()) {69 if (checkValidity.getValue() &&70 !isValidLoopInterchangePermutation(nest, permMap))71 continue;72 permuteLoops(nest, permMap);73 }74 }75}76 77namespace mlir {78void registerTestLoopPermutationPass() {79 PassRegistration<TestLoopPermutation>();80}81} // namespace mlir82