67 lines · cpp
1//===- TestAffineLoopUnswitching.cpp - Test affine if/else hoisting -------===//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 hoist affine if/else structures.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Dialect/Affine/Analysis/Utils.h"14#include "mlir/Dialect/Affine/IR/AffineOps.h"15#include "mlir/Dialect/Affine/Utils.h"16#include "mlir/Pass/Pass.h"17#include "mlir/Transforms/Passes.h"18 19#define PASS_NAME "test-affine-loop-unswitch"20 21using namespace mlir;22using namespace mlir::affine;23 24namespace {25 26/// This pass applies the permutation on the first maximal perfect nest.27struct TestAffineLoopUnswitching28 : public PassWrapper<TestAffineLoopUnswitching, OperationPass<>> {29 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestAffineLoopUnswitching)30 31 StringRef getArgument() const final { return PASS_NAME; }32 StringRef getDescription() const final {33 return "Tests affine loop unswitching / if/else hoisting";34 }35 TestAffineLoopUnswitching() = default;36 TestAffineLoopUnswitching(const TestAffineLoopUnswitching &pass) = default;37 38 void runOnOperation() override;39 40 /// The maximum number of iterations to run this for.41 constexpr static unsigned kMaxIterations = 5;42};43 44} // namespace45 46void TestAffineLoopUnswitching::runOnOperation() {47 // Each hoisting invalidates a lot of IR around. Just stop the walk after the48 // first if/else hoisting, and repeat until no more hoisting can be done, or49 // the maximum number of iterations have been run.50 Operation *op = getOperation();51 unsigned i = 0;52 do {53 auto walkFn = [](AffineIfOp op) {54 return succeeded(hoistAffineIfOp(op)) ? WalkResult::interrupt()55 : WalkResult::advance();56 };57 if (op->walk(walkFn).wasInterrupted())58 break;59 } while (++i < kMaxIterations);60}61 62namespace mlir {63void registerTestAffineLoopUnswitchingPass() {64 PassRegistration<TestAffineLoopUnswitching>();65}66} // namespace mlir67