Fix optimizer crash (#8274)

onnxruntime/core/optimizer/conv_activation_fusion.cc

@@ -121,24 +121,32 @@ Status ConvActivationFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l
         graph_utils::FinalizeNodeFusion(graph, {conv_node, act_node}, fused_conv);
         modified = true;
       } else if (graph_utils::IsSupportedOptypeVersionAndDomain(next_node, "Add", {6, 7, 13, 14})) {
+        if (next_node.GetOutputEdgesCount() != 1) {
+          continue;
+        }
         const auto& last_node = *(next_node.OutputNodesBegin());
         if (last_node.GetExecutionProviderType() != node->GetExecutionProviderType()) {
           continue;
         }
-        if (graph_utils::IsSupportedOptypeVersionAndDomain(last_node, "Relu", {6, 13, 14}) &&
-            next_node.GetOutputEdgesCount() == 1) {
+        if (graph_utils::IsSupportedOptypeVersionAndDomain(last_node, "Relu", {6, 13, 14})) {
           Node& conv_node = *node;
           Node& add_node = *graph.GetNode(next_node.Index());
           Node& act_node = *graph.GetNode(last_node.Index());
           auto conv_inputs = conv_node.MutableInputDefs();
           auto conv_outputs = conv_node.MutableOutputDefs();
           auto add_inputs = add_node.MutableInputDefs();
-          for (auto add_input : add_inputs) {
-            if (add_input->Name() != conv_outputs[0]->Name()) {
+          int32_t dependent = 0, independent = 0;
+          for (auto add_input: add_inputs) {
+            if (add_input->Name() == conv_outputs[0]->Name()) {
+              dependent++;
+            } else {
               conv_inputs.push_back(add_input);
-              break;
+              independent++;
             }
           }
+          if (dependent != 1 || independent != 1) {
+            continue;
+          }
           auto node_name = graph.GenerateNodeName(conv_node.Name() + "_" +
                                                   add_node.Name() + "_" +
                                                   act_node.Name());

onnxruntime/test/optimizer/graph_transform_test.cc

@@ -658,6 +658,7 @@ TEST_F(GraphTransformationTests, NotWhereFusion) {
 }
 
 #if defined(USE_CUDA) && !defined(DISABLE_CONTRIB_OPS)
+// Conv->Add->Relu will be transformed to FusedConv
 TEST_F(GraphTransformationTests, FuseCudaConvAddRelu) {
   auto model_uri = MODEL_FOLDER "fusion/conv_add_relu.onnx";
   std::shared_ptr<Model> p_model;
@@ -673,9 +674,50 @@ TEST_F(GraphTransformationTests, FuseCudaConvAddRelu) {
   graph_transformation_mgr.Register(std::make_unique<ConvActivationFusion>(), TransformerLevel::Level2);
   ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, *logger_));
   op_to_count = CountOpsInGraph(graph);
-  ASSERT_TRUE(op_to_count["Add"] == 0);
-  ASSERT_TRUE(op_to_count["Relu"] == 0);
+  ASSERT_TRUE(op_to_count["Add"] == 0); //Add removed from graph
+  ASSERT_TRUE(op_to_count["Relu"] == 0); //Relu removed from graph
 }
+
+//Conv->Add->Relu will be left intact since there is Identity depend on Add
+TEST_F(GraphTransformationTests, FuseCudaConvAddReluIdentity) {
+  auto model_uri = MODEL_FOLDER "fusion/conv_add_relu_identity.onnx";
+  std::shared_ptr<Model> p_model;
+  ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_));
+  Graph& graph = p_model->MainGraph();
+  for (auto& node : p_model->MainGraph().Nodes()) {
+    node.SetExecutionProviderType(kCudaExecutionProvider);
+  }
+  std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
+  ASSERT_TRUE(op_to_count["Add"] == 1);
+  ASSERT_TRUE(op_to_count["Relu"] == 1);
+  ASSERT_TRUE(op_to_count["Identity"] == 1);
+  onnxruntime::GraphTransformerManager graph_transformation_mgr{5};
+  graph_transformation_mgr.Register(std::make_unique<ConvActivationFusion>(), TransformerLevel::Level2);
+  ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, *logger_));
+  op_to_count = CountOpsInGraph(graph);
+  ASSERT_TRUE(op_to_count["Add"] == 1); //Add remains
+  ASSERT_TRUE(op_to_count["Relu"] == 1); //Relu remains
+  ASSERT_TRUE(op_to_count["Identity"] == 1); //Identity remains
+}
+
+//Conv->Add will be left intact since there is no Relu follows
+TEST_F(GraphTransformationTests, FuseCudaConvAdd) {
+  auto model_uri = MODEL_FOLDER "fusion/conv_add.onnx";
+  std::shared_ptr<Model> p_model;
+  ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_));
+  Graph& graph = p_model->MainGraph();
+  for (auto& node : p_model->MainGraph().Nodes()) {
+    node.SetExecutionProviderType(kCudaExecutionProvider);
+  }
+  std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
+  ASSERT_TRUE(op_to_count["Add"] == 1);
+  onnxruntime::GraphTransformerManager graph_transformation_mgr{5};
+  graph_transformation_mgr.Register(std::make_unique<ConvActivationFusion>(), TransformerLevel::Level2);
+  ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, *logger_));
+  op_to_count = CountOpsInGraph(graph);
+  ASSERT_TRUE(op_to_count["Add"] == 1); //Add remains, no transform applied to the graph
+}
+
 #endif
 
 #ifndef DISABLE_CONTRIB_OPS

onnxruntime/test/testdata/transform/fusion/conv_add.onnx

@@ -0,0 +1,44 @@
+:°

+
+
X
+
W
+
B
C"Conv
+
+
C
+
A
Y"AddgraphZ
+
X
+

+

+

+
+Z
+
W
+

+
+

+
+Z
+
B
+
+

+Z
+
A
+

+

+
+
+b
+
Y
+

+

+
+
+BB
+com.microsoft.experimental
B
+ai.onnx.preview.training
B
+com.microsoft.nchwc
B
+com.microsoft.mlfeaturizers
B
+
+ai.onnx.mlB
+
com.microsoft
B
+ai.onnx.training

onnxruntime/test/testdata/transform/fusion/conv_add_relu_identity.onnx

@@ -0,0 +1,46 @@
+:í

+
+
X
+
W
+
B
C"Conv
+
+
C
+
A
S"Add
+
+
S
Y"Relu
+
+
S
I"IdentitygraphZ
+
X
+

+

+

+
+Z
+
W
+

+
+

+
+Z
+
B
+
+

+Z
+
A
+

+

+
+
+b
+
Y
+

+

+
+
+b
+
I
+

+

+
+
+B