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