[Target] Add support for target object with host field compatible with previous api

include/tvm/target/target.h

@@ -35,6 +35,7 @@
 namespace tvm {
 
 class TargetInternal;
+class Target;
 
 /*!
  * \brief Compilation target.
@@ -60,6 +61,8 @@ class TargetNode : public Object {
   TVM_DLL const std::string& str() const;
   /*! \return Export target to JSON-like configuration */
   TVM_DLL Map<String, ObjectRef> Export() const;
+  /*! \return The Optional<Target> typed target host of the TargetNode */
+  TVM_DLL Optional<Target> GetHost() const;
 
   void VisitAttrs(AttrVisitor* v) {
     v->Visit("kind", &kind);

python/tvm/auto_scheduler/measure.py

@@ -223,8 +223,8 @@ def recover_measure_input(inp, rebuild_state=False):
     task = inp.task
     new_task = SearchTask(
         workload_key=task.workload_key,
-        target=task.target,
-        target_host=task.target_host,
+        target=tvm.target.Target(task.target, task.target_host),
+        target_host=None,
         hardware_params=task.hardware_params,
         layout_rewrite_option=task.layout_rewrite_option,
         task_inputs=list(task.task_input_names),

python/tvm/auto_scheduler/relay_integration.py

@@ -108,10 +108,8 @@ def extract_tasks(
     """
     # pylint: disable=import-outside-toplevel
 
-    if isinstance(target, str):
-        target = tvm.target.Target(target)
-    if isinstance(target_host, str):
-        target_host = tvm.target.Target(target_host)
+    target = tvm.target.Target(target, target_host)
+    target_host = target.host
 
     # Run the compiler to collect all TOPI calls during compilation.
     env = TracingEnvironment(

python/tvm/auto_scheduler/search_task.py

@@ -398,6 +398,9 @@ def __init__(
         if isinstance(target_host, str):
             target_host = Target(target_host)
 
+        target = Target(target, target_host)
+        target_host = target.host
+
         if layout_rewrite_option is None:
             layout_rewrite_option = LayoutRewriteOption.get_target_default(target)
 
@@ -509,8 +512,8 @@ def __getstate__(self):
         return {
             "compute_dag": self.compute_dag,
             "workload_key": self.workload_key,
-            "target": self.target,
-            "target_host": self.target_host,
+            "target": Target(self.target, self.target_host),
+            "target_host": Target(self.target, self.target_host).host,
             "hardware_params": self.hardware_params,
             "layout_rewrite_option": self.layout_rewrite_option,
             "task_input_names": self.task_input_names,
@@ -534,8 +537,8 @@ def __setstate__(self, state):
             _ffi_api.SearchTask,
             state["compute_dag"],
             state["workload_key"],
-            state["target"],
-            state["target_host"],
+            Target(state["target"], state["target_host"]),
+            Target(state["target"], state["target_host"]).host,
             state["hardware_params"],
             state["layout_rewrite_option"],
             state["task_input_names"],

python/tvm/autotvm/graph_tuner/base_graph_tuner.py

@@ -525,6 +525,8 @@ def _callback(_, inputs, results):
                 continue
 
             records = []
+            target = Target(target, target_host)
+            target_host = target.host
             task = autotvm.task.create(
                 "layout_transform", args=args, target=self._target, target_host=target_host
             )

python/tvm/autotvm/measure/measure_methods.py

@@ -418,6 +418,10 @@ def set_task(self, task):
 def _build_func_common(measure_input, check_gpu=None, cuda_arch=None, build_option=None):
     """Common part for building a configuration"""
     target, task, config = measure_input
+
+    target = tvm.target.Target(target, task.target_host)
+    task.target_host = target.host
+
     with target:
         s, args = task.instantiate(config)
 

python/tvm/autotvm/task/relay_integration.py

@@ -25,6 +25,7 @@
 
 import tvm
 from tvm.autotvm.task.dispatcher import DispatchContext, FallbackContext
+from tvm.target import Target
 from .task import create
 from .topi_integration import TaskExtractEnv
 
@@ -89,6 +90,8 @@ def extract_from_program(mod, params, target, target_host=None, ops=None):
     task: Array of autotvm.task.Task
         collected tasks
     """
+    target = Target(target, target_host)
+    target_host = target.host
     return extract_from_multiple_program([mod], [params], target, target_host, ops)
 
 
@@ -148,6 +151,10 @@ def extract_from_multiple_program(mods, params, target, target_host=None, ops=No
 
         logger.disabled = old_state
 
+    # merge target and target host
+    target = Target(target, target_host)
+    target_host = target.host
+
     # create tasks for target
     tasks = []
     for task_name, args in env.get_tasks():

python/tvm/autotvm/task/task.py

@@ -175,14 +175,15 @@ def __getstate__(self):
         # and restore the function by name when unpickling it.
         import cloudpickle  # pylint: disable=import-outside-toplevel
 
+        self.target = Target(self.target, self.target_host)
         return {
             "name": self.name,
             "args": self.args,
             "kwargs": self.kwargs,
             "config_space": self.config_space,
             "flop": self.flop,
             "target": self.target,
-            "target_host": self.target_host,
+            "target_host": self.target.host,
             "func": cloudpickle.dumps(self.func),
         }
 
@@ -195,8 +196,8 @@ def __setstate__(self, state):
         self.config_space = state["config_space"]
         self.func = cloudpickle.loads(state["func"])
         self.flop = state["flop"]
-        self.target = state["target"]
-        self.target_host = state["target_host"]
+        self.target = Target(state["target"], state["target_host"])
+        self.target_host = self.target.host
 
     def __repr__(self):
         return "Task(func_name=%s, args=%s, kwargs=%s, workload=%s)" % (
@@ -448,6 +449,9 @@ def create(task_name, args, target, target_host=None):
     if isinstance(target, str):
         target = Target(target)
 
+    target = Target(target, target_host)
+    target_host = target.host
+
     # init config space
     ret.config_space = ConfigSpace()
 

python/tvm/contrib/peak.py

@@ -106,6 +106,9 @@ def measure_bandwidth_sum(
     s[y].bind(yi, te.thread_axis("threadIdx.x"))
     s[y].unroll(k)
 
+    target = tvm.target.Target(target, target_host)
+    target_host = target.host
+
     try:
         func = tvm.build(s, [x, y], target, target_host=target_host)
 
@@ -153,6 +156,9 @@ def measure_bandwidth_all_types(
     """
     max_threads = target.max_num_threads
 
+    target = tvm.target.Target(target, target_host)
+    target_host = target.host
+
     result = []
     for base_type in ["float"]:
         for bits in [32]:
@@ -229,6 +235,9 @@ def measure_compute_mad(
 
     max_threads = target.max_num_threads
 
+    target = tvm.target.Target(target, target_host)
+    target_host = target.host
+
     base_type = str(base_type) + str(bits)
     dtype = base_type if lanes == 1 else base_type + "x" + str(lanes)
 
@@ -313,6 +322,9 @@ def measure_compute_all_types(
     result: list
         a list of (type_name, GFLOPS/GIOPS) pairs
     """
+    target = tvm.target.Target(target, target_host)
+    target_host = target.host
+
     result = []
     for base_type in ["float", "int"]:
         for bits in [16, 32, 64]:
@@ -357,7 +369,8 @@ def measure_peak_all(target, target_host, host, port):
     port: int
     """
 
-    target = tvm.target.Target(target)
+    target = tvm.target.Target(target, target_host)
+    target_host = target.host
     remote = rpc.connect(host, port)
     n_times = 20
 

python/tvm/driver/build_module.py

@@ -231,8 +231,8 @@ def _build_for_device(input_mod, target, target_host):
     mdev : tvm.module
         A module that contains device code.
     """
-    target = Target(target)
-    target_host = Target(target_host)
+    target = Target(target, target_host)
+    target_host = target.host
     device_type = ndarray.context(target.kind.name, 0).device_type
 
     mod_mixed = input_mod
@@ -386,10 +386,12 @@ def build(inputs, args=None, target=None, target_host=None, name="default_functi
             f"but got {type(inputs)}."
         )
 
+    flag_target_inputs = False
     if not isinstance(inputs, (dict, container.Map)):
         target = Target.current() if target is None else target
         target = target if target else "llvm"
         target_input_mod = {target: input_mod}
+        flag_target_inputs = True
     else:
         target_input_mod = inputs
 
@@ -399,6 +401,11 @@ def build(inputs, args=None, target=None, target_host=None, name="default_functi
         if not isinstance(mod, tvm.IRModule):
             raise ValueError("inputs must be Schedule, IRModule," "or dict of str to IRModule.")
 
+    target = Target(target, target_host)
+    target_host = target.host
+    if flag_target_inputs:
+        target_input_mod = {target: input_mod}
+
     if not target_host:
         for tar, _ in target_input_mod.items():
             tar = Target(tar)
@@ -409,6 +416,11 @@ def build(inputs, args=None, target=None, target_host=None, name="default_functi
     if not target_host:
         target_host = "llvm" if tvm.runtime.enabled("llvm") else "stackvm"
 
+    target = Target(target, target_host)
+    target_host = target.host
+    if flag_target_inputs:
+        target_input_mod = {target: input_mod}
+
     mod_host_all = tvm.IRModule({})
 
     device_modules = []
@@ -427,6 +439,8 @@ def build(inputs, args=None, target=None, target_host=None, name="default_functi
 
     if not isinstance(target_host, Target):
         target_host = Target(target_host)
+        target = Target(target, target_host)
+        target_host = target.host
     if (
         target_host.attrs.get("runtime", tvm.runtime.String("c++")) == "c"
         and target_host.attrs.get("system-lib", 0).value == 1

python/tvm/driver/tvmc/autotuner.py

@@ -23,6 +23,8 @@
 
 from urllib.parse import urlparse
 
+import tvm
+
 from tvm import autotvm, auto_scheduler
 from tvm.autotvm.tuner import GATuner
 from tvm.autotvm.tuner import GridSearchTuner
@@ -242,6 +244,8 @@ def drive_tune(args):
             )
 
     target, extra_targets = common.target_from_cli(args.target)
+    target = tvm.target.Target(target, args.target_host)
+    target_host = target.host
     mod, params = frontends.load_model(args.FILE, args.model_format, shape_dict=args.input_shapes)
 
     for codegen_from_cli in extra_targets:
@@ -298,7 +302,7 @@ def drive_tune(args):
             mod=mod,
             params=params,
             target=target,
-            target_host=args.target_host,
+            target_host=target_host,
             alter_layout=args.desired_layout,
             hardware_params=hardware_params,
             include_simple_tasks=args.include_simple_tasks,
@@ -321,7 +325,7 @@ def drive_tune(args):
             mod=mod,
             params=params,
             target=target,
-            target_host=args.target_host,
+            target_host=target_host,
             alter_layout=args.desired_layout,
         )
 
@@ -365,6 +369,9 @@ def autotvm_get_tuning_tasks(mod, params, target, target_host=None, alter_layout
     if alter_layout:
         mod = common.convert_graph_layout(mod, alter_layout)
 
+    target = tvm.target.Target(target, target_host)
+    target_host = target.host
+
     tasks = autotvm.task.extract_from_program(
         mod["main"],
         target=target,
@@ -413,6 +420,9 @@ def autoscheduler_get_tuning_tasks(
     if alter_layout:
         mod = common.convert_graph_layout(mod, alter_layout)
 
+    target = tvm.target.Target(target, target_host)
+    target_host = target.host
+
     # Extract the tasks
     tasks, task_weights = auto_scheduler.extract_tasks(
         mod["main"],

python/tvm/driver/tvmc/compiler.py

@@ -191,7 +191,8 @@ def compile_model(
         mod = common.convert_graph_layout(mod, alter_layout)
 
     tvm_target, extra_targets = common.target_from_cli(target)
-    target_host = tvm_target if not target_host else target_host
+    tvm_target = tvm.target.Target(tvm_target, tvm_target if not target_host else target_host)
+    target_host = tvm_target.host
 
     for codegen_from_cli in extra_targets:
         codegen = composite_target.get_codegen_by_target(codegen_from_cli["name"])

python/tvm/relay/backend/_backend.py

@@ -80,6 +80,8 @@ def build(mod, target, target_host=None):
     """
     if target_host == "":
         target_host = None
+    target = tvm.target.Target(target, target_host)
+    target_host = target.host
     return tvm.driver.build(mod, target=target, target_host=target_host)
 
 

python/tvm/relay/backend/vm.py

@@ -65,6 +65,13 @@ def compile(mod, target=None, target_host=None, params=None):
     compiler = VMCompiler()
     if params:
         compiler.set_params(params)
+    if isinstance(target, dict):
+        for k in target:
+            target[k] = tvm.target.Target(target[k], target_host)
+            target_host = target[k].host
+    else:
+        target = tvm.target.Target(target, target_host)
+        target_host = target.host
     compiler.lower(mod, target, target_host)
     compiler.codegen()
     return compiler.get_exec()
@@ -130,6 +137,15 @@ def lower(self, mod, target=None, target_host=None):
         """
         target = self._update_target(target)
         target_host = self._update_target_host(target, target_host)
+
+        if isinstance(target, dict):
+            for k in target:
+                target[k] = tvm.target.Target(target[k], target_host)
+                target_host = target[k].host
+        else:
+            target = tvm.target.Target(target, target_host)
+            target_host = target.host
+
         tophub_context = self._tophub_context(target)
         with tophub_context:
             self._lower(mod, target, target_host)
@@ -167,6 +183,15 @@ def optimize(self, mod, target=None, target_host=None, params=None):
         """
         target = self._update_target(target)
         target_host = self._update_target_host(target, target_host)
+
+        if isinstance(target, dict):
+            for k in target:
+                target[k] = tvm.target.Target(target[k], target_host)
+                target_host = target[k].host
+        else:
+            target = tvm.target.Target(target, target_host)
+            target_host = target.host
+
         if params:
             self.set_params(params)
         return self._optimize(mod, target, target_host), self.get_params()