diff --git a/include/tvm/target/target.h b/include/tvm/target/target.h index 2abdb558baf8..64bd251c0ded 100644 --- a/include/tvm/target/target.h +++ b/include/tvm/target/target.h @@ -44,6 +44,8 @@ class TargetNode : public Object { public: /*! \brief The kind of the target device */ TargetKind kind; + /*! \brief Target host information, must be Target type */ + Optional host; /*! \brief Tag of the the target, can be empty */ String tag; /*! \brief Keys for this target */ @@ -64,6 +66,7 @@ class TargetNode : public Object { v->Visit("tag", &tag); v->Visit("keys", &keys); v->Visit("attrs", &attrs); + v->Visit("host", &host); } /*! @@ -122,12 +125,12 @@ class Target : public ObjectRef { TVM_DLL explicit Target(std::nullptr_t) { data_ = nullptr; } /*! * \brief Construct a Target given a string - * \param tag_or_config_or_target_str the string to parse + * \param tag_or_config_or_target_str the string to parse for target */ TVM_DLL explicit Target(const String& tag_or_config_or_target_str); /*! * \brief Construct a Target using a JSON-like configuration - * \param config The JSON-like configuration + * \param config The JSON-like configuration for target */ TVM_DLL explicit Target(const Map& config); /*! @@ -139,7 +142,13 @@ class Target : public ObjectRef { * allow_not_defined is true. */ TVM_DLL static tvm::Target Current(bool allow_not_defined = true); - + /*! + * \brief Construct a Target given target and host + * \param target The Target typed object with host field undefined for target + * \param host The Target typed object for target host + * \return The Target with given target and host context information + */ + TVM_DLL explicit Target(Target target, Target host); TVM_DEFINE_OBJECT_REF_METHODS(Target, ObjectRef, TargetNode); private: diff --git a/include/tvm/target/target_kind.h b/include/tvm/target/target_kind.h index 72c41c6f4647..e7da2dd413a0 100644 --- a/include/tvm/target/target_kind.h +++ b/include/tvm/target/target_kind.h @@ -376,7 +376,8 @@ inline TargetKindRegEntry& TargetKindRegEntry::set_name() { .add_attr_option("tag") \ .add_attr_option("device") \ .add_attr_option("model") \ - .add_attr_option>("libs") + .add_attr_option>("libs") \ + .add_attr_option("host") } // namespace tvm diff --git a/python/tvm/target/target.py b/python/tvm/target/target.py index 8942957d32c9..8c60260e640a 100644 --- a/python/tvm/target/target.py +++ b/python/tvm/target/target.py @@ -46,7 +46,7 @@ class Target(Object): - :py:func:`tvm.target.intel_graphics` create Intel Graphics target """ - def __init__(self, tag_or_str_or_dict): + def __init__(self, tag_or_str_or_dict, host_tag_or_str_or_dict=None): """Construct a TVM target object from 1) Raw target string 2) Target config dict @@ -86,10 +86,22 @@ def __init__(self, tag_or_str_or_dict): mfloat-abi : str (optional) An llvm setting that is one of 'hard' or 'soft' indicating whether to use hardware or software floating-point operations. + host : Union[str, Dict[str, Any]] (optional) + Description for target host. Can be recursive. Similar to tag_or_str_or_dict. + host_tag_or_str_or_dict : Optional[Union[str, Dict[str, Any]]] + Similar to tag_or_str_or_dict but for target host. Can be one of a literal + target host string, a json string describing a configuration, or a dictionary of + configuration options. When using a dictionary or json string to configure target, + the possible values are same as tag_or_str_or_dict. """ if not isinstance(tag_or_str_or_dict, (dict, str, Target)): raise ValueError("target has to be a string or dictionary.") - self.__init_handle_by_constructor__(_ffi_api.Target, tag_or_str_or_dict) + if host_tag_or_str_or_dict is not None: + self.__init_handle_by_constructor__( + _ffi_api.Target, Target(tag_or_str_or_dict), Target(host_tag_or_str_or_dict) + ) + else: + self.__init_handle_by_constructor__(_ffi_api.Target, tag_or_str_or_dict) def __enter__(self): _ffi_api.TargetEnterScope(self) diff --git a/src/target/target.cc b/src/target/target.cc index e44a15c3ff59..b5ca4c38bbb9 100644 --- a/src/target/target.cc +++ b/src/target/target.cc @@ -373,6 +373,15 @@ Target::Target(const Map& config) { data_ = std::move(target); } +Target::Target(Target target, Target host) { + ObjectPtr n = make_object(*target.get()); + CHECK(!n->host.defined()) + << "ValueError: Adding a host to a target whose host field has been defined"; + // add target host into host field + n->host = std::move(host); + data_ = std::move(n); +} + std::vector TargetNode::GetKeys() const { std::vector result; for (auto& expr : keys) { @@ -456,8 +465,18 @@ void TargetInternal::ConstructorDispatcher(TVMArgs args, TVMRetValue* rv) { << runtime::ArgTypeCode2Str(arg.type_code()); } return; + } else if (args.num_args == 2) { + if (args[0].IsObjectRef() && args[1].IsObjectRef()) { + Target target = args[0]; + Target host = args[1]; + *rv = Target(target, host); + } else { + LOG(FATAL) << "ValueError: Invalid type of arguments. Expect 2 Target arguments."; + } + return; } - LOG(FATAL) << "ValueError: Invalid number of arguments. Expect 1, but gets: " << args.num_args; + LOG(FATAL) << "ValueError: Invalid number of arguments. Expect 1 or 2, but gets: " + << args.num_args; } ObjectPtr TargetInternal::FromString(const String& tag_or_config_or_target_str) { @@ -527,6 +546,7 @@ ObjectPtr TargetInternal::FromConfig(std::unordered_map target = make_object(); // parse 'kind' if (config.count(kKind)) { @@ -599,6 +619,13 @@ ObjectPtr TargetInternal::FromConfig(std::unordered_maphost = PackedFunc(ConstructorDispatcher)(config[kHost]).AsObjectRef(); + config.erase(kHost); + } else { + target->host = NullOpt; + } // set default attribute values if they do not exist for (const auto& kv : target->kind->key2default_) { if (!attrs.count(kv.first)) { diff --git a/src/target/target_kind.cc b/src/target/target_kind.cc index a3b1b207f290..863d99993f4a 100644 --- a/src/target/target_kind.cc +++ b/src/target/target_kind.cc @@ -309,7 +309,6 @@ TVM_REGISTER_TARGET_KIND("hybrid", kDLCPU) // line break .add_attr_option("system-lib"); TVM_REGISTER_TARGET_KIND("composite", kDLCPU) - .add_attr_option("target_host") .add_attr_option>("devices"); /********** Registry **********/ diff --git a/tests/python/unittest/test_target_target.py b/tests/python/unittest/test_target_target.py index a0a60cb0c4fd..7b998bef34a5 100644 --- a/tests/python/unittest/test_target_target.py +++ b/tests/python/unittest/test_target_target.py @@ -16,6 +16,7 @@ # under the License. import json import tvm +import pytest from tvm import te from tvm.target import cuda, rocm, mali, intel_graphics, arm_cpu, vta, bifrost, hexagon @@ -113,18 +114,14 @@ def test_config_map(): attributes fails as expected. """ target_config = {"kind": "llvm", "libs": {"a": "b", "c": "d"}} - failed = False - try: + with pytest.raises(ValueError): tvm.target.Target(target_config) - except ValueError: - failed = True - assert failed def test_composite_target(): - tgt = tvm.target.Target("composite --target_host=llvm --devices=cuda,opencl") + tgt = tvm.target.Target("composite --host=llvm --devices=cuda,opencl") assert tgt.kind.name == "composite" - assert tgt.attrs["target_host"].kind.name == "llvm" + assert tgt.attrs["host"].kind.name == "llvm" assert len(tgt.attrs["devices"]) == 2 cuda_device, opencl_device = tgt.attrs["devices"] assert cuda_device.kind.name == "cuda" @@ -158,6 +155,70 @@ def test_list_kinds(): assert all(isinstance(target_name, str) for target_name in targets) +def test_target_host_tags(): + tgt = tvm.target.Target("nvidia/jetson-nano", "nvidia/geforce-rtx-2080-ti") + assert tgt.kind.name == "cuda" + assert tgt.attrs["arch"] == "sm_53" + assert tgt.attrs["shared_memory_per_block"] == 49152 + assert tgt.attrs["max_threads_per_block"] == 1024 + assert tgt.attrs["thread_warp_size"] == 32 + assert tgt.attrs["registers_per_block"] == 32768 + assert tgt.host.kind.name == "cuda" + assert tgt.host.attrs["arch"] == "sm_75" + assert tgt.host.attrs["shared_memory_per_block"] == 49152 + assert tgt.host.attrs["max_threads_per_block"] == 1024 + assert tgt.host.attrs["thread_warp_size"] == 32 + assert tgt.host.attrs["registers_per_block"] == 65536 + + +def test_target_host_tag_dict(): + tgt = tvm.target.Target("nvidia/jetson-nano", {"kind": "llvm"}) + assert tgt.kind.name == "cuda" + assert tgt.attrs["arch"] == "sm_53" + assert tgt.attrs["shared_memory_per_block"] == 49152 + assert tgt.attrs["max_threads_per_block"] == 1024 + assert tgt.attrs["thread_warp_size"] == 32 + assert tgt.attrs["registers_per_block"] == 32768 + assert tgt.host.kind.name == "llvm" + + +def test_target_host_single_dict(): + tgt = tvm.target.Target({"kind": "llvm", "host": "nvidia/jetson-nano"}) + assert tgt.kind.name == "llvm" + assert tgt.host.kind.name == "cuda" + assert tgt.host.attrs["arch"] == "sm_53" + assert tgt.host.attrs["shared_memory_per_block"] == 49152 + assert tgt.host.attrs["max_threads_per_block"] == 1024 + assert tgt.host.attrs["thread_warp_size"] == 32 + assert tgt.host.attrs["registers_per_block"] == 32768 + + +def test_target_host_single_string(): + tgt = tvm.target.Target("cuda --host llvm") + assert tgt.kind.name == "cuda" + assert tgt.host.kind.name == "llvm" + + +def test_target_host_single_string_with_tag(): + tgt = tvm.target.Target("cuda --host nvidia/jetson-nano") + assert tgt.kind.name == "cuda" + assert tgt.host.kind.name == "cuda" + assert tgt.host.attrs["arch"] == "sm_53" + assert tgt.host.attrs["shared_memory_per_block"] == 49152 + assert tgt.host.attrs["max_threads_per_block"] == 1024 + assert tgt.host.attrs["thread_warp_size"] == 32 + assert tgt.host.attrs["registers_per_block"] == 32768 + + +def test_target_host_warning(): + """ + Confirm that constructing a target with invalid + attributes fails as expected. + """ + with pytest.raises(ValueError): + tgt = tvm.target.Target("cuda --host nvidia/jetson-nano", "llvm") + + if __name__ == "__main__": test_target_dispatch() test_target_string_parse()