1 //===- ControlFlowSinkUtils.cpp - Code to perform control-flow sinking ----===//
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 utilities for control-flow sinking. Control-flow
10 // sinking moves operations whose only uses are in conditionally-executed blocks
11 // into those blocks so that they aren't executed on paths where their results
12 // are not needed.
13 //
14 // Control-flow sinking is not implemented on BranchOpInterface because
15 // sinking ops into the successors of branch operations may move ops into loops.
16 // It is idiomatic MLIR to perform optimizations at IR levels that readily
17 // provide the necessary information.
18 //
19 //===----------------------------------------------------------------------===//
20 
21 #include "mlir/Transforms/ControlFlowSinkUtils.h"
22 #include "mlir/IR/Dominance.h"
23 #include "mlir/IR/Matchers.h"
24 #include "mlir/Interfaces/ControlFlowInterfaces.h"
25 #include <vector>
26 
27 #define DEBUG_TYPE "cf-sink"
28 
29 using namespace mlir;
30 
31 namespace {
32 /// A helper struct for control-flow sinking.
33 class Sinker {
34 public:
35   /// Create an operation sinker with given dominance info.
36   Sinker(function_ref<bool(Operation *, Region *)> shouldMoveIntoRegion,
37          DominanceInfo &domInfo)
38       : shouldMoveIntoRegion(shouldMoveIntoRegion), domInfo(domInfo) {}
39 
40   /// Given a list of regions, find operations to sink and sink them. Return the
41   /// number of operations sunk.
42   size_t sinkRegions(ArrayRef<Region *> regions);
43 
44 private:
45   /// Given a region and an op which dominates the region, returns true if all
46   /// users of the given op are dominated by the entry block of the region, and
47   /// thus the operation can be sunk into the region.
48   bool allUsersDominatedBy(Operation *op, Region *region);
49 
50   /// Given a region and a top-level op (an op whose parent region is the given
51   /// region), determine whether the defining ops of the op's operands can be
52   /// sunk into the region.
53   ///
54   /// Add moved ops to the work queue.
55   void tryToSinkPredecessors(Operation *user, Region *region,
56                              std::vector<Operation *> &stack);
57 
58   /// Iterate over all the ops in a region and try to sink their predecessors.
59   /// Recurse on subgraphs using a work queue.
60   void sinkRegion(Region *region);
61 
62   /// The callback to determine whether an op should be moved in to a region.
63   function_ref<bool(Operation *, Region *)> shouldMoveIntoRegion;
64   /// Dominance info to determine op user dominance with respect to regions.
65   DominanceInfo &domInfo;
66   /// The number of operations sunk.
67   size_t numSunk = 0;
68 };
69 } // end anonymous namespace
70 
71 bool Sinker::allUsersDominatedBy(Operation *op, Region *region) {
72   assert(region->findAncestorOpInRegion(*op) == nullptr &&
73          "expected op to be defined outside the region");
74   return llvm::all_of(op->getUsers(), [&](Operation *user) {
75     // The user is dominated by the region if its containing block is dominated
76     // by the region's entry block.
77     return domInfo.dominates(&region->front(), user->getBlock());
78   });
79 }
80 
81 void Sinker::tryToSinkPredecessors(Operation *user, Region *region,
82                                    std::vector<Operation *> &stack) {
83   LLVM_DEBUG(user->print(llvm::dbgs() << "\nContained op:\n"));
84   for (Value value : user->getOperands()) {
85     Operation *op = value.getDefiningOp();
86     // Ignore block arguments and ops that are already inside the region.
87     if (!op || op->getParentRegion() == region)
88       continue;
89     LLVM_DEBUG(op->print(llvm::dbgs() << "\nTry to sink:\n"));
90 
91     // If the op's users are all in the region and it can be moved, then do so.
92     if (allUsersDominatedBy(op, region) && shouldMoveIntoRegion(op, region)) {
93       // Move the op into the region's entry block. If the op is part of a
94       // subgraph, dependee ops would have been moved first, so inserting before
95       // the start of the block will ensure SSA dominance is preserved locally
96       // in the subgraph. Ops can only be safely moved into the entry block as
97       // the region's other blocks may for a loop.
98       op->moveBefore(&region->front(), region->front().begin());
99       ++numSunk;
100       // Add the op to the work queue.
101       stack.push_back(op);
102     }
103   }
104 }
105 
106 void Sinker::sinkRegion(Region *region) {
107   // Initialize the work queue with all the ops in the region.
108   std::vector<Operation *> stack;
109   for (Operation &op : region->getOps())
110     stack.push_back(&op);
111 
112   // Process all the ops depth-first. This ensures that nodes of subgraphs are
113   // sunk in the correct order.
114   while (!stack.empty()) {
115     Operation *op = stack.back();
116     stack.pop_back();
117     tryToSinkPredecessors(op, region, stack);
118   }
119 }
120 
121 size_t Sinker::sinkRegions(ArrayRef<Region *> regions) {
122   for (Region *region : regions)
123     if (!region->empty())
124       sinkRegion(region);
125   return numSunk;
126 }
127 
128 size_t mlir::controlFlowSink(
129     ArrayRef<Region *> regions, DominanceInfo &domInfo,
130     function_ref<bool(Operation *, Region *)> shouldMoveIntoRegion) {
131   return Sinker(shouldMoveIntoRegion, domInfo).sinkRegions(regions);
132 }
133 
134 void mlir::getSinglyExecutedRegionsToSink(RegionBranchOpInterface branch,
135                                           SmallVectorImpl<Region *> &regions) {
136   // Collect constant operands.
137   SmallVector<Attribute> operands(branch->getNumOperands(), Attribute());
138   for (auto &it : llvm::enumerate(branch->getOperands()))
139     (void)matchPattern(it.value(), m_Constant(&operands[it.index()]));
140 
141   // Get the invocation bounds.
142   SmallVector<InvocationBounds> bounds;
143   branch.getRegionInvocationBounds(operands, bounds);
144 
145   // For a simple control-flow sink, only consider regions that are executed at
146   // most once.
147   for (auto it : llvm::zip(branch->getRegions(), bounds)) {
148     const InvocationBounds &bound = std::get<1>(it);
149     if (bound.getUpperBound() && *bound.getUpperBound() <= 1)
150       regions.push_back(&std::get<0>(it));
151   }
152 }
153