From 3725d0211fb4ac26197de71f7f7f05443dbb58a0 Mon Sep 17 00:00:00 2001
From: Yufeng Li <liyufeng1987@gmail.com>
Date: Thu, 13 May 2021 23:58:15 -0700
Subject: [PATCH] support maxpool QDQ fusion from opset 12 (#7693)

---
 .../qdq_transformer/qdq_simple_ops.cc         | 17 +++++++++++++-
 .../optimizer/qdq_transformer/registry.cc     |  4 ++--
 .../test/optimizer/qdq_transformer_test.cc    | 23 ++++++++++++-------
 3 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/onnxruntime/core/optimizer/qdq_transformer/qdq_simple_ops.cc b/onnxruntime/core/optimizer/qdq_transformer/qdq_simple_ops.cc
index e45328ca77..96d8b0b0d9 100644
--- a/onnxruntime/core/optimizer/qdq_transformer/qdq_simple_ops.cc
+++ b/onnxruntime/core/optimizer/qdq_transformer/qdq_simple_ops.cc
@@ -4,6 +4,7 @@
 #include <vector>
 
 #include "core/graph/graph.h"
+#include "core/graph/graph_utils.h"
 #include "core/optimizer/qdq_transformer/qdq_op_transformer.h"
 #include "core/optimizer/qdq_transformer/qdq_util.h"
 #include "core/optimizer/qdq_transformer/registry.h"
@@ -39,7 +40,21 @@ class QDQSimpleTransformer : public QDQOperatorTransformer {
   }
 };
 
-DEFINE_QDQ_CREATOR(MaxPool, QDQSimpleTransformer)
+class QDQMaxPoolTransformer : public QDQSimpleTransformer {
+ public:
+  QDQMaxPoolTransformer(Node& node, Graph& graph) : QDQSimpleTransformer(node, graph) {}
+
+ protected:
+  bool Check(const std::vector<const Node*>& dq_nodes, const std::vector<const Node*>& q_nodes) const override {
+    if (!QDQSimpleTransformer::Check(dq_nodes, q_nodes)) {
+      return false;
+    }
+
+    return graph_utils::IsSupportedOptypeVersionAndDomain(node_, "MaxPool", {12});
+  }
+};
+
+DEFINE_QDQ_CREATOR(MaxPool, QDQMaxPoolTransformer)
 DEFINE_QDQ_CREATOR(Reshape, QDQSimpleTransformer)
 DEFINE_QDQ_CREATOR(Gather, QDQSimpleTransformer)
 DEFINE_QDQ_CREATOR(Transpose, QDQSimpleTransformer)
diff --git a/onnxruntime/core/optimizer/qdq_transformer/registry.cc b/onnxruntime/core/optimizer/qdq_transformer/registry.cc
index 49f6e4e0ad..8af6122897 100644
--- a/onnxruntime/core/optimizer/qdq_transformer/registry.cc
+++ b/onnxruntime/core/optimizer/qdq_transformer/registry.cc
@@ -5,7 +5,7 @@
 
 namespace onnxruntime {
 DECLARE_QDQ_CREATOR(Conv, QDQConvTransformer);
-DECLARE_QDQ_CREATOR(MaxPool, QDQSimpleTransformer);
+DECLARE_QDQ_CREATOR(MaxPool, QDQMaxPoolTransformer);
 DECLARE_QDQ_CREATOR(Reshape, QDQSimpleTransformer);
 DECLARE_QDQ_CREATOR(Gather, QDQSimpleTransformer);
 DECLARE_QDQ_CREATOR(Transpose, QDQSimpleTransformer);
@@ -17,7 +17,7 @@ DECLARE_QDQ_CREATOR(Concat, QDQConcatTransformer);
 
 std::unordered_map<std::string, QDQRegistry::QDQTransformerCreator> QDQRegistry::qdqtransformer_creators_{
     REGISTER_QDQ_CREATOR(Conv, QDQConvTransformer),
-    REGISTER_QDQ_CREATOR(MaxPool, QDQSimpleTransformer),
+    REGISTER_QDQ_CREATOR(MaxPool, QDQMaxPoolTransformer),
     REGISTER_QDQ_CREATOR(Reshape, QDQSimpleTransformer),
     REGISTER_QDQ_CREATOR(Gather, QDQSimpleTransformer),
     REGISTER_QDQ_CREATOR(Transpose, QDQSimpleTransformer),
diff --git a/onnxruntime/test/optimizer/qdq_transformer_test.cc b/onnxruntime/test/optimizer/qdq_transformer_test.cc
index 32aa1b9da8..c35b31b043 100644
--- a/onnxruntime/test/optimizer/qdq_transformer_test.cc
+++ b/onnxruntime/test/optimizer/qdq_transformer_test.cc
@@ -81,7 +81,7 @@ TEST(QDQTransformerTests, Conv) {
 }
 
 TEST(QDQTransformerTests, ConvMaxPoolReshape_UInt8) {
-  auto test_case = [&](const std::vector<int64_t>& input_shape, const std::vector<int64_t>& weights_shape) {
+  auto test_case = [&](const std::vector<int64_t>& input_shape, const std::vector<int64_t>& weights_shape, int opset_version) {
     auto build_test_case = [&](ModelTestBuilder& builder) {
       auto* input_arg = builder.MakeInput<float>(input_shape, -1.f, 1.f);
       auto* output_arg = builder.MakeOutput();
@@ -118,16 +118,23 @@ TEST(QDQTransformerTests, ConvMaxPoolReshape_UInt8) {
       EXPECT_EQ(op_to_count["QLinearConv"], 1);
       EXPECT_EQ(op_to_count["MaxPool"], 1);
       EXPECT_EQ(op_to_count["Reshape"], 1);
-      EXPECT_EQ(op_to_count["QuantizeLinear"], 1);
-      EXPECT_EQ(op_to_count["DequantizeLinear"], 0);
+      EXPECT_EQ(op_to_count["QuantizeLinear"], opset_version < 12 ? 2 : 1);
+      EXPECT_EQ(op_to_count["DequantizeLinear"], opset_version < 12 ? 1 : 0);
     };
 
-    TransformerTester(build_test_case, check_mp_reshape_graph, TransformerLevel::Level1, TransformerLevel::Level2);
+    TransformerTester(build_test_case,
+                      check_mp_reshape_graph,
+                      TransformerLevel::Level1,
+                      TransformerLevel::Level2,
+                      opset_version);
   };
 
-  test_case({1, 12, 37}, {32, 12, 5});
-  test_case({1, 23, 13, 13}, {30, 23, 3, 3});
-  test_case({1, 22, 11, 13, 15}, {30, 22, 5, 3, 3});
+  test_case({1, 12, 37}, {32, 12, 5}, 11);
+  test_case({1, 12, 37}, {32, 12, 5}, 12);
+  test_case({1, 23, 13, 13}, {30, 23, 3, 3}, 11);
+  test_case({1, 23, 13, 13}, {30, 23, 3, 3}, 12);
+  test_case({1, 22, 11, 13, 15}, {30, 22, 5, 3, 3}, 11);
+  test_case({1, 22, 11, 13, 15}, {30, 22, 5, 3, 3}, 12);
 }
 
 TEST(QDQTransformerTests, ConvMaxPoolReshape_Int8) {
@@ -317,7 +324,7 @@ TEST(QDQTransformerTests, Transpose_No_Fusion) {
       builder.AddDequantizeLinearNode<int8_t>(input1_arg, .003f, 1, dq_output);
 
       // add Transpose
-      auto* transpose_output = builder.MakeOutput(); // transpose output is graph output
+      auto* transpose_output = builder.MakeOutput();  // transpose output is graph output
       Node& transpose_node = builder.AddNode("Transpose", {dq_output}, {transpose_output});
       transpose_node.AddAttribute("perm", perms);