[BYOC][VitisAI] Fix issue in Vitis AI codegen out tensor names matching & update docs and docker

docker/Dockerfile.demo_vitis_ai

@@ -20,7 +20,6 @@ FROM xilinx/vitis-ai:latest
 
 RUN apt-get update --fix-missing
 
-
 COPY install/ubuntu_install_core.sh /install/ubuntu_install_core.sh
 RUN bash /install/ubuntu_install_core.sh
 

docker/install/ubuntu_install_vitis_ai_core.sh

@@ -21,9 +21,9 @@ set -u
 set -o pipefail
 
 # install libraries for building Vitis-AI on ubuntu
-apt-get update && apt-get install -y --no-install-recommends \
-    graphviz\
-    gnupg2
-
-apt-get update && apt-get install -y gcc-aarch64-linux-gnu
-
+apt-get update && apt-get install -y \
+    graphviz \
+    gnupg2 \
+    gpg-agent \
+    gcc-aarch64-linux-gnu \
+    && rm -rf /var/lib/apt/lists/*

docs/deploy/vitis_ai.rst

@@ -309,7 +309,7 @@ Edge hardware setup
    https://github.com/Xilinx/PYNQ/releases/tag/v2.5
 2. Follow Pynq instructions for setting up the board: `pynq
    setup <https://pynq.readthedocs.io/en/latest/getting_started.html>`__
-3. After connecting to the board, make sure to run as root. Execute
+3. After connecting to the board, make sure to run as root. **Execute**
    ``su``
 4. Set up DPU on Pynq by following the steps here: `DPU Pynq
    setup <https://github.com/Xilinx/DPU-PYNQ>`__
@@ -441,7 +441,7 @@ TVM.
    import tvm
    import tvm.relay as relay
    from tvm.contrib.target import vitis_ai
-   from tvm.contrib import util, graph_runtime
+   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
 
@@ -524,6 +524,8 @@ model in TVM with Vitis-AI at the edge. The first couple of steps will
 have to be run on the host machine and take care of quantization and
 compilation for deployment at the edge.
 
+A complete ResNet 18 example can be found `here <https://github.com/Xilinx/pyxir/tree/master/examples/tvm>`__.
+
 Host steps
 ^^^^^^^^^^
 
@@ -541,20 +543,50 @@ TVM.
    import tvm
    import tvm.relay as relay
    from tvm.contrib.target import vitis_ai
-   from tvm.contrib import util, graph_runtime
+   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
 
 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.
 
+.. note::
+
+    We recommend switching DPU convolutions' data layouts to NHWC and CPU comvolutions'
+    data layouts to NCHW for best DPU and CPU performance. You can use the ConvertLayout
+    transformation pass two times to achieve this as demonstrated in the code block
+    underneath.
+
 .. code:: python
 
    mod["main"] = bind_params_by_name(mod["main"], params)
+
+   # For edge DPU we recommend switching the convolutions'data layout
+   #    to NHWC for best performance. Therefore, we first convert the layouts
+   #    of all convolutions to NHWC before partitioning. Afterwards, we can
+   #    convert any remaining convolutions (to be executed on CPU) back to NCHW.
+   desired_layouts = {'nn.conv2d': ['NHWC', 'default']}
+   seq = tvm.transform.Sequential([relay.transform.RemoveUnusedFunctions(),
+                                   relay.transform.ConvertLayout(desired_layouts),
+                                   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)
+
+   # After partitioning we recommend transforming the remaining convolutions
+   #    (that will be executed on CPU, if any) back to NCHW data layout
+   #    for best CPU performance
+   desired_layouts = {'nn.conv2d': ['NCHW', 'default']}
+   seq = tvm.transform.Sequential([relay.transform.RemoveUnusedFunctions(),
+                                   relay.transform.ConvertLayout(desired_layouts),
+                                   relay.transform.FoldConstant()])
+   with tvm.transform.PassContext(opt_level=3):
+       mod = seq(mod)
 
 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
@@ -572,13 +604,9 @@ can be included.
 
 .. code:: python
 
-   from tvm.contrib import util
-
-   temp = util.tempdir()
-
    tvm_target = 'llvm'
    target='DPUCZDX8G-zcu104'
-   export_rt_mod_file = temp.relpath("vitis_ai.rtmod")
+   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}):
@@ -604,9 +632,9 @@ Save the TVM lib module so that the Vitis-AI runtime module will also be exporte
 
 .. code:: python
 
-   from tvm.contrib import util
+   from tvm.contrib import utils
 
-   temp = util.tempdir()
+   temp = utils.tempdir()
    lib.export_library(temp.relpath("tvm_lib.so"))
 
 After quantizing and compiling the model for Vitis-AI acceleration using the
@@ -638,15 +666,31 @@ Edge steps
 ^^^^^^^^^^
 
 After setting up TVM with Vitis-AI on the edge device, you can now load
-the TVM runtime module into memory and feed inputs for inference.
+the TVM runtime module into memory and feed inputs for inference. A nearly
+complete runtiem script can be found underneath. Make sure to run the script
+as root (execute ``su`` in terminal to log into root).
+
+
+.. note::
+
+    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``).
 
 .. code:: python
 
+   import pyxir
+   import tvm
+   from tvm.contrib import graph_runtime
+
    ctx = tvm.cpu()
+
+   # input_name = ...
+   # input_data = ...
 
    # load the module into memory
    lib = tvm.runtime.load_module("tvm_dpu_arm.so")
 
    module = graph_runtime.GraphModule(lib["default"](tvm.cpu()))
-   module.set_input(name, data)
+   module.set_input(input_name, input_data)
    module.run()

python/tvm/contrib/target/vitis_ai.py

@@ -132,12 +132,12 @@ def vitis_ai_compiler(ref):
         layers = xgraph.get_layers()
 
         # Get the output tensor names using XGraph and output Relay ids
-        out_tensor_names = []
+        out_tensor_names = [1] * len(output_relay_ids)
         for layer in layers:
             if not layer.internal:
                 for relay_id in layer.attrs["relay_id"]:
                     if relay_id in output_relay_ids:
-                        out_tensor_names.append(layer.name)
+                        out_tensor_names[output_relay_ids.index(relay_id)] = layer.name
                         break
         if not out_tensor_names:
             raise ValueError(

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

@@ -85,6 +85,10 @@ def visit_call(self, call):
 
 def annotation(mod, params, target):
     """Annotate Relay expression for Vitis-AI DPU accelerators"""
+    # We need type information for supporting models that contain operations that don't
+    #   have a Relay to XLayer translation
+    mod = relay.transform.InferType()(mod)
+
     xgraph = pyxir.frontend.tvm.from_relay(mod, params, postprocessing=None)
     xgraph = pyxir.partition(xgraph, targets=[target])