[Target] Add target host field for target specification (#7462)

* Add target host field in Target

* Add host as a config field to target

* Add target host support for Python api

* Add unit tests

* Adjust format for cpplint

* Remove unnecessary  after  in Python file

* Remove redundancy and add param description

* Fix format issue

* Fix param description

* Add unit test for duplicate target hosts

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<ObjectRef> 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<String, ObjectRef>& 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:

include/tvm/target/target_kind.h

@@ -376,7 +376,8 @@ inline TargetKindRegEntry& TargetKindRegEntry::set_name() {
           .add_attr_option<String>("tag")                         \
           .add_attr_option<String>("device")                      \
           .add_attr_option<String>("model")                       \
-          .add_attr_option<Array<String>>("libs")
+          .add_attr_option<Array<String>>("libs")                 \
+          .add_attr_option<Target>("host")
 
 }  // namespace tvm
 

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)

src/target/target.cc

@@ -373,6 +373,15 @@ Target::Target(const Map<String, ObjectRef>& config) {
   data_ = std::move(target);
 }
 
+Target::Target(Target target, Target host) {
+  ObjectPtr<TargetNode> n = make_object<TargetNode>(*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<std::string> TargetNode::GetKeys() const {
   std::vector<std::string> 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<Target>() && args[1].IsObjectRef<Target>()) {
+      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<Object> TargetInternal::FromString(const String& tag_or_config_or_target_str) {
@@ -527,6 +546,7 @@ ObjectPtr<Object> TargetInternal::FromConfig(std::unordered_map<String, ObjectRe
   const String kTag = "tag";
   const String kKeys = "keys";
   const String kDeviceName = "device";
+  const String kHost = "host";
   ObjectPtr<TargetNode> target = make_object<TargetNode>();
   // parse 'kind'
   if (config.count(kKind)) {
@@ -599,6 +619,13 @@ ObjectPtr<Object> TargetInternal::FromConfig(std::unordered_map<String, ObjectRe
       throw dmlc::Error(": Error when parsing target[\"" + key + "\"]" + e.what());
     }
   }
+  // parse host
+  if (config.count(kHost)) {
+    target->host = PackedFunc(ConstructorDispatcher)(config[kHost]).AsObjectRef<Target>();
+    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)) {

src/target/target_kind.cc

@@ -309,7 +309,6 @@ TVM_REGISTER_TARGET_KIND("hybrid", kDLCPU)  // line break
     .add_attr_option<Bool>("system-lib");
 
 TVM_REGISTER_TARGET_KIND("composite", kDLCPU)
-    .add_attr_option<Target>("target_host")
     .add_attr_option<Array<Target>>("devices");
 
 /**********  Registry  **********/

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()