MLBEDSW-3268: Refactor mark_tensors
- Refactored mark_tensor_purpose
- Initial weight compression is now always done in insert_dma
- Removed mark_tensor_format
Change-Id: Ic719b9bcd1d27e1390d7b9ce8cd21795139ec814
Signed-off-by: Louis Verhaard <louis.verhaard@arm.com>
diff --git a/ethosu/vela/rewrite_graph.py b/ethosu/vela/rewrite_graph.py
index e71b228..42acaf9 100644
--- a/ethosu/vela/rewrite_graph.py
+++ b/ethosu/vela/rewrite_graph.py
@@ -82,14 +82,16 @@
return sg
-def visit_graph_post_order(sg, arch, tensor_visit_list, op_visit_list):
-
+def visit_graph_post_order(start_tensors, arch, tensor_visit_list, op_visit_list):
+ # Depth-first graph traversal, starting from the given list of tensors
+ # (typically a subgraph's output_tensors).
+ # Visits ops and tensors in input to output order.
op_visit_dict = dict()
tens_visit_dict = dict()
def visit_op(op):
if op in op_visit_dict:
- return op_visit_dict[op]
+ return
op_visit_dict[op] = op
for tens in op.inputs:
@@ -101,11 +103,9 @@
for tens in op.outputs:
visit_tens(tens)
- return op
-
def visit_tens(tens):
- if tens in tens_visit_dict:
- return tens_visit_dict[tens]
+ if tens is None or tens in tens_visit_dict:
+ return
tens_visit_dict[tens] = tens
@@ -115,15 +115,9 @@
for visit in tensor_visit_list:
visit(tens, arch)
- return tens
-
- for tens in sg.output_tensors:
+ for tens in start_tensors:
visit_tens(tens)
- sg.refresh_after_modification()
-
- return sg
-
def verify_graph_health(nng):