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 tvm/relay/feature.h |
22 | * \brief Detect features used in Expr/Module. |
23 | */ |
24 | #ifndef TVM_RELAY_FEATURE_H_ |
25 | #define TVM_RELAY_FEATURE_H_ |
26 | |
27 | #include <tvm/ir/module.h> |
28 | #include <tvm/relay/expr.h> |
29 | |
30 | #include <bitset> |
31 | #include <string> |
32 | |
33 | namespace tvm { |
34 | namespace relay { |
35 | |
36 | /*! \brief Different kinds of relay feature a program might use. */ |
37 | enum Feature : int { |
38 | fVar = 0, |
39 | fGlobalVar = 1, |
40 | fConstant = 2, |
41 | fTuple = 3, |
42 | fTupleGetItem = 4, |
43 | fFunction = 5, |
44 | fOp = 6, |
45 | fCall = 7, |
46 | fLet = 8, |
47 | fIf = 9, |
48 | fRefCreate = 10, |
49 | fRefRead = 11, |
50 | fRefWrite = 12, |
51 | fConstructor = 13, |
52 | fMatch = 14, |
53 | /*! \brief Whether any non-atom fragment of the program is shared, making the program a graph. */ |
54 | fGraph = 15, |
55 | /*! \brief Whether there is local fixpoint in the program. */ |
56 | fLetRec = 16 |
57 | }; |
58 | |
59 | constexpr size_t feature_count = 17; |
60 | |
61 | /*! |
62 | * \brief A finite set of Feature. |
63 | */ |
64 | class FeatureSet { |
65 | public: |
66 | FeatureSet(const FeatureSet&) = default; |
67 | /*! \brief A singleton set containing a single Feature. */ |
68 | explicit FeatureSet(Feature ft) { bs_.set(static_cast<size_t>(ft)); } |
69 | explicit FeatureSet(const tvm::Array<tvm::Integer>& ft) { |
70 | for (Integer i : ft) { |
71 | *this += Feature(i.IntValue()); |
72 | } |
73 | } |
74 | explicit operator Array<Integer>() const { |
75 | Array<Integer> ret; |
76 | for (size_t i = 0; i < feature_count; ++i) { |
77 | if (bs_[i]) { |
78 | ret.push_back(Integer(i)); |
79 | } |
80 | } |
81 | return ret; |
82 | } |
83 | /*! \brief A set that contain all the Feature. */ |
84 | static FeatureSet All() { |
85 | FeatureSet fs; |
86 | fs.bs_.flip(); |
87 | return fs; |
88 | } |
89 | /*! \brief The empty set. Contain no Feature. */ |
90 | static FeatureSet No() { |
91 | FeatureSet fs; |
92 | return fs; |
93 | } |
94 | template <typename T> |
95 | FeatureSet& operator+=(const T& rhs) { |
96 | bs_ |= FeatureSet(rhs).bs_; |
97 | return *this; |
98 | } |
99 | /*! \brief Set union. */ |
100 | template <typename T> |
101 | FeatureSet operator+(const T& rhs) const { |
102 | FeatureSet fs(*this); |
103 | fs += rhs; |
104 | return fs; |
105 | } |
106 | template <typename T> |
107 | FeatureSet& operator-=(const T& rhs) { |
108 | bs_ &= ~(FeatureSet(rhs)).bs_; |
109 | return *this; |
110 | } |
111 | /*! \brief Set difference. */ |
112 | template <typename T> |
113 | FeatureSet operator-(const T& rhs) const { |
114 | FeatureSet fs(*this); |
115 | fs -= rhs; |
116 | return fs; |
117 | } |
118 | /*! |
119 | * \brief Is this a subset of rhs? |
120 | * |
121 | * \param rhs another FeatureSet. |
122 | * |
123 | * \return true only if this is a subset of rhs. |
124 | */ |
125 | bool is_subset_of(const FeatureSet& rhs) const { return ((*this) - rhs).bs_.none(); } |
126 | |
127 | /*! |
128 | * \brief return a string representation. |
129 | */ |
130 | std::string ToString() const; |
131 | |
132 | private: |
133 | std::bitset<feature_count> bs_; |
134 | FeatureSet() = default; |
135 | explicit FeatureSet(const std::bitset<feature_count>& bs) : bs_(bs) {} |
136 | }; |
137 | |
138 | /*! |
139 | * \brief Calculate the feature of the program. |
140 | * |
141 | * \param expr The expression. |
142 | * |
143 | * \return The FeatureSet. |
144 | */ |
145 | FeatureSet DetectFeature(const RelayExpr& expr); |
146 | |
147 | /*! |
148 | * \brief Calculate the feature of the program. |
149 | * |
150 | * \param mod The module. |
151 | * |
152 | * \return The FeatureSet. |
153 | */ |
154 | FeatureSet DetectFeature(const IRModule& mod); |
155 | |
156 | /*! |
157 | * \brief Calculate the feature of the program. |
158 | * |
159 | * \param expr The expression. |
160 | * \param mod The module. |
161 | * |
162 | * \return The FeatureSet. |
163 | */ |
164 | inline FeatureSet DetectFeature(const Expr& expr, const IRModule& mod) { |
165 | return DetectFeature(expr) + DetectFeature(mod); |
166 | } |
167 | |
168 | /*! |
169 | * \brief Check the feature of the program. |
170 | * |
171 | * \param expr The expression. |
172 | * \param fs The feature set of the program. |
173 | */ |
174 | void CheckFeature(const RelayExpr& expr, const FeatureSet& fs); |
175 | |
176 | /*! |
177 | * \brief Check the feature of the program. |
178 | * |
179 | * \param mod The module. |
180 | * \param fs The feature set of the program. |
181 | */ |
182 | void CheckFeature(const IRModule& mod, const FeatureSet& fs); |
183 | |
184 | /*! |
185 | * \brief Check the feature of the program. |
186 | * |
187 | * \param expr The expression. |
188 | * \param mod The module. |
189 | * \param fs The feature set of the program. |
190 | */ |
191 | inline void CheckFeature(const RelayExpr& expr, const IRModule& mod, const FeatureSet& fs) { |
192 | CheckFeature(expr, fs); |
193 | CheckFeature(mod, fs); |
194 | } |
195 | |
196 | } // namespace relay |
197 | } // namespace tvm |
198 | |
199 | #endif // TVM_RELAY_FEATURE_H_ |
200 | |