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<FuncOp>> {
27   /// Get the command-line argument of the test pass.
28   StringRef getArgument() const final { return "test-control-flow-sink"; }
29   /// Get the description of the test pass.
30   StringRef getDescription() const final {
31     return "Test control-flow sink pass";
32   }
33 
34   /// Runs the pass on the function.
35   void runOnOperation() override {
36     auto &domInfo = getAnalysis<DominanceInfo>();
37     auto shouldMoveIntoRegion = [](Operation *op, Region *region) {
38       return region->getRegionNumber() == 0 &&
39              op->getName().getStringRef() == "test.sink_me";
40     };
41     auto moveIntoRegion = [](Operation *op, Region *region) {
42       Block &entry = region->front();
43       op->moveBefore(&entry, entry.begin());
44       op->setAttr("was_sunk",
45                   Builder(op).getI32IntegerAttr(region->getRegionNumber()));
46     };
47 
48     getOperation()->walk([&](Operation *op) {
49       if (op->getName().getStringRef() != "test.sink_target")
50         return;
51       SmallVector<Region *> regions =
52           llvm::to_vector(RegionRange(op->getRegions()));
53       controlFlowSink(regions, domInfo, shouldMoveIntoRegion, moveIntoRegion);
54     });
55   }
56 };
57 } // end anonymous namespace
58 
59 namespace mlir {
60 namespace test {
61 void registerTestControlFlowSink() {
62   PassRegistration<TestControlFlowSinkPass>();
63 }
64 } // end namespace test
65 } // end namespace mlir
66