saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add check of graph output in Bert Fusions (#4126)

Browse source 
* Refine node output edge checking in bert related fusions
This commit is contained in:
Tianlei Wu 2020-06-06 00:06:07 -07:00 committed by GitHub
parent ffed43e9b8
commit 7c8e1580a1
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
12 changed files with 104 additions and 76 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

33
onnxruntime/core/optimizer/attention_fusion.cc
View file

@ -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) {

3
onnxruntime/core/optimizer/bias_gelu_fusion.cc
View file

@ -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 <deque>
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;
}

24
onnxruntime/core/optimizer/embed_layer_norm_fusion.cc
View file

@ -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.

35
onnxruntime/core/optimizer/fast_gelu_fusion.cc
View file

@ -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<Node&>(*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<Node&>(*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",

10
onnxruntime/core/optimizer/gelu_fusion.cc
View file

@ -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;
}

33
onnxruntime/core/optimizer/layer_norm_fusion.cc
View file

@ -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 <deque>
@ -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() ||

2
onnxruntime/core/optimizer/reshape_fusion.cc
View file

@ -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;
}

9
onnxruntime/core/optimizer/utils.cc
View file

@ -227,5 +227,14 @@ bool IsSupportedDataType(const Node& node, const std::vector<std::string>& 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

7
onnxruntime/core/optimizer/utils.h
View file

@ -66,5 +66,12 @@ int32_t IndexOfNodeInput(const Node& node, const NodeArg& node_arg);
*/
bool IsSupportedDataType(const Node& node, const std::vector<std::string>& 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

17
onnxruntime/test/optimizer/graph_transform_test.cc
View file

@ -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<Model> 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<GeluFusion>(), TransformerLevel::Level2);
graph_transformation_mgr.Register(onnxruntime::make_unique<BiasGelu>(), TransformerLevel::Level2);
auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, *logger_);
ASSERT_TRUE(ret.IsOK());
std::map<std::string, int> 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<Model> p_model;

BIN
onnxruntime/test/testdata/transform/fusion/gelu_format2_0_with_bias_use_graph_output.onnx vendored Normal file
View file

Binary file not shown.

7
onnxruntime/test/testdata/transform/fusion/gelu_gen.py vendored
View file

@ -14,10 +14,12 @@ Generate test model for Gelu subgraph pattern 2:
has_bias = True # change it to True to generate gelu_format2_*_with_bias.onnx
gelu_use_graph_input = False # change it to False to let Gelu don't have graph inputs as inputs.
node_has_graph_output = True # change it to False to let Gelu don't have graph output
switch_order = True # switch order of inputs for Mul and Add
X = helper.make_tensor_value_info('input', TensorProto.FLOAT, ["batch", "seqlen", 64])
Y = helper.make_tensor_value_info('output', TensorProto.FLOAT, ["batch", "seqlen", 64])
Z = helper.make_tensor_value_info('div', TensorProto.FLOAT, ["batch", "seqlen", 64])
value = (0.01 * np.arange(64)).astype(np.float32).reshape((64))
bias_initializer = numpy_helper.from_array(value, "input_bias")
@ -74,7 +76,7 @@ if has_bias:
initializers.extend([initializer_sqrt_2, initializer_1, initializer_0_5])
# Create the graph (GraphProto)
graph_def = helper.make_graph(nodes, 'gelu_pattern_2', [X], [Y], initializers)
graph_def = helper.make_graph(nodes, 'gelu_pattern_2', [X], [Y, Z] if node_has_graph_output else [Y], initializers)
opsets = []
onnxdomain = OperatorSetIdProto()
@ -99,6 +101,9 @@ if has_bias:
if gelu_use_graph_input:
file_name += "_use_graph_input"
if node_has_graph_output:
file_name += "_use_graph_output"
file_name += ".onnx"
onnx.save(model_def, file_name)
print(file_name)