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