Add nhwc transformer support and unittest for qlinear concat. (#7587)

onnxruntime/core/optimizer/nhwc_transformer.cc

@@ -40,6 +40,25 @@ class NhwcTransformerImpl {
     return (it != nhwc_args_.end()) ? it->second.get() : nullptr;
   }
 
+  static bool NchwAxisToNhwc(int64_t& axis, int rank) {
+    if (axis < -rank || axis >= rank) {
+      return false;
+    }
+    bool is_negative_axis = (axis < 0);
+    if (is_negative_axis) {
+      axis = axis + rank;
+    }
+    if (axis == 1) {
+      axis = rank - 1;
+    } else if (axis > 1) {
+      axis = axis - 1;
+    }
+    if (is_negative_axis) {
+      axis = axis - rank;
+    }
+    return true;
+  }
+
   size_t RemoveOutputEdge(Node& node, size_t output_index);
   void CreateNhwcArgument(Node& node, Node& nhwc_node, int rank, size_t output_index);
   void CreateNhwcArgument(Node& node, Node& nhwc_node, int rank);
@@ -52,6 +71,7 @@ class NhwcTransformerImpl {
   void TransformMaxPool(Node& node);
   void TransformSplit(Node& node);
   void TransformPad(Node& node);
+  void TransformQLinearConcat(Node& node);
 
   Graph& graph_;
 
@@ -228,6 +248,45 @@ void NhwcTransformerImpl::TransformQLinearActivation(Node& node) {
   CreateNhwcArgument(node, node, nhwc_input->rank_);
 }
 
+void NhwcTransformerImpl::TransformQLinearConcat(Node& node) {
+  auto& input_defs = node.MutableInputDefs();
+
+  int rank = 0;
+  for (size_t def_index = 2, def_count = input_defs.size(); def_index < def_count; def_index += 3) {
+    auto* nhwc_input = LookupNhwcArgument(input_defs[def_index]);
+    if (nhwc_input == nullptr || (def_index > 2 && nhwc_input->rank_ != rank)) {
+      return;
+    }
+    if (def_index == 2) {
+      rank = nhwc_input->rank_;
+    }
+  }
+
+  // Change the axis attribute accordingly for NCHW to NHWC model.
+  const auto* axis_attr = graph_utils::GetNodeAttribute(node, "axis");
+  if (axis_attr != nullptr && utils::HasInt(*axis_attr)) {
+    int64_t axis = axis_attr->i();
+    if (!NchwAxisToNhwc(axis, rank)) {
+      // direct return on invalid axis
+      return;
+    }
+    node.AddAttribute("axis", axis);
+  } else {
+    // direct return on invalid node
+    return;
+  }
+
+  for (size_t def_index = 2, def_count = input_defs.size(); def_index < def_count; def_index += 3) {
+    // Update the node to directly use the NHWC inputs and decrement the original
+    // use counts of the NHWC inputs.
+    auto* nhwc_input = LookupNhwcArgument(input_defs[def_index]);
+    input_defs[def_index] = nhwc_input->nhwc_arg_;
+    nhwc_input->remaining_original_uses_--;
+  }
+
+  CreateNhwcArgument(node, node, rank);
+}
+
 void NhwcTransformerImpl::TransformQLinearGlobalAveragePool(Node& node) {
   auto& input_defs = node.MutableInputDefs();
 
@@ -300,20 +359,13 @@ void NhwcTransformerImpl::TransformSplit(Node& node) {
   const auto* axis_attr = graph_utils::GetNodeAttribute(node, "axis");
   if (axis_attr != nullptr && utils::HasInt(*axis_attr)) {
     int64_t axis = axis_attr->i();
-    if (axis < -nhwc_input->rank_ || axis >= nhwc_input->rank_) {
+    if (!NchwAxisToNhwc(axis, nhwc_input->rank_)) {
       // direct return on invalid axis
       return;
     }
-    if (axis < 0) {
-      axis = axis + nhwc_input->rank_;
-    }
-    if (axis == 1) {
-      axis = nhwc_input->rank_ - 1;
-    } else if (axis > 1) {
-      axis = axis - 1;
-    }
     node.AddAttribute("axis", axis);
   }
+  // default axis is 0 when attribute not exists, which do not need update
 
   // Update the node to directly use the NHWC inputs and decrement the original
   // use counts of the NHWC inputs.
@@ -377,6 +429,8 @@ void NhwcTransformerImpl::Transform(Node& node) {
     TransformQLinearActivation(node);
   } else if (graph_utils::IsSupportedOptypeVersionAndDomain(node, "QLinearGlobalAveragePool", {1}, kMSDomain)) {
     TransformQLinearGlobalAveragePool(node);
+  } else if (graph_utils::IsSupportedOptypeVersionAndDomain(node, "QLinearConcat", {1}, kMSDomain)) {
+    TransformQLinearConcat(node);
   } else if (graph_utils::IsSupportedOptypeVersionAndDomain(node, "MaxPool", {12})) {
     TransformMaxPool(node);
   } else if (graph_utils::IsSupportedOptypeVersionAndDomain(node, "Split", {2, 11, 13})) {

onnxruntime/test/optimizer/graph_transform_test_builder.h

@@ -213,6 +213,22 @@ class ModelTestBuilder {
     return AddNode(op_type, input_args, {output_arg}, kMSDomain);
   }
 
+  Node& AddQLinearConcatLike(const std::string& op_type,
+                             NodeArg* output_arg,
+                             float output_scale,
+                             uint8_t output_zero_point,
+                             std::vector<std::tuple<NodeArg*, float, uint8_t>> quantized_inputs) {
+    std::vector<NodeArg*> input_args;
+    input_args.push_back(MakeScalarInitializer<float>(output_scale));
+    input_args.push_back(MakeScalarInitializer<uint8_t>(output_zero_point));
+    for (size_t input_index = 0; input_index < quantized_inputs.size(); ++input_index) {
+      input_args.push_back(std::get<0>(quantized_inputs[input_index]));
+      input_args.push_back(MakeScalarInitializer<float>(std::get<1>(quantized_inputs[input_index])));
+      input_args.push_back(MakeScalarInitializer<uint8_t>(std::get<2>(quantized_inputs[input_index])));
+    }
+    return AddNode(op_type, input_args, {output_arg}, kMSDomain);
+  }
+
   Node& AddQLinearActivationNode(const std::string& op_type,
                                  NodeArg* input_arg,
                                  float input_scale,

onnxruntime/test/optimizer/nhwc_transformer_test.cc

@@ -290,6 +290,50 @@ TEST(NhwcTransformerTests, ConvSplit) {
   }
 }
 
+TEST(NhwcTransformerTests, ConvSplitQLinearConcat) {
+  for (int64_t axis = -4LL; axis < 4; axis++) {
+    auto build_test_case = [&, axis](ModelTestBuilder& builder) {
+      auto* input_arg = builder.MakeInput<uint8_t>({2, 23, 16, 16}, 0, 31);
+      auto* conv_output_arg = builder.MakeIntermediate();
+      auto* split_output1_arg = builder.MakeIntermediate();
+      auto* split_output2_arg = builder.MakeIntermediate();
+      auto* qlconcat_output_arg = builder.MakeIntermediate();
+      auto* output_arg = builder.MakeOutput();
+
+      const int64_t conv1_output_channels = 32;
+      auto* conv1_weight_arg = NhwcMakeInitializer<uint8_t>(builder, {conv1_output_channels, 23, 3, 3});
+      Node& conv_node = builder.AddQLinearConvNode<uint8_t>(input_arg, .01f, 135,
+                                                            conv1_weight_arg, .02f, 126,
+                                                            conv_output_arg, .37f, 131);
+      conv_node.AddAttribute("pads", std::vector<int64_t>{1, 1, 1, 1});
+
+      Node& split_node = builder.AddNode("Split", {conv_output_arg}, {split_output1_arg, split_output2_arg});
+      split_node.AddAttribute("axis", static_cast<int64_t>(axis));
+
+      Node& qlconcat_node = builder.AddQLinearConcatLike(
+          "QLinearConcat", qlconcat_output_arg, .37f, 131,
+          {{split_output1_arg, .37f, uint8_t(131)}, {split_output2_arg, .37f, uint8_t(131)}});
+      qlconcat_node.AddAttribute("axis", static_cast<int64_t>(axis));
+
+      auto* conv2_weight_arg = NhwcMakeInitializer<uint8_t>(builder, {17, conv1_output_channels, 3, 3});
+      builder.AddQLinearConvNode<uint8_t>(qlconcat_output_arg, .43f, 126,
+                                          conv2_weight_arg, .02f, 126,
+                                          output_arg, .37f, 131);
+    };
+
+    auto check_nhwc_graph = [&](InferenceSessionWrapper& session) {
+      auto op_to_count = CountOpsInGraph(session.GetGraph());
+      EXPECT_EQ(op_to_count["com.microsoft.QLinearConv"], 2);
+      EXPECT_EQ(op_to_count["Transpose"], 2);
+    };
+
+    TransformerTester(build_test_case,
+                      check_nhwc_graph,
+                      TransformerLevel::Level2,
+                      TransformerLevel::Level3);
+  }
+}
+
 TEST(NhwcTransformerTests, ConvPad) {
   std::vector<std::string> pad_modes{"constant", "reflect", "edge"};
   for (const auto& mode : pad_modes) {