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/target/target.h |
22 | * \brief Compilation target object. |
23 | */ |
24 | #ifndef TVM_TARGET_TARGET_H_ |
25 | #define TVM_TARGET_TARGET_H_ |
26 | |
27 | #include <tvm/ir/expr.h> |
28 | #include <tvm/ir/module.h> |
29 | #include <tvm/node/node.h> |
30 | #include <tvm/support/with.h> |
31 | #include <tvm/target/target_kind.h> |
32 | |
33 | #include <string> |
34 | #include <unordered_set> |
35 | #include <vector> |
36 | |
37 | namespace tvm { |
38 | |
39 | class TargetInternal; |
40 | class Target; |
41 | |
42 | /*! |
43 | * \brief Compilation target. |
44 | * \sa Target |
45 | */ |
46 | class TargetNode : public Object { |
47 | public: |
48 | /*! \brief The kind of the target device */ |
49 | TargetKind kind; |
50 | /*! \brief Target host information, must be Target type */ |
51 | Optional<ObjectRef> host; |
52 | /*! \brief Tag of the target, can be empty */ |
53 | String tag; |
54 | /*! \brief Keys for this target */ |
55 | Array<String> keys; |
56 | /*! \brief Collection of attributes */ |
57 | Map<String, ObjectRef> attrs; |
58 | /*! \brief Target features */ |
59 | Map<String, ObjectRef> features; |
60 | |
61 | /*! |
62 | * \brief The raw string representation of the target |
63 | * \return the full device string to pass to codegen::Build |
64 | * \note It will be deprecated after the Target RFC is fully landed. |
65 | */ |
66 | TVM_DLL const std::string& str() const; |
67 | /*! \return Export target to JSON-like configuration */ |
68 | TVM_DLL Map<String, ObjectRef> Export() const; |
69 | /*! \return The Optional<Target> typed target host of the TargetNode */ |
70 | TVM_DLL Optional<Target> GetHost() const; |
71 | /*! \return The device type for this target */ |
72 | TVM_DLL int GetTargetDeviceType() const; |
73 | |
74 | /*! |
75 | * \brief Returns a human readable representation of \p Target which includes all fields, |
76 | * especially the host. Useful for diagnostic messages and debugging. |
77 | * |
78 | * TODO(mbs): The ReprPrinter version should perhaps switch to this form, however currently |
79 | * code depends on str() and << being the same. |
80 | */ |
81 | String ToDebugString() const; |
82 | |
83 | void VisitAttrs(AttrVisitor* v) { |
84 | v->Visit("kind" , &kind); |
85 | v->Visit("tag" , &tag); |
86 | v->Visit("keys" , &keys); |
87 | v->Visit("attrs" , &attrs); |
88 | v->Visit("features" , &features); |
89 | v->Visit("host" , &host); |
90 | } |
91 | |
92 | /*! |
93 | * \brief Get an entry from attrs of the target |
94 | * \tparam TObjectRef Type of the attribute |
95 | * \param attr_key The name of the attribute key |
96 | * \param default_value The value returned if the key is not present |
97 | * \return An optional, NullOpt if not found, otherwise the value found |
98 | */ |
99 | template <typename TObjectRef> |
100 | Optional<TObjectRef> GetAttr( |
101 | const std::string& attr_key, |
102 | Optional<TObjectRef> default_value = Optional<TObjectRef>(nullptr)) const { |
103 | static_assert(std::is_base_of<ObjectRef, TObjectRef>::value, |
104 | "Can only call GetAttr with ObjectRef types." ); |
105 | auto it = attrs.find(attr_key); |
106 | if (it != attrs.end()) { |
107 | return Downcast<Optional<TObjectRef>>((*it).second); |
108 | } else { |
109 | return default_value; |
110 | } |
111 | } |
112 | /*! |
113 | * \brief Get an entry from attrs of the target |
114 | * \tparam TObjectRef Type of the attribute |
115 | * \param attr_key The name of the attribute key |
116 | * \param default_value The value returned if the key is not present |
117 | * \return An optional, NullOpt if not found, otherwise the value found |
118 | */ |
119 | template <typename TObjectRef> |
120 | Optional<TObjectRef> GetAttr(const std::string& attr_key, TObjectRef default_value) const { |
121 | return GetAttr<TObjectRef>(attr_key, Optional<TObjectRef>(default_value)); |
122 | } |
123 | |
124 | /*! |
125 | * \brief Get a Target feature |
126 | * |
127 | * \param feature_key The feature key. |
128 | * \param default_value The default value if the key does not exist, defaults to nullptr. |
129 | * |
130 | * \return The result |
131 | * |
132 | * \tparam TOBjectRef the expected object type. |
133 | * \throw Error if the key exists but the value does not match TObjectRef |
134 | * |
135 | * \code |
136 | * |
137 | * void GetTargetFeature(const Target& target) { |
138 | * Bool has_feature = target->GetFeature<Bool>("has_feature", false).value(); |
139 | * } |
140 | * |
141 | * \endcode |
142 | */ |
143 | template <typename TObjectRef> |
144 | Optional<TObjectRef> GetFeature( |
145 | const std::string& feature_key, |
146 | Optional<TObjectRef> default_value = Optional<TObjectRef>(nullptr)) const { |
147 | Optional<TObjectRef> feature = Downcast<Optional<TObjectRef>>(features.Get(feature_key)); |
148 | if (!feature) { |
149 | return default_value; |
150 | } |
151 | return feature; |
152 | } |
153 | // variant that uses TObjectRef to enable implicit conversion to default value. |
154 | template <typename TObjectRef> |
155 | Optional<TObjectRef> GetFeature(const std::string& attr_key, TObjectRef default_value) const { |
156 | return GetFeature<TObjectRef>(attr_key, Optional<TObjectRef>(default_value)); |
157 | } |
158 | |
159 | /*! \brief Get the keys for this target as a vector of string */ |
160 | TVM_DLL std::vector<std::string> GetKeys() const; |
161 | /*! \brief Get the keys for this target as an unordered_set of string */ |
162 | TVM_DLL std::unordered_set<std::string> GetLibs() const; |
163 | |
164 | bool SEqualReduce(const TargetNode* other, SEqualReducer equal) const; |
165 | void SHashReduce(SHashReducer hash_reduce) const; |
166 | |
167 | static constexpr const char* _type_key = "Target" ; |
168 | static constexpr const bool _type_has_method_sequal_reduce = true; |
169 | static constexpr const bool _type_has_method_shash_reduce = true; |
170 | TVM_DECLARE_FINAL_OBJECT_INFO(TargetNode, Object); |
171 | |
172 | private: |
173 | /*! \brief Internal string repr. */ |
174 | mutable std::string str_repr_; |
175 | |
176 | friend class TargetInternal; |
177 | }; |
178 | |
179 | /*! |
180 | * \brief Managed reference class to TargetNode. |
181 | * \sa TargetNode |
182 | */ |
183 | class Target : public ObjectRef { |
184 | public: |
185 | /*! \brief Construct a null Target */ |
186 | TVM_DLL explicit Target(std::nullptr_t) { data_ = nullptr; } |
187 | /*! |
188 | * \brief Construct a Target given a string |
189 | * \param tag_or_config_or_target_str the string to parse for target |
190 | */ |
191 | TVM_DLL explicit Target(const String& tag_or_config_or_target_str); |
192 | /*! |
193 | * \brief Construct a Target using a JSON-like configuration |
194 | * \param config The JSON-like configuration for target |
195 | */ |
196 | TVM_DLL explicit Target(const Map<String, ObjectRef>& config); |
197 | /*! |
198 | * \brief Get the current target context from thread local storage. |
199 | * \param allow_not_defined If the context stack is empty and this is set to true, an |
200 | * undefined Target will be returned. Otherwise, an empty context stack will cause a |
201 | * runtime error. |
202 | * \return The target that is the current context. The target may not be defined if |
203 | * allow_not_defined is true. |
204 | */ |
205 | TVM_DLL static tvm::Target Current(bool allow_not_defined = true); |
206 | /*! |
207 | * \brief Construct a Target given target and host |
208 | * \param target The Target typed object with host field undefined for target |
209 | * \param host The Target typed object for target host |
210 | * \return The Target with given target and host context information |
211 | */ |
212 | TVM_DLL explicit Target(Target target, Target host); |
213 | TVM_DEFINE_OBJECT_REF_METHODS(Target, ObjectRef, TargetNode); |
214 | /*! |
215 | * \brief Create a new Target object with given target (w.o host) and target host. |
216 | * \param target The current Target typed object target, with or without host field. |
217 | * \param host The given Target typed object target host |
218 | * \return The new Target object with the given target and host field of given host. |
219 | */ |
220 | static Target WithHost(const Target& target, const Target& host); |
221 | |
222 | /*! |
223 | * \brief Returns true if \p this target represents an external codegen. If so, |
224 | * \p this->kind->name can be used as the "Compiler" attribute on partitioned functions, |
225 | * and can be used to retrieve a partitioning pattern table using |
226 | * \p get_pattern_table. |
227 | */ |
228 | bool IsExternalCodegen() const; |
229 | |
230 | /*! |
231 | * \brief Returns true if \p this target represents an external codegen which is compatible |
232 | * with \p that target. In particular: |
233 | * - \p this has a true ::tvm::attr::kIsExternalCodegen attribute |
234 | * - \p that does not have a true ::tvm::attr::kIsExternalCodegen attribute |
235 | * - \p this and \p that have the same GetTargetDeviceType() |
236 | * |
237 | * After partitioning, the external codegen compilation path may use \p that to guide it's |
238 | * compilation to a \p runtime::Module. Given \p this, an appropriate \p that can be |
239 | * found using \p CompilationConfig::FindPrimitiveTargetOrFail(this->GetTargetDeviceType()). |
240 | * |
241 | * The \p CollagePartition pass uses this method to guide it's search over candidate partitions |
242 | * using external codegen. |
243 | */ |
244 | bool IsExternalCodegenFor(const Target& that) const; |
245 | |
246 | private: |
247 | Target(TargetKind kind, Optional<ObjectRef> host, String tag, Array<String> keys, |
248 | Map<String, ObjectRef> attrs); |
249 | |
250 | // enable with syntax. |
251 | friend class TargetInternal; |
252 | friend class With<Target>; |
253 | /*! |
254 | * \brief Push a new target context onto the thread local stack. |
255 | * The Target on top of the stack is used to determine which |
256 | * specialization to use when invoking a GenericFunc. |
257 | */ |
258 | TVM_DLL void EnterWithScope(); |
259 | /*! |
260 | * \brief Pop a target off the thread local context stack, |
261 | * restoring the previous target as the current context. |
262 | */ |
263 | TVM_DLL void ExitWithScope(); |
264 | }; |
265 | |
266 | /*! |
267 | * \brief Check and update host field of the given legacy target and target host pair. |
268 | * Note that this function is for legacy target api compatibility issue only, not |
269 | * recommended for other use. |
270 | * \param target The pointer to a Target typed object with host field to be updated |
271 | * \param host The pointer to a Target typed object for target host to be updated |
272 | */ |
273 | void CheckAndUpdateHostConsistency(Target* target, Target* host); |
274 | |
275 | /*! |
276 | * \brief Check and update host field of the given legacy heterogeneous targets and |
277 | * target host.Note that this function is for legacy target api compatibility issue only, |
278 | * not recommended for other use. |
279 | * \param ir_modules The pointer to a Map objects with keys being Target objects |
280 | * \param host The Target typed object for target host to be updated |
281 | */ |
282 | void CheckAndUpdateHostConsistency(Map<Target, IRModule>* ir_modules, Target* host); |
283 | |
284 | } // namespace tvm |
285 | #endif // TVM_TARGET_TARGET_H_ |
286 | |