[TVMC][VitisAI] Enable Vitis AI target through TVMC

docs/deploy/vitis_ai.rst

@@ -196,7 +196,7 @@ Hardware setup and docker build
       pip3 install -e . --user
 
 Edge (DPUCZDX8G)
-^^^^^^^^^^^^^^^^
+~~~~~~~~~~~~~~~~~
 
 
 For edge deployment we make use of two systems referred to as host and
@@ -435,8 +435,8 @@ Cloud usage
 This section shows how to accelerate a convolutional neural network
 model in TVM with Vitis-AI on the cloud.
 
-To be able to target the Vitis-AI cloud DPUCADX8G target we first have
-to import the target in PyXIR. This PyXIR package is the interface being
+To be able to target the Vitis-AI cloud DPUCADX8G we first have
+to import the DPU target in PyXIR. This PyXIR package is the interface being
 used by TVM to integrate with the Vitis-AI stack. Additionaly, import
 the typical TVM and Relay modules and the Vitis-AI contrib module inside
 TVM.
@@ -451,32 +451,29 @@ TVM.
    from tvm.contrib.target import vitis_ai
    from tvm.contrib import utils, graph_runtime
    from tvm.relay.build_module import bind_params_by_name
-   from tvm.relay.op.contrib.vitis_ai import annotation
+   from tvm.relay.op.contrib.vitis_ai import partition_for_vitis_ai
 
 After importing a convolutional neural network model using the usual
 Relay API's, annotate the Relay expression for the given Vitis-AI DPU
 target and partition the graph.
 
 .. code:: python
-
-   mod["main"] = bind_params_by_name(mod["main"], params)
-   mod = annotation(mod, params, target)
-   mod = relay.transform.MergeCompilerRegions()(mod)
-   mod = relay.transform.PartitionGraph()(mod)
+
+   dpu = 'DPUCADX8G'
+   mod = partition_for_vitis_ai(mod, params, dpu)
 
 Now, we can build the TVM runtime library for executing the model. The
 TVM target is 'llvm' as the operations that can't be handled by the DPU
-are executed on the CPU. The Vitis-AI target is DPUCADX8G as we are
-targeting the cloud DPU and this target is passed as a config to the TVM
+are executed on the CPU. The Vitis-AI DPU is DPUCADX8G as we are
+targeting the cloud DPU and this DPU identifier is passed as a config to the TVM
 build call.
 
 .. code:: python
 
-   tvm_target = 'llvm'
-   target='DPUCADX8G'
+   target = 'llvm'
 
-   with tvm.transform.PassContext(opt_level=3, config= {'relay.ext.vitis_ai.options.target': target}):
-      lib = relay.build(mod, tvm_target, params=params)
+   with tvm.transform.PassContext(opt_level=3, config= {'relay.ext.vitis_ai.options': {'dpu': dpu}}):
+      lib = relay.build(mod, target, params=params)
 
 As one more step before we can accelerate a model with Vitis-AI in TVM
 we have to quantize and compile the model for execution on the DPU. We
@@ -537,8 +534,8 @@ A complete ResNet 18 example can be found `here <https://github.com/Xilinx/pyxir
 Host steps
 ^^^^^^^^^^
 
-To be able to target the Vitis-AI cloud DPUCZDX8G target we first have
-to import the target in PyXIR. This PyXIR package is the interface being
+To be able to target the Vitis-AI cloud DPUCZDX8G we first have
+to import the DPU target in PyXIR. This PyXIR package is the interface being
 used by TVM to integrate with the Vitis-AI stack. Additionaly, import
 the typical TVM and Relay modules and the Vitis-AI contrib module inside
 TVM.
@@ -553,11 +550,11 @@ TVM.
    from tvm.contrib.target import vitis_ai
    from tvm.contrib import utils, graph_runtime
    from tvm.relay.build_module import bind_params_by_name
-   from tvm.relay.op.contrib.vitis_ai import annotation
+   from tvm.relay.op.contrib.vitis_ai import partition_for_vitis_ai
 
 After importing a convolutional neural network model using the usual
 Relay API's, annotate the Relay expression for the given Vitis-AI DPU
-target and partition the graph.
+and partition the graph.
 
 .. note::
 
@@ -585,11 +582,10 @@ target and partition the graph.
                                    relay.transform.FoldConstant()])
    with tvm.transform.PassContext(opt_level=3):
        mod = seq(mod)
-
-   # Annotate and partition the Relay expression for the given target
-   mod = annotation(mod, params, target)
-   mod = relay.transform.MergeCompilerRegions()(mod)
-   mod = relay.transform.PartitionGraph()(mod)
+
+   dpu = 'DPUCZDX8G-zcu104'
+   # Annotate and partition the Relay expression for the given DPU
+   mod = partition_for_vitis_ai(mod, params, dpu)
 
    # After partitioning we recommend transforming the remaining convolutions
    #    (that will be executed on CPU, if any) back to NCHW data layout
@@ -604,10 +600,10 @@ target and partition the graph.
 Now, we can build the TVM runtime library for executing the model. The
 TVM target is 'llvm' as the operations that can't be handled by the DPU
 are executed on the CPU. At this point that means the CPU on the host machine.
-The Vitis-AI target is DPUCZDX8G-zcu104 as we are targeting the edge DPU
-on the ZCU104 board and this target is passed as a config to the TVM
+The Vitis-AI DPU identifier is DPUCZDX8G-zcu104 as we are targeting the edge DPU
+on the ZCU104 board and this identifier is passed as a config to the TVM
 build call. Note that different identifiers can be passed for different
-targets, see `edge targets info <#edge-requirements>`__. Additionally, we
+DPU's, see `edge DPU's info <#edge-requirements>`__. Additionally, we
 provide the 'export_runtime_module' config that points to a file to which we
 can export the Vitis-AI runtime module. We have to do this because we will
 first be compiling and quantizing the model on the host machine before building
@@ -617,13 +613,15 @@ can be included.
 
 .. code:: python
 
-   tvm_target = 'llvm'
-   target='DPUCZDX8G-zcu104'
+   target = 'llvm'
    export_rt_mod_file = "vitis_ai.rtmod"
-
-   with tvm.transform.PassContext(opt_level=3, config= {'relay.ext.vitis_ai.options.target': target,
-   						        'relay.ext.vitis_ai.options.export_runtime_module': export_rt_mod_file}):
-      lib = relay.build(mod, tvm_target, params=params)
+
+   build_options = {
+      'dpu': dpu,
+      'export_runtime_module': export_rt_mod_file
+   }
+   with tvm.transform.PassContext(opt_level=3, config= {'relay.ext.vitis_ai.options': build_options}):
+      lib = relay.build(mod, target, params=params)
 
 We will quantize and compile the model for execution on the DPU using on-the-fly
 quantization on the host machine. This makes use of TVM inference calls
@@ -658,15 +656,17 @@ in the TVM build.
 .. code:: python
 
    # Export lib for aarch64 target
-   tvm_target = tvm.target.arm_cpu('ultra96')
+   target = tvm.target.arm_cpu('ultra96')
    lib_kwargs = {
         'fcompile': contrib.cc.create_shared,
         'cc': "/usr/aarch64-linux-gnu/bin/ld"
    }
-
-   with tvm.transform.PassContext(opt_level=3,
-                                  config={'relay.ext.vitis_ai.options.load_runtime_module': export_rt_mod_file}):
-        lib_arm = relay.build(mod, tvm_target, params=params)
+
+   build_options = {
+        'load_runtime_module': export_rt_mod_file
+   }
+   with tvm.transform.PassContext(opt_level=3, config={'relay.ext.vitis_ai.options': build_options}):
+        lib_arm = relay.build(mod, target, params=params)
 
    lib_dpuv2.export_library('tvm_dpu_arm.so', **lib_kwargs)
 
@@ -688,7 +688,7 @@ as root (execute ``su`` in terminal to log into root).
 
     You will see a warning about the 'cpu-tf' runtime not being found. This warning is
     expected on the board and can be ignored. Note also that you **shouldn't** import the
-    PyXIR targets in the run script (``import pyxir.contrib.target.DPUCZDX8G``).
+    PyXIR DPU targets in the run script (``import pyxir.contrib.target.DPUCZDX8G``).
 
 .. code:: python
 

python/tvm/contrib/target/vitis_ai.py

@@ -36,12 +36,12 @@ def __init__(self, model_name, function):
         self.function = function
         self.params = {}
 
-    def convert_pyxir(self, target):
+    def convert_pyxir(self, dpu_target):
         """Convert Relay expression to PyXIR XGraph"""
         xgraph = pyxir.frontend.tvm.from_relay(
             self.function, params=self.params, postprocessing=None
         )
-        xgraph = pyxir.partition(xgraph, targets=[target])
+        xgraph = pyxir.partition(xgraph, targets=[dpu_target])
         return xgraph
 
     def get_output_names(self):
@@ -71,44 +71,69 @@ def vitis_ai_compiler(ref):
 
     pass_context = tvm.get_global_func("transform.GetCurrentPassContext")()
 
-    # The target Vitis-AI accelerator device
-    target = (
-        str(pass_context.config["relay.ext.vitis_ai.options.target"])
-        if "relay.ext.vitis_ai.options.target" in pass_context.config
+    cfg = (
+        pass_context.config["relay.ext.vitis_ai.options"]
+        if "relay.ext.vitis_ai.options" in pass_context.config
         else None
     )
 
-    # (Optional configs) The build and work directories to be used by Vitis-AI
-    vai_build_dir = (
-        str(pass_context.config["relay.ext.vitis_ai.options.build_dir"])
-        if "relay.ext.vitis_ai.options.build_dir" in pass_context.config
-        else tvm.contrib.utils.tempdir().relpath("")
-    )
-    vai_work_dir = (
-        str(pass_context.config["relay.ext.vitis_ai.options.work_dir"])
-        if "relay.ext.vitis_ai.options.work_dir" in pass_context.config
-        else tvm.contrib.utils.tempdir().relpath("")
-    )
+    # Backward compatibility with old pass context configs
+    if cfg is None:
+        warnings.warn(
+            "You are using a deprecated way of passing build configs (e.g."
+            " `relay.ext.vitis_ai.options.target`). Check out the Vitis AI "
+            " documentation here: https://tvm.apache.org/docs/deploy/vitis_ai.html"
+            " to switch to recommended way for passing build configs."
+        )
 
-    # (Optional configs) Export and load PyXIR runtime module to file if provided. This is used to
-    #   compile and quantize a model on the host and deploy it at the edge
-    export_runtime_module = (
-        str(pass_context.config["relay.ext.vitis_ai.options.export_runtime_module"])
-        if "relay.ext.vitis_ai.options.export_runtime_module" in pass_context.config
-        else ""
-    )
-    load_runtime_module = (
-        str(pass_context.config["relay.ext.vitis_ai.options.load_runtime_module"])
-        if "relay.ext.vitis_ai.options.load_runtime_module" in pass_context.config
-        else ""
-    )
+        # The target Vitis-AI accelerator device
+        dpu_target = (
+            str(pass_context.config["relay.ext.vitis_ai.options.target"])
+            if "relay.ext.vitis_ai.options.target" in pass_context.config
+            else None
+        )
+
+        # (Optional configs) The build and work directories to be used by Vitis-AI
+        vai_build_dir = (
+            str(pass_context.config["relay.ext.vitis_ai.options.build_dir"])
+            if "relay.ext.vitis_ai.options.build_dir" in pass_context.config
+            else tvm.contrib.utils.tempdir().relpath("")
+        )
+        vai_work_dir = (
+            str(pass_context.config["relay.ext.vitis_ai.options.work_dir"])
+            if "relay.ext.vitis_ai.options.work_dir" in pass_context.config
+            else tvm.contrib.utils.tempdir().relpath("")
+        )
+
+        # (Optional configs) Export and load PyXIR runtime module to file if provided. This is
+        #   used to compile and quantize a model on the host and deploy it at the edge
+        export_runtime_module = (
+            str(pass_context.config["relay.ext.vitis_ai.options.export_runtime_module"])
+            if "relay.ext.vitis_ai.options.export_runtime_module" in pass_context.config
+            else ""
+        )
+        load_runtime_module = (
+            str(pass_context.config["relay.ext.vitis_ai.options.load_runtime_module"])
+            if "relay.ext.vitis_ai.options.load_runtime_module" in pass_context.config
+            else ""
+        )
+    else:
+        dpu_target = cfg.dpu if cfg.dpu else None
+        # (Optional configs) The build and work directories to be used by Vitis AI
+        vai_build_dir = cfg.build_dir if cfg.build_dir else tvm.contrib.utils.tempdir().relpath("")
+
+        # (Optional configs) Export and load PyXIR runtime module to file if provided. This is
+        #   used to compile and quantize a model on the host and deploy it at the edge
+        vai_work_dir = cfg.work_dir if cfg.work_dir else tvm.contrib.utils.tempdir().relpath("")
+        export_runtime_module = cfg.export_runtime_module
+        load_runtime_module = cfg.load_runtime_module
 
     # Config checks
-    if load_runtime_module and target is not None:
+    if load_runtime_module and dpu_target is not None:
         warnings.warn(
-            "Both `load_runtime_module` and `target` configs were specified."
+            "Both `load_runtime_module` and `dpu` configs were specified."
             " The `load_runtime_module` points to a prebuilt runtime module with"
-            " an internal target so the `target` config will be ignored"
+            " an internal DPU target so the `dpu` config will be ignored"
         )
     if load_runtime_module and "relay.ext.vitis_ai.options.build_dir" in pass_context.config:
         warnings.warn(
@@ -127,7 +152,7 @@ def vitis_ai_compiler(ref):
     if load_runtime_module == "":
         # Convert Relay expression into XGraph and do partitioning inside PyXIR
         builder = CodegenVitisAI(name, ref)
-        xgraph = builder.convert_pyxir(target)
+        xgraph = builder.convert_pyxir(dpu_target)
         output_relay_ids = builder.get_output_names()
         layers = xgraph.get_layers()
 
@@ -141,15 +166,17 @@ def vitis_ai_compiler(ref):
                         break
         if any([name == "unkown_name" for name in out_tensor_names]):
             raise ValueError(
-                "During codegeneration the loading of subexpression \
-                             failed due to output tensor name mismatch in Relay PyXIR interface."
+                "During codegeneration the loading of subexpression"
+                " failed due to output tensor name mismatch in Relay PyXIR interface."
             )
         xgraph.meta_attrs["tvm_out_tensors"] = out_tensor_names
         xgraph_str = pyxir.get_xgraph_str(xgraph)
 
         runtime_func = "tvm.vitis_ai_runtime.from_xgraph"
         fcreate = tvm._ffi.get_global_func(runtime_func)
-        return fcreate(name, xgraph_str, target, vai_build_dir, vai_work_dir, export_runtime_module)
+        return fcreate(
+            name, xgraph_str, dpu_target, vai_build_dir, vai_work_dir, export_runtime_module
+        )
 
     runtime_func = "tvm.vitis_ai_runtime.from_rt_mod"
     fcreate = tvm._ffi.get_global_func(runtime_func)

python/tvm/driver/tvmc/autotuner.py

@@ -247,7 +247,7 @@ def drive_tune(args):
     for codegen_from_cli in extra_targets:
         codegen = composite_target.get_codegen_by_target(codegen_from_cli["name"])
         partition_function = codegen["pass_pipeline"]
-        mod = partition_function(mod, params)
+        mod = partition_function(mod, params, **codegen_from_cli["opts"])
 
     # min_repeat_ms should be:
     # a. the value provided by the user, if any, or

python/tvm/driver/tvmc/common.py

@@ -133,7 +133,7 @@ def tokenize_target(target):
 
     target_pattern = (
         r"(\-{0,2}[\w\-]+\=?"
-        r"(?:[\w\+\-]+(?:,[\w\+\-])*|[\'][\w\+\-,\s]+[\']|[\"][\w\+\-,\s]+[\"])*|,)"
+        r"(?:[\w\+\-.]+(?:,[\w\+\-])*|[\'][\w\+\-,\s]+[\']|[\"][\w\+\-,\s]+[\"])*|,)"
     )
 
     return re.findall(target_pattern, target)

python/tvm/driver/tvmc/compiler.py

@@ -196,7 +196,7 @@ def compile_model(
     for codegen_from_cli in extra_targets:
         codegen = composite_target.get_codegen_by_target(codegen_from_cli["name"])
         partition_function = codegen["pass_pipeline"]
-        mod = partition_function(mod, params)
+        mod = partition_function(mod, params, **codegen_from_cli["opts"])
         if codegen["config_key"] is not None:
             config[codegen["config_key"]] = codegen_from_cli["opts"]
 

python/tvm/driver/tvmc/composite_target.py

@@ -21,6 +21,8 @@
 
 from tvm.relay.op.contrib.arm_compute_lib import partition_for_arm_compute_lib
 from tvm.relay.op.contrib.ethosn import partition_for_ethosn
+from tvm.relay.op.contrib.vitis_ai import partition_for_vitis_ai
+from tvm.contrib.target import vitis_ai  # pylint: disable=unused-import
 
 from .common import TVMCException
 
@@ -40,6 +42,10 @@
         "config_key": "relay.ext.ethos-n.options",
         "pass_pipeline": partition_for_ethosn,
     },
+    "vitis-ai": {
+        "config_key": "relay.ext.vitis_ai.options",
+        "pass_pipeline": partition_for_vitis_ai,
+    },
 }
 
 

python/tvm/relay/op/contrib/arm_compute_lib.py

@@ -43,7 +43,7 @@ def is_arm_compute_runtime_enabled():
     return False
 
 
-def partition_for_arm_compute_lib(mod, params=None):
+def partition_for_arm_compute_lib(mod, params=None, **opts):
     """Partition the graph greedily offloading supported
     operators to Arm Compute Library.
 

python/tvm/relay/op/contrib/ethosn.py

@@ -46,7 +46,7 @@ def ethosn_available():
     return Available.SW_AND_HW if hw else Available.SW_ONLY
 
 
-def partition_for_ethosn(mod, params=None):
+def partition_for_ethosn(mod, params=None, **opts):
     """Partition the graph greedily offloading supported
     operators to Arm Ethos-N NPU.