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
37namespace tvm {
38
39class TargetInternal;
40class Target;
41
42/*!
43 * \brief Compilation target.
44 * \sa Target
45 */
46class 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 */
183class 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 */
273void 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 */
282void CheckAndUpdateHostConsistency(Map<Target, IRModule>* ir_modules, Target* host);
283
284} // namespace tvm
285#endif // TVM_TARGET_TARGET_H_
286