From 7c8e1580a13ce333e47a41146bccfc90b3a70db5 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 6 Jun 2020 00:06:07 -0700 Subject: [PATCH] Add check of graph output in Bert Fusions (#4126) * Refine node output edge checking in bert related fusions --- .../core/optimizer/attention_fusion.cc | 33 ++++++++--------- .../core/optimizer/bias_gelu_fusion.cc | 3 +- .../core/optimizer/embed_layer_norm_fusion.cc | 24 ++++++------ .../core/optimizer/fast_gelu_fusion.cc | 35 ++++++++---------- onnxruntime/core/optimizer/gelu_fusion.cc | 10 ++--- .../core/optimizer/layer_norm_fusion.cc | 33 +++++++---------- onnxruntime/core/optimizer/reshape_fusion.cc | 2 +- onnxruntime/core/optimizer/utils.cc | 9 +++++ onnxruntime/core/optimizer/utils.h | 7 ++++ .../test/optimizer/graph_transform_test.cc | 17 +++++++++ ..._format2_0_with_bias_use_graph_output.onnx | Bin 0 -> 847 bytes .../testdata/transform/fusion/gelu_gen.py | 7 +++- 12 files changed, 104 insertions(+), 76 deletions(-) create mode 100644 onnxruntime/test/testdata/transform/fusion/gelu_format2_0_with_bias_use_graph_output.onnx diff --git a/onnxruntime/core/optimizer/attention_fusion.cc b/onnxruntime/core/optimizer/attention_fusion.cc index 68284f3e6b..bffe2372c6 100644 --- a/onnxruntime/core/optimizer/attention_fusion.cc +++ b/onnxruntime/core/optimizer/attention_fusion.cc @@ -413,16 +413,14 @@ bool AttentionFusion::FuseSubGraph(Node& layer_norm, const Node& add_after_layer return false; } - if (add.GetOutputEdgesCount() != 1 || - matmul.GetOutputEdgesCount() != 1 || - reshape.GetOutputEdgesCount() != 1 || - transpose.GetOutputEdgesCount() != 1 || - qkv_matmul.GetOutputEdgesCount() != 1 || - v_transpose.GetOutputEdgesCount() != 1 || - v_reshape.GetOutputEdgesCount() != 1 || - v_add.GetOutputEdgesCount() != 1 || - v_matmul.GetOutputEdgesCount() != 1 || - v_root.GetOutputEdgesCount() != 4) { + // Internal nodes of attention subgraph only allow edges within the subgraph, and no graph output is allowed. + // No constraints for four nodes: reshape node is last node of Attention; and add, matmul and v_root are not in attention subgraph. + if (!optimizer_utils::CheckOutputEdges(graph, transpose, 1) || + !optimizer_utils::CheckOutputEdges(graph, qkv_matmul, 1) || + !optimizer_utils::CheckOutputEdges(graph, v_transpose, 1) || + !optimizer_utils::CheckOutputEdges(graph, v_reshape, 1) || + !optimizer_utils::CheckOutputEdges(graph, v_add, 1) || + !optimizer_utils::CheckOutputEdges(graph, v_matmul, 1)) { DEBUG_LOG("Output edge count not expected for nodes in path v"); return false; } @@ -513,12 +511,13 @@ bool AttentionFusion::FuseSubGraph(Node& layer_norm, const Node& add_after_layer const Node& mask_unsqueeze_2 = *p_mask_unsqueeze_2; const Node& mask_unsqueeze_1 = *p_mask_unsqueeze_1; - if (softmax.GetOutputEdgesCount() != 1 || - mask_add.GetOutputEdgesCount() != 1 || - mask_sub.GetOutputEdgesCount() != 1 || - (p_mask_cast != nullptr && (*p_mask_cast).GetOutputEdgesCount() != 1) || - mask_unsqueeze_2.GetOutputEdgesCount() != 1 || - mask_unsqueeze_1.GetOutputEdgesCount() != 1) { + + if (!optimizer_utils::CheckOutputEdges(graph, softmax, 1) || + !optimizer_utils::CheckOutputEdges(graph, mask_add, 1) || + !optimizer_utils::CheckOutputEdges(graph, mask_sub, 1) || + (p_mask_cast != nullptr && !optimizer_utils::CheckOutputEdges(graph, *p_mask_cast, 1)) || + !optimizer_utils::CheckOutputEdges(graph, mask_unsqueeze_2, 1) || + !optimizer_utils::CheckOutputEdges(graph, mask_unsqueeze_1, 1)) { DEBUG_LOG("Output edge count not expected for mask nodes"); return false; } @@ -713,7 +712,7 @@ bool AttentionFusion::FuseSubGraph(Node& layer_norm, const Node& add_after_layer k_matmul.Index()}; // When the last Attention node is fused. Original mask processing nodes can be removed safely. - if (mask_mul.GetOutputEdgesCount() == 1) { + if (optimizer_utils::CheckOutputEdges(graph, mask_mul, 1)) { nodes_to_remove.push_back(mask_mul.Index()); nodes_to_remove.push_back(mask_sub.Index()); if (p_mask_cast != nullptr) { diff --git a/onnxruntime/core/optimizer/bias_gelu_fusion.cc b/onnxruntime/core/optimizer/bias_gelu_fusion.cc index 5bdad45b9c..9906848f35 100644 --- a/onnxruntime/core/optimizer/bias_gelu_fusion.cc +++ b/onnxruntime/core/optimizer/bias_gelu_fusion.cc @@ -4,6 +4,7 @@ #include "core/optimizer/initializer.h" #include "core/optimizer/bias_gelu_fusion.h" #include "core/graph/graph_utils.h" +#include "core/optimizer/utils.h" #include using namespace ONNX_NAMESPACE; @@ -25,7 +26,7 @@ Status BiasGelu::ApplyImpl(Graph& graph, bool& modified, int graph_level, const if (!graph_utils::IsSupportedOptypeVersionAndDomain(node, "Add", {7}) || !graph_utils::IsSupportedProvider(node, GetCompatibleExecutionProviders()) || - node.GetOutputEdgesCount() != 1) { + !optimizer_utils::CheckOutputEdges(graph, node, 1)) { continue; } diff --git a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc index d4afa07b85..8136e39ed7 100644 --- a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc +++ b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc @@ -125,7 +125,7 @@ static bool MatchPositionSubgraph( return false; } for (size_t i = 0; i < edges.size(); i++) { - if (edges[i]->GetNode().GetOutputEdgesCount() != 1) { + if (!optimizer_utils::CheckOutputEdges(graph, edges[i]->GetNode(), 1)) { DEBUG_LOG("Output edge count not expected for nodes in path 1 of position shape."); return false; } @@ -151,9 +151,9 @@ static bool MatchPositionSubgraph( return false; } - if (edges[0]->GetNode().GetOutputEdgesCount() != 1 || - edges[1]->GetNode().GetOutputEdgesCount() != 2 || - edges[2]->GetNode().GetOutputEdgesCount() != 1) { + if (!optimizer_utils::CheckOutputEdges(graph, edges[0]->GetNode(), 1) || + !optimizer_utils::CheckOutputEdges(graph, edges[1]->GetNode(), 2) || + !optimizer_utils::CheckOutputEdges(graph, edges[2]->GetNode(), 1)) { DEBUG_LOG("Output edge count not expected for nodes in path 2 of position shape."); return false; } @@ -238,11 +238,10 @@ static bool MatchPositionEmbeddingSubgraph1( return false; } const size_t gather_index = 8; - auto gather_output_edges_count = pg_edges[gather_index]->GetNode().GetOutputEdgesCount(); // All nodes in Path 1 must have only 1 output edge, except the gather node allowed 1 or 2 output edges for (size_t i = 0; i < pg_edges.size(); i++) { - if (pg_edges[i]->GetNode().GetOutputEdgesCount() != 1) { - if (i == gather_index && gather_output_edges_count == 2) { + if (!optimizer_utils::CheckOutputEdges(graph, pg_edges[i]->GetNode(), 1)) { + if (i == gather_index && optimizer_utils::CheckOutputEdges(graph, pg_edges[i]->GetNode(), 2)) { continue; } DEBUG_LOG("Output edge count not expected for nodes in path1."); @@ -253,7 +252,7 @@ static bool MatchPositionEmbeddingSubgraph1( Node& expand_node = *graph.GetNode(pg_edges[0]->GetNode().Index()); Node& gather_node = *graph.GetNode(pg_edges[gather_index]->GetNode().Index()); - if (gather_output_edges_count == 1) { + if (gather_node.GetOutputEdgesCount() == 1) { // Check if the second input of the Gather node in the path has a constant input of 1 // For gather_output_edges_count == 2, such checks are in MatchPositionSubgraph function. if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(gather_node.InputDefs()[1]), int64_t(1), true)) { @@ -347,7 +346,7 @@ static bool MatchPositionEmbeddingSubgraph2( size_t last_edge = edges.size() - 1; for (size_t i = 0; i < edges.size(); i++) { - if (edges[i]->GetNode().GetOutputEdgesCount() != (i == last_edge ? 2u : 1u)) { + if (!optimizer_utils::CheckOutputEdges(graph, edges[i]->GetNode(), (i == last_edge ? 2u : 1u))) { DEBUG_LOG("Output edge count not expected for nodes in path 1."); return false; } @@ -390,7 +389,7 @@ static bool MatchPositionEmbeddingSubgraph( return false; } Node& position_gather_node = *graph.GetNode(edges[0]->GetNode().Index()); - if (position_gather_node.GetOutputEdgesCount() != 1) { + if (!optimizer_utils::CheckOutputEdges(graph, position_gather_node, 1)) { return false; } @@ -544,7 +543,7 @@ Status EmbedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l continue; } Node& segment_gather_node = *graph.GetNode(edges[0]->GetNode().Index()); - if (segment_gather_node.GetOutputEdgesCount() != 1) { + if (!optimizer_utils::CheckOutputEdges(graph, segment_gather_node, 1)) { continue; } // The first input of segment_gather_node must be 2d. @@ -566,7 +565,8 @@ Status EmbedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l } Node& add_node = *graph.GetNode(edges[0]->GetNode().Index()); Node& word_gather_node = *graph.GetNode(edges[1]->GetNode().Index()); - if (add_node.GetOutputEdgesCount() != 1 || word_gather_node.GetOutputEdgesCount() != 1) { + if (!optimizer_utils::CheckOutputEdges(graph, add_node, 1) || + !optimizer_utils::CheckOutputEdges(graph, word_gather_node, 1)) { continue; } // The first input of word_gather_node must be 2d. diff --git a/onnxruntime/core/optimizer/fast_gelu_fusion.cc b/onnxruntime/core/optimizer/fast_gelu_fusion.cc index c71555bb83..a2e8528ff1 100644 --- a/onnxruntime/core/optimizer/fast_gelu_fusion.cc +++ b/onnxruntime/core/optimizer/fast_gelu_fusion.cc @@ -24,12 +24,13 @@ static bool IsSupportedDataType(const Node& node) { } } -static bool CheckNode(const Node& node, const std::string& op_name, int32_t opset_version, ProviderType provider, - bool require_single_output = false) { +static bool CheckNode(Graph& graph, const Node& node, const std::string& op_name, int32_t opset_version, ProviderType provider, + bool require_single_output) { return graph_utils::IsSupportedOptypeVersionAndDomain(node, op_name, {opset_version}) && node.GetExecutionProviderType() == provider && IsSupportedDataType(node) && - (!require_single_output || node.GetOutputEdgesCount() == 1); + (!require_single_output || node.GetOutputEdgesCount() == 1) && + graph.GetNodeOutputsInGraphOutputs(node).empty(); } MatchResult FastGeluFusion::CheckFirstFormula(Graph& graph, Node& mul1_node, @@ -58,7 +59,7 @@ MatchResult FastGeluFusion::CheckFirstFormula(Graph& graph, Node& mul1_node, Node& mul2_node = *graph.GetNode(mul1_node.OutputNodesBegin()->Index()); input_index = optimizer_utils::IndexOfNodeInput(mul2_node, *mul1_node.MutableOutputDefs()[0]); - if (!CheckNode(mul2_node, "Mul", 7, mul1_node.GetExecutionProviderType(), true) || + if (!CheckNode(graph, mul2_node, "Mul", 7, mul1_node.GetExecutionProviderType(), true) || mul2_node.MutableInputDefs()[(input_index + 1) % 2]->Name() != gelu_without_bias_input_arg->Name()) { return matchResult; } @@ -66,14 +67,14 @@ MatchResult FastGeluFusion::CheckFirstFormula(Graph& graph, Node& mul1_node, Node& add1_node = *graph.GetNode(mul2_node.OutputNodesBegin()->Index()); input_index = optimizer_utils::IndexOfNodeInput(add1_node, *mul2_node.MutableOutputDefs()[0]); - if (!CheckNode(add1_node, "Add", 7, mul1_node.GetExecutionProviderType(), true) || + if (!CheckNode(graph, add1_node, "Add", 7, mul1_node.GetExecutionProviderType(), true) || !optimizer_utils::IsInitializerWithExpectedValue(graph, *(add1_node.InputDefs()[(input_index + 1) % 2]), 1.0f, true)) { return matchResult; } nodes_to_fuse.push_back(add1_node); Node& mul3_node = *graph.GetNode(add1_node.OutputNodesBegin()->Index()); - if (!CheckNode(mul3_node, "Mul", 7, mul1_node.GetExecutionProviderType(), true)) { + if (!CheckNode(graph, mul3_node, "Mul", 7, mul1_node.GetExecutionProviderType(), true)) { return matchResult; } nodes_to_fuse.push_back(mul3_node); @@ -82,7 +83,7 @@ MatchResult FastGeluFusion::CheckFirstFormula(Graph& graph, Node& mul1_node, const Node* p_mul3_input_node = graph_utils::GetInputNode(mul3_node, (input_index + 1) % 2); if (p_mul3_input_node == nullptr) return matchResult; Node& mul4_node = const_cast(*p_mul3_input_node); - if (!CheckNode(mul4_node, "Mul", 7, mul1_node.GetExecutionProviderType(), true)) { + if (!CheckNode(graph, mul4_node, "Mul", 7, mul1_node.GetExecutionProviderType(), true)) { return matchResult; } @@ -124,7 +125,7 @@ MatchResult FastGeluFusion::CheckSecondFormula(Graph& graph, Node& pow1_node, Node& mul1_node = *graph.GetNode(pow1_node.OutputNodesBegin()->Index()); auto input_index = optimizer_utils::IndexOfNodeInput(mul1_node, *pow1_node.MutableOutputDefs()[0]); - if (!CheckNode(mul1_node, "Mul", 7, pow1_node.GetExecutionProviderType(), true) || + if (!CheckNode(graph, mul1_node, "Mul", 7, pow1_node.GetExecutionProviderType(), true) || !optimizer_utils::IsInitializerWithExpectedValue(graph, *(mul1_node.InputDefs()[(input_index + 1) % 2]), 0.044714998453855515f, true)) { return matchResult; @@ -133,7 +134,7 @@ MatchResult FastGeluFusion::CheckSecondFormula(Graph& graph, Node& pow1_node, Node& add1_node = *graph.GetNode(mul1_node.OutputNodesBegin()->Index()); input_index = optimizer_utils::IndexOfNodeInput(add1_node, *mul1_node.MutableOutputDefs()[0]); - if (!CheckNode(add1_node, "Add", 7, pow1_node.GetExecutionProviderType(), true) || + if (!CheckNode(graph, add1_node, "Add", 7, pow1_node.GetExecutionProviderType(), true) || add1_node.MutableInputDefs()[(input_index + 1) % 2]->Name() != pow_input_arg->Name()) { return matchResult; } @@ -141,7 +142,7 @@ MatchResult FastGeluFusion::CheckSecondFormula(Graph& graph, Node& pow1_node, Node& mul2_node = *graph.GetNode(add1_node.OutputNodesBegin()->Index()); input_index = optimizer_utils::IndexOfNodeInput(mul2_node, *add1_node.MutableOutputDefs()[0]); - if (!CheckNode(mul2_node, "Mul", 7, pow1_node.GetExecutionProviderType(), true) || + if (!CheckNode(graph, mul2_node, "Mul", 7, pow1_node.GetExecutionProviderType(), true) || !optimizer_utils::IsInitializerWithExpectedValue(graph, *(mul2_node.InputDefs()[(input_index + 1) % 2]), 0.7978845834732056f, true)) { return matchResult; @@ -176,12 +177,12 @@ Status FastGeluFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, }; Node& tanh_node = *graph.GetNode(matchRet.tanh_input_node->OutputNodesBegin()->Index()); - if (!CheckNode(tanh_node, "Tanh", 6, node.GetExecutionProviderType(), true)) { + if (!CheckNode(graph, tanh_node, "Tanh", 6, node.GetExecutionProviderType(), true)) { continue; } Node& add2_node = *graph.GetNode(tanh_node.OutputNodesBegin()->Index()); - if (!CheckNode(add2_node, "Add", 7, node.GetExecutionProviderType(), true)) { + if (!CheckNode(graph, add2_node, "Add", 7, node.GetExecutionProviderType(), true)) { continue; } @@ -192,12 +193,7 @@ Status FastGeluFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, Node& mul5_node = *graph.GetNode(add2_node.OutputNodesBegin()->Index()); // This is the output of the Gelu subgraph, we don't need check it has single edge. - if (!CheckNode(mul5_node, "Mul", 7, node.GetExecutionProviderType(), false)) { - continue; - } - - // ingnore the transformer if Gelu's output is the graph's output. - if (!graph.GetNodeOutputsInGraphOutputs(mul5_node).empty()) { + if (!CheckNode(graph, mul5_node, "Mul", 7, node.GetExecutionProviderType(), false)) { continue; } @@ -205,7 +201,7 @@ Status FastGeluFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, const Node* p_mul5_input_node = graph_utils::GetInputNode(mul5_node, (input_index + 1) % 2); if (p_mul5_input_node == nullptr) continue; Node& mul6_node = const_cast(*p_mul5_input_node); - if (!CheckNode(mul6_node, "Mul", 7, node.GetExecutionProviderType(), false)) { + if (!CheckNode(graph, mul6_node, "Mul", 7, node.GetExecutionProviderType(), false)) { continue; } @@ -224,6 +220,7 @@ Status FastGeluFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, nodes_to_fuse.insert(nodes_to_fuse.end(), {tanh_node, add2_node, mul6_node, mul5_node}); auto type_info = *node.MutableOutputDefs()[0]->TypeAsProto(); + // TODO: re-use node arg of mul5 so that it is allowed to be graph output (Need modify CheckNode as well). auto& shape_output = graph.GetOrCreateNodeArg(graph.GenerateNodeArgName("fast_gelu_output"), &type_info); Node& fast_gelu_node = graph.AddNode(graph.GenerateNodeName("GPT2Gelu"), "FastGelu", diff --git a/onnxruntime/core/optimizer/gelu_fusion.cc b/onnxruntime/core/optimizer/gelu_fusion.cc index 6697c5117d..a93338f6c3 100644 --- a/onnxruntime/core/optimizer/gelu_fusion.cc +++ b/onnxruntime/core/optimizer/gelu_fusion.cc @@ -57,7 +57,7 @@ Status GeluFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, cons if (!graph_utils::IsSupportedOptypeVersionAndDomain(div, "Div", {7}) || !graph_utils::IsSupportedProvider(div, GetCompatibleExecutionProviders()) || - div.GetOutputEdgesCount() != 1 || + !optimizer_utils::CheckOutputEdges(graph, div, 1) || !IsSupportedDataType(div)) { continue; } @@ -73,7 +73,7 @@ Status GeluFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, cons Node& erf_node = *graph.GetNode(div.OutputNodesBegin()->Index()); if (!graph_utils::IsSupportedOptypeVersionAndDomain(erf_node, "Erf", {9}) || erf_node.GetExecutionProviderType() != div.GetExecutionProviderType() || - erf_node.GetOutputEdgesCount() != 1 || + !optimizer_utils::CheckOutputEdges(graph, erf_node, 1) || !IsSupportedDataType(erf_node)) { continue; } @@ -81,7 +81,7 @@ Status GeluFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, cons Node& add_node = *graph.GetNode(erf_node.OutputNodesBegin()->Index()); if (!graph_utils::IsSupportedOptypeVersionAndDomain(add_node, "Add", {7}) || add_node.GetExecutionProviderType() != div.GetExecutionProviderType() || - add_node.GetOutputEdgesCount() != 1 || + !optimizer_utils::CheckOutputEdges(graph, add_node, 1) || !IsSupportedDataType(add_node)) { continue; } @@ -108,7 +108,7 @@ Status GeluFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, cons Node& mul2_node = *graph.GetNode(p_mul2_node->Index()); if (!graph_utils::IsSupportedOptypeVersionAndDomain(mul2_node, "Mul", {7}) || mul2_node.GetExecutionProviderType() != div.GetExecutionProviderType() || - mul2_node.GetOutputEdgesCount() != 1 || + !optimizer_utils::CheckOutputEdges(graph, mul2_node, 1) || !IsSupportedDataType(mul2_node)) { continue; } @@ -129,7 +129,7 @@ Status GeluFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, cons is_pattern_1 = false; // Match subgraph pattern 2 - if (mul_node.GetOutputEdgesCount() != 1) { + if (!optimizer_utils::CheckOutputEdges(graph, mul_node, 1)) { continue; } diff --git a/onnxruntime/core/optimizer/layer_norm_fusion.cc b/onnxruntime/core/optimizer/layer_norm_fusion.cc index 39a1907452..f2ce33b645 100644 --- a/onnxruntime/core/optimizer/layer_norm_fusion.cc +++ b/onnxruntime/core/optimizer/layer_norm_fusion.cc @@ -3,6 +3,7 @@ #include "core/optimizer/initializer.h" #include "core/optimizer/layer_norm_fusion.h" #include "core/graph/graph_utils.h" +#include "core/optimizer/utils.h" #include "float.h" #include @@ -61,6 +62,7 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, if (!graph_utils::IsSupportedOptypeVersionAndDomain(reduce_mean_node, "ReduceMean", {1, 11}) || !graph_utils::IsSupportedProvider(reduce_mean_node, GetCompatibleExecutionProviders()) || (reduce_mean_node.GetOutputEdgesCount() != 1 && reduce_mean_node.GetOutputEdgesCount() != 2) || + !graph.GetNodeOutputsInGraphOutputs(reduce_mean_node).empty() || !IsSupportedDataType(reduce_mean_node)) { continue; } @@ -93,6 +95,7 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, Node& sub_node = *graph.GetNode(p_sub_node->Index()); if (!graph_utils::IsSupportedOptypeVersionAndDomain(sub_node, "Sub", {7}) || sub_node.GetExecutionProviderType() != reduce_mean_node.GetExecutionProviderType() || + !optimizer_utils::CheckOutputEdges(graph, sub_node, subCnt == 1 ? 2u: 1u) || !IsSupportedDataType(sub_node)) { continue; } @@ -107,7 +110,7 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, Node& sub_node_dup = *graph.GetNode(p_sub_node_dup->Index()); if (!graph_utils::IsSupportedOptypeVersionAndDomain(sub_node_dup, "Sub", {7}) || sub_node_dup.GetExecutionProviderType() != reduce_mean_node.GetExecutionProviderType() || - sub_node_dup.GetOutputEdgesCount() != 1 || + !optimizer_utils::CheckOutputEdges(graph, sub_node, 1) || !IsSupportedDataType(sub_node_dup)) { continue; } @@ -124,7 +127,7 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, Node& div_node = *graph.GetNode(p_div->Index()); if (!graph_utils::IsSupportedOptypeVersionAndDomain(div_node, "Div", {7}) || div_node.GetExecutionProviderType() != reduce_mean_node.GetExecutionProviderType() || - div_node.GetOutputEdgesCount() != 1 || + !optimizer_utils::CheckOutputEdges(graph, div_node, 1) || !IsSupportedDataType(div_node)) { continue; } @@ -139,7 +142,7 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, if (!graph_utils::IsSupportedOptypeVersionAndDomain(sqrt_node, "Sqrt", {6}) || sqrt_node.GetExecutionProviderType() != reduce_mean_node.GetExecutionProviderType() || - sqrt_node.GetOutputEdgesCount() != 1 || + !optimizer_utils::CheckOutputEdges(graph, sqrt_node, 1) || !IsSupportedDataType(sqrt_node) || sqrt_node.GetInputEdgesCount() == 0) { continue; @@ -150,7 +153,7 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, Node& add2_node = *graph.GetNode(sqrt_node.InputNodesBegin()->Index()); if (!graph_utils::IsSupportedOptypeVersionAndDomain(add2_node, "Add", {7}) || add2_node.GetExecutionProviderType() != reduce_mean_node.GetExecutionProviderType() || - add2_node.GetOutputEdgesCount() != 1 || + !optimizer_utils::CheckOutputEdges(graph, add2_node, 1) || !IsSupportedDataType(add2_node)) { continue; } @@ -165,7 +168,7 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, Node& reduce_mean2_node = *graph.GetNode(p_reduce_mean2->Index()); if (!graph_utils::IsSupportedOptypeVersionAndDomain(reduce_mean2_node, "ReduceMean", {1, 11}) || reduce_mean2_node.GetExecutionProviderType() != reduce_mean_node.GetExecutionProviderType() || - reduce_mean2_node.GetOutputEdgesCount() != 1 || + !optimizer_utils::CheckOutputEdges(graph, reduce_mean2_node, 1) || !IsSupportedDataType(reduce_mean2_node) || reduce_mean2_node.GetInputEdgesCount() == 0) { continue; @@ -176,7 +179,7 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, Node& pow_node = *graph.GetNode(reduce_mean2_node.InputNodesBegin()->Index()); if (!graph_utils::IsSupportedOptypeVersionAndDomain(pow_node, "Pow", {7, 12}) || pow_node.GetExecutionProviderType() != reduce_mean_node.GetExecutionProviderType() || - pow_node.GetOutputEdgesCount() != 1 || + !optimizer_utils::CheckOutputEdges(graph, pow_node, 1) || !IsSupportedDataType(pow_node)) { continue; } @@ -184,20 +187,8 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, // Traceback the pow node to find sub --> pow const Node* p_sub2_node = graph_utils::FirstParentByType(pow_node, "Sub"); - if (p_sub2_node == nullptr) { - continue; - } - Node& sub2_node = *graph.GetNode(p_sub2_node->Index()); - if (!graph_utils::IsSupportedOptypeVersionAndDomain(sub2_node, "Sub", {7}) || - sub2_node.GetExecutionProviderType() != reduce_mean_node.GetExecutionProviderType() || - !IsSupportedDataType(sub2_node)) { - continue; - } - - // Traceback the sub node to find reduceMean --> sub - const Node* p_reduce_mean_check = graph_utils::FirstParentByType(sub2_node, "ReduceMean"); - // Check if the reduceMean node after traceback is the same node as the reduceMean node from the beginning. - if (p_reduce_mean_check == nullptr || p_reduce_mean_check != &reduce_mean_node) { + if (p_sub2_node == nullptr || + (p_sub2_node != p_sub_node && p_sub2_node != p_sub_node_dup)) { continue; } @@ -205,12 +196,14 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, Node& mul_node = *graph.GetNode(div_node.OutputNodesBegin()->Index()); if (!graph_utils::IsSupportedOptypeVersionAndDomain(mul_node, "Mul", {7}) || mul_node.GetExecutionProviderType() != reduce_mean_node.GetExecutionProviderType() || + !optimizer_utils::CheckOutputEdges(graph, mul_node, 1) || !IsSupportedDataType(mul_node)) { continue; } nodes_to_remove.push_back(mul_node); // mul --> add + // Need not check output edges of last node since they will be moved to fused node. Node& last_add_node = *graph.GetNode(mul_node.OutputNodesBegin()->Index()); if (!graph_utils::IsSupportedOptypeVersionAndDomain(last_add_node, "Add", {7}) || last_add_node.GetExecutionProviderType() != reduce_mean_node.GetExecutionProviderType() || diff --git a/onnxruntime/core/optimizer/reshape_fusion.cc b/onnxruntime/core/optimizer/reshape_fusion.cc index 3b00e55e6e..6fc37df9bc 100644 --- a/onnxruntime/core/optimizer/reshape_fusion.cc +++ b/onnxruntime/core/optimizer/reshape_fusion.cc @@ -116,7 +116,7 @@ bool ReshapeFusion::Fuse_Subgraph1(Node& reshape, Graph& graph, const logging::L } auto concat_input_count = concat.InputArgCount().front(); - if (concat.GetOutputEdgesCount() > 1) { + if (!optimizer_utils::CheckOutputEdges(graph, concat, 1)) { return false; } diff --git a/onnxruntime/core/optimizer/utils.cc b/onnxruntime/core/optimizer/utils.cc index 052f5ff67f..b4c7fcb224 100644 --- a/onnxruntime/core/optimizer/utils.cc +++ b/onnxruntime/core/optimizer/utils.cc @@ -227,5 +227,14 @@ bool IsSupportedDataType(const Node& node, const std::vector& suppo return true; } + +bool CheckOutputEdges(const Graph& graph, const Node& node, size_t expected_output_edges) { + if (!graph.GetNodeOutputsInGraphOutputs(node).empty()) { + return false; + } + + return node.GetOutputEdgesCount() == expected_output_edges; +} + } // namespace optimizer_utils } // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/utils.h b/onnxruntime/core/optimizer/utils.h index 0cafe0d49c..0f1317efa5 100644 --- a/onnxruntime/core/optimizer/utils.h +++ b/onnxruntime/core/optimizer/utils.h @@ -66,5 +66,12 @@ int32_t IndexOfNodeInput(const Node& node, const NodeArg& node_arg); */ bool IsSupportedDataType(const Node& node, const std::vector& supported_data_types); +/** Check whether node's output edges count is expected. +@remarks graph output is not included in output edges, and this node shall not have graph output. + A node with graph output cannot be fused unless the graph output also exists in outputs of fused node. +@returns false when the node has graph output, or number of output edges are not expected. +*/ +bool CheckOutputEdges(const Graph& graph, const Node& node, size_t expected_output_edges); + } // namespace optimizer_utils } // namespace onnxruntime diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 1c3e2e9d73..2a9b1bb686 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -1684,6 +1684,23 @@ TEST_F(GraphTransformationTests, GeluFusionTestFormat2GraphInput) { ASSERT_TRUE(op_to_count["Gelu"] == 1); } +TEST_F(GraphTransformationTests, GeluFusionTestFormat2GraphOutput) { + auto model_uri = MODEL_FOLDER "fusion/gelu_format2_0_with_bias_use_graph_output.onnx"; + std::shared_ptr p_model; + ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_)); + Graph& graph = p_model->MainGraph(); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + graph_transformation_mgr.Register(onnxruntime::make_unique(), TransformerLevel::Level2); + graph_transformation_mgr.Register(onnxruntime::make_unique(), TransformerLevel::Level2); + auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, *logger_); + ASSERT_TRUE(ret.IsOK()); + + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_TRUE(op_to_count["Gelu"] == 0); + ASSERT_TRUE(op_to_count["BiasGelu"] == 0); +} + TEST_F(GraphTransformationTests, BiasGeluTest) { auto model_uri = MODEL_FOLDER "fusion/bias_gelu_fusion.onnx"; std::shared_ptr p_model; diff --git a/onnxruntime/test/testdata/transform/fusion/gelu_format2_0_with_bias_use_graph_output.onnx b/onnxruntime/test/testdata/transform/fusion/gelu_format2_0_with_bias_use_graph_output.onnx new file mode 100644 index 0000000000000000000000000000000000000000..146e24ef1818a8843d1cb74ee42fa83d9f2c3c43 GIT binary patch literal 847 zcmaiyJ#5op5Qg*PB<`uu)I_RkAkY>{Q;CWuAR$!W2N5U`(V<A=J}F-@gP;NdL4-+RxU?>;^xOcqUZ z=eo9|6iSA6rx}D>5YWw1%@V@8s+pEWZjqrR_^Fi0Q*a^cH3JsGQHNb5MqP099T%Op;wOAD80tc> zShHwE;?(hoG}D>Tcpk#pEzPLOCB?F|iYdpJU;FsGVQ=!pNYZ;g4?WObSc7hJtE)(uuQ|ps5cgAg;-evcF>t*)F-R= zr&(_|!X%hs)H6F;y1+K=C~KwW*z!3V$$0l^?mPF%6_59s=I9T`&e47x(ZilD{=KUg z;XLmZW*{U4vWk^^2D}igYGp$+f##mQ2JXg${}``~F516cNe(<=NKV6