coordinate pre and late rewrite

Author: wayuanho

tf2onnx/tfonnx.py

@@ -2419,19 +2419,6 @@ def tensorflow_onnx_mapping(g, continue_on_error, custom_op_handlers):
     return mapped_op, unmapped_op
 
 
-def run_pre_rewriters(g, rewriters):
-    ops = g.get_nodes()
-    for rewrite in rewriters:
-        ops = rewrite(g, ops)
-        g.set_nodes(ops)
-    for node in ops:
-        body_graphs = node.get_body_graphs()
-        if body_graphs:
-            for attr, b_g in body_graphs.items():
-                log.debug("start rewriting subgraph of %s's attribute %s", node.name, attr)
-                run_pre_rewriters(b_g, rewriters)
-
-
 def transpose_inputs(ctx, inputs_as_nchw):
     """Insert a transpose from NHWC to NCHW on model input on users request."""
     ops = []
@@ -2488,17 +2475,18 @@ def topological_sort(g, continue_on_error):
             pass
 
 
-def run_late_rewriters(g, funcs, continue_on_error):
-    if g.contained_graphs:
-        for dict_val in g.contained_graphs.values():
-            for attr_name, b_g in dict_val.items():
-                run_late_rewriters(b_g, funcs, attr_name)
-
-    topological_sort(g, continue_on_error)
+def run_rewriters(g, funcs, continue_on_error=False, need_sort=True):
+    if need_sort:
+        topological_sort(g, continue_on_error)
     for func in funcs:
         ops = func(g, g.get_nodes())
         g.set_nodes(ops)
 
+    if g.contained_graphs:
+        for dict_val in g.contained_graphs.values():
+            for attr_name, b_g in dict_val.items():
+                run_rewriters(b_g, funcs, attr_name)
+
 
 def process_tf_graph(tf_graph, continue_on_error=False, verbose=False, target=None,
                      opset=None, custom_op_handlers=None, custom_rewriter=None,
@@ -2566,7 +2554,7 @@ def process_tf_graph(tf_graph, continue_on_error=False, verbose=False, target=No
     if custom_rewriter is not None:
         rewriters.extend(custom_rewriter)
 
-    run_pre_rewriters(g, rewriters)
+    run_rewriters(g, rewriters, need_sort=False)
 
     # some nodes may already copied into inner Graph, so remove them from main Graph.
     g.delete_unused_nodes(output_names)
@@ -2583,7 +2571,7 @@ def process_tf_graph(tf_graph, continue_on_error=False, verbose=False, target=No
     if TARGET_RS6 in target:
         late_rewriters.append(rewrite_incomplete_type_support_rs6)
     if late_rewriters:
-        run_late_rewriters(g, late_rewriters, continue_on_error)
+        run_rewriters(g, late_rewriters, continue_on_error, True)
 
     # onnx requires topological sorting
     topological_sort(g, continue_on_error)