1 //===- TestControlFlowSink.cpp - Test control-flow sink pass --------------===// 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 pass tests the control-flow sink utilities by implementing an example 10 // control-flow sink pass. 11 // 12 //===----------------------------------------------------------------------===// 13 14 #include "mlir/Dialect/Func/IR/FuncOps.h" 15 #include "mlir/IR/Dominance.h" 16 #include "mlir/Pass/Pass.h" 17 #include "mlir/Transforms/ControlFlowSinkUtils.h" 18 19 using namespace mlir; 20 21 namespace { 22 /// An example control-flow sink pass to test the control-flow sink utilites. 23 /// This pass will sink ops named `test.sink_me` and tag them with an attribute 24 /// `was_sunk` into the first region of `test.sink_target` ops. 25 struct TestControlFlowSinkPass 26 : public PassWrapper<TestControlFlowSinkPass, OperationPass<func::FuncOp>> { MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID__anon204f40370111::TestControlFlowSinkPass27 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestControlFlowSinkPass) 28 29 /// Get the command-line argument of the test pass. 30 StringRef getArgument() const final { return "test-control-flow-sink"; } 31 /// Get the description of the test pass. getDescription__anon204f40370111::TestControlFlowSinkPass32 StringRef getDescription() const final { 33 return "Test control-flow sink pass"; 34 } 35 36 /// Runs the pass on the function. runOnOperation__anon204f40370111::TestControlFlowSinkPass37 void runOnOperation() override { 38 auto &domInfo = getAnalysis<DominanceInfo>(); 39 auto shouldMoveIntoRegion = [](Operation *op, Region *region) { 40 return region->getRegionNumber() == 0 && 41 op->getName().getStringRef() == "test.sink_me"; 42 }; 43 auto moveIntoRegion = [](Operation *op, Region *region) { 44 Block &entry = region->front(); 45 op->moveBefore(&entry, entry.begin()); 46 op->setAttr("was_sunk", 47 Builder(op).getI32IntegerAttr(region->getRegionNumber())); 48 }; 49 50 getOperation()->walk([&](Operation *op) { 51 if (op->getName().getStringRef() != "test.sink_target") 52 return; 53 SmallVector<Region *> regions = 54 llvm::to_vector(RegionRange(op->getRegions())); 55 controlFlowSink(regions, domInfo, shouldMoveIntoRegion, moveIntoRegion); 56 }); 57 } 58 }; 59 } // end anonymous namespace 60 61 namespace mlir { 62 namespace test { registerTestControlFlowSink()63void registerTestControlFlowSink() { 64 PassRegistration<TestControlFlowSinkPass>(); 65 } 66 } // end namespace test 67 } // end namespace mlir 68