From 00c1da5e0a019c644f83481e849d29c363635d2f Mon Sep 17 00:00:00 2001
From: Jian Chen <cjian@microsoft.com>
Date: Mon, 15 May 2023 22:34:41 +0000
Subject: [PATCH] Fixing NhwcFusedConv fp16 (#15950)

### Description
<!-- Describe your changes. -->

This should produced fused Resnet50.fp16.onnx

### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
---
 .../core/optimizer/conv_add_act_fusion.cc     | 27 ++++++++++++++-----
 .../selector_action_transformer.cc            |  7 ++---
 2 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/onnxruntime/core/optimizer/conv_add_act_fusion.cc b/onnxruntime/core/optimizer/conv_add_act_fusion.cc
index b9016927de..7c8bfeaec5 100644
--- a/onnxruntime/core/optimizer/conv_add_act_fusion.cc
+++ b/onnxruntime/core/optimizer/conv_add_act_fusion.cc
@@ -175,17 +175,27 @@ class ConvAddActivationSelector : public NodeSelector {
       // Check if this is a single use convolution that hasn't already
       // been fused with another Add/Sum node. The Add/Sum can also only be
       // fused if the convolution isn't itself fused with an activation.
-      if (
-          (inputs_node[n]->OpType() == "Conv" || inputs_node[n]->OpType() == "NhwcFusedConv") && (pre_input_defs_count < 4) && (producer_input_args_count.size() < 4) && (graph_utils::GetNodeAttribute(*inputs_node[n], "activation") == nullptr) && (inputs_node[n]->GetOutputEdgesCount() == 1)) {
-        if (pre_input_defs_count < 3) {
-          // The optional bias parameter is empty so set to an empty string.
+      if ((inputs_node[n]->OpType() == "Conv") && (pre_input_defs_count < 4) &&
+          (producer_input_args_count.size() < 4) &&
+          (graph_utils::GetNodeAttribute(*inputs_node[n], "activation") == nullptr) &&
+          (inputs_node[n]->GetOutputEdgesCount() == 1)) {
+        if (pre_input_defs_count < 3) {  // The optional bias parameter is empty so set to an empty string.
+          // TODO, add a new null arguments for bias
+          continue;
+        }
+        return inputs_node[n];
+      }
+      if (inputs_node[n]->OpType() == "NhwcFusedConv" && (pre_input_defs_count < 4) &&
+          (producer_input_args_count.size() < 5) &&
+          (graph_utils::GetNodeAttribute(*inputs_node[n], "activation") == nullptr) &&
+          (inputs_node[n]->GetOutputEdgesCount() == 1)) {
+        if (pre_input_defs_count < 3) {  // The optional bias parameter is empty so set to an empty string.
           // TODO, add a new null arguments for bias
           continue;
         }
         return inputs_node[n];
       }
     }
-
     return nullptr;
   }
 };
@@ -275,11 +285,14 @@ class FuseConvAddActivationAction : public ReplaceWithNew {
 }  // namespace actions
 
 void RegisterConvAddActivationFusionRules(SelectorActionRegistry& registry) {
-  const auto name = "ConvAddAct";
   auto action = std::make_unique<actions::FuseConvAddActivationAction>();
   auto selector = std::make_unique<selectors::ConvAddActivationSelector>();
-  registry.RegisterSelectorAndAction(name, {{"Conv", {1, 11}}},
+  registry.RegisterSelectorAndAction("ConvAddAct", {{"Conv", {1, 11}}},
                                      std::move(selector), std::move(action));
+  auto action_nhwc = std::make_unique<actions::FuseConvAddActivationAction>();
+  auto selector_nhwc = std::make_unique<selectors::ConvAddActivationSelector>();
+  registry.RegisterSelectorAndAction("NhwcFusedConvAct", {{"NhwcFusedConv", {1, 11}}},
+                                     std::move(selector_nhwc), std::move(action_nhwc));
 }
 
 SelectorActionRegistry CreateSelectorActionRegistry() {
diff --git a/onnxruntime/core/optimizer/selectors_actions/selector_action_transformer.cc b/onnxruntime/core/optimizer/selectors_actions/selector_action_transformer.cc
index 540e0e92d3..e182b6c695 100644
--- a/onnxruntime/core/optimizer/selectors_actions/selector_action_transformer.cc
+++ b/onnxruntime/core/optimizer/selectors_actions/selector_action_transformer.cc
@@ -93,9 +93,10 @@ static Status MatchAndProcess(
   Status status = Status::OK();
 
   do {
-    // TODO: for now this just needs to support ONNX ops. If we ever had a transformer that was going to
-    // target non-ONNX ops we'd need to rework a few things to include the op domain in the matches
-    if (node.Domain() != kOnnxDomain) {
+    // TODO: for now this just needs to support ONNX and Micrsoft Domain ops.
+    // If we ever had a transformer that was going to target non-ONNX ops,
+    // we'd need to rework a few things to include the op domain in the matches
+    if (node.Domain() != kOnnxDomain && node.Domain() != kMSDomain) {
       break;
     }