[Target] Add target host field for target specification

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 of the target device */
+  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,14 +125,28 @@ 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 given a string
+   * \param tag_or_config_or_target_str the string to parse for target
+   * \param host_tag_or_config_or_host_str the string to parse for target host
+   */
+  TVM_DLL explicit Target(const String& tag_or_config_or_target_str,
+                          const String& host_tag_or_config_or_host_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);
+  /*!
+   * \brief Construct a Target using a JSON-like configuration
+   * \param config The JSON-like configuration for target
+   * \param host_config The JSON-like configuration for target host
+   */
+  TVM_DLL explicit Target(const Map<String, ObjectRef>& config,
+                          const Map<String, ObjectRef>& host_config);
   /*!
    * \brief Get the current target context from thread local storage.
    * \param allow_not_defined If the context stack is empty and this is set to true, an
@@ -138,6 +155,11 @@ class Target : public ObjectRef {
    * \return The target that is the current context. The target may not be defined if
    * allow_not_defined is true.
    */
+  TVM_DLL explicit Target(Target target, Target host);
+  /*!
+   * \brief Construct a Target given target and host
+   * \return The Target with given target and host context information
+   */
   TVM_DLL static tvm::Target Current(bool allow_not_defined = true);
 
   TVM_DEFINE_OBJECT_REF_METHODS(Target, ObjectRef, TargetNode);

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,24 @@ 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 : 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:
+            if not isinstance(host_tag_or_str_or_dict, (dict, str, Target)):
+                raise ValueError("target host has to be a string or dictionary.")
+            self.__init_handle_by_constructor__(
+                _ffi_api.Target, tag_or_str_or_dict, 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

@@ -362,6 +362,13 @@ Target::Target(const String& tag_or_config_or_target_str) {
   data_ = std::move(target);
 }
 
+Target::Target(const String& tag_or_config_or_target_str,
+               const String& host_tag_or_config_or_host_str) {
+  ObjectPtr<TargetNode> n = make_object<TargetNode>(
+      *Target(Target(tag_or_config_or_target_str), Target(host_tag_or_config_or_host_str)).get());
+  data_ = std::move(n);
+}
+
 Target::Target(const Map<String, ObjectRef>& config) {
   ObjectPtr<Object> target;
   try {
@@ -373,6 +380,21 @@ Target::Target(const Map<String, ObjectRef>& config) {
   data_ = std::move(target);
 }
 
+Target::Target(const Map<String, ObjectRef>& config, const Map<String, ObjectRef>& host_config) {
+  ObjectPtr<TargetNode> n =
+      make_object<TargetNode>(*Target(Target(config), Target(host_config)).get());
+  data_ = std::move(n);
+}
+
+Target::Target(Target target, Target host) {
+  ObjectPtr<TargetNode> n = make_object<TargetNode>(*target.get());
+  if (n->host.defined())
+    throw dmlc::Error(": Error when adding target host to target already with host");
+  // 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,6 +478,11 @@ void TargetInternal::ConstructorDispatcher(TVMArgs args, TVMRetValue* rv) {
                  << runtime::ArgTypeCode2Str(arg.type_code());
     }
     return;
+  } else if (args.num_args == 2) {
+    const auto &argt = args[0], &argh = args[1];
+    auto func = PackedFunc(ConstructorDispatcher);
+    *rv = Target(func(argt).AsObjectRef<Target>(), func(argh).AsObjectRef<Target>());
+    return;
   }
   LOG(FATAL) << "ValueError: Invalid number of arguments. Expect 1, but gets: " << args.num_args;
 }
@@ -527,6 +554,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 +627,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

@@ -219,6 +219,7 @@ TVM_REGISTER_TARGET_KIND("llvm", kDLCPU)
     .add_attr_option<Bool>("system-lib")
     .add_attr_option<String>("runtime")
     .add_attr_option<Bool>("link-params", Bool(false))
+    .add_attr_option<Target>("host")
     .set_default_keys({"cpu"});
 
 TVM_REGISTER_TARGET_KIND("c", kDLCPU)
@@ -227,6 +228,7 @@ TVM_REGISTER_TARGET_KIND("c", kDLCPU)
     .add_attr_option<String>("runtime")
     .add_attr_option<String>("mcpu")
     .add_attr_option<String>("march")
+    .add_attr_option<Target>("host")
     .set_default_keys({"cpu"});
 
 TVM_REGISTER_TARGET_KIND("cuda", kDLGPU)
@@ -238,6 +240,7 @@ TVM_REGISTER_TARGET_KIND("cuda", kDLGPU)
     .add_attr_option<Integer>("shared_memory_per_block")
     .add_attr_option<Integer>("registers_per_block")
     .add_attr_option<Integer>("max_threads_per_block")
+    .add_attr_option<Target>("host")
     .set_default_keys({"cuda", "gpu"});
 
 TVM_REGISTER_TARGET_KIND("nvptx", kDLGPU)
@@ -246,6 +249,7 @@ TVM_REGISTER_TARGET_KIND("nvptx", kDLGPU)
     .add_attr_option<Bool>("system-lib")
     .add_attr_option<Integer>("max_num_threads", Integer(1024))
     .add_attr_option<Integer>("thread_warp_size", Integer(32))
+    .add_attr_option<Target>("host")
     .set_default_keys({"cuda", "gpu"})
     .set_attrs_preprocessor(UpdateNVPTXAttrs);
 
@@ -255,40 +259,48 @@ TVM_REGISTER_TARGET_KIND("rocm", kDLROCM)
     .add_attr_option<Bool>("system-lib")
     .add_attr_option<Integer>("max_num_threads", Integer(256))
     .add_attr_option<Integer>("thread_warp_size", Integer(64))
+    .add_attr_option<Target>("host")
     .set_default_keys({"rocm", "gpu"})
     .set_attrs_preprocessor(UpdateROCmAttrs);
 
 TVM_REGISTER_TARGET_KIND("opencl", kDLOpenCL)
     .add_attr_option<Bool>("system-lib")
     .add_attr_option<Integer>("max_num_threads", Integer(256))
     .add_attr_option<Integer>("thread_warp_size")
+    .add_attr_option<Target>("host")
     .set_default_keys({"opencl", "gpu"});
 
 TVM_REGISTER_TARGET_KIND("metal", kDLMetal)
     .add_attr_option<Bool>("system-lib")
     .add_attr_option<Integer>("max_num_threads", Integer(256))
+    .add_attr_option<Target>("host")
     .set_default_keys({"metal", "gpu"});
 
 TVM_REGISTER_TARGET_KIND("vulkan", kDLVulkan)
     .add_attr_option<Bool>("system-lib")
     .add_attr_option<Integer>("max_num_threads", Integer(256))
+    .add_attr_option<Target>("host")
     .set_default_keys({"vulkan", "gpu"});
 
 TVM_REGISTER_TARGET_KIND("webgpu", kDLWebGPU)
     .add_attr_option<Bool>("system-lib")
     .add_attr_option<Integer>("max_num_threads", Integer(256))
+    .add_attr_option<Target>("host")
     .set_default_keys({"webgpu", "gpu"});
 
 TVM_REGISTER_TARGET_KIND("sdaccel", kDLOpenCL)
     .add_attr_option<Bool>("system-lib")
+    .add_attr_option<Target>("host")
     .set_default_keys({"sdaccel", "hls"});
 
 TVM_REGISTER_TARGET_KIND("aocl", kDLAOCL)
     .add_attr_option<Bool>("system-lib")
+    .add_attr_option<Target>("host")
     .set_default_keys({"aocl", "hls"});
 
 TVM_REGISTER_TARGET_KIND("aocl_sw_emu", kDLAOCL)
     .add_attr_option<Bool>("system-lib")
+    .add_attr_option<Target>("host")
     .set_default_keys({"aocl", "hls"});
 
 TVM_REGISTER_TARGET_KIND("hexagon", kDLHexagon)
@@ -297,19 +309,23 @@ TVM_REGISTER_TARGET_KIND("hexagon", kDLHexagon)
     .add_attr_option<String>("mtriple")
     .add_attr_option<Bool>("system-lib")
     .add_attr_option<Array<String>>("llvm-options")
+    .add_attr_option<Target>("host")
     .set_default_keys({"hexagon"});
 
 TVM_REGISTER_TARGET_KIND("stackvm", kDLCPU)  // line break
-    .add_attr_option<Bool>("system-lib");
+    .add_attr_option<Bool>("system-lib")
+    .add_attr_option<Target>("host");
 
 TVM_REGISTER_TARGET_KIND("ext_dev", kDLExtDev)  // line break
-    .add_attr_option<Bool>("system-lib");
+    .add_attr_option<Bool>("system-lib")
+    .add_attr_option<Target>("host");
 
 TVM_REGISTER_TARGET_KIND("hybrid", kDLCPU)  // line break
-    .add_attr_option<Bool>("system-lib");
+    .add_attr_option<Bool>("system-lib")
+    .add_attr_option<Target>("host");
 
 TVM_REGISTER_TARGET_KIND("composite", kDLCPU)
-    .add_attr_option<Target>("target_host")
+    .add_attr_option<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,44 @@ 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
+
+
 if __name__ == "__main__":
     test_target_dispatch()
     test_target_string_parse()