From 3044aa874353f9d8d4327941d762a92d09460cf0 Mon Sep 17 00:00:00 2001
From: Adrian Lizarraga <adlizarraga@microsoft.com>
Date: Tue, 28 May 2024 21:27:54 -0700
Subject: [PATCH] [Quant tool] Extend support for QDQ type conversion at graph
 output (#20841)

### Description
Allows mixed-precision overrides that adds a QDQ quantization type
conversion sequence at a graph output that **is not** consumed by other
nodes. This is not a common use-case but should handle it instead of
raising an error.

#### Example
Original model

![image](https://github.com/microsoft/onnxruntime/assets/19691973/4c9c3bb0-4ca1-4213-9259-9d0506ed22f2)

mixed-precision overrides:
```python
        mixed_prec_overrides = {
            "input_0": [{"quant_type": QuantType.QUInt16}],
            "op_0_out": [
                {
                    "quant_type": QuantType.QUInt16,
                    "convert": {"quant_type": QuantType.QUInt8},
                }
            ],
        }
        quantize_static(
            float_model_path,
            qdq_model_path,
            data_reader,
            quant_format=QuantFormat.QDQ,
            activation_type=QuantType.QUInt8,
            op_types_to_quantize=[node.op_type for node in float_model.graph.node],
            extra_options={
                "TensorQuantOverrides": mixed_prec_overrides,
            },
        )
```

QDQ model:

![image](https://github.com/microsoft/onnxruntime/assets/19691973/804fc89b-4a00-43bc-a4ff-21edd6f27e98)

### Motivation and Context
This scenario is arising for certain quantization configurations. Should
handle it gracefully.
---
 .../tools/quantization/qdq_quantizer.py       |  6 ++-
 .../test/python/quantization/test_qdq.py      | 52 +++++++++++++++++++
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/onnxruntime/python/tools/quantization/qdq_quantizer.py b/onnxruntime/python/tools/quantization/qdq_quantizer.py
index b7f4b06904..c368d887fd 100644
--- a/onnxruntime/python/tools/quantization/qdq_quantizer.py
+++ b/onnxruntime/python/tools/quantization/qdq_quantizer.py
@@ -636,8 +636,12 @@ class QDQQuantizer(BaseQuantizer):
                                 |        +-> <Graph output>
                                 |
                                 +-> DQ1' ---> Q2 ---> DQ2 ---> <Consumers of converted type>
+
+        5) Tensor T is a graph output that is not consumed by any other nodes.
+
+            <Producer> ---> Q1 ---> DQ1 ---> Q2 ---> DQ2 ---> <Graph output>
         """
-        tensor_recv_nodes = set([node.name for node in self.tensor_to_its_receiving_nodes[tensor_name]])
+        tensor_recv_nodes = set([node.name for node in self.tensor_to_its_receiving_nodes.get(tensor_name, [])])
 
         if (
             self.dedicated_qdq_pair
diff --git a/onnxruntime/test/python/quantization/test_qdq.py b/onnxruntime/test/python/quantization/test_qdq.py
index 8a7cfda8f0..0cde8c31a5 100644
--- a/onnxruntime/test/python/quantization/test_qdq.py
+++ b/onnxruntime/test/python/quantization/test_qdq.py
@@ -1335,6 +1335,58 @@ class TestQDQMixedPrecision(TestQDQFormat):
 
                 self.assertIn("op_0_out", graph_outputs)
 
+    def test_add_tensor_qdq_ops_case_5(self):
+        """
+        Tensor T is a graph output without any consumers.
+        <Producer> ---> Q1 ---> DQ1 ---> Q2 ---> DQ2 ---> <Graph output>
+        """
+        float_model_path = os.path.join(self._tmp_dir_path, "case_5.onnx")
+        qdq_model_path = os.path.join(self._tmp_dir_path, "case_5.qdq.onnx")
+
+        # Build model with input_0 -> op_0 -> op_0_out
+        # The graph output has no consumers.
+        float_model = self.build_test_model_for_add_qdq_ops(0, True)
+        onnx.save_model(float_model, float_model_path)
+
+        data_reader = self.input_feeds(3, {"input_0": (1, 2, 3)}, np.float32)
+
+        mixed_prec_overrides = {
+            "input_0": [{"quant_type": QuantType.QUInt16}],
+            "op_0_out": [
+                {
+                    "quant_type": QuantType.QUInt16,
+                    "convert": {"quant_type": QuantType.QUInt8},
+                }
+            ],
+        }
+        quantize_static(
+            float_model_path,
+            qdq_model_path,
+            data_reader,
+            quant_format=QuantFormat.QDQ,
+            activation_type=QuantType.QUInt8,
+            op_types_to_quantize=[node.op_type for node in float_model.graph.node],
+            extra_options={
+                "TensorQuantOverrides": mixed_prec_overrides,
+            },
+        )
+
+        # Expect the following QDQ model:
+        # input_0 --> Q_16 --> DQ_16 --> op_0 --> Q_16 --> DQ_16 --> Q_8 --> DQ_8 --> output_0
+        qdq_node_counts = {"QuantizeLinear": 3, "DequantizeLinear": 3}
+        check_op_type_count(self, qdq_model_path, **qdq_node_counts)
+
+        qdq_model = onnx.load_model(qdq_model_path)
+        onnx.checker.check_model(qdq_model, True)
+
+        initializers = {init.name: init for init in qdq_model.graph.initializer}
+
+        # Check zero-point data types
+        orig_zp_init = initializers["op_0_out_zero_point"]
+        self.assertEqual(orig_zp_init.data_type, onnx.TensorProto.UINT16)
+        convert_zp_init = initializers["op_0_out_zero_point_convert"]
+        self.assertEqual(convert_zp_init.data_type, onnx.TensorProto.UINT8)
+
     def build_test_model_1(self, shape):
         """
         Returns the following float32 model.