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/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, FunctionPass> { 28 TestAffineLoopUnswitching() = default; TestAffineLoopUnswitching__anone8bd92820111::TestAffineLoopUnswitching29 TestAffineLoopUnswitching(const TestAffineLoopUnswitching &pass) {} 30 31 void runOnFunction() override; 32 33 /// The maximum number of iterations to run this for. 34 constexpr static unsigned kMaxIterations = 5; 35 }; 36 37 } // end anonymous namespace 38 runOnFunction()39void TestAffineLoopUnswitching::runOnFunction() { 40 // Each hoisting invalidates a lot of IR around. Just stop the walk after the 41 // first if/else hoisting, and repeat until no more hoisting can be done, or 42 // the maximum number of iterations have been run. 43 auto func = getFunction(); 44 unsigned i = 0; 45 do { 46 auto walkFn = [](AffineIfOp op) { 47 return succeeded(hoistAffineIfOp(op)) ? WalkResult::interrupt() 48 : WalkResult::advance(); 49 }; 50 if (func.walk(walkFn).wasInterrupted()) 51 break; 52 } while (++i < kMaxIterations); 53 } 54 55 namespace mlir { registerTestAffineLoopUnswitchingPass()56void registerTestAffineLoopUnswitchingPass() { 57 PassRegistration<TestAffineLoopUnswitching>( 58 PASS_NAME, "Tests affine loop unswitching / if/else hoisting"); 59 } 60 } // namespace mlir 61