1 | /* |
2 | * Licensed to the Apache Software Foundation (ASF) under one |
3 | * or more contributor license agreements. See the NOTICE file |
4 | * distributed with this work for additional information |
5 | * regarding copyright ownership. The ASF licenses this file |
6 | * to you under the Apache License, Version 2.0 (the |
7 | * "License"); you may not use this file except in compliance |
8 | * with the License. You may obtain a copy of the License at |
9 | * |
10 | * http://www.apache.org/licenses/LICENSE-2.0 |
11 | * |
12 | * Unless required by applicable law or agreed to in writing, |
13 | * software distributed under the License is distributed on an |
14 | * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
15 | * KIND, either express or implied. See the License for the |
16 | * specific language governing permissions and limitations |
17 | * under the License. |
18 | */ |
19 | |
20 | /*! |
21 | * \file src/relay/backend/liveness_analysis.h |
22 | * \brief Analysis that collects the live variables before and after each node. |
23 | * NOTE: the input IR should be in ANF. |
24 | */ |
25 | |
26 | #ifndef TVM_RELAY_BACKEND_LIVENESS_ANALYSIS_H_ |
27 | #define TVM_RELAY_BACKEND_LIVENESS_ANALYSIS_H_ |
28 | |
29 | #include <tvm/relay/transform.h> |
30 | |
31 | #include <unordered_map> |
32 | #include <unordered_set> |
33 | #include <vector> |
34 | |
35 | #include "../../support/arena.h" |
36 | #include "../op/memory/device_copy.h" |
37 | #include "../transforms/device_aware_visitors.h" |
38 | #include "../transforms/let_list.h" |
39 | |
40 | namespace tvm { |
41 | namespace relay { |
42 | namespace transform { |
43 | |
44 | using support::Arena; |
45 | using VarSet = std::unordered_set<Var, ObjectPtrHash, ObjectPtrEqual>; |
46 | |
47 | // TODO(@altanh, @mbs, @mbrookhart): we should do a survey of all "*-flow graphs" in the codebase |
48 | // to see what can be deduplicated. |
49 | |
50 | // TODO(@altanh): support Relay Refs once/if they are supported by the VM. |
51 | |
52 | /*! |
53 | * \brief A representation of an input expression (typically a Function) as a directed graph of |
54 | * basic blocks, with edges between basic blocks corresponding to control flow branching. |
55 | */ |
56 | class ControlFlowGraph { |
57 | public: |
58 | struct Node; |
59 | struct BasicBlock; |
60 | |
61 | using NodePtr = Node*; |
62 | using BasicBlockPtr = BasicBlock*; |
63 | |
64 | /*! |
65 | * \brief A chunk of IR that does not have any control flow branching. At this stage in the IR, |
66 | * basic blocks correspond to: |
67 | * (1) a sequence of nested Let expressions, where each node in the block corresponds to a |
68 | * binding and the last node is either the (non-Let) body or a binding that branches |
69 | * (e.g. "let %x = if (%c) { true_block } else { false_block }"). |
70 | * (2) an atomic expression representing the target expression of a control flow branch, e.g. |
71 | * %v and %u in "let %x = if (%c) { %v } else { %u }". |
72 | */ |
73 | struct BasicBlock { |
74 | // The nodes of the basic block. |
75 | std::vector<NodePtr> nodes; |
76 | // The predecessor basic blocks. |
77 | std::vector<BasicBlockPtr> pred; |
78 | // The successor basic blocks. |
79 | std::vector<BasicBlockPtr> succ; |
80 | |
81 | static BasicBlockPtr Make(support::Arena* arena) { return arena->make<BasicBlock>(); } |
82 | }; |
83 | |
84 | /*! |
85 | * \brief Roughly corresponds to a "statement" in the IR, such as an individual binding in a |
86 | * basic block or the "return value" of a block. Each node maps to a single corresponding expr in |
87 | * the IR, but the converse is not true (e.g. in the case of variables). |
88 | */ |
89 | struct Node { |
90 | /*! \brief The basic block this node belongs to. */ |
91 | BasicBlockPtr parent; |
92 | /*! \brief The index into the parent basic block where this node is. */ |
93 | size_t index; |
94 | /*! \brief The expr this node corresponds to. */ |
95 | Expr expr; |
96 | |
97 | /*! \brief Returns whether or not this node is the first one in the parent basic block. */ |
98 | bool IsFirst() const { return index == 0; } |
99 | |
100 | /*! \brief Returns whether or not this node is the last one in the parent basic block. */ |
101 | bool IsLast() const { return index == parent->nodes.size() - 1; } |
102 | |
103 | /*! \brief Returns the predecessor nodes of this node. */ |
104 | std::vector<NodePtr> GetPred() const { |
105 | std::vector<NodePtr> pred; |
106 | if (IsFirst()) { |
107 | for (const BasicBlockPtr& pred_block : parent->pred) { |
108 | pred.push_back(pred_block->nodes.back()); |
109 | } |
110 | } else { |
111 | pred.push_back(parent->nodes[index - 1]); |
112 | } |
113 | return pred; |
114 | } |
115 | |
116 | /*! \brief Returns the successor nodes of this node. */ |
117 | std::vector<NodePtr> GetSucc() const { |
118 | std::vector<NodePtr> succ; |
119 | if (IsLast()) { |
120 | for (const BasicBlockPtr& succ_block : parent->succ) { |
121 | succ.push_back(succ_block->nodes.front()); |
122 | } |
123 | } else { |
124 | succ.push_back(parent->nodes[index + 1]); |
125 | } |
126 | return succ; |
127 | } |
128 | |
129 | /*! \brief Creates a node with the given expr and appends it to the parent basic block. */ |
130 | static NodePtr Make(Arena* arena, BasicBlockPtr parent, Expr expr) { |
131 | NodePtr n = arena->make<Node>(); |
132 | n->parent = parent; |
133 | n->expr = expr; |
134 | n->index = parent->nodes.size(); |
135 | parent->nodes.push_back(n); |
136 | return n; |
137 | } |
138 | }; |
139 | |
140 | /*! \brief The basic block where control flow begins. */ |
141 | BasicBlockPtr entry; |
142 | |
143 | /*! |
144 | * \brief Mapping from Let expressions to their corresponding nodes. Note that Let expressions |
145 | * are never shared in ANF (unlike vars), so this is an injection. |
146 | */ |
147 | std::unordered_map<Expr, NodePtr, ObjectPtrHash, ObjectPtrEqual> let_map; |
148 | |
149 | /*! \brief The nodes of the CFG in reverse post order. */ |
150 | std::vector<NodePtr> reverse_post_order; |
151 | |
152 | /*! \brief Creates and returns the CFG of the given expression. */ |
153 | static ControlFlowGraph Create(Arena* arena, const Expr& body); |
154 | |
155 | private: |
156 | class Creator; |
157 | }; |
158 | |
159 | /*! \brief Helper class for building CFGs. */ |
160 | class ControlFlowGraph::Creator : private ExprFunctor<void(const Expr&, BasicBlockPtr)> { |
161 | public: |
162 | Creator() {} |
163 | |
164 | ControlFlowGraph Create(Arena* arena, const Expr& body); |
165 | |
166 | private: |
167 | /*! \brief The arena allocator. */ |
168 | Arena* arena_; |
169 | |
170 | /*! \brief The CFG being built. */ |
171 | ControlFlowGraph cfg_; |
172 | /*! |
173 | * \brief Whether or not we are in a function. CFGs do not support nested functions so this is |
174 | * used to error out in such a case. |
175 | */ |
176 | bool in_func_ = false; |
177 | |
178 | /*! |
179 | * \brief Link \p to as a successor block to \p from. |
180 | */ |
181 | void Succ(BasicBlockPtr from, BasicBlockPtr to); |
182 | |
183 | #define DEFAULT_CFG(OP) \ |
184 | void VisitExpr_(const OP* op, BasicBlockPtr parent) final { \ |
185 | NodePtr n = Node::Make(arena_, parent, GetRef<Expr>(op)); \ |
186 | cfg_.reverse_post_order.push_back(n); \ |
187 | } |
188 | |
189 | void VisitExpr_(const FunctionNode* f, BasicBlockPtr parent) final; |
190 | void VisitExpr_(const LetNode* let_node, BasicBlockPtr parent) final; |
191 | void VisitExpr_(const IfNode* if_node, BasicBlockPtr parent); |
192 | void VisitExpr_(const MatchNode* match_node, BasicBlockPtr parent); |
193 | |
194 | DEFAULT_CFG(VarNode); |
195 | DEFAULT_CFG(GlobalVarNode); |
196 | DEFAULT_CFG(ConstantNode); |
197 | DEFAULT_CFG(CallNode); |
198 | DEFAULT_CFG(OpNode); |
199 | DEFAULT_CFG(TupleNode); |
200 | DEFAULT_CFG(TupleGetItemNode); |
201 | }; |
202 | |
203 | /*! |
204 | * \brief Helper class for collecting the variables used/read by an expression. NOTE: for If exprs, |
205 | * only the condition is included (not the branches). Similarly, for Match exprs only the value |
206 | * being deconstructed is included. |
207 | */ |
208 | class VarUseCollector : public ExprFunctor<VarSet(const Expr& e)> { |
209 | public: |
210 | VarSet VisitExpr_(const VarNode* var_node); |
211 | VarSet VisitExpr_(const CallNode* call_node); |
212 | VarSet VisitExpr_(const TupleNode* tuple_node); |
213 | VarSet VisitExpr_(const TupleGetItemNode* get_node); |
214 | VarSet VisitExpr_(const IfNode* if_node); |
215 | VarSet VisitExpr_(const MatchNode* match_node); |
216 | |
217 | VarSet VisitExpr_(const ConstructorNode* cons_node) { return {}; } |
218 | VarSet VisitExpr_(const GlobalVarNode* gvar_node) { return {}; } |
219 | VarSet VisitExpr_(const ConstantNode* const_node) { return {}; } |
220 | VarSet VisitExpr_(const OpNode* op_node) { return {}; } |
221 | VarSet VisitExpr_(const FunctionNode* func_node) { return {}; } |
222 | }; |
223 | |
224 | /*! |
225 | * \brief Analysis that collects the variables used and defined at each node. |
226 | */ |
227 | struct UseDefAnalysis { |
228 | using CFG = ControlFlowGraph; |
229 | |
230 | /*! \brief Mapping of node -> variables used/read by node. */ |
231 | std::unordered_map<CFG::NodePtr, VarSet> use; |
232 | |
233 | /*! \brief Mapping of node -> variable defined/written by node. */ |
234 | std::unordered_map<CFG::NodePtr, Var> def; |
235 | |
236 | VarUseCollector use_collector; |
237 | |
238 | static UseDefAnalysis Analyze(const CFG& cfg); |
239 | }; |
240 | |
241 | /*! \brief Returns whether \p a and \p b are the same set of vars. */ |
242 | bool SetEqual(const VarSet& a, const VarSet& b); |
243 | |
244 | /*! |
245 | * \brief Analysis that collects the live variables before and after each node. |
246 | */ |
247 | struct LivenessAnalysis { |
248 | using CFG = ControlFlowGraph; |
249 | |
250 | /*! \brief Mapping of node -> set of variables live before node. */ |
251 | std::unordered_map<CFG::NodePtr, VarSet> live_in; |
252 | |
253 | /*! \brief Mapping of node -> set of variables live after node. */ |
254 | std::unordered_map<CFG::NodePtr, VarSet> live_out; |
255 | |
256 | /*! |
257 | * \brief Analyze the input \p cfg (using info from \p use_def). |
258 | * |
259 | * \param cfg The input control flow graph. |
260 | * \param use_def Use-def analysis of \p cfg. |
261 | * \return LivenessAnalysis |
262 | */ |
263 | static LivenessAnalysis Analyze(const ControlFlowGraph& cfg, const UseDefAnalysis& use_def); |
264 | }; |
265 | |
266 | } // namespace transform |
267 | } // namespace relay |
268 | } // namespace tvm |
269 | |
270 | #endif // TVM_RELAY_BACKEND_LIVENESS_ANALYSIS_H_ |
271 | |