Add unit tests for DoubleQDQPairsRemover

adrianlizarraga · adrianlizarraga · commit fd4f5a59d020 · 2024-05-22T20:19:07.000-07:00
diff --git a/onnxruntime/core/optimizer/double_qdq_pairs_remover.cc b/onnxruntime/core/optimizer/double_qdq_pairs_remover.cc
@@ -2,6 +2,8 @@
 // Licensed under the MIT License.
 #include "core/optimizer/double_qdq_pairs_remover.h"
 #include <cassert>
+#include <string>
+#include <vector>
 
 #include "core/common/gsl.h"
 #include "core/graph/graph_utils.h"
@@ -88,7 +90,8 @@ static bool FindNewZeroPointAndScale(const Graph& graph, const Node& node1, cons
 // After removing the middle two nodes, the zero point and scale of the final (outer) ops must be recomputed
 // for correctness.
 template <typename ZeroPointType>
-static bool RecomputeOuterQDQZeroPointAndScale(Graph& graph, Node& q1, const Node& dq1, const Node& q2, gsl::span<Node*> dq2s) {
+static bool RecomputeOuterQDQZeroPointAndScale(Graph& graph, Node& q1, const Node& dq1, const Node& q2,
+                                               gsl::span<Node*> dq2s) {
   if (dq2s.empty()) {
     return false;
   }
diff --git a/onnxruntime/test/optimizer/graph_transform_test_builder.h b/onnxruntime/test/optimizer/graph_transform_test_builder.h
@@ -6,6 +6,7 @@
 #include <type_traits>
 #include <vector>
 
+#include "core/common/gsl.h"
 #include "core/common/type_utils.h"
 #include "core/graph/graph.h"
 #include "core/framework/framework_common.h"
@@ -195,13 +196,14 @@ class ModelTestBuilder {
     return &graph_.GetOrCreateNodeArg(name, &type_proto);
   }
 
-  template <typename T>
-  NodeArg* MakeInitializer(const std::vector<int64_t>& shape, const std::vector<T>& data) {
+  // Makes an initializer from the provided shape, element type, and raw_data bytes.
+  NodeArg* MakeInitializer(gsl::span<const int64_t> shape, ONNX_NAMESPACE::TensorProto_DataType elem_type,
+                           gsl::span<const std::byte> raw_data) {
     std::string name = graph_.GenerateNodeArgName("constant");
     ONNX_NAMESPACE::TensorProto tensor_proto;
     tensor_proto.set_name(name);
-    tensor_proto.set_data_type(utils::ToTensorProtoElementType<T>());
-    tensor_proto.set_raw_data(data.data(), data.size() * sizeof(T));
+    tensor_proto.set_data_type(elem_type);
+    tensor_proto.set_raw_data(raw_data.data(), raw_data.size());
 
     for (auto& dim : shape) {
       tensor_proto.add_dims(dim);
@@ -212,6 +214,12 @@ class ModelTestBuilder {
     return &graph_.GetOrCreateNodeArg(name, nullptr);
   }
 
+  template <typename T>
+  NodeArg* MakeInitializer(const std::vector<int64_t>& shape, const std::vector<T>& data) {
+    gsl::span<const std::byte> raw_data = ReinterpretAsSpan<const std::byte, const T>(data);
+    return MakeInitializer(shape, utils::ToTensorProtoElementType<T>(), raw_data);
+  }
+
   // Special handle for std::vector<bool>.
   NodeArg* MakeInitializerBool(const std::vector<int64_t>& shape, const std::vector<bool>& data) {
     std::string name = graph_.GenerateNodeArgName("constant");
@@ -342,6 +350,57 @@ class ModelTestBuilder {
     return AddNode("QuantizeLinear", input_args, {output_arg}, domain, attributes);
   }
 
+  static std::vector<std::byte> GetZeroPointBytes(int64_t zero_point, ONNX_NAMESPACE::TensorProto_DataType type) {
+    switch (type) {
+      case ONNX_NAMESPACE::TensorProto_DataType_INT8: {
+        int8_t val = static_cast<int8_t>(zero_point);
+        auto span = ReinterpretAsSpan<const std::byte, const int8_t>(gsl::make_span(&val, 1));
+        return std::vector<std::byte>(span.begin(), span.end());
+      }
+      case ONNX_NAMESPACE::TensorProto_DataType_UINT8: {
+        uint8_t val = static_cast<uint8_t>(zero_point);
+        auto span = ReinterpretAsSpan<const std::byte, const uint8_t>(gsl::make_span(&val, 1));
+        return std::vector<std::byte>(span.begin(), span.end());
+      }
+      case ONNX_NAMESPACE::TensorProto_DataType_INT16: {
+        int16_t val = static_cast<int16_t>(zero_point);
+        auto span = ReinterpretAsSpan<const std::byte, const int16_t>(gsl::make_span(&val, 1));
+        return std::vector<std::byte>(span.begin(), span.end());
+      }
+      case ONNX_NAMESPACE::TensorProto_DataType_UINT16: {
+        uint16_t val = static_cast<uint16_t>(zero_point);
+        auto span = ReinterpretAsSpan<const std::byte, const uint16_t>(gsl::make_span(&val, 1));
+        return std::vector<std::byte>(span.begin(), span.end());
+      }
+      case ONNX_NAMESPACE::TensorProto_DataType_INT32: {
+        int32_t val = static_cast<int32_t>(zero_point);
+        auto span = ReinterpretAsSpan<const std::byte, const int32_t>(gsl::make_span(&val, 1));
+        return std::vector<std::byte>(span.begin(), span.end());
+      }
+      default:
+        ORT_THROW("Unhandled zero-point type ", type, ".");
+    }
+  }
+
+  // Adds a Q node with a runtime configurable zero-point type.
+  // Takes in an int64_t zero_point value, which is large enough to represent all ONNX zero-point types.
+  Node& AddQuantizeLinearNode(NodeArg* input_arg,
+                              float input_scale,
+                              int64_t input_zero_point,
+                              ONNX_NAMESPACE::TensorProto_DataType zero_point_type,
+                              NodeArg* output_arg,
+                              bool use_ms_domain = false) {
+    std::vector<NodeArg*> input_args;
+    input_args.push_back(input_arg);
+    input_args.push_back(MakeScalarInitializer<float>(input_scale));
+
+    std::vector<std::byte> zp_bytes = GetZeroPointBytes(input_zero_point, zero_point_type);
+    input_args.push_back(MakeInitializer({}, zero_point_type, zp_bytes));
+
+    std::string domain = use_ms_domain ? kMSDomain : "";
+    return AddNode("QuantizeLinear", input_args, {output_arg}, domain);
+  }
+
   template <typename T>
   typename std::enable_if<IsTypeDequantLinearCompatible<T>::value, Node&>::type
   AddDequantizeLinearNode(NodeArg* input_arg,
@@ -400,6 +459,25 @@ class ModelTestBuilder {
     return AddNode("DequantizeLinear", input_args, {output_arg}, domain, attributes);
   }
 
+  // Adds a DQ node with a runtime configurable zero-point type.
+  // Takes in an int64_t zero_point value, which is large enough to represent all ONNX zero-point types.
+  Node& AddDequantizeLinearNode(NodeArg* input_arg,
+                                float input_scale,
+                                int64_t input_zero_point,
+                                ONNX_NAMESPACE::TensorProto_DataType zero_point_type,
+                                NodeArg* output_arg,
+                                bool use_ms_domain = false) {
+    std::vector<NodeArg*> input_args;
+    input_args.push_back(input_arg);
+    input_args.push_back(MakeScalarInitializer<float>(input_scale));
+
+    std::vector<std::byte> zp_bytes = GetZeroPointBytes(input_zero_point, zero_point_type);
+    input_args.push_back(MakeInitializer({}, zero_point_type, zp_bytes));
+
+    std::string domain = use_ms_domain ? kMSDomain : "";
+    return AddNode("DequantizeLinear", input_args, {output_arg}, domain);
+  }
+
   template <typename TWeight>
   Node& AddQLinearConvNode(NodeArg* input_arg,
                            float input_scale,
diff --git a/onnxruntime/test/optimizer/qdq_test_utils.cc b/onnxruntime/test/optimizer/qdq_test_utils.cc
@@ -164,5 +164,48 @@ std::vector<std::string> GetNodeOpTypesInTopologicalOrder(const Graph& graph, bo
   return op_types;
 }
 
+GetQDQTestCaseFn BuildDoubleQDQTestCaseWithDuplicateLastDQs(
+    const std::vector<int64_t>& input_shape,
+    const std::vector<float>& input_data,
+    const std::vector<int64_t>& zero_points,
+    const std::vector<ONNX_NAMESPACE::TensorProto_DataType>& zero_point_types,
+    const std::vector<float>& scales,
+    int graph_output_index,
+    bool use_contrib_qdq) {
+  const size_t num_nodes = zero_points.size();
+  bool valid_inputs = (num_nodes >= 4) &&
+                      (zero_point_types.size() == num_nodes) &&
+                      (scales.size() == num_nodes) &&
+                      (graph_output_index >= 0 && graph_output_index < 4);
+  if (!valid_inputs) {
+    ORT_THROW("Invalid inputs for call to BuildDoubleQDQTestCaseWithDuplicateLastDQs()");
+  }
+
+  return [=](ModelTestBuilder& builder) {
+    auto* input_arg = builder.MakeInput<float>(input_shape, input_data);
+    std::vector<NodeArg*> node_outputs(num_nodes);
+
+    for (size_t i = 0; i < num_nodes; i++) {
+      if (i == graph_output_index || i >= 3) {
+        node_outputs[i] = builder.MakeOutput();
+      } else {
+        node_outputs[i] = builder.MakeIntermediate();
+      }
+    }
+
+    builder.AddQuantizeLinearNode(input_arg, scales[0], zero_points[0], zero_point_types[0], node_outputs[0],
+                                  use_contrib_qdq);
+    builder.AddDequantizeLinearNode(node_outputs[0], scales[1], zero_points[1], zero_point_types[1], node_outputs[1],
+                                    use_contrib_qdq);
+    builder.AddQuantizeLinearNode(node_outputs[1], scales[2], zero_points[2], zero_point_types[2], node_outputs[2],
+                                  use_contrib_qdq);
+
+    for (size_t i = 3; i < num_nodes; i++) {
+      builder.AddDequantizeLinearNode(node_outputs[2], scales[i], zero_points[i], zero_point_types[i],
+                                      node_outputs[i], use_contrib_qdq);
+    }
+  };
+}
+
 }  // namespace test
 }  // namespace onnxruntime
diff --git a/onnxruntime/test/optimizer/qdq_test_utils.h b/onnxruntime/test/optimizer/qdq_test_utils.h
@@ -460,6 +460,27 @@ GetQDQTestCaseFn BuildDoubleQDQTestCases(Type1 zp_1, Type2 zp_2, Type3 zp_3, Typ
   };
 }
 
+/// <summary>
+/// Returns a function that builds a model with a double QDQ sequence (Q1 -> DQ1 -> Q2 -> DQ2*),
+/// where DQ2 can be repeated. Must provide at least 4 zero-point and scale values.
+/// </summary>
+/// <param name="input_shape">Shape of input float data.</param>
+/// <param name="input_data">Input float data.</param>
+/// <param name="zero_points">Ordered list of zero-point values for each node in the sequence.</param>
+/// <param name="zero_point_types">Ordered list of zero-point types for each node in the sequence.</param>
+/// <param name="zero_points">Ordered list of scale values for each node in the sequence.</param>
+/// <param name="graph_output_index">Index of the node that provides a graph output.</param>
+/// <param name="use_contrib_qdq">Set to true to use the 'com.microsoft' domain for Q and DQ ops.</param>
+/// <returns>A function for building the model</returns>
+GetQDQTestCaseFn BuildDoubleQDQTestCaseWithDuplicateLastDQs(
+    const std::vector<int64_t>& input_shape,
+    const std::vector<float>& input_data,
+    const std::vector<int64_t>& zero_points,
+    const std::vector<ONNX_NAMESPACE::TensorProto_DataType>& zero_point_types,
+    const std::vector<float>& scales,
+    int graph_output_index,
+    bool use_contrib_qdq = false);
+
 template <typename T>
 GetQDQTestCaseFn BuildDoubleQDQWithoutLastOutput(int output_index, bool use_contrib_qdq = false) {
   return [=](ModelTestBuilder& builder) {
diff --git a/onnxruntime/test/optimizer/qdq_transformer_test.cc b/onnxruntime/test/optimizer/qdq_transformer_test.cc
@@ -7,6 +7,7 @@
 #include "core/graph/model.h"
 #include "core/graph/onnx_protobuf.h"
 #include "core/mlas/inc/mlas.h"
+#include "core/optimizer/double_qdq_pairs_remover.h"
 #include "core/optimizer/qdq_transformer/qdq_final_cleanup.h"
 #include "core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.h"
 #include "core/optimizer/qdq_transformer/selectors_actions/qdq_selector_action_transformer.h"
@@ -1233,6 +1234,96 @@ TEST(QDQTransformerTests, DoubleQDQ_Without_Last_Node_Being_Output) {
   RunDoubleQDQWithoutLastNodeBeingOutput<uint16_t>(3, 1, 1, !use_ms_qdq, 21);
 }
 
+static void RunDoubleQDQWithDuplicateLastDQs(int expected_Q_count, int expected_DQ_count,
+                                             const std::vector<int64_t>& input_shape,
+                                             const std::vector<float>& input_data,
+                                             const std::vector<int64_t>& zero_points,
+                                             const std::vector<ONNX_NAMESPACE::TensorProto_DataType>& zero_point_types,
+                                             const std::vector<float>& scales,
+                                             int graph_output_index,
+                                             bool use_contrib_qdq = false,
+                                             int opset = 19) {
+  auto graph_checker = [&](InferenceSessionWrapper& session) {
+    auto op_to_count = CountOpsInGraph(session.GetGraph());
+    const QDQOpKeys qdq_keys = GetQDQOpKeys(use_contrib_qdq);
+    EXPECT_EQ(op_to_count[qdq_keys.quantize_linear], expected_Q_count);
+    EXPECT_EQ(op_to_count[qdq_keys.dequantize_linear], expected_DQ_count);
+  };
+
+  auto model_build_fn = BuildDoubleQDQTestCaseWithDuplicateLastDQs(input_shape, input_data, zero_points,
+                                                                   zero_point_types, scales, graph_output_index,
+                                                                   use_contrib_qdq);
+  TransformerTester(model_build_fn,
+                    graph_checker,
+                    TransformerLevel::Default,
+                    TransformerLevel::Level1,
+                    opset,
+                    /*per_sample_tolerance*/ 0.0,
+                    /*relative_per_sample_tolerance*/ 0.0,
+                    std::make_unique<DoubleQDQPairsRemover>());
+}
+
+// Test QDQDoublePairsRemover when the sequence ends with duplicate DQs.
+TEST(QDQTransformerTests, DoubleQDQPairsRemover_DuplicateLastDQs) {
+  std::vector<int64_t> shape = {1, 2, 2, 2};
+  std::vector<float> input_data = {-3.0f, -2.0f, -1.0f, 0.0f, 0.5f, 1.0f, 2.0f, 3.0f};
+
+  constexpr auto int8_type = ONNX_NAMESPACE::TensorProto_DataType_INT8;
+  constexpr auto uint8_type = ONNX_NAMESPACE::TensorProto_DataType_UINT8;
+  constexpr auto int16_type = ONNX_NAMESPACE::TensorProto_DataType_INT16;
+  constexpr auto uint16_type = ONNX_NAMESPACE::TensorProto_DataType_UINT16;
+  std::vector<ONNX_NAMESPACE::TensorProto_DataType> quant_types = {int8_type, uint8_type, int16_type, uint16_type};
+
+  // Input graph:
+  // input -> Q1 -> DQ1 -> Q2 --+--> DQ2 -> output0
+  //                            |
+  //                            ...
+  //                            |
+  //                            +--> DQ2'' -> outputN
+  // Expected graph after DoubleQDQPairsRemover:
+  // input -> Q1 --+--> DQ2 -> output0
+  //               |
+  //               ...
+  //               |
+  //               +--> DQ2'' -> outputN
+  for (auto quant_type : quant_types) {
+    for (size_t num_dq2s = 1; num_dq2s <= 1; num_dq2s++) {
+      const size_t num_nodes = 3 + num_dq2s;
+      std::vector<int64_t> zp_vals(num_nodes, 1);
+      std::vector<ONNX_NAMESPACE::TensorProto_DataType> zp_types(num_nodes, quant_type);
+      std::vector<float> scale_vals(num_nodes, 0.1f);
+
+      const int expected_q_nodes = 1;
+      const int expected_dq_nodes = static_cast<int>(num_dq2s);
+      RunDoubleQDQWithDuplicateLastDQs(expected_q_nodes, expected_dq_nodes, shape, input_data, zp_vals, zp_types,
+                                       scale_vals, 3, false, 21);
+      RunDoubleQDQWithDuplicateLastDQs(expected_q_nodes, expected_dq_nodes, shape, input_data, zp_vals, zp_types,
+                                       scale_vals, 3, quant_type == int16_type || quant_type == uint16_type, 19);
+    }
+  }
+
+  // Should not remove QDQ pair because the middle nodes produce a graph output.
+  for (auto quant_type : quant_types) {
+    for (int output_index = 0; output_index < 3; output_index++) {
+      for (size_t num_dq2s = 1; num_dq2s <= 1; num_dq2s++) {
+        const size_t num_nodes = 3 + num_dq2s;
+        std::vector<int64_t> zp_vals(num_nodes, 1);
+        std::vector<ONNX_NAMESPACE::TensorProto_DataType> zp_types(num_nodes, quant_type);
+        std::vector<float> scale_vals(num_nodes, 0.1f);
+
+        const int expected_q_nodes = 2;
+        int expected_dq_nodes = 1 + static_cast<int>(num_dq2s);
+        if (output_index == 1) {
+          // EnsureUniqueDQ pass will create a duplicate DQ if it produces a graph output.
+          expected_dq_nodes += 1;
+        }
+        RunDoubleQDQWithDuplicateLastDQs(expected_q_nodes, expected_dq_nodes, shape, input_data, zp_vals, zp_types,
+                                         scale_vals, output_index, false, 21);
+      }
+    }
+  }
+}
+
 // Runs a test that checks if DQ -> Split -> Q (many) is replaced with just Split.
 template <typename QuantType>
 static void RunDropSplitQDQTestCase(const std::vector<int64_t>& input_shape, int64_t axis,
