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
33namespace tvm {
34namespace relay {
35
36/*! \brief Different kinds of relay feature a program might use. */
37enum 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
59constexpr size_t feature_count = 17;
60
61/*!
62 * \brief A finite set of Feature.
63 */
64class 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 */
145FeatureSet 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 */
154FeatureSet 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 */
164inline 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 */
174void 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 */
182void 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 */
191inline 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