comaniac commented on a change in pull request #7577:
URL: https://github.com/apache/tvm/pull/7577#discussion_r586649354



##########
File path: 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

Review comment:
       We don't need this line?

##########
File path: python/tvm/contrib/target/vitis_ai.py
##########
@@ -71,37 +71,66 @@ 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
+        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:
+        target = cfg.target if cfg.target 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("")

Review comment:
       ```suggestion
               cfg.build_dir if cfg.build_dir else 
tvm.contrib.utils.tempdir().relpath("")
   ```

##########
File path: python/tvm/relay/op/contrib/vitis_ai.py
##########
@@ -80,25 +82,61 @@ def visit_call(self, call):
                 else:
                     return super().visit_call(call)
 
+        xgraph = pyxir.frontend.tvm.from_relay(mod, self.params, 
postprocessing=None)
+        xgraph = pyxir.partition(xgraph, targets=[self.target])
+
+        layers = xgraph.get_layers()
+        relay_ids = [
+            list(np.array(layer.attrs["relay_id"]).flatten())
+            for layer in layers
+            if layer.target == self.target
+        ]
+        self.relay_ids = [item for sublist in relay_ids for item in sublist]
+
         return Annotator().visit(func)
 
 
 def annotation(mod, params, target):
-    """Annotate Relay expression for Vitis-AI DPU accelerators"""
+    """Annotate Relay expression for offloading operators to Vitis AI DPU 
accelerators
+    NOTE: This function does the same as the next one 
(`partition_for_vitis_ai`) but is
+    still here for backward compatibility"""

Review comment:
       For this case, we usually log a deprecate warning, and we remove the 
implementation in this or next release cycle.
   You could refer to the following cases:
   - 
https://github.com/apache/tvm/blob/main/python/tvm/auto_scheduler/search_task.py#L403
   - https://github.com/apache/tvm/blob/main/python/tvm/target/target.py#L481

##########
File path: python/tvm/relay/op/contrib/vitis_ai.py
##########
@@ -80,25 +82,61 @@ def visit_call(self, call):
                 else:
                     return super().visit_call(call)
 
+        xgraph = pyxir.frontend.tvm.from_relay(mod, self.params, 
postprocessing=None)
+        xgraph = pyxir.partition(xgraph, targets=[self.target])
+
+        layers = xgraph.get_layers()
+        relay_ids = [
+            list(np.array(layer.attrs["relay_id"]).flatten())
+            for layer in layers
+            if layer.target == self.target
+        ]
+        self.relay_ids = [item for sublist in relay_ids for item in sublist]
+
         return Annotator().visit(func)
 
 
 def annotation(mod, params, target):
-    """Annotate Relay expression for Vitis-AI DPU accelerators"""
+    """Annotate Relay expression for offloading operators to Vitis AI DPU 
accelerators
+    NOTE: This function does the same as the next one 
(`partition_for_vitis_ai`) but is
+    still here for backward compatibility"""
     # We need type information for supporting models that contain operations 
that don't
     #   have a Relay to XLayer translation
     mod = relay.transform.InferType()(mod)
+    mod = VitisAIAnnotationPass("vitis_ai", target, params)(mod)
+    return mod
 
-    xgraph = pyxir.frontend.tvm.from_relay(mod, params, postprocessing=None)
-    xgraph = pyxir.partition(xgraph, targets=[target])
 
-    layers = xgraph.get_layers()
-    relay_ids = [
-        list(np.array(layer.attrs["relay_id"]).flatten())
-        for layer in layers
-        if layer.target == target
-    ]
-    relay_ids_flatten = [item for sublist in relay_ids for item in sublist]
-    mod = VitisAIAnnotationPass("vitis_ai", relay_ids_flatten)(mod)
+def partition_for_vitis_ai(mod, params=None, target=None, **opts):
+    """Partition the Relay expression for offloading operators to Vitis AI DPU
 
-    return mod
+    Parameters
+    ----------
+    mod : Module
+        The module to run passes on.
+    params : Optional[Dict[str, NDArray]]
+        Constant input parameters.
+    target : str
+        The DPU identifier (e.g. DPUCZDX8G-zcu104, DPUCADX8G)
+
+    Returns
+    -------
+    ret : annotated and partitioned module.

Review comment:
       ```suggestion
       ret : Module
           annotated and partitioned module.
   ```

##########
File path: python/tvm/contrib/target/vitis_ai.py
##########
@@ -71,37 +71,66 @@ 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
+        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:
+        target = cfg.target if cfg.target 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("")

Review comment:
       ```suggestion
               cfg.work_dir if cfg.work_dir else 
tvm.contrib.utils.tempdir().relpath("")
   ```




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
[email protected]


Reply via email to