tmoreau89 commented on a change in pull request #6125:
URL: https://github.com/apache/incubator-tvm/pull/6125#discussion_r477013567
##########
File path: vta/python/vta/top/graphpack.py
##########
@@ -449,18 +549,46 @@ def graph_pack(expr,
'expr.astext(show_meta_data=False)'. When count_meta is True, the
operator increase
logic would count the meta.
+ device_annot: boolean, optional
+ if we want to annoate the device_type
+
+ annot_start_name: str, optional
+ device annotation start node, from which we mark the nodes as `ext_dev`
+
+ annot_end_name: str, optional
+ device annotation end node, after which we mark the nodes as 'cpu'
+
Returns
-------
expr : Expr
The transformed expression.
"""
assert isinstance(expr, relay.Function)
- assert ((start_name != stop_name) or (start_name_idx < stop_name_idx))
+ assert ((start_name != stop_name) or (start_name_idx is None !=
stop_name_idx is None) or \
+ (not (start_name_idx is None and stop_name_idx is None)) \
+ or (start_name_idx < stop_name_idx))
expr = get_subgraph(expr, start_name, stop_name, start_name_idx,
stop_name_idx, count_meta)
expr = run_opt_pass(expr, transform.InferType())
packer = ExprPack(
bfactor, cfactor,
weight_bits)
expr = packer.visit(expr)
assert not packer.start_pack
- return run_opt_pass(expr, transform.InferType())
+ expr = run_opt_pass(expr, transform.InferType())
+
+ if device_annot:
+ expr_locator = ExprLocater()
Review comment:
ExprLocator
----------------------------------------------------------------
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]
