diff --git a/onnxruntime/core/graph/graph_utils.cc b/onnxruntime/core/graph/graph_utils.cc index 0334597549609..f083297c16fea 100644 --- a/onnxruntime/core/graph/graph_utils.cc +++ b/onnxruntime/core/graph/graph_utils.cc @@ -8,6 +8,7 @@ #include "core/common/logging/logging.h" #include +#include #include #include #include @@ -411,6 +412,14 @@ const ONNX_NAMESPACE::AttributeProto* GetNodeAttribute(const Node& node, const s return iter == attrs.end() ? nullptr : &iter->second; } +bool IsFullShapeNode(const Node& node) { + const auto* start_attr = GetNodeAttribute(node, "start"); + const auto* end_attr = GetNodeAttribute(node, "end"); + // end=INT64_MAX is the runtime default meaning "all dimensions" (full shape). + return (!start_attr || start_attr->i() == 0) && + (!end_attr || end_attr->i() == std::numeric_limits::max()); +} + static NodeArg& GetOrCreateNodeArg(Graph& graph, const ONNX_NAMESPACE::TensorProto& new_initializer) { ONNX_NAMESPACE::TypeProto new_type; auto* typeproto_tensor = new_type.mutable_tensor_type(); diff --git a/onnxruntime/core/graph/graph_utils.h b/onnxruntime/core/graph/graph_utils.h index 2106da1a96327..5681d4e1d08f0 100644 --- a/onnxruntime/core/graph/graph_utils.h +++ b/onnxruntime/core/graph/graph_utils.h @@ -8,6 +8,7 @@ #include "core/graph/onnx_protobuf.h" #include "core/graph/graph.h" +#include #include #include @@ -31,6 +32,10 @@ bool IsSupportedOptypeVersionAndDomain(const Node& node, /** Returns the attribute of a Node with a given name. */ const ONNX_NAMESPACE::AttributeProto* GetNodeAttribute(const Node& node, const std::string& attr_name); +/** Checks whether a Shape node returns the full tensor shape (all dimensions). + * Returns false if start/end attributes restrict the output to a subset of dimensions. */ +bool IsFullShapeNode(const Node& node); + /** Add a new initializer to 'graph'. Checks that new_initializer does not already exist in 'graph' before adding it. @returns The NodeArg for the new initializer. diff --git a/onnxruntime/core/optimizer/attention_fusion.cc b/onnxruntime/core/optimizer/attention_fusion.cc index 7fe7c914fa796..30eaaafccca82 100644 --- a/onnxruntime/core/optimizer/attention_fusion.cc +++ b/onnxruntime/core/optimizer/attention_fusion.cc @@ -349,7 +349,7 @@ static bool TryFuseMobileClipMHA(Node& qkv_matmul, const Node* sequence_transpose = graph_utils::GetInputNode(qkv_matmul, 0); if (sequence_transpose == nullptr || - !graph_utils::IsSupportedOptypeVersionAndDomain(*sequence_transpose, "Transpose", {1, 13}, kOnnxDomain) || + !graph_utils::IsSupportedOptypeVersionAndDomain(*sequence_transpose, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain) || !HasExpectedPerm(*sequence_transpose, {0, 2, 1}) || !optimizer_utils::CheckOutputEdges(graph, *sequence_transpose, 1)) { return false; @@ -357,14 +357,14 @@ static bool TryFuseMobileClipMHA(Node& qkv_matmul, const Node* input_reshape = graph_utils::GetInputNode(*sequence_transpose, 0); if (input_reshape == nullptr || - !graph_utils::IsSupportedOptypeVersionAndDomain(*input_reshape, "Reshape", {5, 13, 14}, kOnnxDomain) || + !graph_utils::IsSupportedOptypeVersionAndDomain(*input_reshape, "Reshape", {5, 13, 14, 19, 21, 23, 24, 25}, kOnnxDomain) || !optimizer_utils::CheckOutputEdges(graph, *input_reshape, 1)) { return fail("missing input Reshape before sequence transpose"); } Node* qkv_reshape = GetOnlyChildByOutputIndex(graph, qkv_matmul, 0, "Reshape"); if (qkv_reshape == nullptr || - !graph_utils::IsSupportedOptypeVersionAndDomain(*qkv_reshape, "Reshape", {5, 13, 14}, kOnnxDomain) || + !graph_utils::IsSupportedOptypeVersionAndDomain(*qkv_reshape, "Reshape", {5, 13, 14, 19, 21, 23, 24, 25}, kOnnxDomain) || !optimizer_utils::CheckOutputEdges(graph, *qkv_reshape, 1)) { return fail("qkv Reshape after MatMul not matched"); } @@ -379,9 +379,9 @@ static bool TryFuseMobileClipMHA(Node& qkv_matmul, Node* k_squeeze = GetOnlyChildByOutputIndex(graph, *split, 1, "Squeeze"); Node* v_transpose = GetOnlyChildByOutputIndex(graph, *split, 2, "Transpose"); if (q_transpose == nullptr || k_squeeze == nullptr || v_transpose == nullptr || - !graph_utils::IsSupportedOptypeVersionAndDomain(*q_transpose, "Transpose", {1, 13}, kOnnxDomain) || - !graph_utils::IsSupportedOptypeVersionAndDomain(*k_squeeze, "Squeeze", {13}, kOnnxDomain) || - !graph_utils::IsSupportedOptypeVersionAndDomain(*v_transpose, "Transpose", {1, 13}, kOnnxDomain) || + !graph_utils::IsSupportedOptypeVersionAndDomain(*q_transpose, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain) || + !graph_utils::IsSupportedOptypeVersionAndDomain(*k_squeeze, "Squeeze", {13, 21, 23, 24, 25}, kOnnxDomain) || + !graph_utils::IsSupportedOptypeVersionAndDomain(*v_transpose, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain) || !HasExpectedPerm(*q_transpose, {2, 0, 3, 1, 4}) || !HasExpectedPerm(*v_transpose, {2, 0, 3, 1, 4}) || !HasExpectedAxesInput(graph, *k_squeeze, {2})) { @@ -391,8 +391,8 @@ static bool TryFuseMobileClipMHA(Node& qkv_matmul, Node* q_squeeze = GetOnlyChildByOutputIndex(graph, *q_transpose, 0, "Squeeze"); Node* v_squeeze = GetOnlyChildByOutputIndex(graph, *v_transpose, 0, "Squeeze"); if (q_squeeze == nullptr || v_squeeze == nullptr || - !graph_utils::IsSupportedOptypeVersionAndDomain(*q_squeeze, "Squeeze", {13}, kOnnxDomain) || - !graph_utils::IsSupportedOptypeVersionAndDomain(*v_squeeze, "Squeeze", {13}, kOnnxDomain) || + !graph_utils::IsSupportedOptypeVersionAndDomain(*q_squeeze, "Squeeze", {13, 21, 23, 24, 25}, kOnnxDomain) || + !graph_utils::IsSupportedOptypeVersionAndDomain(*v_squeeze, "Squeeze", {13, 21, 23, 24, 25}, kOnnxDomain) || !HasExpectedAxesInput(graph, *q_squeeze, {0}) || !HasExpectedAxesInput(graph, *v_squeeze, {0})) { return fail("q/v squeeze pattern not matched"); @@ -402,7 +402,7 @@ static bool TryFuseMobileClipMHA(Node& qkv_matmul, Node* k_transpose = GetOnlyChildByOutputIndex(graph, *k_squeeze, 0, "Transpose"); if (q_scale_mul == nullptr || k_transpose == nullptr || !graph_utils::IsSupportedOptypeVersionAndDomain(*q_scale_mul, "Mul", {7, 13, 14}, kOnnxDomain) || - !graph_utils::IsSupportedOptypeVersionAndDomain(*k_transpose, "Transpose", {1, 13}, kOnnxDomain) || + !graph_utils::IsSupportedOptypeVersionAndDomain(*k_transpose, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain) || !HasExpectedPerm(*k_transpose, {0, 2, 3, 1})) { return fail("q scale Mul or k Transpose(0,2,3,1) not matched"); } @@ -460,7 +460,7 @@ static bool TryFuseMobileClipMHA(Node& qkv_matmul, Node* transpose_3 = GetOnlyChildByOutputIndex(graph, *qkv_matmul_1, 0, "Transpose"); if (transpose_3 == nullptr || - !graph_utils::IsSupportedOptypeVersionAndDomain(*transpose_3, "Transpose", {1, 13}, kOnnxDomain) || + !graph_utils::IsSupportedOptypeVersionAndDomain(*transpose_3, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain) || !HasExpectedPerm(*transpose_3, {0, 2, 1, 3}) || !optimizer_utils::CheckOutputEdges(graph, *transpose_3, 1)) { return fail("output Transpose(0,2,1,3) not matched"); @@ -468,7 +468,7 @@ static bool TryFuseMobileClipMHA(Node& qkv_matmul, Node* reshape_2 = GetOnlyChildByOutputIndex(graph, *transpose_3, 0, "Reshape"); if (reshape_2 == nullptr || - !graph_utils::IsSupportedOptypeVersionAndDomain(*reshape_2, "Reshape", {5, 13, 14}, kOnnxDomain) || + !graph_utils::IsSupportedOptypeVersionAndDomain(*reshape_2, "Reshape", {5, 13, 14, 19, 21, 23, 24, 25}, kOnnxDomain) || !optimizer_utils::CheckOutputEdges(graph, *reshape_2, 1)) { return fail("output Reshape not matched"); } @@ -497,7 +497,7 @@ static bool TryFuseMobileClipMHA(Node& qkv_matmul, if (proj_gemm == nullptr) { proj_gemm_input_reshape = GetOnlyChildByOutputIndex(graph, *reshape_2, 0, "Reshape"); if (proj_gemm_input_reshape == nullptr || - !graph_utils::IsSupportedOptypeVersionAndDomain(*proj_gemm_input_reshape, "Reshape", {5, 13, 14}, kOnnxDomain) || + !graph_utils::IsSupportedOptypeVersionAndDomain(*proj_gemm_input_reshape, "Reshape", {5, 13, 14, 19, 21, 23, 24, 25}, kOnnxDomain) || !optimizer_utils::CheckOutputEdges(graph, *proj_gemm_input_reshape, 1)) { return fail("projection MatMul/Gemm not matched"); } @@ -511,7 +511,7 @@ static bool TryFuseMobileClipMHA(Node& qkv_matmul, proj_gemm_output_reshape = GetOnlyChildByOutputIndex(graph, *proj_gemm, 0, "Reshape"); if (proj_gemm_output_reshape == nullptr || - !graph_utils::IsSupportedOptypeVersionAndDomain(*proj_gemm_output_reshape, "Reshape", {5, 13, 14}, kOnnxDomain) || + !graph_utils::IsSupportedOptypeVersionAndDomain(*proj_gemm_output_reshape, "Reshape", {5, 13, 14, 19, 21, 23, 24, 25}, kOnnxDomain) || !optimizer_utils::CheckOutputEdges(graph, *proj_gemm_output_reshape, 1)) { return fail("normalized projection Gemm output Reshape not matched"); } @@ -920,11 +920,11 @@ static bool FuseSubGraphQKImpl(Node& layer_norm, } std::vector q_path{ - {0, 0, "Transpose", {1, 13}, kOnnxDomain}, - {0, 0, "Reshape", {5, 13}, kOnnxDomain}, - {0, 0, "Add", {7, 13}, kOnnxDomain}, + {0, 0, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "Reshape", {5, 13, 14, 19, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "Add", {7, 13, 14}, kOnnxDomain}, {0, 0, "MatMul", {1, 9, 13}, kOnnxDomain}, - {0, 0, "LayerNormalization", {1}, kOnnxDomain}}; + {0, 0, "LayerNormalization", {1, 17}, kOnnxDomain}}; if (!graph_utils::FindPath(edges[edges.size() - 1]->GetNode(), true, q_path, edges, logger)) { DEBUG_LOG("Failed to find path for q"); return false; @@ -953,9 +953,9 @@ static bool FuseSubGraphQKImpl(Node& layer_norm, } std::vector k_path{ - {0, 1, "Transpose", {1, 13}, kOnnxDomain}, - {0, 0, "Reshape", {5, 13}, kOnnxDomain}, - {0, 0, "Add", {7, 13}, kOnnxDomain}, + {0, 1, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "Reshape", {5, 13, 14, 19, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "Add", {7, 13, 14}, kOnnxDomain}, {0, 0, "MatMul", {1, 9, 13}, kOnnxDomain}, {0, 0, "LayerNormalization", {1, 17}, kOnnxDomain}}; @@ -1070,8 +1070,8 @@ static bool FuseSubGraphQK(Node& layer_norm, const logging::Logger& logger) { // path to q std::vector q_varience_path{ - {0, 0, "Div", {7, 13}, kOnnxDomain}, - {0, 0, "MatMul", {1, 9}, kOnnxDomain}}; + {0, 0, "Div", {7, 13, 14}, kOnnxDomain}, + {0, 0, "MatMul", {1, 9, 13}, kOnnxDomain}}; std::vector edges; if (!graph_utils::FindPath(*(mask_nodes.add), true, q_varience_path, edges, logger)) { DEBUG_LOG("Failed to find path for q"); @@ -1163,7 +1163,7 @@ static bool FuseSubGraphQKDistilBert(Node& layer_norm, // path to q std::vector q_varience_path{ {0, 2, "MatMul", {1, 9, 13}, kOnnxDomain}, - {0, 0, "Div", {7, 13}, kOnnxDomain}}; + {0, 0, "Div", {7, 13, 14}, kOnnxDomain}}; std::vector edges; if (!graph_utils::FindPath(*(mask_nodes.where), true, q_varience_path, edges, logger)) { DEBUG_LOG("Failed to find path for q"); @@ -1265,14 +1265,14 @@ bool AttentionFusion::FuseSubGraph(Node& layer_norm, std::map& mask_int32_map, const logging::Logger& logger) { std::vector parent_path{ - {0, 0, "Add", {7, 13}, kOnnxDomain}, + {0, 0, "Add", {7, 13, 14}, kOnnxDomain}, {0, 0, "MatMul", {1, 9, 13}, kOnnxDomain}, - {0, 0, "Reshape", {5, 13}, kOnnxDomain}, - {0, 0, "Transpose", {1, 13}, kOnnxDomain}, + {0, 0, "Reshape", {5, 13, 14, 19, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain}, {0, 0, "MatMul", {1, 9, 13}, kOnnxDomain}, - {0, 1, "Transpose", {1, 13}, kOnnxDomain}, - {0, 0, "Reshape", {5, 13}, kOnnxDomain}, - {0, 0, "Add", {7, 13}, kOnnxDomain}, + {0, 1, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "Reshape", {5, 13, 14, 19, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "Add", {7, 13, 14}, kOnnxDomain}, {0, 0, "MatMul", {1, 9, 13}, kOnnxDomain}, {0, 0, "LayerNormalization", {1, 17}, kOnnxDomain}}; diff --git a/onnxruntime/core/optimizer/attention_fusion_helper.h b/onnxruntime/core/optimizer/attention_fusion_helper.h index a328a6d451a89..b7bd36c8df9e4 100644 --- a/onnxruntime/core/optimizer/attention_fusion_helper.h +++ b/onnxruntime/core/optimizer/attention_fusion_helper.h @@ -74,14 +74,14 @@ bool MatchGemmSubgraph(Graph& graph, DEBUG_LOG("Start MatchGemmSubgraph"); // GPT Attention fusion supports opset version 9 or later. std::vector parent_path{ - {0, dst_arg_index, "Reshape", {5, 13}, kOnnxDomain}, + {0, dst_arg_index, "Reshape", {5, 13, 14, 19, 21, 23, 24, 25}, kOnnxDomain}, {0, 0, "Gemm", {9, 11, 13}, kOnnxDomain}, - {0, 0, "Reshape", {5, 13}, kOnnxDomain}, + {0, 0, "Reshape", {5, 13, 14, 19, 21, 23, 24, 25}, kOnnxDomain}, {0, 1, "Concat", {4, 11, 13}, kOnnxDomain}, - {0, 1, "Unsqueeze", {1, 11, 13}, kOnnxDomain}, - {0, 0, "Squeeze", {1, 11, 13}, kOnnxDomain}, + {0, 1, "Unsqueeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "Squeeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, {0, 0, "Slice", {1, 10, 11, 13}, kOnnxDomain}, - {0, 0, "Shape", {1, 13}, kOnnxDomain}}; + {0, 0, "Shape", {1, 13, 15, 19, 21, 23, 24, 25}, kOnnxDomain}}; std::vector edges; if (!graph_utils::FindPath(node_after_gemm_reshape, true, parent_path, edges, logger)) { @@ -98,6 +98,14 @@ bool MatchGemmSubgraph(Graph& graph, const Node& slice = edges[6]->GetNode(); const Node& shape_before_slice = edges[7]->GetNode(); + // The downstream Slice/Squeeze/Gather nodes assume Shape returns the full tensor shape so + // that indices map directly to tensor dimensions. A partial shape (opset 15+ start/end + // attributes) would produce incorrect index mapping. + if (!graph_utils::IsFullShapeNode(shape_before_slice)) { + DEBUG_LOG("Shape node has non-default start/end attributes"); + return false; + } + const auto& subgraph_input = shape_before_slice.InputDefs()[0]; if (reshape_before_gemm.InputDefs()[0]->Name() != subgraph_input->Name()) { DEBUG_LOG("Input of reshape_before_gemm is not the input of subgraph"); @@ -166,9 +174,9 @@ bool MatchGemmSubgraph(Graph& graph, // Match: [Input] ----> Shape --> Gather (indices=0 or 1) --> Unsqueeze (axes=0) ----> Concat ( , , ) for (int i = 0; i < 2; i++) { std::vector gather_path1{ - {0, i, "Unsqueeze", {1, 11, 13}, kOnnxDomain}, + {0, i, "Unsqueeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, {0, 0, "Gather", {1, 11, 13}, kOnnxDomain}, - {0, 0, "Shape", {1, 13}, kOnnxDomain}}; + {0, 0, "Shape", {1, 13, 15, 19, 21, 23, 24, 25}, kOnnxDomain}}; if (!graph_utils::FindPath(concat_after_gather, true, gather_path1, edges, logger)) { DEBUG_LOG("Faild to match gemm gather path"); @@ -375,8 +383,8 @@ bool ValidateUnidirMask(const Graph& graph, const NodeArg& mask, bool& is_unidir bool MatchUnidirMaskSubgraph(const Graph& graph, const Node& add_node, MatchUnidirMaskResult& result, bool use_shared_node, const logging::Logger& logger) { DEBUG_LOG("Start MatchUnidirMaskSubgraph"); std::vector root_path{ - {0, 0, "Where", {9}, kOnnxDomain}, - {0, 1, "Div", {7, 13}, kOnnxDomain}}; + {0, 0, "Where", {9, 16}, kOnnxDomain}, + {0, 1, "Div", {7, 13, 14}, kOnnxDomain}}; std::vector edges; if (!graph_utils::FindPath(add_node, true, root_path, edges, logger)) { @@ -392,14 +400,14 @@ bool MatchUnidirMaskSubgraph(const Graph& graph, const Node& add_node, MatchUnid } std::vector path1{ - {0, 0, "Cast", {9, 13}, kOnnxDomain}, + {0, 0, "Cast", {9, 13, 19, 21, 23, 24, 25}, kOnnxDomain}, {0, 0, "Slice", {10, 11, 13}, kOnnxDomain}, // Last Slice {0, 0, "Slice", {10, 11, 13}, kOnnxDomain}, // Mask Slice - {0, 1, "Unsqueeze", {9, 11, 13}, kOnnxDomain}, - {0, 0, "Sub", {7, 13}, kOnnxDomain}, - {0, 0, "Squeeze", {1, 11, 13}, kOnnxDomain}, + {0, 1, "Unsqueeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "Sub", {7, 13, 14}, kOnnxDomain}, + {0, 0, "Squeeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, {0, 0, "Slice", {10, 11, 13}, kOnnxDomain}, // Slice 1 - {0, 0, "Shape", {1, 13}, kOnnxDomain}}; + {0, 0, "Shape", {1, 13, 15, 19, 21, 23, 24, 25}, kOnnxDomain}}; if (!graph_utils::FindPath(where_node, true, path1, edges, logger)) { DEBUG_LOG("Faild to match path 1 for unidirectional mask"); @@ -454,8 +462,8 @@ bool MatchUnidirMaskSubgraph(const Graph& graph, const Node& add_node, MatchUnid } std::vector slice_ends_path{ - {0, 2, "Unsqueeze", {9, 11, 13}, kOnnxDomain}, - {0, 0, "Squeeze", {1, 11, 13}, kOnnxDomain}}; + {0, 2, "Unsqueeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "Squeeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}}; if (!graph_utils::FindPath(last_slice, true, slice_ends_path, edges, logger) || edges[1]->GetNode().Index() != squeeze1.Index()) { @@ -482,9 +490,9 @@ bool MatchUnidirMaskSubgraph(const Graph& graph, const Node& add_node, MatchUnid } std::vector path4{ - {0, 1, "Squeeze", {1, 11, 13}, kOnnxDomain}, + {0, 1, "Squeeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, {0, 0, "Slice", {10, 11, 13}, kOnnxDomain}, // Slice 2 - {0, 0, "Shape", {1, 13}, kOnnxDomain}}; + {0, 0, "Shape", {1, 13, 15, 19, 21, 23, 24, 25}, kOnnxDomain}}; if (!graph_utils::FindPath(sub, true, path4, edges, logger)) { DEBUG_LOG("Faild to match path 4 for unidirectional mask"); @@ -632,9 +640,9 @@ bool MatchInputMaskSubgraph(const Graph& graph, const Node& qkv_matmul, Attentio } std::vector mask_path{ - {0, 0, "Add", {7, 13}, kOnnxDomain}, - {0, 1, "Mul", {7, 13}, kOnnxDomain}, - {0, 0, "Sub", {7, 13}, kOnnxDomain}}; + {0, 0, "Add", {7, 13, 14}, kOnnxDomain}, + {0, 1, "Mul", {7, 13, 14}, kOnnxDomain}, + {0, 0, "Sub", {7, 13, 14}, kOnnxDomain}}; if (!graph_utils::FindPath(softmax, true, mask_path, edges, logger)) { DEBUG_LOG("Failed to find path for mask"); @@ -766,7 +774,7 @@ bool MatchInputMaskSubgraph(const Graph& graph, const Node& layer_norm, const No // expand has another input Shape <-- qk_MatMul std::vector shape_path{ - {0, 1, "Shape", {1, 13}, kOnnxDomain}, + {0, 1, "Shape", {1, 13, 15, 19, 21, 23, 24, 25}, kOnnxDomain}, {0, 0, "MatMul", {1, 9, 13}, kOnnxDomain}}; if (!graph_utils::FindPath(expand, true, shape_path, edges, logger)) { DEBUG_LOG("Failed to find shape path"); @@ -791,9 +799,9 @@ bool MatchInputMaskSubgraph(const Graph& graph, const Node& layer_norm, const No // reshape node's shape input std::vector reshape_shape_path_1{ {0, 1, "Concat", {4, 11, 13}, kOnnxDomain}, - {0, 0, "Unsqueeze", {1, 11, 13}, kOnnxDomain}, + {0, 0, "Unsqueeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, {0, 0, "Gather", {1, 11, 13}, kOnnxDomain}, - {0, 0, "Shape", {1, 13}, kOnnxDomain}}; + {0, 0, "Shape", {1, 13, 15, 19, 21, 23, 24, 25}, kOnnxDomain}}; if (!graph_utils::FindPath(reshape, true, reshape_shape_path_1, edges, logger)) { DEBUG_LOG("Failed to find reshape shape path 1"); return false; @@ -809,9 +817,9 @@ bool MatchInputMaskSubgraph(const Graph& graph, const Node& layer_norm, const No } std::vector reshape_shape_path_2{ - {0, 3, "Unsqueeze", {1, 11, 13}, kOnnxDomain}, + {0, 3, "Unsqueeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, {0, 0, "Gather", {1, 11, 13}, kOnnxDomain}, - {0, 0, "Shape", {1, 13}, kOnnxDomain}}; + {0, 0, "Shape", {1, 13, 15, 19, 21, 23, 24, 25}, kOnnxDomain}}; if (!graph_utils::FindPath(concat, true, reshape_shape_path_2, edges, logger)) { DEBUG_LOG("Failed to find reshape shape path 2"); return false; @@ -887,7 +895,7 @@ bool MatchPastSubgraph(Graph& graph, const Node& k_concat, const Node& v_concat, MatchPastResult& result, const logging::Logger& logger) { DEBUG_LOG("Start MatchPastSubgraph"); std::vector past_k_path{ - {0, 0, "Transpose", {1, 13}, kOnnxDomain}, + {0, 0, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain}, {0, 0, "Gather", {1, 11, 13}, kOnnxDomain}}; if (transpose_optimized_pattern) { @@ -904,13 +912,13 @@ bool MatchPastSubgraph(Graph& graph, const Node& k_concat, const Node& v_concat, const Node& past_k_gather = edges[i++]->GetNode(); std::vector present_k_path{ - {0, 0, "Transpose", {1, 13}, kOnnxDomain}, - {0, 0, "Unsqueeze", {1, 11, 13}, kOnnxDomain}, + {0, 0, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "Unsqueeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, {0, 0, "Concat", {4, 11, 13}, kOnnxDomain}}; if (transpose_optimized_pattern) { present_k_path = { - {0, 0, "Unsqueeze", {1, 11, 13}, kOnnxDomain}, + {0, 0, "Unsqueeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, {0, 0, "Concat", {4, 11, 13}, kOnnxDomain}}; } @@ -924,7 +932,7 @@ bool MatchPastSubgraph(Graph& graph, const Node& k_concat, const Node& v_concat, const Node& present_concat = edges[i++]->GetNode(); std::vector present_past_v_path{ - {0, 1, "Unsqueeze", {1, 11, 13}, kOnnxDomain}, + {0, 1, "Unsqueeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, {0, 0, "Concat", {4, 11, 13}, kOnnxDomain}, {0, 0, "Gather", {1, 11, 13}, kOnnxDomain}}; if (!graph_utils::FindPath(present_concat, true, present_past_v_path, edges, logger)) { @@ -1024,7 +1032,7 @@ bool CheckDistilBertReshapeShape(const Graph& graph, const Node& reshape, int64_ // lazy check: record unqueeze first and then check in the mask path std::vector shape_path{ {0, 1, "Concat", {4, 11, 13}, kOnnxDomain}, - {0, 0, "Unsqueeze", {1, 11, 13}, kOnnxDomain}}; + {0, 0, "Unsqueeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}}; std::vector edges; if (!graph_utils::FindPath(reshape, true, shape_path, edges, logger)) { DEBUG_LOG("Failed to find shape path"); @@ -1339,9 +1347,9 @@ bool FuseGptAttention(Node& layer_norm, Graph& graph, int64_t hidden_size, std:: } std::vector path1{ - {0, 0, "Reshape", {5, 13}, kOnnxDomain}, - {0, 0, "Transpose", {1, 13}, kOnnxDomain}, - {0, 0, "MatMul", {1, 9}, kOnnxDomain}}; + {0, 0, "Reshape", {5, 13, 14, 19, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "MatMul", {1, 9, 13}, kOnnxDomain}}; std::vector edges; if (!graph_utils::FindPath(*gemm1_result.input_node, true, path1, edges, logger)) { @@ -1361,9 +1369,9 @@ bool FuseGptAttention(Node& layer_norm, Graph& graph, int64_t hidden_size, std:: bool has_past = graph_utils::IsSupportedOptypeVersionAndDomain(*v_concat, "Concat", {4, 11, 13}, kOnnxDomain); std::vector path2{ - {0, 1, "Transpose", {1, 13}, kOnnxDomain}, - {0, 0, "Reshape", {5, 13}, kOnnxDomain}, - {2, 0, "Split", {2, 11, 13}, kOnnxDomain}}; + {0, 1, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "Reshape", {5, 13, 14, 19, 21, 23, 24, 25}, kOnnxDomain}, + {2, 0, "Split", {2, 11, 13, 18}, kOnnxDomain}}; if (!graph_utils::FindPath(has_past ? *v_concat : qkv_matmul, true, path2, edges, logger)) { DEBUG_LOG("Faild to find path v to Split"); @@ -1413,9 +1421,9 @@ bool FuseGptAttention(Node& layer_norm, Graph& graph, int64_t hidden_size, std:: // path to q std::vector q_path{ {0, 0, "MatMul", {1, 9, 13}, kOnnxDomain}, - {0, 0, "Transpose", {1, 13}, kOnnxDomain}, - {0, 0, "Reshape", {5, 13}, kOnnxDomain}, - {0, 0, "Split", {2, 11, 13}, kOnnxDomain}}; + {0, 0, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "Reshape", {5, 13, 14, 19, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "Split", {2, 11, 13, 18}, kOnnxDomain}}; const Node* qk_div = unidir_mask_result.div_node; if (!graph_utils::FindPath(*qk_div, true, q_path, edges, logger)) { @@ -1447,7 +1455,7 @@ bool FuseGptAttention(Node& layer_norm, Graph& graph, int64_t hidden_size, std:: return false; } - if (graph_utils::IsSupportedOptypeVersionAndDomain(*k_concat, "Transpose", {1, 13, 21}, kOnnxDomain)) { + if (graph_utils::IsSupportedOptypeVersionAndDomain(*k_concat, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain)) { transpose_optimized_pattern = true; DEBUG_LOG("Using transpose optimized pattern"); opt_k_transpose = k_concat; @@ -1471,9 +1479,9 @@ bool FuseGptAttention(Node& layer_norm, Graph& graph, int64_t hidden_size, std:: // path to k std::vector k_path{ - {0, 1, "Transpose", {1, 13}, kOnnxDomain}, - {0, 0, "Reshape", {5, 13}, kOnnxDomain}, - {1, 0, "Split", {2, 11, 13}, kOnnxDomain}}; + {0, 1, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "Reshape", {5, 13, 14, 19, 21, 23, 24, 25}, kOnnxDomain}, + {1, 0, "Split", {2, 11, 13, 18}, kOnnxDomain}}; if (!graph_utils::FindPath(has_past ? *k_concat : qk_matmul, true, k_path, edges, logger)) { DEBUG_LOG("Failed to find path for k"); diff --git a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc index 606e91ce91bbb..e2a3ef2a4fe2d 100644 --- a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc +++ b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc @@ -115,9 +115,9 @@ static bool MatchInputToConcatSubgraph( const NodeIndex expected_gather_node_1_index) { std::vector expand_parent_path1{ {0, index, "Concat", {4, 11, 13}, kOnnxDomain}, - {0, 0, "Unsqueeze", {1, 11, 13}, kOnnxDomain}, + {0, 0, "Unsqueeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, {0, 0, "Gather", {1, 11, 13}, kOnnxDomain}, - {0, 0, "Shape", {1, 13}, kOnnxDomain}, + {0, 0, "Shape", {1, 13, 15, 19, 21, 23, 24, 25}, kOnnxDomain}, }; std::vector edges; @@ -138,6 +138,14 @@ static bool MatchInputToConcatSubgraph( } } + // The Gather(index=0) below assumes Shape returns the full tensor shape. A partial shape + // (opset 15+ start/end attributes) would cause Gather to pick the wrong dimension. + const Node& shape_node_path1 = edges[shape_index]->GetNode(); + if (!graph_utils::IsFullShapeNode(shape_node_path1)) { + DEBUG_LOG("Shape node in path 1 has non-default start/end attributes."); + return false; + } + Node& concat_node = *graph.GetNode(edges[0]->GetNode().Index()); Node& gather_node_0 = *graph.GetNode(edges[2]->GetNode().Index()); Node& shape_node_0 = *graph.GetNode(edges[3]->GetNode().Index()); @@ -147,9 +155,9 @@ static bool MatchInputToConcatSubgraph( } std::vector concat_parent_path{ - {0, 1, "Unsqueeze", {1, 11, 13}, kOnnxDomain}, + {0, 1, "Unsqueeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, {0, 0, "Gather", {1, 11, 13}, kOnnxDomain}, - {0, 0, "Shape", {1, 13}, kOnnxDomain}}; + {0, 0, "Shape", {1, 13, 15, 19, 21, 23, 24, 25}, kOnnxDomain}}; if (!graph_utils::FindPath(concat_node, true, concat_parent_path, edges, logger)) { DEBUG_LOG("Failed to find path 2 of position shape."); @@ -167,6 +175,13 @@ static bool MatchInputToConcatSubgraph( Node& gather_node_1 = *graph.GetNode(edges[1]->GetNode().Index()); Node& shape_node_1 = *graph.GetNode(edges[2]->GetNode().Index()); + // The Gather(index=1) below assumes Shape returns the full tensor shape. A partial shape + // (opset 15+ start/end attributes) would cause Gather to pick the wrong dimension. + if (!graph_utils::IsFullShapeNode(shape_node_1)) { + DEBUG_LOG("Shape node in path 2 has non-default start/end attributes."); + return false; + } + // The gather node (with second input indices==1) is also shared by other subgraph if (expected_gather_node_1_index != gather_node_1.Index()) { DEBUG_LOG("Gather node in path 2 is not linked to another subgraph."); @@ -231,42 +246,42 @@ static bool MatchPositionEmbeddingSubgraphsFromGather( // --> Cast --> Unsqueeze --> Expand --> Gather std::vector parent_path_1{ {0, 1, "Expand", {8, 13}, kOnnxDomain}, - {0, 0, "Unsqueeze", {1, 11, 13}, kOnnxDomain}, - {0, 0, "Cast", {9, 13}, kOnnxDomain}, - {0, 0, "Squeeze", {1, 11, 13}, kOnnxDomain}, - {0, 0, "Transpose", {1, 13}, kOnnxDomain}, + {0, 0, "Unsqueeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "Cast", {9, 13, 19, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "Squeeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain}, {0, 0, "NonZero", {9, 13}, kOnnxDomain}, - {0, 0, "ConstantOfShape", {9}, kOnnxDomain}, - {0, 0, "Unsqueeze", {1, 11, 13}, kOnnxDomain}, + {0, 0, "ConstantOfShape", {9, 20, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "Unsqueeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, {0, 0, "Gather", {1, 11, 13}, kOnnxDomain}, - {0, 0, "Shape", {1, 13}, kOnnxDomain}}; + {0, 0, "Shape", {1, 13, 15, 19, 21, 23, 24, 25}, kOnnxDomain}}; // Look for Path 2 (Path 1 with no cast): std::vector parent_path_2{ {0, 1, "Expand", {8, 13}, kOnnxDomain}, - {0, 0, "Unsqueeze", {1, 11, 13}, kOnnxDomain}, - {0, 0, "Squeeze", {1, 11, 13}, kOnnxDomain}, - {0, 0, "Transpose", {1, 13}, kOnnxDomain}, + {0, 0, "Unsqueeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "Squeeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain}, {0, 0, "NonZero", {9, 13}, kOnnxDomain}, - {0, 0, "ConstantOfShape", {9}, kOnnxDomain}, - {0, 0, "Unsqueeze", {1, 11, 13}, kOnnxDomain}, + {0, 0, "ConstantOfShape", {9, 20, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "Unsqueeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, {0, 0, "Gather", {1, 11, 13}, kOnnxDomain}, - {0, 0, "Shape", {1, 13}, kOnnxDomain}}; + {0, 0, "Shape", {1, 13, 15, 19, 21, 23, 24, 25}, kOnnxDomain}}; // Path 3 Pattern: // Shape -> Gather -> Cast (to=7) -> Range (start=0, delta=1) -> Unsqueeze -> Expand std::vector parent_path_3{ {0, 1, "Expand", {8, 13}, kOnnxDomain}, - {0, 0, "Unsqueeze", {1, 11, 13}, kOnnxDomain}, + {0, 0, "Unsqueeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, {0, 0, "Range", {1, 11}, kOnnxDomain}, - {0, 1, "Cast", {9, 13}, kOnnxDomain}, + {0, 1, "Cast", {9, 13, 19, 21, 23, 24, 25}, kOnnxDomain}, {0, 0, "Gather", {1, 11, 13}, kOnnxDomain}, - {0, 0, "Shape", {1, 13}, kOnnxDomain}}; + {0, 0, "Shape", {1, 13, 15, 19, 21, 23, 24, 25}, kOnnxDomain}}; // Path 4 pattern (Path 3 with no "Cast"): std::vector parent_path_4{ {0, 1, "Expand", {8, 13}, kOnnxDomain}, - {0, 0, "Unsqueeze", {1, 11, 13}, kOnnxDomain}, + {0, 0, "Unsqueeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, {0, 0, "Range", {1, 11}, kOnnxDomain}, {0, 1, "Gather", {1, 11, 13}, kOnnxDomain}, - {0, 0, "Shape", {1, 13}, kOnnxDomain}}; + {0, 0, "Shape", {1, 13, 15, 19, 21, 23, 24, 25}, kOnnxDomain}}; // Match one of the three path patterns. if (!graph_utils::FindPath(position_gather_node, true, parent_path_1, pg_edges, logger) && !graph_utils::FindPath(position_gather_node, true, parent_path_2, pg_edges, logger) && @@ -318,7 +333,7 @@ static bool MatchPositionEmbeddingSubgraphsFromGather( // Match Shape --> Expand path. std::vector pg_edges_2; - if (!graph_utils::FindPath(expand_node, true, {graph_utils::EdgeEndToMatch{0, 1, "Shape", {1, 13}, kOnnxDomain}}, pg_edges_2, logger)) { + if (!graph_utils::FindPath(expand_node, true, {graph_utils::EdgeEndToMatch{0, 1, "Shape", {1, 13, 15, 19, 21, 23, 24, 25}, kOnnxDomain}}, pg_edges_2, logger)) { DEBUG_LOG("Failed to match Shape node. "); return false; } @@ -338,9 +353,9 @@ static bool MatchPositionEmbeddingSubgraphsFromGather( // -------------------- std::vector pg_edges_2; std::vector path_to_match_1{ - {0, 1, "Where", {9}, kOnnxDomain}, - {0, 0, "Equal", {1, 7, 11, 13}, kOnnxDomain}, - {0, 0, "Reshape", {5, 13}, kOnnxDomain}}; + {0, 1, "Where", {9, 16}, kOnnxDomain}, + {0, 0, "Equal", {1, 7, 11, 13, 19}, kOnnxDomain}, + {0, 0, "Reshape", {5, 13, 14, 19, 21, 23, 24, 25}, kOnnxDomain}}; if (graph_utils::FindPath(expand_node, true, path_to_match_1, pg_edges_2, logger)) { if (!optimizer_utils::CheckOutputEdges(graph, pg_edges_2[0]->GetNode(), 1) || !optimizer_utils::CheckOutputEdges(graph, pg_edges_2[1]->GetNode(), 1) || @@ -559,7 +574,7 @@ static bool FuseSubGraph(Graph& graph, // Trace back to find Gather --> Add --> LayerNormalization std::vector word_embedding_path{ - {0, 0, "Add", {7, 13}, kOnnxDomain}, + {0, 0, "Add", {7, 13, 14}, kOnnxDomain}, {0, 0, "Gather", {1, 11, 13}, kOnnxDomain}}; if (!graph_utils::FindPath(layer_norm_add_node, true, word_embedding_path, edges, logger)) { return false; @@ -843,7 +858,7 @@ Status EmbedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l std::vector edges; // Find Add --> LayerNormalization - if (!graph_utils::FindPath(layer_norm_node, true, {graph_utils::EdgeEndToMatch{0, 0, "Add", {7, 13}, kOnnxDomain}}, edges, logger)) { + if (!graph_utils::FindPath(layer_norm_node, true, {graph_utils::EdgeEndToMatch{0, 0, "Add", {7, 13, 14}, kOnnxDomain}}, edges, logger)) { continue; } Node& layer_norm_add_node = *graph.GetNode(edges[0]->GetNode().Index()); diff --git a/onnxruntime/core/optimizer/group_query_attention_pre_norm_fusion.cc b/onnxruntime/core/optimizer/group_query_attention_pre_norm_fusion.cc index 909229822f134..3271c4cc9a22f 100644 --- a/onnxruntime/core/optimizer/group_query_attention_pre_norm_fusion.cc +++ b/onnxruntime/core/optimizer/group_query_attention_pre_norm_fusion.cc @@ -79,7 +79,7 @@ bool MatchPreNormReshapeChain(Graph& graph, Node* reshape_outer = graph.GetMutableProducerNode(consumer_input->Name()); if (reshape_outer == nullptr || - !graph_utils::IsSupportedOptypeVersionAndDomain(*reshape_outer, "Reshape", {5, 13, 14, 19, 21, 23})) { + !graph_utils::IsSupportedOptypeVersionAndDomain(*reshape_outer, "Reshape", {5, 13, 14, 19, 21, 23, 24, 25})) { return false; } if (reshape_outer->GetOutputEdgesCount() != 1) { @@ -174,7 +174,7 @@ bool MatchPreNormReshapeChain(Graph& graph, } Node* reshape_inner = graph.GetMutableProducerNode(sln->InputDefs()[0]->Name()); if (reshape_inner == nullptr || - !graph_utils::IsSupportedOptypeVersionAndDomain(*reshape_inner, "Reshape", {5, 13, 14, 19, 21, 23})) { + !graph_utils::IsSupportedOptypeVersionAndDomain(*reshape_inner, "Reshape", {5, 13, 14, 19, 21, 23, 24, 25})) { return false; } if (reshape_inner->GetOutputEdgesCount() != 1) { diff --git a/onnxruntime/core/optimizer/reshape_fusion.cc b/onnxruntime/core/optimizer/reshape_fusion.cc index f50c0e2e635bc..65f477d9c7479 100644 --- a/onnxruntime/core/optimizer/reshape_fusion.cc +++ b/onnxruntime/core/optimizer/reshape_fusion.cc @@ -15,7 +15,10 @@ namespace onnxruntime { bool GetAxesFromUnsqueezeNode(const Graph& graph, const Node& unsqueeze, InlinedVector& axes) { if (graph_utils::MatchesOpSinceVersion(unsqueeze, {1, 11})) { return graph_utils::GetRepeatedNodeAttributeValues(unsqueeze, "axes", axes); - } else if (graph_utils::MatchesOpSinceVersion(unsqueeze, {13})) { + } + + // Opset 13+ moved axes from attribute to input[1]. + if (unsqueeze.InputDefs().size() > 1) { const NodeArg* axes_node_arg = unsqueeze.InputDefs()[1]; return optimizer_utils::AppendTensorFromInitializer(graph, *axes_node_arg, axes, true); } @@ -169,12 +172,11 @@ bool ReshapeFusion::Match_One_Element_Output_Subgraph_1(Graph& graph, const Node const Node& gather = edges[1]->GetNode(); const Node& shape = edges[2]->GetNode(); - if (graph_utils::MatchesOpSinceVersion(shape, {15})) { - const ONNX_NAMESPACE::AttributeProto* start_attr = graph_utils::GetNodeAttribute(shape, "start"); - const ONNX_NAMESPACE::AttributeProto* end_attr = graph_utils::GetNodeAttribute(shape, "end"); - if (!((!start_attr || static_cast(start_attr->i()) == 0) && (!end_attr))) { - return false; - } + // The fusion assumes Shape returns the full tensor shape so that Gather indices correspond + // directly to tensor dimensions. A partial shape (opset 15+ start/end attributes) would shift + // the index mapping and produce incorrect results. + if (!graph_utils::IsFullShapeNode(shape)) { + return false; } InlinedVector axes; diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index f14327af4b8dd..5c76074d6b119 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -15,6 +15,7 @@ #include "onnx/defs/parser.h" #include "onnx/defs/printer.h" +#include "onnx/defs/schema.h" #include "core/common/span_utils.h" #include "core/framework/data_types.h" @@ -6597,6 +6598,26 @@ TEST_F(GraphTransformationTests, AttentionFusionDistilBertTest) { EXPECT_EQ(op_to_count["Shape"], 0); } +// These tests verify that attention fusions fire at the CURRENT max ONNX opset. +// When the ONNX opset advances (e.g., via submodule update), nodes will report a new +// SinceVersion(). If the optimizer's version lists are not updated, the fusion will +// fail to match and these tests will fail. +// +// To fix: update the EdgeEndToMatch version lists in the indicated optimizer file +// to include the new opset version for each affected op type. +// +// NOTE: GPT-2 and DistilBert attention models use attribute-based Unsqueeze/Squeeze +// in their mask subgraphs (opset <= 12 only). Those patterns cannot be converted to +// current opset without rewriting the mask matching logic. The MobileCLIP test +// (programmatically built) covers the current-opset attention fusion paths. + +static int GetCurrentOnnxOpset() { + const auto& map = ONNX_NAMESPACE::OpSchemaRegistry::DomainToVersionRange::Instance().Map(); + auto it = map.find(ONNX_NAMESPACE::ONNX_DOMAIN); + EXPECT_TRUE(it != map.end()) << "ONNX domain not found in OpSchemaRegistry"; + return it != map.end() ? it->second.second : 0; +} + enum class MobileClipProjectionType { MatMulAdd, GemmWithReshapes, @@ -6883,6 +6904,21 @@ TEST_F(GraphTransformationTests, AttentionFusionMobileClipMhaTest) { std::make_unique()); } +TEST_F(GraphTransformationTests, AttentionFusionMobileClipMhaCurrentOpsetTest) { + auto build_test_case = [](ModelTestBuilder& builder) { + BuildMobileClipAttentionTestCase(builder, MobileClipProjectionType::MatMulAdd); + }; + + TransformerTester(build_test_case, + CheckMobileClipAttentionFusedSession, + TransformerLevel::Level1, + TransformerLevel::Level2, + GetCurrentOnnxOpset(), + 1e-3, + 0.0, + std::make_unique()); +} + TEST_F(GraphTransformationTests, AttentionFusionMobileClipMhaProjectionGemmTest) { auto build_test_case = [](ModelTestBuilder& builder) { BuildMobileClipAttentionTestCase(builder, MobileClipProjectionType::GemmWithReshapes); @@ -6975,6 +7011,255 @@ TEST_F(GraphTransformationTests, AttentionFusionMobileClipMhaProjectionRewriteFa CheckMobileClipAttentionUnfusedMatMulGraph)); } +// Current-opset regression tests for fusion optimizers. +// These construct minimal graphs at the current ONNX opset and verify the optimizer fires. +// If ONNX bumps an op's since_version, the optimizer's version list will miss it and the test fails. + +TEST_F(GraphTransformationTests, GeluFusionCurrentOpsetTest) { + // Pattern: x -> Div(sqrt2) -> Erf -> Add(1) -> Mul(x) -> Mul(0.5) -> output + int current_opset = GetCurrentOnnxOpset(); + auto build_test_case = [](ModelTestBuilder& builder) { + auto* input = builder.MakeInput({{2, 3, 64}}); + auto* sqrt2 = builder.MakeInitializer({}, {1.4142135f}); + auto* one = builder.MakeInitializer({}, {1.0f}); + auto* half = builder.MakeInitializer({}, {0.5f}); + + auto* div_out = builder.MakeIntermediate(); + auto* erf_out = builder.MakeIntermediate(); + auto* add_out = builder.MakeIntermediate(); + auto* mul1_out = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + + builder.AddNode("Div", {input, sqrt2}, {div_out}); + builder.AddNode("Erf", {div_out}, {erf_out}); + builder.AddNode("Add", {erf_out, one}, {add_out}); + builder.AddNode("Mul", {input, add_out}, {mul1_out}); + builder.AddNode("Mul", {mul1_out, half}, {output}); + }; + + auto post_graph_checker = [current_opset](Graph& graph) { + auto op_to_count = CountOpsInGraph(graph); + if (op_to_count["Gelu"] == 1 || op_to_count["com.microsoft.Gelu"] == 1) { + TEST_RETURN_IF_NOT(op_to_count["Div"] == 0); + TEST_RETURN_IF_NOT(op_to_count["Erf"] == 0); + TEST_RETURN_IF_NOT(op_to_count["Add"] == 0); + TEST_RETURN_IF_NOT(op_to_count["Mul"] == 0); + return Status::OK(); + } + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, + "Gelu fusion failed at opset ", current_opset, + ". Remaining ops: Div=", op_to_count["Div"], + " Erf=", op_to_count["Erf"], + " Add=", op_to_count["Add"], + " Mul=", op_to_count["Mul"], + ". Either update version lists in " + "onnxruntime/core/optimizer/gelu_fusion.cc" + " or skip this opset in the test if the fusion is not expected to apply."); + }; + + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, current_opset, *logger_, + std::make_unique(), + TransformerLevel::Level1, 1, nullptr, post_graph_checker)); +} + +TEST_F(GraphTransformationTests, FastGeluFusionCurrentOpsetTest) { + // FastGelu pattern: x*x*0.044715 + x -> mul(sqrt(2/pi)) -> tanh -> add(1) -> mul(0.5) -> mul(x) + int current_opset = GetCurrentOnnxOpset(); + auto build_test_case = [](ModelTestBuilder& builder) { + auto* input = builder.MakeInput({{2, 3, 64}}); + auto* coeff = builder.MakeInitializer({}, {0.044715f}); + auto* sqrt2pi = builder.MakeInitializer({}, {0.7978845f}); // sqrt(2/pi) + auto* one = builder.MakeInitializer({}, {1.0f}); + auto* half = builder.MakeInitializer({}, {0.5f}); + auto* three = builder.MakeInitializer({}, {3.0f}); + + auto* pow_out = builder.MakeIntermediate(); + auto* mul1_out = builder.MakeIntermediate(); + auto* add1_out = builder.MakeIntermediate(); + auto* mul2_out = builder.MakeIntermediate(); + auto* tanh_out = builder.MakeIntermediate(); + auto* add2_out = builder.MakeIntermediate(); + auto* mul_half_out = builder.MakeIntermediate(); + auto* mul_final_out = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + + builder.AddNode("Pow", {input, three}, {pow_out}); + builder.AddNode("Mul", {pow_out, coeff}, {mul1_out}); + builder.AddNode("Add", {input, mul1_out}, {add1_out}); + builder.AddNode("Mul", {add1_out, sqrt2pi}, {mul2_out}); + builder.AddNode("Tanh", {mul2_out}, {tanh_out}); + builder.AddNode("Add", {tanh_out, one}, {add2_out}); + builder.AddNode("Mul", {input, half}, {mul_half_out}); + builder.AddNode("Mul", {mul_half_out, add2_out}, {mul_final_out}); + builder.AddNode("Identity", {mul_final_out}, {output}); + }; + + auto post_graph_checker = [current_opset](Graph& graph) { + auto op_to_count = CountOpsInGraph(graph); + if (op_to_count["com.microsoft.FastGelu"] == 1) { + TEST_RETURN_IF_NOT(op_to_count["Tanh"] == 0); + TEST_RETURN_IF_NOT(op_to_count["Pow"] == 0); + return Status::OK(); + } + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, + "FastGelu fusion failed at opset ", current_opset, + ". Remaining ops: Tanh=", op_to_count["Tanh"], + " Pow=", op_to_count["Pow"], + " Mul=", op_to_count["Mul"], + " Add=", op_to_count["Add"], + ". Either update version lists in " + "onnxruntime/core/optimizer/fast_gelu_fusion.cc" + " or skip this opset in the test if the fusion is not expected to apply."); + }; + + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, current_opset, *logger_, + std::make_unique(), + TransformerLevel::Level2, 1, nullptr, post_graph_checker)); +} + +TEST_F(GraphTransformationTests, BiasGeluFusionCurrentOpsetTest) { + // BiasGelu pattern: Add(input, bias) -> Gelu -> output (requires opset >= 20 for ONNX Gelu) + int current_opset = GetCurrentOnnxOpset(); + if (current_opset < 20) { + GTEST_SKIP() << "BiasGelu with ONNX Gelu requires opset >= 20"; + } + auto build_test_case = [](ModelTestBuilder& builder) { + auto* input = builder.MakeInput({{2, 3, 64}}); + auto* bias = builder.MakeInitializer({64}, -0.5f, 0.5f); + auto* add_out = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + + builder.AddNode("Add", {input, bias}, {add_out}); + builder.AddNode("Gelu", {add_out}, {output}); + }; + + auto post_graph_checker = [current_opset](Graph& graph) { + auto op_to_count = CountOpsInGraph(graph); + if (op_to_count["com.microsoft.BiasGelu"] == 1) { + TEST_RETURN_IF_NOT(op_to_count["Add"] == 0); + TEST_RETURN_IF_NOT(op_to_count["Gelu"] == 0); + return Status::OK(); + } + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, + "BiasGelu fusion failed at opset ", current_opset, + ". Remaining ops: Add=", op_to_count["Add"], + " Gelu=", op_to_count["Gelu"], + ". Either update version lists in " + "onnxruntime/core/optimizer/bias_gelu_fusion.cc" + " or skip this opset in the test if the fusion is not expected to apply."); + }; + + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, current_opset, *logger_, + std::make_unique(), + TransformerLevel::Level2, 1, nullptr, post_graph_checker)); +} + +TEST_F(GraphTransformationTests, MatMulAddFusionCurrentOpsetTest) { + // MatMul + Add -> Gemm fusion + int current_opset = GetCurrentOnnxOpset(); + auto build_test_case = [](ModelTestBuilder& builder) { + auto* input = builder.MakeInput({{2, 64}}); + auto* weights = builder.MakeInitializer({64, 32}, -1.0f, 1.0f); + auto* bias = builder.MakeInitializer({32}, -0.5f, 0.5f); + auto* matmul_out = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + + builder.AddNode("MatMul", {input, weights}, {matmul_out}); + builder.AddNode("Add", {matmul_out, bias}, {output}); + }; + + auto post_graph_checker = [current_opset](Graph& graph) { + auto op_to_count = CountOpsInGraph(graph); + if (op_to_count["Gemm"] == 1) { + TEST_RETURN_IF_NOT(op_to_count["MatMul"] == 0); + TEST_RETURN_IF_NOT(op_to_count["Add"] == 0); + return Status::OK(); + } + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, + "MatMulAdd fusion failed at opset ", current_opset, + ". Remaining ops: MatMul=", op_to_count["MatMul"], + " Add=", op_to_count["Add"], + ". Either update version lists in " + "onnxruntime/core/optimizer/matmul_add_fusion.cc" + " or skip this opset in the test if the fusion is not expected to apply."); + }; + + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, current_opset, *logger_, + std::make_unique(), + TransformerLevel::Level1, 1, nullptr, post_graph_checker)); +} + +TEST_F(GraphTransformationTests, DivMulFusionCurrentOpsetTest) { + // 1/x * y -> y/x fusion + int current_opset = GetCurrentOnnxOpset(); + auto build_test_case = [](ModelTestBuilder& builder) { + auto* input_x = builder.MakeInput({{2, 64}}); + auto* input_y = builder.MakeInput({{2, 64}}); + auto* one = builder.MakeInitializer({}, {1.0f}); + auto* div_out = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + + builder.AddNode("Div", {one, input_x}, {div_out}); + builder.AddNode("Mul", {div_out, input_y}, {output}); + }; + + auto post_graph_checker = [current_opset](Graph& graph) { + auto op_to_count = CountOpsInGraph(graph); + if (op_to_count["Div"] == 1 && op_to_count["Mul"] == 0) { + return Status::OK(); + } + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, + "DivMul fusion failed at opset ", current_opset, + ". Remaining ops: Div=", op_to_count["Div"], + " Mul=", op_to_count["Mul"], + ". Either update version lists in " + "onnxruntime/core/optimizer/div_mul_fusion.cc" + " or skip this opset in the test if the fusion is not expected to apply."); + }; + + auto rule_transformer = std::make_unique("DivMulFusionCurrentOpset"); + ASSERT_STATUS_OK(rule_transformer->Register(std::make_unique())); + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, current_opset, *logger_, + std::move(rule_transformer), + TransformerLevel::Level1, 1, nullptr, post_graph_checker)); +} + +TEST_F(GraphTransformationTests, QuickGeluFusionCurrentOpsetTest) { + // x * Sigmoid(alpha * x) -> QuickGelu(x, alpha) fusion + int current_opset = GetCurrentOnnxOpset(); + auto build_test_case = [](ModelTestBuilder& builder) { + auto* input = builder.MakeInput({{2, 3, 64}}); + auto* alpha = builder.MakeInitializer({}, {1.702f}); + auto* mul1_out = builder.MakeIntermediate(); + auto* sigmoid_out = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + + builder.AddNode("Mul", {input, alpha}, {mul1_out}); + builder.AddNode("Sigmoid", {mul1_out}, {sigmoid_out}); + builder.AddNode("Mul", {input, sigmoid_out}, {output}); + }; + + auto post_graph_checker = [current_opset](Graph& graph) { + auto op_to_count = CountOpsInGraph(graph); + if (op_to_count["com.microsoft.QuickGelu"] == 1) { + TEST_RETURN_IF_NOT(op_to_count["Mul"] == 0); + TEST_RETURN_IF_NOT(op_to_count["Sigmoid"] == 0); + return Status::OK(); + } + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, + "QuickGelu fusion failed at opset ", current_opset, + ". Remaining ops: Mul=", op_to_count["Mul"], + " Sigmoid=", op_to_count["Sigmoid"], + ". Either update version lists in " + "onnxruntime/core/optimizer/quick_gelu_fusion.cc" + " or skip this opset in the test if the fusion is not expected to apply."); + }; + + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, current_opset, *logger_, + std::make_unique(), + TransformerLevel::Level2, 1, nullptr, post_graph_checker)); +} + TEST_F(GraphTransformationTests, GeluFusionTest) { constexpr const ORTCHAR_T* model_uri = MODEL_FOLDER "fusion/gelu.onnx"; std::shared_ptr p_model; @@ -8219,8 +8504,14 @@ TEST_F(GraphTransformationTests, ReshapeFusionOpsetTest) { return Status::OK(); }; - const std::vector opsets{11, 12, 13, 14, 15, 18}; - bool shape_test_for_opset15 = false; + // Include the current max opset to ensure the fusion stays up-to-date. + // If this test fails at the current opset, update version lists in + // onnxruntime/core/optimizer/reshape_fusion.cc. + std::vector opsets{11, 12, 13, 14, 15, 18, 19, 21, 23, 24, 25}; + int current_opset = GetCurrentOnnxOpset(); + if (std::find(opsets.begin(), opsets.end(), current_opset) == opsets.end()) { + opsets.push_back(current_opset); + } for (auto& opset : opsets) { auto build_test_case = [&](ModelTestBuilder& builder) { @@ -8245,14 +8536,7 @@ TEST_F(GraphTransformationTests, ReshapeFusionOpsetTest) { builder.AddNode("Add", {input_arg0, input_arg1}, {add_out}); if (opset_version >= 15) { - if (shape_test_for_opset15) { - auto& shape_1 = builder.AddNode("Shape", {add_out}, {shape_out}); - shape_1.AddAttribute("start", (int64_t)1); - shape_1.AddAttribute("end", (int64_t)2); - } else { - builder.AddNode("Shape", {add_out}, {shape_out}).AddAttribute("start", (int64_t)0); - shape_test_for_opset15 = true; - } + builder.AddNode("Shape", {add_out}, {shape_out}).AddAttribute("start", (int64_t)0); } else { builder.AddNode("Shape", {add_out}, {shape_out}); } @@ -8271,13 +8555,48 @@ TEST_F(GraphTransformationTests, ReshapeFusionOpsetTest) { builder.AddNode("Reshape", {add_out, concattraining1_out}, {out}); }; + // Test that the fusion fires for every opset. std::unique_ptr transformer = std::make_unique(); - if (opset >= 15 && shape_test_for_opset15) { - ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, opset, *logger_, std::move(transformer), TransformerLevel::Level1, 1, - pre_graph_checker, pre_graph_checker)); - } else { - ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, opset, *logger_, std::move(transformer), TransformerLevel::Level1, 1, - pre_graph_checker, post_graph_checker)); + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, opset, *logger_, std::move(transformer), TransformerLevel::Level1, 1, + pre_graph_checker, post_graph_checker)); + + // For opset >= 15, also test that partial Shape (start=1, end=2) prevents fusion. + if (opset >= 15) { + auto build_partial_shape_case = [&](ModelTestBuilder& builder) { + auto* input_arg0 = builder.MakeInput({{batch_size, seq_lenth, hidden_size}}); + auto* input_arg1 = builder.MakeInput({{hidden_size}}); + auto* scalar_int_0 = builder.MakeInitializer({}, {0}); + auto* scalar_int_1 = builder.MakeInitializer({}, {1}); + auto* single_value_1d_int_0 = builder.MakeInitializer({1}, {0}); + auto* single_value_1d_int_16 = builder.MakeInitializer({1}, {16}); + auto* single_value_1d_int_64 = builder.MakeInitializer({1}, {64}); + auto* add_out = builder.MakeIntermediate(); + auto* shape_out = builder.MakeIntermediate(); + auto* gather_out_0 = builder.MakeIntermediate(); + auto* gather_out_1 = builder.MakeIntermediate(); + auto* unsqueeze_out_0 = builder.MakeIntermediate(); + auto* unsqueeze_out_1 = builder.MakeIntermediate(); + auto* concattraining1_out = builder.MakeIntermediate(); + auto* concattraining1_length = builder.MakeIntermediate(); + auto* out = builder.MakeOutput(); + + builder.AddNode("Add", {input_arg0, input_arg1}, {add_out}); + auto& shape_1 = builder.AddNode("Shape", {add_out}, {shape_out}); + shape_1.AddAttribute("start", (int64_t)1); + shape_1.AddAttribute("end", (int64_t)2); + builder.AddNode("Gather", {shape_out, scalar_int_0}, {gather_out_0}); + builder.AddNode("Gather", {shape_out, scalar_int_1}, {gather_out_1}); + builder.AddNode("Unsqueeze", {gather_out_0, single_value_1d_int_0}, {unsqueeze_out_0}); + builder.AddNode("Unsqueeze", {gather_out_1, single_value_1d_int_0}, {unsqueeze_out_1}); + builder.AddNode("ConcatTraining", {unsqueeze_out_0, unsqueeze_out_1, single_value_1d_int_16, single_value_1d_int_64}, + {concattraining1_out, concattraining1_length}, "com.microsoft") + .AddAttribute("axis", static_cast(0)); + builder.AddNode("Reshape", {add_out, concattraining1_out}, {out}); + }; + + std::unique_ptr transformer_no_fuse = std::make_unique(); + ASSERT_STATUS_OK(TestGraphTransformer(build_partial_shape_case, opset, *logger_, std::move(transformer_no_fuse), + TransformerLevel::Level1, 1, pre_graph_checker, pre_graph_checker)); } } } diff --git a/onnxruntime/test/optimizer/graph_transform_test_layernorm.cc b/onnxruntime/test/optimizer/graph_transform_test_layernorm.cc index 6ce2357f648d8..263003d212b3a 100644 --- a/onnxruntime/test/optimizer/graph_transform_test_layernorm.cc +++ b/onnxruntime/test/optimizer/graph_transform_test_layernorm.cc @@ -9,6 +9,8 @@ #include "gtest/gtest.h" +#include "onnx/defs/schema.h" + #include "core/graph/graph_utils.h" #include "core/graph/graph_viewer.h" #include "core/graph/model.h" @@ -36,6 +38,13 @@ namespace test { #ifndef DISABLE_CONTRIB_OPS +static int GetCurrentOnnxOpset() { + const auto& map = ONNX_NAMESPACE::OpSchemaRegistry::DomainToVersionRange::Instance().Map(); + auto it = map.find(ONNX_NAMESPACE::ONNX_DOMAIN); + EXPECT_TRUE(it != map.end()) << "ONNX domain not found in OpSchemaRegistry"; + return it != map.end() ? it->second.second : 0; +} + TEST_F(GraphTransformationTests, LayerNormFusionTest) { constexpr const ORTCHAR_T* model_uri = MODEL_FOLDER "fusion/layer_norm.onnx"; std::shared_ptr p_model; @@ -625,6 +634,111 @@ static void TestSkipLayerNormFusion(const std::basic_string& file_pat ASSERT_TRUE(op_to_count["Cast"] == cast_count); } +// Current-opset regression tests for LayerNorm and SkipLayerNorm fusions. +// These construct minimal graphs at the current ONNX opset and verify the optimizer fires. + +TEST_F(GraphTransformationTests, LayerNormFusionCurrentOpsetTest) { + // LayerNorm pattern: ReduceMean -> Sub -> Pow(2) -> ReduceMean -> Add(eps) -> Sqrt -> Div -> Mul(gamma) -> Add(beta) + int current_opset = GetCurrentOnnxOpset(); + auto build_test_case = [](ModelTestBuilder& builder) { + constexpr int64_t hidden_size = 64; + auto* input = builder.MakeInput({{2, 3, hidden_size}}); + auto* gamma = builder.MakeInitializer({hidden_size}, -1.0f, 1.0f); + auto* beta = builder.MakeInitializer({hidden_size}, -0.5f, 0.5f); + auto* two = builder.MakeInitializer({}, {2.0f}); + auto* eps = builder.MakeInitializer({}, {1e-5f}); + + auto* axes = builder.MakeInitializer({1}, {-1}); + auto* mean1_out = builder.MakeIntermediate(); + auto* sub_out = builder.MakeIntermediate(); + auto* pow_out = builder.MakeIntermediate(); + auto* mean2_out = builder.MakeIntermediate(); + auto* add_eps_out = builder.MakeIntermediate(); + auto* sqrt_out = builder.MakeIntermediate(); + auto* div_out = builder.MakeIntermediate(); + auto* mul_out = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + + builder.AddNode("ReduceMean", {input, axes}, {mean1_out}); + builder.AddNode("Sub", {input, mean1_out}, {sub_out}); + builder.AddNode("Pow", {sub_out, two}, {pow_out}); + builder.AddNode("ReduceMean", {pow_out, axes}, {mean2_out}); + builder.AddNode("Add", {mean2_out, eps}, {add_eps_out}); + builder.AddNode("Sqrt", {add_eps_out}, {sqrt_out}); + builder.AddNode("Div", {sub_out, sqrt_out}, {div_out}); + builder.AddNode("Mul", {div_out, gamma}, {mul_out}); + builder.AddNode("Add", {mul_out, beta}, {output}); + }; + + auto post_graph_checker = [current_opset](Graph& graph) { + auto op_to_count = CountOpsInGraph(graph); + if (op_to_count["LayerNormalization"] == 1) { + TEST_RETURN_IF_NOT(op_to_count["ReduceMean"] == 0); + TEST_RETURN_IF_NOT(op_to_count["Sub"] == 0); + TEST_RETURN_IF_NOT(op_to_count["Pow"] == 0); + TEST_RETURN_IF_NOT(op_to_count["Sqrt"] == 0); + TEST_RETURN_IF_NOT(op_to_count["Div"] == 0); + return Status::OK(); + } + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, + "LayerNorm fusion failed at opset ", current_opset, + ". Remaining ops: ReduceMean=", op_to_count["ReduceMean"], + " Sub=", op_to_count["Sub"], + " Pow=", op_to_count["Pow"], + " Sqrt=", op_to_count["Sqrt"], + " Div=", op_to_count["Div"], + ". Either update version lists in " + "onnxruntime/core/optimizer/layer_norm_fusion.cc" + " or skip this opset in the test if the fusion is not expected to apply."); + }; + + const InlinedHashSet no_limit_empty_ep_list = {}; + // LayerNorm fusion at Level1 when opset >= 17 (ONNX LayerNormalization available). + // At Level2, it skips if fuse_in_level_1 is true. + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, current_opset, *logger_, + std::make_unique(no_limit_empty_ep_list, TransformerLevel::Level1), + TransformerLevel::Level1, 1, nullptr, post_graph_checker)); +} + +TEST_F(GraphTransformationTests, SkipLayerNormFusionCurrentOpsetTest) { + // SkipLayerNorm pattern: Add(input, skip) -> LayerNormalization(gamma, beta) + int current_opset = GetCurrentOnnxOpset(); + auto build_test_case = [](ModelTestBuilder& builder) { + constexpr int64_t hidden_size = 64; + auto* input = builder.MakeInput({{2, 3, hidden_size}}); + auto* skip = builder.MakeInput({{2, 3, hidden_size}}); + auto* gamma = builder.MakeInitializer({hidden_size}, -1.0f, 1.0f); + auto* beta = builder.MakeInitializer({hidden_size}, -0.5f, 0.5f); + + auto* add_out = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + + builder.AddNode("Add", {input, skip}, {add_out}); + builder.AddNode("LayerNormalization", {add_out, gamma, beta}, {output}) + .AddAttribute("axis", static_cast(-1)); + }; + + auto post_graph_checker = [current_opset](Graph& graph) { + auto op_to_count = CountOpsInGraph(graph); + if (op_to_count["com.microsoft.SkipLayerNormalization"] == 1) { + TEST_RETURN_IF_NOT(op_to_count["Add"] == 0); + TEST_RETURN_IF_NOT(op_to_count["LayerNormalization"] == 0); + return Status::OK(); + } + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, + "SkipLayerNorm fusion failed at opset ", current_opset, + ". Remaining ops: Add=", op_to_count["Add"], + " LayerNormalization=", op_to_count["LayerNormalization"], + ". Either update version lists in " + "onnxruntime/core/optimizer/skip_layer_norm_fusion.cc" + " or skip this opset in the test if the fusion is not expected to apply."); + }; + + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, current_opset, *logger_, + std::make_unique(), + TransformerLevel::Level2, 1, nullptr, post_graph_checker)); +} + TEST_F(GraphTransformationTests, SkipLayerNormFusionTest) { TestSkipLayerNormFusion(MODEL_FOLDER "fusion/skip_layer_norm_format1.onnx", 0, 0, 1, 0, logger_.get()); TestSkipLayerNormFusion(MODEL_FOLDER "fusion/skip_layer_norm_format2.onnx", 0, 0, 1, 0, logger_.get()); @@ -1279,6 +1393,135 @@ TEST_F(GraphTransformationTests, EmbedLayerNormFusionFormat2) { ASSERT_TRUE(op_to_count["com.microsoft.EmbedLayerNormalization"] == 1); } +// These tests verify that EmbedLayerNorm fusion fires at the CURRENT max ONNX opset. +// When the ONNX opset advances, nodes will report a new SinceVersion(). If the +// optimizer's version lists are not updated, the fusion will fail to match. +// +// To fix: update the EdgeEndToMatch version lists in +// onnxruntime/core/optimizer/embed_layer_norm_fusion.cc + +// Loads a model and upgrades it to the current ONNX opset. Currently handles converting +// Squeeze/Unsqueeze/Reduce* ops from attribute-based axes to input-based (opset 13+/18+). +// Extend this function if additional op conversions are needed for new test models. +static void LoadModelAtCurrentOpset(const ORTCHAR_T* base_model_uri, + std::shared_ptr& p_model, + const logging::Logger& logger) { + int current_opset = GetCurrentOnnxOpset(); + ONNX_NAMESPACE::ModelProto model_proto; + { + std::shared_ptr base_model; + ASSERT_STATUS_OK(Model::Load(base_model_uri, base_model, nullptr, logger)); + model_proto = base_model->ToProto(); + } + for (auto& opset_import : *model_proto.mutable_opset_import()) { + if (opset_import.domain().empty() || opset_import.domain() == "ai.onnx") { + opset_import.set_version(current_opset); + break; + } + } + + // Convert attribute-based Squeeze/Unsqueeze to input-based for opset 13+ compatibility. + // Also convert ReduceSum/ReduceMean axes attribute to input for opset 18+ compatibility. + auto* graph_proto = model_proto.mutable_graph(); + int next_init_id = 0; + for (auto& node : *graph_proto->mutable_node()) { + bool is_squeeze_unsqueeze = (node.op_type() == "Squeeze" || node.op_type() == "Unsqueeze"); + bool is_reduce = (node.op_type() == "ReduceSum" || node.op_type() == "ReduceMean" || + node.op_type() == "ReduceMax" || node.op_type() == "ReduceMin" || + node.op_type() == "ReduceProd"); + if (!is_squeeze_unsqueeze && !is_reduce) { + continue; + } + int axes_attr_idx = -1; + for (int i = 0; i < node.attribute_size(); ++i) { + if (node.attribute(i).name() == "axes") { + axes_attr_idx = i; + break; + } + } + if (axes_attr_idx < 0) { + continue; + } + const auto& axes_attr = node.attribute(axes_attr_idx); + std::string init_name = "__axes_init_" + std::to_string(next_init_id++); + auto* init = graph_proto->add_initializer(); + init->set_name(init_name); + init->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); + init->add_dims(axes_attr.ints_size()); + for (int i = 0; i < axes_attr.ints_size(); ++i) { + init->add_int64_data(axes_attr.ints(i)); + } + // For Squeeze/Unsqueeze, axes is input[1]. For Reduce ops, axes is also input[1]. + node.add_input(init_name); + node.mutable_attribute()->SwapElements(axes_attr_idx, node.attribute_size() - 1); + node.mutable_attribute()->RemoveLast(); + } + + ASSERT_STATUS_OK(Model::Load(std::move(model_proto), p_model, nullptr, logger)); +} + +TEST_F(GraphTransformationTests, EmbedLayerNormFusionFormat1CurrentOpset) { + constexpr const ORTCHAR_T* model_uri = MODEL_FOLDER "fusion/embed_layer_norm_format1.onnx"; + std::shared_ptr p_model; + ASSERT_NO_FATAL_FAILURE(LoadModelAtCurrentOpset(model_uri, p_model, *logger_)); + Graph& graph = p_model->MainGraph(); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::make_unique(), TransformerLevel::Level2)); + ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, *logger_)); + + std::map op_to_count = CountOpsInGraph(graph); + EXPECT_EQ(op_to_count["com.microsoft.EmbedLayerNormalization"], 1) + << "EmbedLayerNorm fusion (format 1) failed at opset " << GetCurrentOnnxOpset() << ". " + << "Update version lists in onnxruntime/core/optimizer/embed_layer_norm_fusion.cc " + << "(MatchInputToConcatSubgraph, MatchPositionEmbeddingSubgraph)."; + EXPECT_EQ(op_to_count["Gather"], 0); + EXPECT_EQ(op_to_count["Add"], 0); +} + +TEST_F(GraphTransformationTests, EmbedLayerNormFusionFormat2CurrentOpset) { + constexpr const ORTCHAR_T* model_uri = MODEL_FOLDER "fusion/embed_layer_norm_format2.onnx"; + std::shared_ptr p_model; + ASSERT_NO_FATAL_FAILURE(LoadModelAtCurrentOpset(model_uri, p_model, *logger_)); + Graph& graph = p_model->MainGraph(); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::make_unique(), TransformerLevel::Level2)); + ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, *logger_)); + + std::map op_to_count = CountOpsInGraph(graph); + EXPECT_EQ(op_to_count["com.microsoft.EmbedLayerNormalization"], 1) + << "EmbedLayerNorm fusion (format 2: NonZero+Transpose+Squeeze path) failed at opset " + << GetCurrentOnnxOpset() << ". " + << "Update version lists in onnxruntime/core/optimizer/embed_layer_norm_fusion.cc " + << "(parent_path_1, parent_path_2)."; + EXPECT_EQ(op_to_count["Shape"], 0); + EXPECT_EQ(op_to_count["Expand"], 0); + EXPECT_EQ(op_to_count["Gather"], 0); + EXPECT_EQ(op_to_count["Unsqueeze"], 0); +} + +TEST_F(GraphTransformationTests, EmbedLayerNormFusionFormat3CurrentOpset) { + constexpr const ORTCHAR_T* model_uri = MODEL_FOLDER "fusion/embed_layer_norm_format3.onnx"; + std::shared_ptr p_model; + ASSERT_NO_FATAL_FAILURE(LoadModelAtCurrentOpset(model_uri, p_model, *logger_)); + Graph& graph = p_model->MainGraph(); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::make_unique(), TransformerLevel::Level2)); + ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, *logger_)); + + std::map op_to_count = CountOpsInGraph(graph); + EXPECT_EQ(op_to_count["com.microsoft.EmbedLayerNormalization"], 1) + << "EmbedLayerNorm fusion (format 3: Range-based path) failed at opset " + << GetCurrentOnnxOpset() << ". " + << "Update version lists in onnxruntime/core/optimizer/embed_layer_norm_fusion.cc " + << "(parent_path_3, parent_path_4)."; + EXPECT_EQ(op_to_count["Shape"], 0); + EXPECT_EQ(op_to_count["Gather"], 0); + EXPECT_EQ(op_to_count["Unsqueeze"], 0); +} + static void EmbedLayerNormFusionFormat3(const std::basic_string& file_path, logging::Logger* logger) { std::shared_ptr p_model; ASSERT_TRUE(Model::Load(file_path, p_model, nullptr, *logger).IsOK()); diff --git a/onnxruntime/test/optimizer/group_query_attention_pre_norm_fusion_test.cc b/onnxruntime/test/optimizer/group_query_attention_pre_norm_fusion_test.cc index b330475c01486..3efaba78004d6 100644 --- a/onnxruntime/test/optimizer/group_query_attention_pre_norm_fusion_test.cc +++ b/onnxruntime/test/optimizer/group_query_attention_pre_norm_fusion_test.cc @@ -16,6 +16,7 @@ #include "test/optimizer/webgpu_fusion_test_util.h" #include "gtest/gtest.h" +#include "onnx/defs/schema.h" namespace onnxruntime { namespace test { @@ -258,6 +259,20 @@ TEST_F(GraphTransformationTests, GroupQueryAttentionPreNormFusionFusesQwenPatter TransformerLevel::Level2, /*steps=*/1, nullptr, CheckFusedGraph)); } +TEST_F(GraphTransformationTests, GroupQueryAttentionPreNormFusionFusesQwenPatternCurrentOpset) { + // Uses the current max ONNX opset to catch version-list drift. + // If this fails, update version lists in + // onnxruntime/core/optimizer/group_query_attention_pre_norm_fusion.cc. + const auto& map = ONNX_NAMESPACE::OpSchemaRegistry::DomainToVersionRange::Instance().Map(); + auto it = map.find(ONNX_NAMESPACE::ONNX_DOMAIN); + ASSERT_TRUE(it != map.end()) << "ONNX domain not found in OpSchemaRegistry"; + int current_opset = it->second.second; + auto build = [](ModelTestBuilder& builder) { BuildQwenQkPostNormPattern(builder, BuildOptions{}); }; + ASSERT_STATUS_OK(TestGraphTransformer( + build, /*opset_version=*/current_opset, *logger_, MakeWebGpuTransformer(), + TransformerLevel::Level2, /*steps=*/1, nullptr, CheckFusedGraph)); +} + TEST_F(GraphTransformationTests, GroupQueryAttentionPreNormFusionMatchesUnfusedWebGpuResults) { if (!DefaultWebGpuExecutionProvider()) { GTEST_SKIP() << "WebGPU EP unavailable in this build.";