From c4ebc608702c001b53d54ae32d05d3deb194932d Mon Sep 17 00:00:00 2001
From: Yufeng Li <liyufeng1987@gmail.com>
Date: Tue, 30 Mar 2021 09:01:15 -0700
Subject: [PATCH] sort quantized nodes in topo logical order (#7172)

---
 .../python/tools/quantization/onnx_model.py   | 39 +++++++++++++++++++
 .../test/python/quantization/op_test_utils.py |  8 ++++
 .../test/python/quantization/test_qdq.py      |  6 +--
 3 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/onnxruntime/python/tools/quantization/onnx_model.py b/onnxruntime/python/tools/quantization/onnx_model.py
index a38d7e69d7..325a393995 100644
--- a/onnxruntime/python/tools/quantization/onnx_model.py
+++ b/onnxruntime/python/tools/quantization/onnx_model.py
@@ -1,4 +1,5 @@
 import onnx
+import itertools
 from .quant_utils import find_by_name
 from pathlib import Path
 
@@ -197,6 +198,7 @@ class ONNXModel:
         '''
         Save model to external data, which is needed for model size > 2GB
         '''
+        self.topological_sort()
         if use_external_data_format:
             onnx.external_data_helper.convert_model_to_external_data(self.model,
                                                                      all_tensors_to_one_file=True,
@@ -254,3 +256,40 @@ class ONNXModel:
             if output.name == output_name:
                 return True
         return False
+
+    def topological_sort(self):
+        deps_count = [0]*len(self.nodes()) # dependency count of each node
+        deps_to_nodes = {} # input to node indice
+        for node_idx, node in enumerate(self.nodes()):
+            deps_count[node_idx] = len(node.input)
+            for input_name in node.input:
+                if input_name not in deps_to_nodes:
+                    deps_to_nodes[input_name] = [node_idx]
+                else:
+                    deps_to_nodes[input_name].append(node_idx)
+
+
+        # initialize sorted_nodes
+        sorted_nodes = []
+        for input in itertools.chain(self.initializer(), self.model.graph.input):
+            for node_idx in deps_to_nodes[input.name]:
+                deps_count[node_idx] = deps_count[node_idx] - 1
+                if deps_count[node_idx] == 0:
+                    sorted_nodes.append(self.nodes()[node_idx])
+
+        s = 0
+        e = len(sorted_nodes)
+
+        while s < e:
+            for output in sorted_nodes[s].output:
+                if output in deps_to_nodes:
+                    for node_idx in deps_to_nodes[output]:
+                        deps_count[node_idx] = deps_count[node_idx] - 1
+                        if deps_count[node_idx] == 0:
+                            sorted_nodes.append(self.nodes()[node_idx])
+                            e = e + 1
+            s = s + 1
+
+        assert(e == len(self.graph().node)), "Graph is not a DAG"
+        self.graph().ClearField('node')
+        self.graph().node.extend(sorted_nodes)
\ No newline at end of file
diff --git a/onnxruntime/test/python/quantization/op_test_utils.py b/onnxruntime/test/python/quantization/op_test_utils.py
index 31c2a32718..27d5e4364d 100644
--- a/onnxruntime/test/python/quantization/op_test_utils.py
+++ b/onnxruntime/test/python/quantization/op_test_utils.py
@@ -20,6 +20,14 @@ class TestDataFeeds(CalibrationDataReader):
     def rewind(self):
         self.iter_next = iter(self.data_feeds)
 
+def check_op_type_order(testcase, model_path, ops):
+    model = onnx.load(Path(model_path))
+    testcase.assertEqual(len(ops), len(model.graph.node), 'op count is not same')
+    for node_idx, node in enumerate(model.graph.node):
+        testcase.assertEqual(
+            ops[node_idx],
+            node.op_type,
+            'op {} is not in order. Expected: {}, Actual: {}'.format(node_idx, ops[node_idx], node.op_type))
 
 def check_op_type_count(testcase, model_path, **kwargs):
     model = onnx.load(Path(model_path))
diff --git a/onnxruntime/test/python/quantization/test_qdq.py b/onnxruntime/test/python/quantization/test_qdq.py
index 895aed8ece..c7f682ee81 100644
--- a/onnxruntime/test/python/quantization/test_qdq.py
+++ b/onnxruntime/test/python/quantization/test_qdq.py
@@ -11,7 +11,7 @@ import onnx
 import numpy as np
 from onnx import helper, TensorProto
 from onnxruntime.quantization import quantize_static, QuantType, QuantFormat
-from op_test_utils import TestDataFeeds, check_model_correctness, check_op_type_count
+from op_test_utils import TestDataFeeds, check_model_correctness, check_op_type_count, check_op_type_order
 
 class TestQDQFormat(unittest.TestCase):
     def input_feeds(self, n, name2shape):
@@ -163,8 +163,8 @@ class TestQDQFormatConvClip(TestQDQFormat):
                         reduce_range = per_channel
                         )
         data_reader.rewind()
-        qdq_nodes = {'Conv': 1, 'QuantizeLinear': 1, 'DequantizeLinear': 2}
-        check_op_type_count(self, model_int8_qdq_path, **qdq_nodes)
+        #topo sort check
+        check_op_type_order(self, model_int8_qdq_path, ['DequantizeLinear', 'QuantizeLinear', 'DequantizeLinear', 'Conv', 'Clip'])
         check_model_correctness(self, model_fp32_path, model_int8_qdq_path, data_reader.get_next())
 
         data_reader.rewind()