jtuyls commented on a change in pull request #7577:
URL: https://github.com/apache/tvm/pull/7577#discussion_r588548912
##########
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):
Review comment:
Yes, we have been using `target` for some time but I agree it's an
overused term and it would be good to rename it. As we will be deprecating some
API's anyway, I will change it in this PR too then. I think using `dpu` would
be the clearest and it would also avoid other conflicts.
----------------------------------------------------------------
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]
