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
40namespace tvm {
41namespace relay {
42namespace transform {
43
44using support::Arena;
45using 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 */
56class 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. */
160class 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 */
208class 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 */
227struct 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. */
242bool SetEqual(const VarSet& a, const VarSet& b);
243
244/*!
245 * \brief Analysis that collects the live variables before and after each node.
246 */
247struct 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