Skip to content

Commit 0bec7b4

Browse files
tqchentqchen
authored
Use dedicated TargetKind registry and fluent definitions (#20469)
TargetKind only needs canonical named instances and typed configuration options. A small dedicated registry and temporary TargetKindDef builder make registration direct and remove the unused attribute-map API. - Keep TargetKind as a plain Object with its existing reflection and target behavior. - Define common and backend options through ConfigSchema in grouped static initialization blocks. - Remove the generic AttrRegistry dependency and obsolete attribute-map forwarding API.
1 parent cd88be6 commit 0bec7b4

12 files changed

Lines changed: 398 additions & 523 deletions

File tree

‎include/tvm/target/target_kind.h‎

Lines changed: 86 additions & 181 deletions
Original file line numberDiff line numberDiff line change
@@ -24,14 +24,12 @@
2424
#ifndef TVM_TARGET_TARGET_KIND_H_
2525
#define TVM_TARGET_TARGET_KIND_H_
2626

27+
#include <tvm/ffi/container/map.h>
2728
#include <tvm/ffi/function.h>
2829
#include <tvm/ffi/reflection/registry.h>
29-
#include <tvm/ir/attr_registry_map.h>
3030
#include <tvm/ir/config_schema.h>
3131
#include <tvm/runtime/base.h>
3232

33-
#include <memory>
34-
#include <unordered_map>
3533
#include <utility>
3634
#include <vector>
3735

@@ -50,9 +48,6 @@ using FTargetCanonicalizer =
5048

5149
class TargetInternal;
5250

53-
template <typename>
54-
class TargetKindAttrMap;
55-
5651
/*! \brief Target kind, specifies the kind of the target */
5752
class TargetKindNode : public ffi::Object {
5853
public:
@@ -79,20 +74,10 @@ class TargetKindNode : public ffi::Object {
7974
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("target.TargetKind", TargetKindNode, ffi::Object);
8075

8176
private:
82-
/*! \brief Return the index stored in attr registry */
83-
uint32_t AttrRegistryIndex() const { return index_; }
84-
/*! \brief Return the name stored in attr registry */
85-
ffi::String AttrRegistryName() const { return name; }
8677
/*! \brief ConfigSchema for validating and resolving target attributes */
8778
ir::ConfigSchema schema_;
88-
/*! \brief Index used for internal lookup of attribute registry */
89-
uint32_t index_;
90-
91-
template <typename, typename>
92-
friend class AttrRegistry;
93-
template <typename>
94-
friend class AttrRegistryMapContainerMap;
95-
friend class TargetKindRegEntry;
79+
friend class TargetKindRegistry;
80+
friend class TargetKindDef;
9681
friend class TargetInternal;
9782
};
9883

@@ -106,13 +91,10 @@ class TargetKind : public ffi::ObjectRef {
10691
explicit TargetKind(ffi::ObjectPtr<TargetKindNode> data) : ffi::ObjectRef(data) {
10792
TVM_FFI_ICHECK(data != nullptr);
10893
}
109-
/*! \brief Get the attribute map given the attribute name */
110-
template <typename ValueType>
111-
static inline TargetKindAttrMap<ValueType> GetAttrMap(const ffi::String& attr_name);
11294
/*!
11395
* \brief Retrieve the TargetKind given its name
11496
* \param target_kind_name Name of the target kind
115-
* \return The TargetKind requested
97+
* \return The canonical kind, or nullopt when the name is unknown.
11698
*/
11799
TVM_DLL static ffi::Optional<TargetKind> Get(const ffi::String& target_kind_name);
118100
/*! \brief Mutable access to the container class */
@@ -121,198 +103,121 @@ class TargetKind : public ffi::ObjectRef {
121103
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TargetKind, ffi::ObjectRef, TargetKindNode);
122104

123105
private:
124-
TVM_DLL static const AttrRegistryMapContainerMap<TargetKind>& GetAttrMapContainer(
125-
const ffi::String& attr_name);
126-
friend class TargetKindRegEntry;
127106
friend class TargetInternal;
128107
};
129108

130-
/*!
131-
* \brief ffi::Map<TargetKind, ValueType> used to store meta-information about TargetKind
132-
* \tparam ValueType The type of the value stored in map
133-
*/
134-
template <typename ValueType>
135-
class TargetKindAttrMap : public AttrRegistryMap<TargetKind, ValueType> {
136-
public:
137-
using TParent = AttrRegistryMap<TargetKind, ValueType>;
138-
using TParent::count;
139-
using TParent::get;
140-
using TParent::operator[];
141-
explicit TargetKindAttrMap(const AttrRegistryMapContainerMap<TargetKind>& map) : TParent(map) {}
142-
};
143-
144109
/*! \brief Value used with --runtime in target specs to indicate the C++ runtime. */
145110
static constexpr const char* kTvmRuntimeCpp = "c++";
146111

147112
/*! \brief Value used with --runtime in target specs to indicate the C runtime. */
148113
static constexpr const char* kTvmRuntimeCrt = "c";
149114

150115
/*!
151-
* \brief Helper structure to register TargetKind
152-
* \sa TVM_REGISTER_TARGET_KIND
116+
* \brief Process-wide registry of canonical target kinds, keyed by name.
153117
*/
154-
class TargetKindRegEntry {
118+
class TargetKindRegistry {
155119
public:
156120
/*!
157-
* \brief Register additional attributes to target_kind.
158-
* \param attr_name The name of the attribute.
159-
* \param value The value to be set.
160-
* \param plevel The priority level of this attribute,
161-
* an higher priority level attribute
162-
* will replace lower priority level attribute.
163-
* Must be bigger than 0.
164-
*
165-
* Cannot set with same plevel twice in the code.
166-
*
167-
* \tparam ValueType The type of the value to be set.
121+
* \brief Access the process-wide registry.
122+
* \return The singleton registry.
168123
*/
169-
template <typename ValueType>
170-
inline TargetKindRegEntry& set_attr(const ffi::String& attr_name, const ValueType& value,
171-
int plevel = 10);
124+
TVM_DLL static TargetKindRegistry* Global();
172125
/*!
173-
* \brief Set DLPack's device_type the target
174-
* \param device_type Device type
126+
* \brief Register a name or return its existing canonical kind.
127+
* \param name Target kind name.
128+
* \return The canonical target kind for name. Standard target options are
129+
* declared when the name is first registered.
175130
*/
176-
inline TargetKindRegEntry& set_default_device_type(int device_type);
131+
TVM_DLL TargetKind RegisterOrGet(const ffi::String& name);
177132
/*!
178-
* \brief Set DLPack's device_type the target
179-
* \param keys The default keys
133+
* \brief Look up a registered kind.
134+
* \param name Target kind name.
135+
* \return The canonical kind, or nullopt when name is unknown.
180136
*/
181-
inline TargetKindRegEntry& set_default_keys(std::vector<ffi::String> keys);
137+
TVM_DLL ffi::Optional<TargetKind> Get(const ffi::String& name);
182138
/*!
183-
* \brief Set the canonicalizer function applied upon target creation.
184-
* \param canonicalizer The target canonicalizer function.
139+
* \brief List registered kind names.
140+
* \return Names of all registered target kinds.
185141
*/
186-
inline TargetKindRegEntry& set_target_canonicalizer(FTargetCanonicalizer canonicalizer);
142+
TVM_DLL ffi::Array<ffi::String> ListTargetKinds();
187143
/*!
188-
* \brief Register a valid configuration option and its ValueType for validation
189-
* \param key The configuration key
190-
* \param traits Optional traits (e.g. refl::DefaultValue, doc string, or raw default value)
191-
* \tparam ValueType The value type to be registered
192-
* \tparam Traits Optional trait types
144+
* \brief List declared option names and types for a kind.
145+
* \param kind Registered target kind.
146+
* \return Map from option name to type string.
193147
*/
194-
template <typename ValueType, typename... Traits>
195-
inline TargetKindRegEntry& add_attr_option(const ffi::String& key, Traits&&... traits);
196-
/*! \brief Set name of the TargetKind to be the same as registry if it is empty */
197-
inline TargetKindRegEntry& set_name();
148+
TVM_DLL ffi::Map<ffi::String, ffi::String> ListTargetKindOptions(const TargetKind& kind);
149+
150+
private:
151+
ffi::Map<ffi::String, TargetKind> kinds_;
152+
};
153+
154+
/*!
155+
* \brief Temporary fluent builder for registering a target kind.
156+
*
157+
* The registry retains the canonical kind after this builder is destroyed.
158+
* Group related definitions inside TVM_FFI_STATIC_INIT_BLOCK():
159+
* \code
160+
* TVM_FFI_STATIC_INIT_BLOCK() {
161+
* TargetKindDef("llvm")
162+
* .set_default_device_type(kDLCPU)
163+
* .set_default_keys({"cpu"})
164+
* .def_option<ffi::String>("mcpu");
165+
* }
166+
* \endcode
167+
*/
168+
class TargetKindDef {
169+
public:
198170
/*!
199-
* \brief List all the entry names in the registry.
200-
* \return The entry names.
171+
* \brief Register or retrieve the canonical kind with this name.
172+
* \param name Target kind name.
201173
*/
202-
TVM_DLL static ffi::Array<ffi::String> ListTargetKinds();
174+
explicit TargetKindDef(const ffi::String& name)
175+
: kind_(TargetKindRegistry::Global()->RegisterOrGet(name)) {}
203176
/*!
204-
* \brief Get all supported option names and types for a given Target kind.
205-
* \return Map of option name to type
177+
* \brief Set the default DLPack device type.
178+
* \param device_type DLPack device type.
179+
* \return This builder for chaining.
206180
*/
207-
TVM_DLL static ffi::Map<ffi::String, ffi::String> ListTargetKindOptions(const TargetKind& kind);
208-
181+
TargetKindDef& set_default_device_type(int device_type) {
182+
kind_->default_device_type = device_type;
183+
return *this;
184+
}
209185
/*!
210-
* \brief Register or get a new entry.
211-
* \param target_kind_name The name of the TargetKind.
212-
* \return the corresponding entry.
186+
* \brief Set default target keys.
187+
* \param keys Default keys in priority order.
188+
* \return This builder for chaining.
213189
*/
214-
TVM_DLL static TargetKindRegEntry& RegisterOrGet(const ffi::String& target_kind_name);
215-
216-
private:
217-
TargetKind kind_;
218-
ffi::String name;
219-
220-
/*! \brief private constructor */
221-
explicit TargetKindRegEntry(uint32_t reg_index) : kind_(ffi::make_object<TargetKindNode>()) {
222-
kind_->index_ = reg_index;
190+
TargetKindDef& set_default_keys(std::vector<ffi::String> keys) {
191+
kind_->default_keys = keys;
192+
return *this;
223193
}
224194
/*!
225-
* \brief update the attribute TargetKindAttrMap
226-
* \param key The name of the attribute
227-
* \param value The value to be set
228-
* \param plevel The priority level
195+
* \brief Set the canonicalizer used when constructing targets of this kind.
196+
* \param canonicalizer Function from a validated config map to its canonical form.
197+
* \return This builder for chaining. Also updates the kind's ConfigSchema.
229198
*/
230-
TVM_DLL void UpdateAttr(const ffi::String& key, ffi::Any value, int plevel);
231-
template <typename, typename>
232-
friend class AttrRegistry;
233-
friend class TargetKind;
234-
};
235-
236-
template <typename ValueType>
237-
inline TargetKindAttrMap<ValueType> TargetKind::GetAttrMap(const ffi::String& attr_name) {
238-
return TargetKindAttrMap<ValueType>(GetAttrMapContainer(attr_name));
239-
}
240-
241-
template <typename ValueType>
242-
inline TargetKindRegEntry& TargetKindRegEntry::set_attr(const ffi::String& attr_name,
243-
const ValueType& value, int plevel) {
244-
TVM_FFI_ICHECK_GT(plevel, 0) << "plevel in set_attr must be greater than 0";
245-
ffi::Any rv;
246-
rv = value;
247-
UpdateAttr(attr_name, rv, plevel);
248-
return *this;
249-
}
250-
251-
inline TargetKindRegEntry& TargetKindRegEntry::set_default_device_type(int device_type) {
252-
kind_->default_device_type = device_type;
253-
return *this;
254-
}
255-
256-
inline TargetKindRegEntry& TargetKindRegEntry::set_default_keys(std::vector<ffi::String> keys) {
257-
kind_->default_keys = keys;
258-
return *this;
259-
}
260-
261-
inline TargetKindRegEntry& TargetKindRegEntry::set_target_canonicalizer(
262-
FTargetCanonicalizer canonicalizer) {
263-
kind_->target_canonicalizer = canonicalizer;
264-
kind_->schema_.set_canonicalizer(canonicalizer);
265-
return *this;
266-
}
267-
268-
template <typename ValueType, typename... Traits>
269-
inline TargetKindRegEntry& TargetKindRegEntry::add_attr_option(const ffi::String& key,
270-
Traits&&... traits) {
271-
kind_->schema_.def_option<ValueType>(key, std::forward<Traits>(traits)...);
272-
return *this;
273-
}
274-
275-
inline TargetKindRegEntry& TargetKindRegEntry::set_name() {
276-
if (kind_->name.empty()) {
277-
kind_->name = name;
199+
TargetKindDef& set_target_canonicalizer(FTargetCanonicalizer canonicalizer) {
200+
kind_->target_canonicalizer = canonicalizer;
201+
kind_->schema_.set_canonicalizer(canonicalizer);
202+
return *this;
203+
}
204+
/*!
205+
* \brief Declare a typed target option in the kind's ConfigSchema.
206+
* \tparam ValueType Canonical option value type.
207+
* \tparam Traits Optional metadata or validator trait types.
208+
* \param key Option name.
209+
* \param traits Optional traits such as a default value or validator.
210+
* \return This builder for chaining. Duplicate option names raise ValueError.
211+
*/
212+
template <typename ValueType, typename... Traits>
213+
TargetKindDef& def_option(const ffi::String& key, Traits&&... traits) {
214+
kind_->schema_.def_option<ValueType>(key, std::forward<Traits>(traits)...);
215+
return *this;
278216
}
279-
return *this;
280-
}
281-
282-
#define TVM_TARGET_KIND_REGISTER_VAR_DEF \
283-
[[maybe_unused]] static ::tvm::TargetKindRegEntry& __make_##TargetKind
284217

285-
/*!
286-
* \def TVM_REGISTER_TARGET_KIND
287-
* \brief Register a new target kind, or set attribute of the corresponding target kind.
288-
*
289-
* \param TargetKindName The name of target kind
290-
* \param DeviceType The DLDeviceType of the target kind
291-
*
292-
* \code
293-
*
294-
* TVM_REGISTER_TARGET_KIND("llvm")
295-
* .set_attr<TPreCodegenPass>("TPreCodegenPass", a-pre-codegen-pass)
296-
* .add_attr_option<Bool>("system_lib")
297-
* .add_attr_option<ffi::String>("mtriple")
298-
* .add_attr_option<ffi::String>("mattr");
299-
*
300-
* \endcode
301-
*/
302-
#define TVM_REGISTER_TARGET_KIND(TargetKindName, DeviceType) \
303-
TVM_FFI_STR_CONCAT(TVM_TARGET_KIND_REGISTER_VAR_DEF, __COUNTER__) = \
304-
::tvm::TargetKindRegEntry::RegisterOrGet(TargetKindName) \
305-
.set_name() \
306-
.set_default_device_type(DeviceType) \
307-
.add_attr_option<ffi::String>("kind") \
308-
.add_attr_option<ffi::Array<ffi::String>>("keys") \
309-
.add_attr_option<ffi::String>("tag") \
310-
.add_attr_option<ffi::String>("device") \
311-
.add_attr_option<ffi::String>("model") \
312-
.add_attr_option<ffi::Array<ffi::String>>("libs") \
313-
.add_attr_option<Target>("host") \
314-
.add_attr_option<int64_t>("from_device") \
315-
.add_attr_option<int64_t>("target_device_type")
218+
private:
219+
TargetKind kind_;
220+
};
316221

317222
} // namespace tvm
318223

‎python/tvm/target/target.py‎

Lines changed: 0 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -210,21 +210,6 @@ def __getattr__(self, name: str):
210210
return str(value) if isinstance(value, String) else value
211211
raise AttributeError(f"'Target' object has no attribute '{name}'")
212212

213-
def get_kind_attr(self, attr_name):
214-
"""Get additional attribute about the target kind.
215-
216-
Parameters
217-
----------
218-
attr_name : str
219-
The attribute name.
220-
221-
Returns
222-
-------
223-
value : object
224-
The attribute value
225-
"""
226-
return _ffi_api.TargetKindGetAttr(self.kind, attr_name)
227-
228213
def get_target_device_type(self):
229214
"""Returns the device_type for this target."""
230215
return _ffi_api.TargetGetDeviceType(self)

0 commit comments

Comments
 (0)