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 using namespace mlir::affine; 23 24 namespace { 25 26 /// This pass applies the permutation on the first maximal perfect nest. 27 struct TestAffineLoopUnswitching 28 : public PassWrapper<TestAffineLoopUnswitching, OperationPass<>> { MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID__anon8ccd45be0111::TestAffineLoopUnswitching29 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestAffineLoopUnswitching) 30 31 StringRef getArgument() const final { return PASS_NAME; } getDescription__anon8ccd45be0111::TestAffineLoopUnswitching32 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 } // namespace 45 runOnOperation()46void TestAffineLoopUnswitching::runOnOperation() { 47 // Each hoisting invalidates a lot of IR around. Just stop the walk after the 48 // first if/else hoisting, and repeat until no more hoisting can be done, or 49 // 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 62 namespace mlir { registerTestAffineLoopUnswitchingPass()63void registerTestAffineLoopUnswitchingPass() { 64 PassRegistration<TestAffineLoopUnswitching>(); 65 } 66 } // namespace mlir 67