This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 0bec7b439d Use dedicated TargetKind registry and fluent definitions
(#20469)
0bec7b439d is described below
commit 0bec7b439d4a2ed298a41b684c44d43f07c7f46f
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Sep 26 12:06:08 2026 -0400
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.
---
include/tvm/target/target_kind.h | 267 ++++++++++-------------------
python/tvm/target/target.py | 15 --
src/backend/cuda/codegen/target_kind.cc | 50 +++---
src/backend/hexagon/codegen/target_kind.cc | 27 ++-
src/backend/metal/codegen/target_kind.cc | 32 ++--
src/backend/opencl/codegen/target_kind.cc | 37 ++--
src/backend/rocm/codegen/target_kind.cc | 35 ++--
src/backend/trn/codegen/target_kind.cc | 22 +--
src/backend/vulkan/codegen/target_kind.cc | 85 +++++----
src/backend/webgpu/codegen/target_kind.cc | 25 ++-
src/target/target_kind.cc | 254 ++++++++++++++-------------
tests/cpp/target_test.cc | 72 ++++----
12 files changed, 398 insertions(+), 523 deletions(-)
diff --git a/include/tvm/target/target_kind.h b/include/tvm/target/target_kind.h
index 3a3d21ce53..caea01b077 100644
--- a/include/tvm/target/target_kind.h
+++ b/include/tvm/target/target_kind.h
@@ -24,14 +24,12 @@
#ifndef TVM_TARGET_TARGET_KIND_H_
#define TVM_TARGET_TARGET_KIND_H_
+#include <tvm/ffi/container/map.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/attr_registry_map.h>
#include <tvm/ir/config_schema.h>
#include <tvm/runtime/base.h>
-#include <memory>
-#include <unordered_map>
#include <utility>
#include <vector>
@@ -50,9 +48,6 @@ using FTargetCanonicalizer =
class TargetInternal;
-template <typename>
-class TargetKindAttrMap;
-
/*! \brief Target kind, specifies the kind of the target */
class TargetKindNode : public ffi::Object {
public:
@@ -79,20 +74,10 @@ class TargetKindNode : public ffi::Object {
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("target.TargetKind", TargetKindNode,
ffi::Object);
private:
- /*! \brief Return the index stored in attr registry */
- uint32_t AttrRegistryIndex() const { return index_; }
- /*! \brief Return the name stored in attr registry */
- ffi::String AttrRegistryName() const { return name; }
/*! \brief ConfigSchema for validating and resolving target attributes */
ir::ConfigSchema schema_;
- /*! \brief Index used for internal lookup of attribute registry */
- uint32_t index_;
-
- template <typename, typename>
- friend class AttrRegistry;
- template <typename>
- friend class AttrRegistryMapContainerMap;
- friend class TargetKindRegEntry;
+ friend class TargetKindRegistry;
+ friend class TargetKindDef;
friend class TargetInternal;
};
@@ -106,13 +91,10 @@ class TargetKind : public ffi::ObjectRef {
explicit TargetKind(ffi::ObjectPtr<TargetKindNode> data) :
ffi::ObjectRef(data) {
TVM_FFI_ICHECK(data != nullptr);
}
- /*! \brief Get the attribute map given the attribute name */
- template <typename ValueType>
- static inline TargetKindAttrMap<ValueType> GetAttrMap(const ffi::String&
attr_name);
/*!
* \brief Retrieve the TargetKind given its name
* \param target_kind_name Name of the target kind
- * \return The TargetKind requested
+ * \return The canonical kind, or nullopt when the name is unknown.
*/
TVM_DLL static ffi::Optional<TargetKind> Get(const ffi::String&
target_kind_name);
/*! \brief Mutable access to the container class */
@@ -121,26 +103,9 @@ class TargetKind : public ffi::ObjectRef {
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TargetKind, ffi::ObjectRef,
TargetKindNode);
private:
- TVM_DLL static const AttrRegistryMapContainerMap<TargetKind>&
GetAttrMapContainer(
- const ffi::String& attr_name);
- friend class TargetKindRegEntry;
friend class TargetInternal;
};
-/*!
- * \brief ffi::Map<TargetKind, ValueType> used to store meta-information about
TargetKind
- * \tparam ValueType The type of the value stored in map
- */
-template <typename ValueType>
-class TargetKindAttrMap : public AttrRegistryMap<TargetKind, ValueType> {
- public:
- using TParent = AttrRegistryMap<TargetKind, ValueType>;
- using TParent::count;
- using TParent::get;
- using TParent::operator[];
- explicit TargetKindAttrMap(const AttrRegistryMapContainerMap<TargetKind>&
map) : TParent(map) {}
-};
-
/*! \brief Value used with --runtime in target specs to indicate the C++
runtime. */
static constexpr const char* kTvmRuntimeCpp = "c++";
@@ -148,171 +113,111 @@ static constexpr const char* kTvmRuntimeCpp = "c++";
static constexpr const char* kTvmRuntimeCrt = "c";
/*!
- * \brief Helper structure to register TargetKind
- * \sa TVM_REGISTER_TARGET_KIND
+ * \brief Process-wide registry of canonical target kinds, keyed by name.
*/
-class TargetKindRegEntry {
+class TargetKindRegistry {
public:
/*!
- * \brief Register additional attributes to target_kind.
- * \param attr_name The name of the attribute.
- * \param value The value to be set.
- * \param plevel The priority level of this attribute,
- * an higher priority level attribute
- * will replace lower priority level attribute.
- * Must be bigger than 0.
- *
- * Cannot set with same plevel twice in the code.
- *
- * \tparam ValueType The type of the value to be set.
+ * \brief Access the process-wide registry.
+ * \return The singleton registry.
*/
- template <typename ValueType>
- inline TargetKindRegEntry& set_attr(const ffi::String& attr_name, const
ValueType& value,
- int plevel = 10);
+ TVM_DLL static TargetKindRegistry* Global();
/*!
- * \brief Set DLPack's device_type the target
- * \param device_type Device type
+ * \brief Register a name or return its existing canonical kind.
+ * \param name Target kind name.
+ * \return The canonical target kind for name. Standard target options are
+ * declared when the name is first registered.
*/
- inline TargetKindRegEntry& set_default_device_type(int device_type);
+ TVM_DLL TargetKind RegisterOrGet(const ffi::String& name);
/*!
- * \brief Set DLPack's device_type the target
- * \param keys The default keys
+ * \brief Look up a registered kind.
+ * \param name Target kind name.
+ * \return The canonical kind, or nullopt when name is unknown.
*/
- inline TargetKindRegEntry& set_default_keys(std::vector<ffi::String> keys);
+ TVM_DLL ffi::Optional<TargetKind> Get(const ffi::String& name);
/*!
- * \brief Set the canonicalizer function applied upon target creation.
- * \param canonicalizer The target canonicalizer function.
+ * \brief List registered kind names.
+ * \return Names of all registered target kinds.
*/
- inline TargetKindRegEntry& set_target_canonicalizer(FTargetCanonicalizer
canonicalizer);
+ TVM_DLL ffi::Array<ffi::String> ListTargetKinds();
/*!
- * \brief Register a valid configuration option and its ValueType for
validation
- * \param key The configuration key
- * \param traits Optional traits (e.g. refl::DefaultValue, doc string, or
raw default value)
- * \tparam ValueType The value type to be registered
- * \tparam Traits Optional trait types
+ * \brief List declared option names and types for a kind.
+ * \param kind Registered target kind.
+ * \return Map from option name to type string.
*/
- template <typename ValueType, typename... Traits>
- inline TargetKindRegEntry& add_attr_option(const ffi::String& key,
Traits&&... traits);
- /*! \brief Set name of the TargetKind to be the same as registry if it is
empty */
- inline TargetKindRegEntry& set_name();
+ TVM_DLL ffi::Map<ffi::String, ffi::String> ListTargetKindOptions(const
TargetKind& kind);
+
+ private:
+ ffi::Map<ffi::String, TargetKind> kinds_;
+};
+
+/*!
+ * \brief Temporary fluent builder for registering a target kind.
+ *
+ * The registry retains the canonical kind after this builder is destroyed.
+ * Group related definitions inside TVM_FFI_STATIC_INIT_BLOCK():
+ * \code
+ * TVM_FFI_STATIC_INIT_BLOCK() {
+ * TargetKindDef("llvm")
+ * .set_default_device_type(kDLCPU)
+ * .set_default_keys({"cpu"})
+ * .def_option<ffi::String>("mcpu");
+ * }
+ * \endcode
+ */
+class TargetKindDef {
+ public:
/*!
- * \brief List all the entry names in the registry.
- * \return The entry names.
+ * \brief Register or retrieve the canonical kind with this name.
+ * \param name Target kind name.
*/
- TVM_DLL static ffi::Array<ffi::String> ListTargetKinds();
+ explicit TargetKindDef(const ffi::String& name)
+ : kind_(TargetKindRegistry::Global()->RegisterOrGet(name)) {}
/*!
- * \brief Get all supported option names and types for a given Target kind.
- * \return Map of option name to type
+ * \brief Set the default DLPack device type.
+ * \param device_type DLPack device type.
+ * \return This builder for chaining.
*/
- TVM_DLL static ffi::Map<ffi::String, ffi::String>
ListTargetKindOptions(const TargetKind& kind);
-
+ TargetKindDef& set_default_device_type(int device_type) {
+ kind_->default_device_type = device_type;
+ return *this;
+ }
/*!
- * \brief Register or get a new entry.
- * \param target_kind_name The name of the TargetKind.
- * \return the corresponding entry.
+ * \brief Set default target keys.
+ * \param keys Default keys in priority order.
+ * \return This builder for chaining.
*/
- TVM_DLL static TargetKindRegEntry& RegisterOrGet(const ffi::String&
target_kind_name);
-
- private:
- TargetKind kind_;
- ffi::String name;
-
- /*! \brief private constructor */
- explicit TargetKindRegEntry(uint32_t reg_index) :
kind_(ffi::make_object<TargetKindNode>()) {
- kind_->index_ = reg_index;
+ TargetKindDef& set_default_keys(std::vector<ffi::String> keys) {
+ kind_->default_keys = keys;
+ return *this;
}
/*!
- * \brief update the attribute TargetKindAttrMap
- * \param key The name of the attribute
- * \param value The value to be set
- * \param plevel The priority level
+ * \brief Set the canonicalizer used when constructing targets of this kind.
+ * \param canonicalizer Function from a validated config map to its
canonical form.
+ * \return This builder for chaining. Also updates the kind's ConfigSchema.
*/
- TVM_DLL void UpdateAttr(const ffi::String& key, ffi::Any value, int plevel);
- template <typename, typename>
- friend class AttrRegistry;
- friend class TargetKind;
-};
-
-template <typename ValueType>
-inline TargetKindAttrMap<ValueType> TargetKind::GetAttrMap(const ffi::String&
attr_name) {
- return TargetKindAttrMap<ValueType>(GetAttrMapContainer(attr_name));
-}
-
-template <typename ValueType>
-inline TargetKindRegEntry& TargetKindRegEntry::set_attr(const ffi::String&
attr_name,
- const ValueType&
value, int plevel) {
- TVM_FFI_ICHECK_GT(plevel, 0) << "plevel in set_attr must be greater than 0";
- ffi::Any rv;
- rv = value;
- UpdateAttr(attr_name, rv, plevel);
- return *this;
-}
-
-inline TargetKindRegEntry& TargetKindRegEntry::set_default_device_type(int
device_type) {
- kind_->default_device_type = device_type;
- return *this;
-}
-
-inline TargetKindRegEntry&
TargetKindRegEntry::set_default_keys(std::vector<ffi::String> keys) {
- kind_->default_keys = keys;
- return *this;
-}
-
-inline TargetKindRegEntry& TargetKindRegEntry::set_target_canonicalizer(
- FTargetCanonicalizer canonicalizer) {
- kind_->target_canonicalizer = canonicalizer;
- kind_->schema_.set_canonicalizer(canonicalizer);
- return *this;
-}
-
-template <typename ValueType, typename... Traits>
-inline TargetKindRegEntry& TargetKindRegEntry::add_attr_option(const
ffi::String& key,
- Traits&&...
traits) {
- kind_->schema_.def_option<ValueType>(key, std::forward<Traits>(traits)...);
- return *this;
-}
-
-inline TargetKindRegEntry& TargetKindRegEntry::set_name() {
- if (kind_->name.empty()) {
- kind_->name = name;
+ TargetKindDef& set_target_canonicalizer(FTargetCanonicalizer canonicalizer) {
+ kind_->target_canonicalizer = canonicalizer;
+ kind_->schema_.set_canonicalizer(canonicalizer);
+ return *this;
+ }
+ /*!
+ * \brief Declare a typed target option in the kind's ConfigSchema.
+ * \tparam ValueType Canonical option value type.
+ * \tparam Traits Optional metadata or validator trait types.
+ * \param key Option name.
+ * \param traits Optional traits such as a default value or validator.
+ * \return This builder for chaining. Duplicate option names raise
ValueError.
+ */
+ template <typename ValueType, typename... Traits>
+ TargetKindDef& def_option(const ffi::String& key, Traits&&... traits) {
+ kind_->schema_.def_option<ValueType>(key, std::forward<Traits>(traits)...);
+ return *this;
}
- return *this;
-}
-
-#define TVM_TARGET_KIND_REGISTER_VAR_DEF \
- [[maybe_unused]] static ::tvm::TargetKindRegEntry& __make_##TargetKind
-/*!
- * \def TVM_REGISTER_TARGET_KIND
- * \brief Register a new target kind, or set attribute of the corresponding
target kind.
- *
- * \param TargetKindName The name of target kind
- * \param DeviceType The DLDeviceType of the target kind
- *
- * \code
- *
- * TVM_REGISTER_TARGET_KIND("llvm")
- * .set_attr<TPreCodegenPass>("TPreCodegenPass", a-pre-codegen-pass)
- * .add_attr_option<Bool>("system_lib")
- * .add_attr_option<ffi::String>("mtriple")
- * .add_attr_option<ffi::String>("mattr");
- *
- * \endcode
- */
-#define TVM_REGISTER_TARGET_KIND(TargetKindName, DeviceType) \
- TVM_FFI_STR_CONCAT(TVM_TARGET_KIND_REGISTER_VAR_DEF, __COUNTER__) = \
- ::tvm::TargetKindRegEntry::RegisterOrGet(TargetKindName) \
- .set_name() \
- .set_default_device_type(DeviceType) \
- .add_attr_option<ffi::String>("kind") \
- .add_attr_option<ffi::Array<ffi::String>>("keys") \
- .add_attr_option<ffi::String>("tag") \
- .add_attr_option<ffi::String>("device") \
- .add_attr_option<ffi::String>("model") \
- .add_attr_option<ffi::Array<ffi::String>>("libs") \
- .add_attr_option<Target>("host") \
- .add_attr_option<int64_t>("from_device") \
- .add_attr_option<int64_t>("target_device_type")
+ private:
+ TargetKind kind_;
+};
} // namespace tvm
diff --git a/python/tvm/target/target.py b/python/tvm/target/target.py
index a1f7a8091c..0ecfd1d34c 100644
--- a/python/tvm/target/target.py
+++ b/python/tvm/target/target.py
@@ -210,21 +210,6 @@ class Target(Object):
return str(value) if isinstance(value, String) else value
raise AttributeError(f"'Target' object has no attribute '{name}'")
- def get_kind_attr(self, attr_name):
- """Get additional attribute about the target kind.
-
- Parameters
- ----------
- attr_name : str
- The attribute name.
-
- Returns
- -------
- value : object
- The attribute value
- """
- return _ffi_api.TargetKindGetAttr(self.kind, attr_name)
-
def get_target_device_type(self):
"""Returns the device_type for this target."""
return _ffi_api.TargetGetDeviceType(self)
diff --git a/src/backend/cuda/codegen/target_kind.cc
b/src/backend/cuda/codegen/target_kind.cc
index b63f1f6e96..26b8a506fc 100644
--- a/src/backend/cuda/codegen/target_kind.cc
+++ b/src/backend/cuda/codegen/target_kind.cc
@@ -111,35 +111,39 @@ ffi::Map<ffi::String, ffi::Any>
UpdateNVPTXAttrs(ffi::Map<ffi::String, ffi::Any>
return target;
}
-void RegisterTargetKinds() {
+} // namespace cuda
+} // namespace backend
+} // namespace tvm
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ using namespace tvm;
+ using namespace tvm::backend::cuda;
namespace refl = tvm::ffi::reflection;
- TVM_REGISTER_TARGET_KIND("cuda_host", kDLCPU).set_default_keys({"cpu"});
+ TargetKindDef("cuda_host")
+ .set_default_device_type(kDLCPU) // line break
+ .set_default_keys({"cpu"});
- TVM_REGISTER_TARGET_KIND("cuda", kDLCUDA)
- .add_attr_option<ffi::String>("mcpu")
- .add_attr_option<ffi::String>("arch")
- .add_attr_option<int64_t>("max_shared_memory_per_block")
- .add_attr_option<int64_t>("max_threads_per_block")
- .add_attr_option<int64_t>("thread_warp_size", refl::DefaultValue(32))
- .add_attr_option<int64_t>("registers_per_block")
- .add_attr_option<int64_t>("l2_cache_size_bytes")
- .add_attr_option<int64_t>("max_num_threads",
- refl::DefaultValue(1024)) // TODO(@zxybazh):
deprecate it
+ TargetKindDef("cuda")
+ .set_default_device_type(kDLCUDA)
+ .def_option<ffi::String>("mcpu")
+ .def_option<ffi::String>("arch")
+ .def_option<int64_t>("max_shared_memory_per_block")
+ .def_option<int64_t>("max_threads_per_block")
+ .def_option<int64_t>("thread_warp_size", refl::DefaultValue(32))
+ .def_option<int64_t>("registers_per_block")
+ .def_option<int64_t>("l2_cache_size_bytes")
+ .def_option<int64_t>("max_num_threads",
+ refl::DefaultValue(1024)) // TODO(@zxybazh):
deprecate it
.set_default_keys({"cuda", "gpu"})
.set_target_canonicalizer(UpdateCUDAAttrs);
- TVM_REGISTER_TARGET_KIND("nvptx", kDLCUDA)
- .add_attr_option<ffi::String>("mcpu")
- .add_attr_option<ffi::String>("mtriple")
- .add_attr_option<int64_t>("max_num_threads", refl::DefaultValue(1024))
- .add_attr_option<int64_t>("thread_warp_size", refl::DefaultValue(32))
+ TargetKindDef("nvptx")
+ .set_default_device_type(kDLCUDA)
+ .def_option<ffi::String>("mcpu")
+ .def_option<ffi::String>("mtriple")
+ .def_option<int64_t>("max_num_threads", refl::DefaultValue(1024))
+ .def_option<int64_t>("thread_warp_size", refl::DefaultValue(32))
.set_default_keys({"cuda", "gpu"})
.set_target_canonicalizer(UpdateNVPTXAttrs);
}
-
-} // namespace cuda
-} // namespace backend
-} // namespace tvm
-
-TVM_FFI_STATIC_INIT_BLOCK() { tvm::backend::cuda::RegisterTargetKinds(); }
diff --git a/src/backend/hexagon/codegen/target_kind.cc
b/src/backend/hexagon/codegen/target_kind.cc
index c71d734f9f..fd00be35cd 100644
--- a/src/backend/hexagon/codegen/target_kind.cc
+++ b/src/backend/hexagon/codegen/target_kind.cc
@@ -28,22 +28,6 @@
#include <tvm/target/target_kind.h>
namespace tvm {
-namespace backend {
-namespace hexagon {
-
-void RegisterTargetKind() {
- TVM_REGISTER_TARGET_KIND("hexagon", kDLHexagon)
- .add_attr_option<ffi::Array<ffi::String>>("mattr")
- .add_attr_option<ffi::String>("mcpu")
- .add_attr_option<ffi::String>("mtriple")
- .add_attr_option<ffi::Array<ffi::String>>("llvm-options")
- .add_attr_option<int64_t>("num-cores")
- .add_attr_option<int64_t>("vtcm-capacity")
- .set_default_keys({"hexagon", "cpu"});
-}
-
-} // namespace hexagon
-} // namespace backend
#ifdef TVM_LLVM_VERSION
namespace codegen {
@@ -56,7 +40,16 @@ void RegisterHexagonIntrinRules();
} // namespace tvm
TVM_FFI_STATIC_INIT_BLOCK() {
- tvm::backend::hexagon::RegisterTargetKind();
+ using namespace tvm;
+ TargetKindDef("hexagon")
+ .set_default_device_type(kDLHexagon)
+ .def_option<ffi::Array<ffi::String>>("mattr")
+ .def_option<ffi::String>("mcpu")
+ .def_option<ffi::String>("mtriple")
+ .def_option<ffi::Array<ffi::String>>("llvm-options")
+ .def_option<int64_t>("num-cores")
+ .def_option<int64_t>("vtcm-capacity")
+ .set_default_keys({"hexagon", "cpu"});
#ifdef TVM_LLVM_VERSION
tvm::codegen::llvm::RegisterHexagonIntrinRules();
tvm::codegen::RegisterHexagonCodegen();
diff --git a/src/backend/metal/codegen/target_kind.cc
b/src/backend/metal/codegen/target_kind.cc
index 3fda53fe70..f7c35e0a60 100644
--- a/src/backend/metal/codegen/target_kind.cc
+++ b/src/backend/metal/codegen/target_kind.cc
@@ -29,25 +29,6 @@
namespace tvm {
-namespace backend {
-namespace metal {
-
-void RegisterTargetKind() {
- namespace refl = tvm::ffi::reflection;
-
- // Metal limits the number of kernel arguments. `max_function_args`
captures that bound.
- TVM_REGISTER_TARGET_KIND("metal", kDLMetal)
- .add_attr_option<int64_t>("max_num_threads", refl::DefaultValue(256))
- .add_attr_option<int64_t>("max_threads_per_block",
refl::DefaultValue(256))
- .add_attr_option<int64_t>("max_shared_memory_per_block",
refl::DefaultValue(32768))
- .add_attr_option<int64_t>("thread_warp_size", refl::DefaultValue(16))
- .add_attr_option<int64_t>("max_function_args", refl::DefaultValue(31))
- .set_default_keys({"metal", "gpu"});
-}
-
-} // namespace metal
-} // namespace backend
-
namespace codegen {
void RegisterMetalCodegen();
namespace intrin {
@@ -57,7 +38,18 @@ void RegisterMetalIntrinRules();
} // namespace tvm
TVM_FFI_STATIC_INIT_BLOCK() {
- tvm::backend::metal::RegisterTargetKind();
+ using namespace tvm;
+ namespace refl = tvm::ffi::reflection;
+
+ // Metal limits the number of kernel arguments. `max_function_args`
captures that bound.
+ TargetKindDef("metal")
+ .set_default_device_type(kDLMetal)
+ .def_option<int64_t>("max_num_threads", refl::DefaultValue(256))
+ .def_option<int64_t>("max_threads_per_block", refl::DefaultValue(256))
+ .def_option<int64_t>("max_shared_memory_per_block",
refl::DefaultValue(32768))
+ .def_option<int64_t>("thread_warp_size", refl::DefaultValue(16))
+ .def_option<int64_t>("max_function_args", refl::DefaultValue(31))
+ .set_default_keys({"metal", "gpu"});
tvm::codegen::intrin::RegisterMetalIntrinRules();
tvm::codegen::RegisterMetalCodegen();
}
diff --git a/src/backend/opencl/codegen/target_kind.cc
b/src/backend/opencl/codegen/target_kind.cc
index dfdc10e062..70219c7bf3 100644
--- a/src/backend/opencl/codegen/target_kind.cc
+++ b/src/backend/opencl/codegen/target_kind.cc
@@ -28,27 +28,6 @@
#include <tvm/target/target_kind.h>
namespace tvm {
-namespace backend {
-namespace opencl {
-
-void RegisterTargetKind() {
- namespace refl = tvm::ffi::reflection;
-
- TVM_REGISTER_TARGET_KIND("opencl", kDLOpenCL)
- .add_attr_option<int64_t>("max_threads_per_block",
refl::DefaultValue(256))
- .add_attr_option<int64_t>("max_shared_memory_per_block",
refl::DefaultValue(16384))
- .add_attr_option<int64_t>("max_num_threads", refl::DefaultValue(256))
- .add_attr_option<int64_t>("thread_warp_size", refl::DefaultValue(1))
- .add_attr_option<int64_t>("texture_spatial_limit",
refl::DefaultValue(16384))
- .add_attr_option<int64_t>("texture_depth_limit",
refl::DefaultValue(2048))
- // Qualcomm OpenCL runtimes may crash when the number of kernel
arguments is too large.
- .add_attr_option<int64_t>("max_function_args", refl::DefaultValue(128))
- .add_attr_option<int64_t>("image_base_address_alignment",
refl::DefaultValue(64))
- .set_default_keys({"opencl", "gpu"});
-}
-
-} // namespace opencl
-} // namespace backend
namespace codegen {
void RegisterOpenCLCodegen();
@@ -60,7 +39,21 @@ void RegisterOpenCLIntrinRules();
} // namespace tvm
TVM_FFI_STATIC_INIT_BLOCK() {
- tvm::backend::opencl::RegisterTargetKind();
+ using namespace tvm;
+ namespace refl = tvm::ffi::reflection;
+
+ TargetKindDef("opencl")
+ .set_default_device_type(kDLOpenCL)
+ .def_option<int64_t>("max_threads_per_block", refl::DefaultValue(256))
+ .def_option<int64_t>("max_shared_memory_per_block",
refl::DefaultValue(16384))
+ .def_option<int64_t>("max_num_threads", refl::DefaultValue(256))
+ .def_option<int64_t>("thread_warp_size", refl::DefaultValue(1))
+ .def_option<int64_t>("texture_spatial_limit", refl::DefaultValue(16384))
+ .def_option<int64_t>("texture_depth_limit", refl::DefaultValue(2048))
+ // Qualcomm OpenCL runtimes may crash when the number of kernel
arguments is too large.
+ .def_option<int64_t>("max_function_args", refl::DefaultValue(128))
+ .def_option<int64_t>("image_base_address_alignment",
refl::DefaultValue(64))
+ .set_default_keys({"opencl", "gpu"});
tvm::codegen::intrin::RegisterOpenCLIntrinRules();
tvm::codegen::RegisterOpenCLCodegen();
tvm::codegen::RegisterOpenCLDeviceScopeCompatibility();
diff --git a/src/backend/rocm/codegen/target_kind.cc
b/src/backend/rocm/codegen/target_kind.cc
index f0f5898def..cfc24bc0a1 100644
--- a/src/backend/rocm/codegen/target_kind.cc
+++ b/src/backend/rocm/codegen/target_kind.cc
@@ -109,23 +109,6 @@ ffi::Map<ffi::String, ffi::Any>
UpdateROCmAttrs(ffi::Map<ffi::String, ffi::Any>
return target;
}
-void RegisterTargetKind() {
- namespace refl = tvm::ffi::reflection;
-
- TVM_REGISTER_TARGET_KIND("rocm", kDLROCM)
- .add_attr_option<ffi::String>("mcpu")
- .add_attr_option<ffi::String>("mtriple")
- .add_attr_option<ffi::Array<ffi::String>>("mattr")
- // TODO(masahi): Support querying from a target device
- // On RDNA cards, thread_warp_size should be 32
- .add_attr_option<int64_t>("max_num_threads", refl::DefaultValue(256))
- .add_attr_option<int64_t>("max_threads_per_block",
refl::DefaultValue(256))
- .add_attr_option<int64_t>("max_shared_memory_per_block",
refl::DefaultValue(65536))
- .add_attr_option<int64_t>("thread_warp_size", refl::DefaultValue(64))
- .set_default_keys({"rocm", "gpu"})
- .set_target_canonicalizer(UpdateROCmAttrs);
-}
-
} // namespace rocm
} // namespace backend
@@ -140,7 +123,23 @@ void RegisterROCMIntrinRules();
} // namespace tvm
TVM_FFI_STATIC_INIT_BLOCK() {
- tvm::backend::rocm::RegisterTargetKind();
+ using namespace tvm;
+ using namespace tvm::backend::rocm;
+ namespace refl = tvm::ffi::reflection;
+
+ TargetKindDef("rocm")
+ .set_default_device_type(kDLROCM)
+ .def_option<ffi::String>("mcpu")
+ .def_option<ffi::String>("mtriple")
+ .def_option<ffi::Array<ffi::String>>("mattr")
+ // TODO(masahi): Support querying from a target device
+ // On RDNA cards, thread_warp_size should be 32
+ .def_option<int64_t>("max_num_threads", refl::DefaultValue(256))
+ .def_option<int64_t>("max_threads_per_block", refl::DefaultValue(256))
+ .def_option<int64_t>("max_shared_memory_per_block",
refl::DefaultValue(65536))
+ .def_option<int64_t>("thread_warp_size", refl::DefaultValue(64))
+ .set_default_keys({"rocm", "gpu"})
+ .set_target_canonicalizer(UpdateROCmAttrs);
#ifdef TVM_LLVM_VERSION
tvm::codegen::llvm::RegisterROCMIntrinRules();
tvm::codegen::RegisterAMDGPUCodegen();
diff --git a/src/backend/trn/codegen/target_kind.cc
b/src/backend/trn/codegen/target_kind.cc
index 2c1e9676e6..3d29c4aec0 100644
--- a/src/backend/trn/codegen/target_kind.cc
+++ b/src/backend/trn/codegen/target_kind.cc
@@ -29,20 +29,6 @@
namespace tvm {
-namespace backend {
-namespace trn {
-
-void RegisterTargetKind() {
- TVM_REGISTER_TARGET_KIND("trn", kDLTrn)
- .add_attr_option<int64_t>("partition_size", 128)
- .add_attr_option<int64_t>("max_sbuf_size_per_partition", 196608)
- .add_attr_option<int64_t>("max_psum_size_per_partition", 16384)
- .add_attr_option<int64_t>("num-cores");
-}
-
-} // namespace trn
-} // namespace backend
-
namespace codegen {
void RegisterTRNCodegen();
} // namespace codegen
@@ -55,7 +41,13 @@ void RegisterTRNTransforms();
} // namespace tvm
TVM_FFI_STATIC_INIT_BLOCK() {
- tvm::backend::trn::RegisterTargetKind();
+ using namespace tvm;
+ TargetKindDef("trn")
+ .set_default_device_type(kDLTrn)
+ .def_option<int64_t>("partition_size", 128)
+ .def_option<int64_t>("max_sbuf_size_per_partition", 196608)
+ .def_option<int64_t>("max_psum_size_per_partition", 16384)
+ .def_option<int64_t>("num-cores");
tvm::codegen::RegisterTRNCodegen();
tvm::tirx::transform::RegisterTRNTransforms();
}
diff --git a/src/backend/vulkan/codegen/target_kind.cc
b/src/backend/vulkan/codegen/target_kind.cc
index d956196ddd..52ad669bc3 100644
--- a/src/backend/vulkan/codegen/target_kind.cc
+++ b/src/backend/vulkan/codegen/target_kind.cc
@@ -28,51 +28,6 @@
#include <tvm/target/target_kind.h>
namespace tvm {
-namespace backend {
-namespace vulkan {
-
-void RegisterTargetKind() {
- namespace refl = tvm::ffi::reflection;
-
- TVM_REGISTER_TARGET_KIND("vulkan", kDLVulkan)
- .add_attr_option<ffi::Array<ffi::String>>("mattr")
- .add_attr_option<bool>("supports_float16")
- .add_attr_option<bool>("supports_float32", refl::DefaultValue(true))
- .add_attr_option<bool>("supports_float64")
- .add_attr_option<bool>("supports_int8")
- .add_attr_option<bool>("supports_int16")
- .add_attr_option<bool>("supports_int32", refl::DefaultValue(true))
- .add_attr_option<bool>("supports_int64")
- .add_attr_option<bool>("supports_8bit_buffer")
- .add_attr_option<bool>("supports_16bit_buffer")
- .add_attr_option<bool>("supports_storage_buffer_storage_class")
- .add_attr_option<bool>("supports_push_descriptor")
- .add_attr_option<bool>("supports_dedicated_allocation")
- .add_attr_option<bool>("supports_integer_dot_product")
- .add_attr_option<bool>("supports_cooperative_matrix")
- .add_attr_option<int64_t>("supported_subgroup_operations")
- .add_attr_option<int64_t>("max_num_threads", refl::DefaultValue(256))
- .add_attr_option<int64_t>("max_threads_per_block",
refl::DefaultValue(256))
- .add_attr_option<int64_t>("thread_warp_size", refl::DefaultValue(1))
- .add_attr_option<int64_t>("max_block_size_x")
- .add_attr_option<int64_t>("max_block_size_y")
- .add_attr_option<int64_t>("max_block_size_z")
- .add_attr_option<int64_t>("max_push_constants_size")
- .add_attr_option<int64_t>("max_uniform_buffer_range")
- .add_attr_option<int64_t>("max_storage_buffer_range")
- .add_attr_option<int64_t>("max_per_stage_descriptor_storage_buffer")
- .add_attr_option<int64_t>("max_shared_memory_per_block")
- .add_attr_option<ffi::String>("device_type")
- .add_attr_option<ffi::String>("device_name")
- .add_attr_option<ffi::String>("driver_name")
- .add_attr_option<int64_t>("driver_version")
- .add_attr_option<int64_t>("vulkan_api_version")
- .add_attr_option<int64_t>("max_spirv_version")
- .set_default_keys({"vulkan", "gpu"});
-}
-
-} // namespace vulkan
-} // namespace backend
#ifdef TVM_ENABLE_SPIRV
namespace codegen {
@@ -85,7 +40,45 @@ void RegisterVulkanIntrinRules();
} // namespace tvm
TVM_FFI_STATIC_INIT_BLOCK() {
- tvm::backend::vulkan::RegisterTargetKind();
+ using namespace tvm;
+ namespace refl = tvm::ffi::reflection;
+
+ TargetKindDef("vulkan")
+ .set_default_device_type(kDLVulkan)
+ .def_option<ffi::Array<ffi::String>>("mattr")
+ .def_option<bool>("supports_float16")
+ .def_option<bool>("supports_float32", refl::DefaultValue(true))
+ .def_option<bool>("supports_float64")
+ .def_option<bool>("supports_int8")
+ .def_option<bool>("supports_int16")
+ .def_option<bool>("supports_int32", refl::DefaultValue(true))
+ .def_option<bool>("supports_int64")
+ .def_option<bool>("supports_8bit_buffer")
+ .def_option<bool>("supports_16bit_buffer")
+ .def_option<bool>("supports_storage_buffer_storage_class")
+ .def_option<bool>("supports_push_descriptor")
+ .def_option<bool>("supports_dedicated_allocation")
+ .def_option<bool>("supports_integer_dot_product")
+ .def_option<bool>("supports_cooperative_matrix")
+ .def_option<int64_t>("supported_subgroup_operations")
+ .def_option<int64_t>("max_num_threads", refl::DefaultValue(256))
+ .def_option<int64_t>("max_threads_per_block", refl::DefaultValue(256))
+ .def_option<int64_t>("thread_warp_size", refl::DefaultValue(1))
+ .def_option<int64_t>("max_block_size_x")
+ .def_option<int64_t>("max_block_size_y")
+ .def_option<int64_t>("max_block_size_z")
+ .def_option<int64_t>("max_push_constants_size")
+ .def_option<int64_t>("max_uniform_buffer_range")
+ .def_option<int64_t>("max_storage_buffer_range")
+ .def_option<int64_t>("max_per_stage_descriptor_storage_buffer")
+ .def_option<int64_t>("max_shared_memory_per_block")
+ .def_option<ffi::String>("device_type")
+ .def_option<ffi::String>("device_name")
+ .def_option<ffi::String>("driver_name")
+ .def_option<int64_t>("driver_version")
+ .def_option<int64_t>("vulkan_api_version")
+ .def_option<int64_t>("max_spirv_version")
+ .set_default_keys({"vulkan", "gpu"});
#ifdef TVM_ENABLE_SPIRV
tvm::codegen::spirv::RegisterVulkanIntrinRules();
tvm::codegen::RegisterVulkanCodegen();
diff --git a/src/backend/webgpu/codegen/target_kind.cc
b/src/backend/webgpu/codegen/target_kind.cc
index abb40fdc29..658ac4dca7 100644
--- a/src/backend/webgpu/codegen/target_kind.cc
+++ b/src/backend/webgpu/codegen/target_kind.cc
@@ -51,18 +51,6 @@ ffi::Map<ffi::String, ffi::Any>
UpdateWebGPUAttrs(ffi::Map<ffi::String, ffi::Any
return target;
}
-void RegisterTargetKind() {
- namespace refl = tvm::ffi::reflection;
-
- TVM_REGISTER_TARGET_KIND("webgpu", kDLWebGPU)
- .add_attr_option<int64_t>("max_num_threads", refl::DefaultValue(256))
- .add_attr_option<int64_t>("max_shared_memory_per_block",
refl::DefaultValue(32768))
- .add_attr_option<bool>("supports_subgroups", refl::DefaultValue(false))
- .add_attr_option<int64_t>("thread_warp_size", refl::DefaultValue(1))
- .set_target_canonicalizer(UpdateWebGPUAttrs)
- .set_default_keys({"webgpu", "gpu"});
-}
-
} // namespace webgpu
} // namespace backend
@@ -75,7 +63,18 @@ void RegisterWebGPUIntrinRules();
} // namespace tvm
TVM_FFI_STATIC_INIT_BLOCK() {
- tvm::backend::webgpu::RegisterTargetKind();
+ using namespace tvm;
+ using namespace tvm::backend::webgpu;
+ namespace refl = tvm::ffi::reflection;
+
+ TargetKindDef("webgpu")
+ .set_default_device_type(kDLWebGPU)
+ .def_option<int64_t>("max_num_threads", refl::DefaultValue(256))
+ .def_option<int64_t>("max_shared_memory_per_block",
refl::DefaultValue(32768))
+ .def_option<bool>("supports_subgroups", refl::DefaultValue(false))
+ .def_option<int64_t>("thread_warp_size", refl::DefaultValue(1))
+ .set_target_canonicalizer(UpdateWebGPUAttrs)
+ .set_default_keys({"webgpu", "gpu"});
tvm::codegen::intrin::RegisterWebGPUIntrinRules();
tvm::codegen::RegisterWebGPUCodegen();
}
diff --git a/src/target/target_kind.cc b/src/target/target_kind.cc
index 87da6875d1..93c66869dc 100644
--- a/src/target/target_kind.cc
+++ b/src/target/target_kind.cc
@@ -32,7 +32,6 @@
#include <algorithm>
-#include "../ir/attr_registry.h"
#include "../support/utils.h"
#include "./canonicalizer/llvm/canonicalize.h"
@@ -60,40 +59,57 @@ TVM_FFI_STATIC_INIT_BLOCK() {
}
/********** Registry-related code **********/
-using TargetKindRegistry = AttrRegistry<TargetKindRegEntry, TargetKind>;
-
-ffi::Array<ffi::String> TargetKindRegEntry::ListTargetKinds() {
- return TargetKindRegistry::Global()->ListAllNames();
+TargetKindRegistry* TargetKindRegistry::Global() {
+ static TargetKindRegistry registry;
+ return ®istry;
}
-ffi::Map<ffi::String, ffi::String> TargetKindRegEntry::ListTargetKindOptions(
- const TargetKind& target_kind) {
- ffi::Map<ffi::String, ffi::String> options;
- for (const auto& e : target_kind->schema_.ListOptions()) {
- options.Set(e.key, e.type_str);
+TargetKind TargetKindRegistry::RegisterOrGet(const ffi::String& name) {
+ if (auto it = kinds_.find(name); it != kinds_.end()) {
+ return (*it).second;
}
- return options;
+ TargetKind kind(ffi::make_object<TargetKindNode>());
+ kind->name = name;
+ kind->default_device_type = kDLCPU;
+ kind->schema_.def_option<ffi::String>("kind")
+ .def_option<ffi::Array<ffi::String>>("keys")
+ .def_option<ffi::String>("tag")
+ .def_option<ffi::String>("device")
+ .def_option<ffi::String>("model")
+ .def_option<ffi::Array<ffi::String>>("libs")
+ .def_option<Target>("host")
+ .def_option<int64_t>("from_device")
+ .def_option<int64_t>("target_device_type");
+ kinds_.Set(name, kind);
+ return kind;
}
-TargetKindRegEntry& TargetKindRegEntry::RegisterOrGet(const ffi::String&
target_kind_name) {
- return TargetKindRegistry::Global()->RegisterOrGet(target_kind_name);
+ffi::Optional<TargetKind> TargetKindRegistry::Get(const ffi::String& name) {
+ if (auto it = kinds_.find(name); it != kinds_.end()) {
+ return (*it).second;
+ }
+ return std::nullopt;
}
-void TargetKindRegEntry::UpdateAttr(const ffi::String& key, ffi::Any value,
int plevel) {
- TargetKindRegistry::Global()->UpdateAttr(key, kind_, value, plevel);
+ffi::Array<ffi::String> TargetKindRegistry::ListTargetKinds() {
+ ffi::Array<ffi::String> names;
+ for (const auto& entry : kinds_) {
+ names.push_back(entry.first);
+ }
+ return names;
}
-const AttrRegistryMapContainerMap<TargetKind>& TargetKind::GetAttrMapContainer(
- const ffi::String& attr_name) {
- return TargetKindRegistry::Global()->GetAttrMap(attr_name);
+ffi::Map<ffi::String, ffi::String> TargetKindRegistry::ListTargetKindOptions(
+ const TargetKind& target_kind) {
+ ffi::Map<ffi::String, ffi::String> options;
+ for (const auto& e : target_kind->schema_.ListOptions()) {
+ options.Set(e.key, e.type_str);
+ }
+ return options;
}
ffi::Optional<TargetKind> TargetKind::Get(const ffi::String& target_kind_name)
{
- const TargetKindRegEntry* reg =
TargetKindRegistry::Global()->Get(target_kind_name);
- if (reg == nullptr) {
- return std::nullopt;
- }
- return reg->kind_;
+ return TargetKindRegistry::Global()->Get(target_kind_name);
}
/*!
@@ -108,112 +124,112 @@ ffi::Map<ffi::String, ffi::Any>
TestTargetParser(ffi::Map<ffi::String, ffi::Any>
/********** Register Target kinds and attributes **********/
-TVM_REGISTER_TARGET_KIND("llvm", kDLCPU)
- .add_attr_option<ffi::Array<ffi::String>>("mattr")
- .add_attr_option<ffi::String>("mcpu")
- .add_attr_option<ffi::String>("mtriple")
- .add_attr_option<ffi::String>("mfloat-abi")
- .add_attr_option<ffi::String>("mabi")
- .add_attr_option<int64_t>("num-cores")
- // Fast math flags, see https://llvm.org/docs/LangRef.html#fast-math-flags
- .add_attr_option<bool>("fast-math") // implies all the below
- .add_attr_option<bool>("fast-math-nnan")
- .add_attr_option<bool>("fast-math-ninf")
- .add_attr_option<bool>("fast-math-nsz")
- .add_attr_option<bool>("fast-math-arcp")
- .add_attr_option<bool>("fast-math-contract")
- .add_attr_option<bool>("fast-math-reassoc")
- .add_attr_option<int64_t>("opt-level")
- // LLVM command line flags, see below
- .add_attr_option<ffi::Array<ffi::String>>("cl-opt")
- // LLVM JIT engine mcjit/orcjit
- .add_attr_option<ffi::String>("jit")
- // TVM & LLVM custom vector bit width
- .add_attr_option<int64_t>("vector-width")
- .set_default_keys({"cpu"})
- // Force the external codegen kind attribute to be registered, even if no
external
- // codegen targets are enabled by the TVM build.
- .set_target_canonicalizer(tvm::target::canonicalizer::llvm::Canonicalize);
-
-// Note regarding the "cl-opt" attribute:
-// Each string in the array has the format
-// -optionname[[:type]=value]
-// where
-// * optionname is the actual LLVM option (e.g. "unroll-threshold")
-// * type is one of "bool", "int", "uint", or "string"
-// * value is the corresponding option value (for "bool" type is can be 0 or
"false"
-// for false value, or 1 or "true" for true value)
-// If type is omitted, it is assumed to be "bool". If value is omitted, it is
assumed
-// to be "true".
-//
-// The type must match the option type in LLVM. To find the type, search the
LLVM
-// repository (https://github.com/llvm/llvm-project) for optionname, and look
for
-// its definition: it will be a declaration of a variable of type cl::opt<T>
with
-// optionname being an argument to the constructor. The T in the declaration is
-// the type.
-// For example, for unroll-threshold, we get the following declaration:
-// static cl::opt<unsigned>
-// UnrollThreshold("unroll-threshold", cl::Hidden,
-// cl::desc("The cost threshold for loop unrolling"));
-// Hence the type is "uint".
-
-TVM_REGISTER_TARGET_KIND("c", kDLCPU)
- .add_attr_option<ffi::String>("mcpu")
- .add_attr_option<ffi::String>("march")
- .add_attr_option<int64_t>("workspace-byte-alignment")
- .add_attr_option<int64_t>("constants-byte-alignment")
- .set_default_keys({"cpu"})
- .set_target_canonicalizer(tvm::target::canonicalizer::llvm::Canonicalize);
-
-TVM_REGISTER_TARGET_KIND("ext_dev", kDLExtDev);
-
-TVM_REGISTER_TARGET_KIND("composite", kDLCPU) // line break
- .add_attr_option<ffi::Array<Target>>(
- "devices",
-
ir::ConfigSchema::AttrValidator(ffi::TypedFunction<ffi::Any(ffi::Any)>( //
- [](ffi::Any val) -> ffi::Any {
- // Allow elements to be strings or dicts, converting them to
Target objects.
- if (val.try_cast<ffi::Array<Target>>().has_value()) return val;
- auto arr = val.cast<ffi::Array<ffi::Any>>();
- ffi::Array<Target> result;
- for (const auto& elem : arr) {
- if (auto t = elem.try_cast<Target>()) {
- result.push_back(t.value());
- } else if (auto s = elem.try_cast<ffi::String>()) {
- result.push_back(Target(s.value()));
- } else if (auto m = elem.try_cast<ffi::Map<ffi::String,
ffi::Any>>()) {
- result.push_back(Target(m.value()));
- } else {
- TVM_FFI_THROW(TypeError)
- << "Expected Target, string, or dict in 'devices' array,
got '"
- << elem.GetTypeKey() << "'";
+TVM_FFI_STATIC_INIT_BLOCK() {
+ TargetKindDef("llvm")
+ .set_default_device_type(kDLCPU)
+ .def_option<ffi::Array<ffi::String>>("mattr")
+ .def_option<ffi::String>("mcpu")
+ .def_option<ffi::String>("mtriple")
+ .def_option<ffi::String>("mfloat-abi")
+ .def_option<ffi::String>("mabi")
+ .def_option<int64_t>("num-cores")
+ // Fast math flags, see
https://llvm.org/docs/LangRef.html#fast-math-flags
+ .def_option<bool>("fast-math") // implies all the below
+ .def_option<bool>("fast-math-nnan")
+ .def_option<bool>("fast-math-ninf")
+ .def_option<bool>("fast-math-nsz")
+ .def_option<bool>("fast-math-arcp")
+ .def_option<bool>("fast-math-contract")
+ .def_option<bool>("fast-math-reassoc")
+ .def_option<int64_t>("opt-level")
+ // LLVM command line flags, see below
+ .def_option<ffi::Array<ffi::String>>("cl-opt")
+ // LLVM JIT engine mcjit/orcjit
+ .def_option<ffi::String>("jit")
+ // TVM & LLVM custom vector bit width
+ .def_option<int64_t>("vector-width")
+ .set_default_keys({"cpu"})
+ // Force the external codegen kind attribute to be registered, even if
no external
+ // codegen targets are enabled by the TVM build.
+
.set_target_canonicalizer(tvm::target::canonicalizer::llvm::Canonicalize);
+
+ // Note regarding the "cl-opt" attribute:
+ // Each string in the array has the format
+ // -optionname[[:type]=value]
+ // where
+ // * optionname is the actual LLVM option (e.g. "unroll-threshold")
+ // * type is one of "bool", "int", "uint", or "string"
+ // * value is the corresponding option value (for "bool" type is can be 0
or "false"
+ // for false value, or 1 or "true" for true value)
+ // If type is omitted, it is assumed to be "bool". If value is omitted, it
is assumed
+ // to be "true".
+ //
+ // The type must match the option type in LLVM. To find the type, search the
LLVM
+ // repository (https://github.com/llvm/llvm-project) for optionname, and
look for
+ // its definition: it will be a declaration of a variable of type cl::opt<T>
with
+ // optionname being an argument to the constructor. The T in the declaration
is
+ // the type.
+ // For example, for unroll-threshold, we get the following declaration:
+ // static cl::opt<unsigned>
+ // UnrollThreshold("unroll-threshold", cl::Hidden,
+ // cl::desc("The cost threshold for loop unrolling"));
+ // Hence the type is "uint".
+
+ TargetKindDef("c")
+ .set_default_device_type(kDLCPU)
+ .def_option<ffi::String>("mcpu")
+ .def_option<ffi::String>("march")
+ .def_option<int64_t>("workspace-byte-alignment")
+ .def_option<int64_t>("constants-byte-alignment")
+ .set_default_keys({"cpu"})
+
.set_target_canonicalizer(tvm::target::canonicalizer::llvm::Canonicalize);
+
+ TargetKindDef("ext_dev") // line break
+ .set_default_device_type(kDLExtDev);
+
+ TargetKindDef("composite")
+ .set_default_device_type(kDLCPU) // line break
+ .def_option<ffi::Array<Target>>(
+ "devices",
+
ir::ConfigSchema::AttrValidator(ffi::TypedFunction<ffi::Any(ffi::Any)>( //
+ [](ffi::Any val) -> ffi::Any {
+ // Allow elements to be strings or dicts, converting them to
Target objects.
+ if (val.try_cast<ffi::Array<Target>>().has_value()) return val;
+ auto arr = val.cast<ffi::Array<ffi::Any>>();
+ ffi::Array<Target> result;
+ for (const auto& elem : arr) {
+ if (auto t = elem.try_cast<Target>()) {
+ result.push_back(t.value());
+ } else if (auto s = elem.try_cast<ffi::String>()) {
+ result.push_back(Target(s.value()));
+ } else if (auto m = elem.try_cast<ffi::Map<ffi::String,
ffi::Any>>()) {
+ result.push_back(Target(m.value()));
+ } else {
+ TVM_FFI_THROW(TypeError)
+ << "Expected Target, string, or dict in 'devices'
array, got '"
+ << elem.GetTypeKey() << "'";
+ }
}
- }
- return ffi::Any(result);
- })));
+ return ffi::Any(result);
+ })));
-TVM_REGISTER_TARGET_KIND("test", kDLCPU) // line break
- .set_target_canonicalizer(TestTargetParser);
+ TargetKindDef("test")
+ .set_default_device_type(kDLCPU) // line break
+ .set_target_canonicalizer(TestTargetParser);
+}
/********** Registry **********/
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
- .def("target.TargetKindGetAttr",
- [](TargetKind kind, ffi::String attr_name) -> ffi::Any {
- auto target_attr_map =
TargetKind::GetAttrMap<ffi::Any>(attr_name);
- ffi::Any rv;
- if (target_attr_map.count(kind)) {
- rv = target_attr_map[kind];
- }
- return rv;
- })
- .def("target.ListTargetKinds", TargetKindRegEntry::ListTargetKinds)
- .def("target.ListTargetKindOptions",
TargetKindRegEntry::ListTargetKindOptions)
+ .def("target.ListTargetKinds", [] { return
TargetKindRegistry::Global()->ListTargetKinds(); })
+ .def(
+ "target.ListTargetKindOptions",
+ [](TargetKind kind) { return
TargetKindRegistry::Global()->ListTargetKindOptions(kind); })
.def("target.ListTargetKindOptionsFromName", [](ffi::String
target_kind_name) {
TargetKind kind = TargetKind::Get(target_kind_name).value();
- return TargetKindRegEntry::ListTargetKindOptions(kind);
+ return TargetKindRegistry::Global()->ListTargetKindOptions(kind);
});
}
diff --git a/tests/cpp/target_test.cc b/tests/cpp/target_test.cc
index a536d1ca58..0bf728425f 100644
--- a/tests/cpp/target_test.cc
+++ b/tests/cpp/target_test.cc
@@ -27,11 +27,13 @@
using namespace tvm;
-TVM_REGISTER_TARGET_KIND("TestTargetKind", kDLCPU)
- .set_attr<std::string>("Attr1", "Value1")
- .add_attr_option<bool>("my_bool")
- .add_attr_option<ffi::Array<ffi::String>>("your_names")
- .add_attr_option<ffi::Map<ffi::String, int64_t>>("her_maps");
+TVM_FFI_STATIC_INIT_BLOCK() {
+ TargetKindDef("TestTargetKind")
+ .set_default_device_type(kDLCPU)
+ .def_option<bool>("my_bool")
+ .def_option<ffi::Array<ffi::String>>("your_names")
+ .def_option<ffi::Map<ffi::String, int64_t>>("her_maps");
+}
ffi::Map<ffi::String, ffi::Any> TestTargetParser(ffi::Map<ffi::String,
ffi::Any> target) {
ffi::String mcpu = target.at("mcpu").as_or_throw<ffi::String>();
@@ -46,28 +48,26 @@ ffi::Map<ffi::String, ffi::Any>
TestAttrsPreProcessor(ffi::Map<ffi::String, ffi:
return attrs;
}
-TVM_REGISTER_TARGET_KIND("TestTargetParser", kDLCPU)
- .add_attr_option<ffi::String>("mattr")
- .add_attr_option<ffi::String>("mcpu")
- .set_default_keys({"cpu"})
- .set_target_canonicalizer(TestTargetParser);
-
-TVM_REGISTER_TARGET_KIND("TestAttrsPreprocessor", kDLCPU)
- .add_attr_option<ffi::String>("mattr")
- .set_default_keys({"cpu"})
- .set_target_canonicalizer(TestAttrsPreProcessor);
-
-TVM_REGISTER_TARGET_KIND("TestClashingPreprocessor", kDLCPU)
- .add_attr_option<ffi::String>("mattr")
- .add_attr_option<ffi::String>("mcpu")
- .set_default_keys({"cpu"})
- .set_target_canonicalizer(TestTargetParser);
-
-TEST(TargetKind, GetAttrMap) {
- auto map = tvm::TargetKind::GetAttrMap<std::string>("Attr1");
- auto target_kind = tvm::TargetKind::Get("TestTargetKind").value();
- std::string result = map[target_kind];
- TVM_FFI_ICHECK_EQ(result, "Value1");
+TVM_FFI_STATIC_INIT_BLOCK() {
+ TargetKindDef("TestTargetParser")
+ .set_default_device_type(kDLCPU)
+ .def_option<ffi::String>("mattr")
+ .def_option<ffi::String>("mcpu")
+ .set_default_keys({"cpu"})
+ .set_target_canonicalizer(TestTargetParser);
+
+ TargetKindDef("TestAttrsPreprocessor")
+ .set_default_device_type(kDLCPU)
+ .def_option<ffi::String>("mattr")
+ .set_default_keys({"cpu"})
+ .set_target_canonicalizer(TestAttrsPreProcessor);
+
+ TargetKindDef("TestClashingPreprocessor")
+ .set_default_device_type(kDLCPU)
+ .def_option<ffi::String>("mattr")
+ .def_option<ffi::String>("mcpu")
+ .set_default_keys({"cpu"})
+ .set_target_canonicalizer(TestTargetParser);
}
TEST(TargetCreation, NestedConfig) {
@@ -277,11 +277,14 @@ TEST(TargetCreationFail, UnknownNonFeatureKeyStillFails) {
ASSERT_THROW({ Target{config}; }, tvm::ffi::Error);
}
-TVM_REGISTER_TARGET_KIND("TestStringKind", kDLCPU)
- .add_attr_option<ffi::String>("single")
- .add_attr_option<ffi::Array<ffi::String>>("array")
- .add_attr_option<ffi::Array<ffi::Array<ffi::String>>>("nested-array")
-
.add_attr_option<ffi::Array<ffi::Array<ffi::Array<ffi::String>>>>("nested2-array");
+TVM_FFI_STATIC_INIT_BLOCK() {
+ TargetKindDef("TestStringKind")
+ .set_default_device_type(kDLCPU)
+ .def_option<ffi::String>("single")
+ .def_option<ffi::Array<ffi::String>>("array")
+ .def_option<ffi::Array<ffi::Array<ffi::String>>>("nested-array")
+
.def_option<ffi::Array<ffi::Array<ffi::Array<ffi::String>>>>("nested2-array");
+}
TEST(TargetCreation, ProcessStrings) {
// Test single string attribute
@@ -384,13 +387,14 @@ TEST(TargetCreation, DeduplicateKeys) {
}
TEST(TargetKindRegistry, ListTargetKinds) {
- ffi::Array<ffi::String> names = TargetKindRegEntry::ListTargetKinds();
+ ffi::Array<ffi::String> names =
TargetKindRegistry::Global()->ListTargetKinds();
TVM_FFI_ICHECK_EQ(names.empty(), false);
TVM_FFI_ICHECK_EQ(std::count(std::begin(names), std::end(names), "llvm"), 1);
}
TEST(TargetKindRegistry, ListTargetOptions) {
TargetKind llvm = TargetKind::Get("llvm").value();
- ffi::Map<ffi::String, ffi::String> attrs =
TargetKindRegEntry::ListTargetKindOptions(llvm);
+ ffi::Map<ffi::String, ffi::String> attrs =
+ TargetKindRegistry::Global()->ListTargetKindOptions(llvm);
TVM_FFI_ICHECK_EQ(attrs.empty(), false);
}
