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 *
22 * \file tvm/relay/pass/pass_util.h
23 * \brief Utilities for writing passes
24 */
25 #ifndef TVM_RELAY_PASS_PASS_UTIL_H_
26 #define TVM_RELAY_PASS_PASS_UTIL_H_
27
28 #include <tvm/relay/op.h>
29 #include <tvm/relay/expr.h>
30 #include <tvm/relay/attrs/transform.h>
31 #include <memory>
32 #include <unordered_map>
33
34 namespace tvm {
35 namespace relay {
36
37 /*!
38 * \brief Get reference counter of each internal ExprNode in body.
39 * \param body The body expression.
40 * \return The reference count mapping.
41 */
42 std::unordered_map<const Node*, size_t>
43 GetExprRefCount(const Expr& body);
44
45 /*!
46 * \brief Check if expr is positive constant.
47 * \param expr The expression to be checked.
48 * \return Whether all elements of expr is positive constant.
49 */
50 bool IsAllPositiveConstant(const Expr& expr);
51
52 /*!
53 * \brief Substitute var with subst.
54 * \param type The type to be substituted.
55 * \param tvar The type variable to be substituted.
56 * \param subst The target of substitution.
57 * \return The substituted result.
58 */
59 Type TypeSubst(const Type& type, const TypeVar& tvar, const Type& subst);
60
61 /*!
62 * \brief Substitute var with subst.
63 * \param expr The expr to be substituted.
64 * \param tvar The type variable to be substituted.
65 * \param subst The target of substitution.
66 * \return The substituted result.
67 */
68 Expr TypeSubst(const Expr& expr, const TypeVar& tvar, const Type& subst);
69
70 /*!
71 * \brief Substitute type vars in type.
72 * \param type The type to be substituted.
73 * \param subst_map The map of substitution.
74 * \return The substituted result.
75 */
76 Type TypeSubst(const Type& type, const tvm::Map<TypeVar, Type>& subst_map);
77
78 /*!
79 * \brief Substitute type vars in type.
80 * \param expr The expr to be substituted.
81 * \param subst_map The map of substitution.
82 * \return The substituted result.
83 */
84 Expr TypeSubst(const Expr& expr, const tvm::Map<TypeVar, Type>& subst_map);
85
86 /*!
87 * \brief Make arbitrary transformation preserve the out most function.
88 * \param func The transformation.
89 * \param e The expression
90 * \return the transformed expression. If e is a function the return is also a function.
91 */
TransformF(const std::function<Expr (const Expr &)> & func,const Expr & e)92 inline Expr TransformF(const std::function<Expr(const Expr&)>& func, const Expr& e) {
93 if (const FunctionNode* f = e.as<FunctionNode>()) {
94 return FunctionNode::make(f->params, func(f->body), f->ret_type, f->type_params, f->attrs);
95 } else {
96 return func(e);
97 }
98 }
99
100 /*!
101 * \brief Decide whether the expression atomic or not?
102 * \param e the expression
103 * \return
104 * is it atomic?
105 * if so, the compute cost of the expression is bounded so it can be copy without graph mode.
106 */
IsAtomic(const Expr & e)107 inline bool IsAtomic(const Expr& e) {
108 return e.as<VarNode>() || e.as<OpNode>() || e.as<ConstructorNode>() || e.as<GlobalVarNode>();
109 }
110
111 template<typename ConditionNodePtr>
112 struct TreeNode {
113 typedef std::shared_ptr<TreeNode<ConditionNodePtr>> pointer;
~TreeNodeTreeNode114 virtual ~TreeNode() {}
115 };
116
117 template<typename ConditionNodePtr>
118 struct TreeLeafNode : TreeNode<ConditionNodePtr> {
119 using TreeNodePtr = typename TreeNode<ConditionNodePtr>::pointer;
120
121 Expr body;
122
TreeLeafNodeTreeLeafNode123 explicit TreeLeafNode(Expr body): body(body) {}
124
MakeTreeLeafNode125 static TreeNodePtr Make(Expr body) {
126 return std::make_shared<TreeLeafNode>(body);
127 }
128
~TreeLeafNodeTreeLeafNode129 ~TreeLeafNode() {}
130 };
131
132 template<typename ConditionNodePtr>
133 struct TreeLeafFatalNode : TreeNode<ConditionNodePtr> {
134 using TreeNodePtr = typename TreeNode<ConditionNodePtr>::pointer;
135
136 TreeLeafFatalNode() = default;
137
MakeTreeLeafFatalNode138 static TreeNodePtr Make() {
139 return std::make_shared<TreeLeafFatalNode>();
140 }
141
~TreeLeafFatalNodeTreeLeafFatalNode142 ~TreeLeafFatalNode() {}
143 };
144
145 template<typename ConditionNodePtr>
146 struct TreeBranchNode : TreeNode<ConditionNodePtr> {
147 using TreeNodePtr = typename TreeNode<ConditionNodePtr>::pointer;
148
149 ConditionNodePtr cond;
150 TreeNodePtr then_branch;
151 TreeNodePtr else_branch;
152
TreeBranchNodeTreeBranchNode153 TreeBranchNode(ConditionNodePtr cond,
154 TreeNodePtr then_branch,
155 TreeNodePtr else_branch)
156 : cond(cond), then_branch(then_branch), else_branch(else_branch) {}
157
158
MakeTreeBranchNode159 static TreeNodePtr Make(ConditionNodePtr cond,
160 TreeNodePtr then_branch,
161 TreeNodePtr else_branch) {
162 return std::make_shared<TreeBranchNode>(cond, then_branch, else_branch);
163 }
164
~TreeBranchNodeTreeBranchNode165 ~TreeBranchNode() {}
166 };
167
168 } // namespace relay
169 } // namespace tvm
170 #endif // TVM_RELAY_PASS_PASS_UTIL_H_
171