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-exception
6 //
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 
21 using namespace mlir;
22 
23 namespace {
24 
25 /// This pass applies the permutation on the first maximal perfect nest.
26 struct TestAffineLoopUnswitching
27     : public PassWrapper<TestAffineLoopUnswitching, OperationPass<>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID__anon5bb0cbdb0111::TestAffineLoopUnswitching28   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestAffineLoopUnswitching)
29 
30   StringRef getArgument() const final { return PASS_NAME; }
getDescription__anon5bb0cbdb0111::TestAffineLoopUnswitching31   StringRef getDescription() const final {
32     return "Tests affine loop unswitching / if/else hoisting";
33   }
34   TestAffineLoopUnswitching() = default;
35   TestAffineLoopUnswitching(const TestAffineLoopUnswitching &pass) = default;
36 
37   void runOnOperation() override;
38 
39   /// The maximum number of iterations to run this for.
40   constexpr static unsigned kMaxIterations = 5;
41 };
42 
43 } // namespace
44 
runOnOperation()45 void TestAffineLoopUnswitching::runOnOperation() {
46   // Each hoisting invalidates a lot of IR around. Just stop the walk after the
47   // first if/else hoisting, and repeat until no more hoisting can be done, or
48   // the maximum number of iterations have been run.
49   Operation *op = getOperation();
50   unsigned i = 0;
51   do {
52     auto walkFn = [](AffineIfOp op) {
53       return succeeded(hoistAffineIfOp(op)) ? WalkResult::interrupt()
54                                             : WalkResult::advance();
55     };
56     if (op->walk(walkFn).wasInterrupted())
57       break;
58   } while (++i < kMaxIterations);
59 }
60 
61 namespace mlir {
registerTestAffineLoopUnswitchingPass()62 void registerTestAffineLoopUnswitchingPass() {
63   PassRegistration<TestAffineLoopUnswitching>();
64 }
65 } // namespace mlir
66