From 6a1233872a1fdad9957b1c960322284aef1112ec Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Tue, 9 Jun 2026 15:06:56 -0700 Subject: [PATCH 1/8] Update optimizer opset version checks for latest ONNX opset 25 Add newer opset versions (19, 21, 23, 24, 25) to IsSupportedOptypeVersionAndDomain and MatchesOpSinceVersion checks in optimizers where the version bumps are type-constraint widenings only (no semantic changes): - attention_fusion.cc: Reshape, Transpose, Squeeze - attention_fusion_helper.h: Transpose - group_query_attention_pre_norm_fusion.cc: Reshape - reshape_fusion.cc: Unsqueeze, Shape Add corresponding tests at opset 25 for attention fusion, GQA pre-norm fusion, and extend ReshapeFusionOpsetTest to cover opsets 19-25. Fix ReshapeFusionOpsetTest to properly test the fusion path for all opsets including 19+. Previously, a mutable flag caused opsets after 18 to only test the no-fusion (partial Shape) path. --- .../core/optimizer/attention_fusion.cc | 26 +++---- .../core/optimizer/attention_fusion_helper.h | 2 +- .../group_query_attention_pre_norm_fusion.cc | 4 +- onnxruntime/core/optimizer/reshape_fusion.cc | 8 +- .../test/optimizer/graph_transform_test.cc | 74 +++++++++++++++---- ...up_query_attention_pre_norm_fusion_test.cc | 7 ++ 6 files changed, 87 insertions(+), 34 deletions(-) diff --git a/onnxruntime/core/optimizer/attention_fusion.cc b/onnxruntime/core/optimizer/attention_fusion.cc index 7fe7c914fa796..61954c3f7d8c4 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"); } diff --git a/onnxruntime/core/optimizer/attention_fusion_helper.h b/onnxruntime/core/optimizer/attention_fusion_helper.h index a328a6d451a89..77d1ab6c190d9 100644 --- a/onnxruntime/core/optimizer/attention_fusion_helper.h +++ b/onnxruntime/core/optimizer/attention_fusion_helper.h @@ -1447,7 +1447,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; 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..d9e0da0acce1b 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,7 +172,8 @@ 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})) { + // Opset 15+ added start/end attributes to Shape. Reject partial-shape queries. + if (shape.SinceVersion() >= 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))) { diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index f14327af4b8dd..4ff8b43142e6e 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -6883,6 +6883,21 @@ TEST_F(GraphTransformationTests, AttentionFusionMobileClipMhaTest) { std::make_unique()); } +TEST_F(GraphTransformationTests, AttentionFusionMobileClipMhaOpset25Test) { + auto build_test_case = [](ModelTestBuilder& builder) { + BuildMobileClipAttentionTestCase(builder, MobileClipProjectionType::MatMulAdd); + }; + + TransformerTester(build_test_case, + CheckMobileClipAttentionFusedSession, + TransformerLevel::Level1, + TransformerLevel::Level2, + 25, + 1e-3, + 0.0, + std::make_unique()); +} + TEST_F(GraphTransformationTests, AttentionFusionMobileClipMhaProjectionGemmTest) { auto build_test_case = [](ModelTestBuilder& builder) { BuildMobileClipAttentionTestCase(builder, MobileClipProjectionType::GemmWithReshapes); @@ -8219,8 +8234,7 @@ TEST_F(GraphTransformationTests, ReshapeFusionOpsetTest) { return Status::OK(); }; - const std::vector opsets{11, 12, 13, 14, 15, 18}; - bool shape_test_for_opset15 = false; + const std::vector opsets{11, 12, 13, 14, 15, 18, 19, 21, 23, 24, 25}; for (auto& opset : opsets) { auto build_test_case = [&](ModelTestBuilder& builder) { @@ -8245,14 +8259,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 +8278,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/group_query_attention_pre_norm_fusion_test.cc b/onnxruntime/test/optimizer/group_query_attention_pre_norm_fusion_test.cc index b330475c01486..a217e1e88e133 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 @@ -258,6 +258,13 @@ TEST_F(GraphTransformationTests, GroupQueryAttentionPreNormFusionFusesQwenPatter TransformerLevel::Level2, /*steps=*/1, nullptr, CheckFusedGraph)); } +TEST_F(GraphTransformationTests, GroupQueryAttentionPreNormFusionFusesQwenPatternOpset25) { + auto build = [](ModelTestBuilder& builder) { BuildQwenQkPostNormPattern(builder, BuildOptions{}); }; + ASSERT_STATUS_OK(TestGraphTransformer( + build, /*opset_version=*/25, *logger_, MakeWebGpuTransformer(), + TransformerLevel::Level2, /*steps=*/1, nullptr, CheckFusedGraph)); +} + TEST_F(GraphTransformationTests, GroupQueryAttentionPreNormFusionMatchesUnfusedWebGpuResults) { if (!DefaultWebGpuExecutionProvider()) { GTEST_SKIP() << "WebGPU EP unavailable in this build."; From 62264db630e491a6ff49bdf3c777917e0d302ca4 Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Wed, 10 Jun 2026 15:43:09 -0700 Subject: [PATCH 2/8] Update optimizer opset version checks for ONNX opset 26 and add current-opset regression tests - Update version lists in attention_fusion.cc, attention_fusion_helper.h, and embed_layer_norm_fusion.cc to include opset versions up to 25/26. - Add programmatic current-opset regression tests that auto-detect when version lists need updating: Gelu, FastGelu, BiasGelu, LayerNorm, SkipLayerNorm, EmbedLayerNorm (3 formats), MobileClip MHA, GQA PreNorm. - Tests check for fused node first and report remaining op counts with guidance to update version lists or skip the opset. --- .../core/optimizer/attention_fusion.cc | 32 +- .../core/optimizer/attention_fusion_helper.h | 88 +++--- .../core/optimizer/embed_layer_norm_fusion.cc | 56 ++-- .../test/optimizer/graph_transform_test.cc | 284 +++++++++++++++++- .../graph_transform_test_layernorm.cc | 240 +++++++++++++++ ...up_query_attention_pre_norm_fusion_test.cc | 12 +- 6 files changed, 619 insertions(+), 93 deletions(-) diff --git a/onnxruntime/core/optimizer/attention_fusion.cc b/onnxruntime/core/optimizer/attention_fusion.cc index 61954c3f7d8c4..30eaaafccca82 100644 --- a/onnxruntime/core/optimizer/attention_fusion.cc +++ b/onnxruntime/core/optimizer/attention_fusion.cc @@ -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 77d1ab6c190d9..07ab0481b3683 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)) { @@ -166,9 +166,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 +375,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 +392,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 +454,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 +482,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 +632,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 +766,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 +791,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 +809,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 +887,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 +904,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 +924,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 +1024,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 +1339,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 +1361,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 +1413,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)) { @@ -1471,9 +1471,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..8734cc4b2cdad 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; @@ -147,9 +147,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."); @@ -231,42 +231,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 +318,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 +338,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 +559,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 +843,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/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 4ff8b43142e6e..3d7f7995933e0 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,7 +6904,7 @@ TEST_F(GraphTransformationTests, AttentionFusionMobileClipMhaTest) { std::make_unique()); } -TEST_F(GraphTransformationTests, AttentionFusionMobileClipMhaOpset25Test) { +TEST_F(GraphTransformationTests, AttentionFusionMobileClipMhaCurrentOpsetTest) { auto build_test_case = [](ModelTestBuilder& builder) { BuildMobileClipAttentionTestCase(builder, MobileClipProjectionType::MatMulAdd); }; @@ -6892,7 +6913,7 @@ TEST_F(GraphTransformationTests, AttentionFusionMobileClipMhaOpset25Test) { CheckMobileClipAttentionFusedSession, TransformerLevel::Level1, TransformerLevel::Level2, - 25, + GetCurrentOnnxOpset(), 1e-3, 0.0, std::make_unique()); @@ -6990,6 +7011,256 @@ 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) { + TEST_RETURN_IF_NOT(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; @@ -8234,7 +8505,14 @@ TEST_F(GraphTransformationTests, ReshapeFusionOpsetTest) { return Status::OK(); }; - const std::vector opsets{11, 12, 13, 14, 15, 18, 19, 21, 23, 24, 25}; + // 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) { diff --git a/onnxruntime/test/optimizer/graph_transform_test_layernorm.cc b/onnxruntime/test/optimizer/graph_transform_test_layernorm.cc index 6ce2357f648d8..46471486cbd7a 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,132 @@ 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 + +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 a217e1e88e133..1ad3a7a7b1965 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,10 +259,17 @@ TEST_F(GraphTransformationTests, GroupQueryAttentionPreNormFusionFusesQwenPatter TransformerLevel::Level2, /*steps=*/1, nullptr, CheckFusedGraph)); } -TEST_F(GraphTransformationTests, GroupQueryAttentionPreNormFusionFusesQwenPatternOpset25) { +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. + int current_opset = ONNX_NAMESPACE::OpSchemaRegistry::DomainToVersionRange::Instance() + .Map() + .at(ONNX_NAMESPACE::ONNX_DOMAIN) + .second; auto build = [](ModelTestBuilder& builder) { BuildQwenQkPostNormPattern(builder, BuildOptions{}); }; ASSERT_STATUS_OK(TestGraphTransformer( - build, /*opset_version=*/25, *logger_, MakeWebGpuTransformer(), + build, /*opset_version=*/current_opset, *logger_, MakeWebGpuTransformer(), TransformerLevel::Level2, /*steps=*/1, nullptr, CheckFusedGraph)); } From 1e044d75f510c43f080863e4361dedb5e7ea1d47 Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Wed, 10 Jun 2026 18:02:44 -0700 Subject: [PATCH 3/8] Address review comments: fix .at() crash risk, remove redundant check - Replace .at(ONNX_DOMAIN) with find + ASSERT_TRUE in GQA test to avoid potential throw on missing domain (Copilot review, high). - Remove redundant TEST_RETURN_IF_NOT in DivMulFusionCurrentOpsetTest where the condition was already guaranteed by the enclosing if (Copilot review, low). --- onnxruntime/test/optimizer/graph_transform_test.cc | 1 - .../group_query_attention_pre_norm_fusion_test.cc | 8 ++++---- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 3d7f7995933e0..5c76074d6b119 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -7206,7 +7206,6 @@ TEST_F(GraphTransformationTests, DivMulFusionCurrentOpsetTest) { 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) { - TEST_RETURN_IF_NOT(op_to_count["Mul"] == 0); return Status::OK(); } return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, 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 1ad3a7a7b1965..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 @@ -263,10 +263,10 @@ TEST_F(GraphTransformationTests, GroupQueryAttentionPreNormFusionFusesQwenPatter // 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. - int current_opset = ONNX_NAMESPACE::OpSchemaRegistry::DomainToVersionRange::Instance() - .Map() - .at(ONNX_NAMESPACE::ONNX_DOMAIN) - .second; + 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(), From 8ca63206b9c2a9052240cc7e8b8570c3f68327be Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Wed, 10 Jun 2026 18:21:04 -0700 Subject: [PATCH 4/8] Add Shape opset 15+ start/end attribute guards to prevent incorrect fusion with partial-shape queries --- .../core/optimizer/attention_fusion_helper.h | 10 +++++++++ .../core/optimizer/embed_layer_norm_fusion.cc | 21 +++++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/onnxruntime/core/optimizer/attention_fusion_helper.h b/onnxruntime/core/optimizer/attention_fusion_helper.h index 07ab0481b3683..a76ae2027bce4 100644 --- a/onnxruntime/core/optimizer/attention_fusion_helper.h +++ b/onnxruntime/core/optimizer/attention_fusion_helper.h @@ -98,6 +98,16 @@ bool MatchGemmSubgraph(Graph& graph, const Node& slice = edges[6]->GetNode(); const Node& shape_before_slice = edges[7]->GetNode(); + // Opset 15+ added start/end attributes to Shape. Reject partial-shape queries. + if (shape_before_slice.SinceVersion() >= 15) { + const ONNX_NAMESPACE::AttributeProto* start_attr = graph_utils::GetNodeAttribute(shape_before_slice, "start"); + const ONNX_NAMESPACE::AttributeProto* end_attr = graph_utils::GetNodeAttribute(shape_before_slice, "end"); + if (!((!start_attr || static_cast(start_attr->i()) == 0) && (!end_attr))) { + 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"); diff --git a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc index 8734cc4b2cdad..a5d82ec8289eb 100644 --- a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc +++ b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc @@ -138,6 +138,17 @@ static bool MatchInputToConcatSubgraph( } } + // Opset 15+ added start/end attributes to Shape. Reject partial-shape queries. + const Node& shape_node_path1 = edges[shape_index]->GetNode(); + if (shape_node_path1.SinceVersion() >= 15) { + const ONNX_NAMESPACE::AttributeProto* start_attr = graph_utils::GetNodeAttribute(shape_node_path1, "start"); + const ONNX_NAMESPACE::AttributeProto* end_attr = graph_utils::GetNodeAttribute(shape_node_path1, "end"); + if (!((!start_attr || static_cast(start_attr->i()) == 0) && (!end_attr))) { + 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()); @@ -167,6 +178,16 @@ static bool MatchInputToConcatSubgraph( Node& gather_node_1 = *graph.GetNode(edges[1]->GetNode().Index()); Node& shape_node_1 = *graph.GetNode(edges[2]->GetNode().Index()); + // Opset 15+ added start/end attributes to Shape. Reject partial-shape queries. + if (shape_node_1.SinceVersion() >= 15) { + const ONNX_NAMESPACE::AttributeProto* start_attr = graph_utils::GetNodeAttribute(shape_node_1, "start"); + const ONNX_NAMESPACE::AttributeProto* end_attr = graph_utils::GetNodeAttribute(shape_node_1, "end"); + if (!((!start_attr || static_cast(start_attr->i()) == 0) && (!end_attr))) { + 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."); From 033ac6347fd1c75487e1c55c18195967b84591e6 Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Thu, 11 Jun 2026 10:10:47 -0700 Subject: [PATCH 5/8] Fix Shape end-attr guard to accept INT64_MAX as full-shape default --- onnxruntime/core/optimizer/attention_fusion_helper.h | 3 ++- onnxruntime/core/optimizer/embed_layer_norm_fusion.cc | 8 ++++++-- onnxruntime/core/optimizer/reshape_fusion.cc | 4 +++- 3 files changed, 11 insertions(+), 4 deletions(-) diff --git a/onnxruntime/core/optimizer/attention_fusion_helper.h b/onnxruntime/core/optimizer/attention_fusion_helper.h index a76ae2027bce4..cdff78638aada 100644 --- a/onnxruntime/core/optimizer/attention_fusion_helper.h +++ b/onnxruntime/core/optimizer/attention_fusion_helper.h @@ -102,7 +102,8 @@ bool MatchGemmSubgraph(Graph& graph, if (shape_before_slice.SinceVersion() >= 15) { const ONNX_NAMESPACE::AttributeProto* start_attr = graph_utils::GetNodeAttribute(shape_before_slice, "start"); const ONNX_NAMESPACE::AttributeProto* end_attr = graph_utils::GetNodeAttribute(shape_before_slice, "end"); - if (!((!start_attr || static_cast(start_attr->i()) == 0) && (!end_attr))) { + if (!((!start_attr || start_attr->i() == 0) && + (!end_attr || end_attr->i() == std::numeric_limits::max()))) { DEBUG_LOG("Shape node has non-default start/end attributes"); return false; } diff --git a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc index a5d82ec8289eb..90b0e71476ae1 100644 --- a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc +++ b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc @@ -2,6 +2,8 @@ // Licensed under the MIT License. #include "core/optimizer/embed_layer_norm_fusion.h" +#include + #include "core/common/span_utils.h" #include "core/optimizer/initializer.h" #include "core/graph/contrib_ops/contrib_defs.h" @@ -143,7 +145,8 @@ static bool MatchInputToConcatSubgraph( if (shape_node_path1.SinceVersion() >= 15) { const ONNX_NAMESPACE::AttributeProto* start_attr = graph_utils::GetNodeAttribute(shape_node_path1, "start"); const ONNX_NAMESPACE::AttributeProto* end_attr = graph_utils::GetNodeAttribute(shape_node_path1, "end"); - if (!((!start_attr || static_cast(start_attr->i()) == 0) && (!end_attr))) { + if (!((!start_attr || start_attr->i() == 0) && + (!end_attr || end_attr->i() == std::numeric_limits::max()))) { DEBUG_LOG("Shape node in path 1 has non-default start/end attributes."); return false; } @@ -182,7 +185,8 @@ static bool MatchInputToConcatSubgraph( if (shape_node_1.SinceVersion() >= 15) { const ONNX_NAMESPACE::AttributeProto* start_attr = graph_utils::GetNodeAttribute(shape_node_1, "start"); const ONNX_NAMESPACE::AttributeProto* end_attr = graph_utils::GetNodeAttribute(shape_node_1, "end"); - if (!((!start_attr || static_cast(start_attr->i()) == 0) && (!end_attr))) { + if (!((!start_attr || start_attr->i() == 0) && + (!end_attr || end_attr->i() == std::numeric_limits::max()))) { DEBUG_LOG("Shape node in path 2 has non-default start/end attributes."); return false; } diff --git a/onnxruntime/core/optimizer/reshape_fusion.cc b/onnxruntime/core/optimizer/reshape_fusion.cc index d9e0da0acce1b..b503d02c565f3 100644 --- a/onnxruntime/core/optimizer/reshape_fusion.cc +++ b/onnxruntime/core/optimizer/reshape_fusion.cc @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include +#include #include "core/graph/graph_utils.h" #include "core/optimizer/initializer.h" @@ -176,7 +177,8 @@ bool ReshapeFusion::Match_One_Element_Output_Subgraph_1(Graph& graph, const Node if (shape.SinceVersion() >= 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))) { + if (!((!start_attr || start_attr->i() == 0) && + (!end_attr || end_attr->i() == std::numeric_limits::max()))) { return false; } } From 992c85ee80e027ccdf2a2408c8d565dfd09e1afb Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Thu, 11 Jun 2026 10:20:24 -0700 Subject: [PATCH 6/8] Add comments explaining partial-shape guard logic --- onnxruntime/core/optimizer/attention_fusion_helper.h | 6 +++++- .../core/optimizer/embed_layer_norm_fusion.cc | 12 ++++++++++-- onnxruntime/core/optimizer/reshape_fusion.cc | 7 ++++++- 3 files changed, 21 insertions(+), 4 deletions(-) diff --git a/onnxruntime/core/optimizer/attention_fusion_helper.h b/onnxruntime/core/optimizer/attention_fusion_helper.h index cdff78638aada..9bdefa3d1863f 100644 --- a/onnxruntime/core/optimizer/attention_fusion_helper.h +++ b/onnxruntime/core/optimizer/attention_fusion_helper.h @@ -98,10 +98,14 @@ bool MatchGemmSubgraph(Graph& graph, const Node& slice = edges[6]->GetNode(); const Node& shape_before_slice = edges[7]->GetNode(); - // Opset 15+ added start/end attributes to Shape. Reject partial-shape queries. + // Opset 15+ added optional start/end attributes to Shape, allowing it to return only a + // subset of dimensions ("partial shape"). The downstream Slice/Squeeze/Gather nodes assume + // Shape returns the full tensor shape. If start != 0 or end is not the default (INT64_MAX), + // the fusion would produce incorrect index mapping. Reject such cases. if (shape_before_slice.SinceVersion() >= 15) { const ONNX_NAMESPACE::AttributeProto* start_attr = graph_utils::GetNodeAttribute(shape_before_slice, "start"); const ONNX_NAMESPACE::AttributeProto* end_attr = graph_utils::GetNodeAttribute(shape_before_slice, "end"); + // end=INT64_MAX is the runtime default, meaning "all dimensions" (i.e. full shape). if (!((!start_attr || start_attr->i() == 0) && (!end_attr || end_attr->i() == std::numeric_limits::max()))) { DEBUG_LOG("Shape node has non-default start/end attributes"); diff --git a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc index 90b0e71476ae1..1abf95148056e 100644 --- a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc +++ b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc @@ -140,11 +140,15 @@ static bool MatchInputToConcatSubgraph( } } - // Opset 15+ added start/end attributes to Shape. Reject partial-shape queries. + // Opset 15+ added optional start/end attributes to Shape, allowing it to return only a + // subset of dimensions ("partial shape"). The Gather(index=0) below assumes Shape returns + // the full tensor shape. If start != 0 or end is not the default (INT64_MAX), the Gather + // would pick the wrong dimension and fusion would be incorrect. Reject such cases. const Node& shape_node_path1 = edges[shape_index]->GetNode(); if (shape_node_path1.SinceVersion() >= 15) { const ONNX_NAMESPACE::AttributeProto* start_attr = graph_utils::GetNodeAttribute(shape_node_path1, "start"); const ONNX_NAMESPACE::AttributeProto* end_attr = graph_utils::GetNodeAttribute(shape_node_path1, "end"); + // end=INT64_MAX is the runtime default, meaning "all dimensions" (i.e. full shape). if (!((!start_attr || start_attr->i() == 0) && (!end_attr || end_attr->i() == std::numeric_limits::max()))) { DEBUG_LOG("Shape node in path 1 has non-default start/end attributes."); @@ -181,10 +185,14 @@ static bool MatchInputToConcatSubgraph( Node& gather_node_1 = *graph.GetNode(edges[1]->GetNode().Index()); Node& shape_node_1 = *graph.GetNode(edges[2]->GetNode().Index()); - // Opset 15+ added start/end attributes to Shape. Reject partial-shape queries. + // Opset 15+ added optional start/end attributes to Shape, allowing it to return only a + // subset of dimensions ("partial shape"). The Gather(index=1) below assumes Shape returns + // the full tensor shape. If start != 0 or end is not the default (INT64_MAX), the Gather + // would pick the wrong dimension and fusion would be incorrect. Reject such cases. if (shape_node_1.SinceVersion() >= 15) { const ONNX_NAMESPACE::AttributeProto* start_attr = graph_utils::GetNodeAttribute(shape_node_1, "start"); const ONNX_NAMESPACE::AttributeProto* end_attr = graph_utils::GetNodeAttribute(shape_node_1, "end"); + // end=INT64_MAX is the runtime default, meaning "all dimensions" (i.e. full shape). if (!((!start_attr || start_attr->i() == 0) && (!end_attr || end_attr->i() == std::numeric_limits::max()))) { DEBUG_LOG("Shape node in path 2 has non-default start/end attributes."); diff --git a/onnxruntime/core/optimizer/reshape_fusion.cc b/onnxruntime/core/optimizer/reshape_fusion.cc index b503d02c565f3..eb352d1389f4c 100644 --- a/onnxruntime/core/optimizer/reshape_fusion.cc +++ b/onnxruntime/core/optimizer/reshape_fusion.cc @@ -173,10 +173,15 @@ bool ReshapeFusion::Match_One_Element_Output_Subgraph_1(Graph& graph, const Node const Node& gather = edges[1]->GetNode(); const Node& shape = edges[2]->GetNode(); - // Opset 15+ added start/end attributes to Shape. Reject partial-shape queries. + // Opset 15+ added optional start/end attributes to Shape, allowing it to return only a + // subset of dimensions ("partial shape"). The fusion below assumes Shape returns the full + // tensor shape so that Gather indices correspond directly to tensor dimensions. If start != 0 + // or end is set to something other than the default (INT64_MAX = all dims), the Gather index + // semantics change and the fusion would produce incorrect results. Reject such cases. if (shape.SinceVersion() >= 15) { const ONNX_NAMESPACE::AttributeProto* start_attr = graph_utils::GetNodeAttribute(shape, "start"); const ONNX_NAMESPACE::AttributeProto* end_attr = graph_utils::GetNodeAttribute(shape, "end"); + // end=INT64_MAX is the runtime default, meaning "all dimensions" (i.e. full shape). if (!((!start_attr || start_attr->i() == 0) && (!end_attr || end_attr->i() == std::numeric_limits::max()))) { return false; From 79412631aaa9503f9320d10bac7da1865c6f3052 Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Thu, 11 Jun 2026 14:41:33 -0700 Subject: [PATCH 7/8] Add scope comment to LoadModelAtCurrentOpset helper --- onnxruntime/test/optimizer/graph_transform_test_layernorm.cc | 3 +++ 1 file changed, 3 insertions(+) diff --git a/onnxruntime/test/optimizer/graph_transform_test_layernorm.cc b/onnxruntime/test/optimizer/graph_transform_test_layernorm.cc index 46471486cbd7a..263003d212b3a 100644 --- a/onnxruntime/test/optimizer/graph_transform_test_layernorm.cc +++ b/onnxruntime/test/optimizer/graph_transform_test_layernorm.cc @@ -1400,6 +1400,9 @@ TEST_F(GraphTransformationTests, EmbedLayerNormFusionFormat2) { // 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) { From 06ef3dbf408030e80ea03833f6ee8cc441609e87 Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Thu, 11 Jun 2026 14:50:03 -0700 Subject: [PATCH 8/8] Extract IsFullShapeNode helper into graph_utils --- onnxruntime/core/graph/graph_utils.cc | 9 +++++ onnxruntime/core/graph/graph_utils.h | 5 +++ .../core/optimizer/attention_fusion_helper.h | 19 +++------- .../core/optimizer/embed_layer_norm_fusion.cc | 38 +++++-------------- onnxruntime/core/optimizer/reshape_fusion.cc | 19 +++------- 5 files changed, 35 insertions(+), 55 deletions(-) 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_helper.h b/onnxruntime/core/optimizer/attention_fusion_helper.h index 9bdefa3d1863f..b7bd36c8df9e4 100644 --- a/onnxruntime/core/optimizer/attention_fusion_helper.h +++ b/onnxruntime/core/optimizer/attention_fusion_helper.h @@ -98,19 +98,12 @@ bool MatchGemmSubgraph(Graph& graph, const Node& slice = edges[6]->GetNode(); const Node& shape_before_slice = edges[7]->GetNode(); - // Opset 15+ added optional start/end attributes to Shape, allowing it to return only a - // subset of dimensions ("partial shape"). The downstream Slice/Squeeze/Gather nodes assume - // Shape returns the full tensor shape. If start != 0 or end is not the default (INT64_MAX), - // the fusion would produce incorrect index mapping. Reject such cases. - if (shape_before_slice.SinceVersion() >= 15) { - const ONNX_NAMESPACE::AttributeProto* start_attr = graph_utils::GetNodeAttribute(shape_before_slice, "start"); - const ONNX_NAMESPACE::AttributeProto* end_attr = graph_utils::GetNodeAttribute(shape_before_slice, "end"); - // end=INT64_MAX is the runtime default, meaning "all dimensions" (i.e. full shape). - if (!((!start_attr || start_attr->i() == 0) && - (!end_attr || end_attr->i() == std::numeric_limits::max()))) { - DEBUG_LOG("Shape node has non-default start/end attributes"); - return false; - } + // 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]; diff --git a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc index 1abf95148056e..e2a3ef2a4fe2d 100644 --- a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc +++ b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc @@ -2,8 +2,6 @@ // Licensed under the MIT License. #include "core/optimizer/embed_layer_norm_fusion.h" -#include - #include "core/common/span_utils.h" #include "core/optimizer/initializer.h" #include "core/graph/contrib_ops/contrib_defs.h" @@ -140,20 +138,12 @@ static bool MatchInputToConcatSubgraph( } } - // Opset 15+ added optional start/end attributes to Shape, allowing it to return only a - // subset of dimensions ("partial shape"). The Gather(index=0) below assumes Shape returns - // the full tensor shape. If start != 0 or end is not the default (INT64_MAX), the Gather - // would pick the wrong dimension and fusion would be incorrect. Reject such cases. + // 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 (shape_node_path1.SinceVersion() >= 15) { - const ONNX_NAMESPACE::AttributeProto* start_attr = graph_utils::GetNodeAttribute(shape_node_path1, "start"); - const ONNX_NAMESPACE::AttributeProto* end_attr = graph_utils::GetNodeAttribute(shape_node_path1, "end"); - // end=INT64_MAX is the runtime default, meaning "all dimensions" (i.e. full shape). - if (!((!start_attr || start_attr->i() == 0) && - (!end_attr || end_attr->i() == std::numeric_limits::max()))) { - DEBUG_LOG("Shape node in path 1 has non-default start/end attributes."); - return false; - } + 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()); @@ -185,19 +175,11 @@ static bool MatchInputToConcatSubgraph( Node& gather_node_1 = *graph.GetNode(edges[1]->GetNode().Index()); Node& shape_node_1 = *graph.GetNode(edges[2]->GetNode().Index()); - // Opset 15+ added optional start/end attributes to Shape, allowing it to return only a - // subset of dimensions ("partial shape"). The Gather(index=1) below assumes Shape returns - // the full tensor shape. If start != 0 or end is not the default (INT64_MAX), the Gather - // would pick the wrong dimension and fusion would be incorrect. Reject such cases. - if (shape_node_1.SinceVersion() >= 15) { - const ONNX_NAMESPACE::AttributeProto* start_attr = graph_utils::GetNodeAttribute(shape_node_1, "start"); - const ONNX_NAMESPACE::AttributeProto* end_attr = graph_utils::GetNodeAttribute(shape_node_1, "end"); - // end=INT64_MAX is the runtime default, meaning "all dimensions" (i.e. full shape). - if (!((!start_attr || start_attr->i() == 0) && - (!end_attr || end_attr->i() == std::numeric_limits::max()))) { - DEBUG_LOG("Shape node in path 2 has non-default start/end attributes."); - return false; - } + // 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 diff --git a/onnxruntime/core/optimizer/reshape_fusion.cc b/onnxruntime/core/optimizer/reshape_fusion.cc index eb352d1389f4c..65f477d9c7479 100644 --- a/onnxruntime/core/optimizer/reshape_fusion.cc +++ b/onnxruntime/core/optimizer/reshape_fusion.cc @@ -2,7 +2,6 @@ // Licensed under the MIT License. #include -#include #include "core/graph/graph_utils.h" #include "core/optimizer/initializer.h" @@ -173,19 +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(); - // Opset 15+ added optional start/end attributes to Shape, allowing it to return only a - // subset of dimensions ("partial shape"). The fusion below assumes Shape returns the full - // tensor shape so that Gather indices correspond directly to tensor dimensions. If start != 0 - // or end is set to something other than the default (INT64_MAX = all dims), the Gather index - // semantics change and the fusion would produce incorrect results. Reject such cases. - if (shape.SinceVersion() >= 15) { - const ONNX_NAMESPACE::AttributeProto* start_attr = graph_utils::GetNodeAttribute(shape, "start"); - const ONNX_NAMESPACE::AttributeProto* end_attr = graph_utils::GetNodeAttribute(shape, "end"); - // end=INT64_MAX is the runtime default, meaning "all dimensions" (i.e. full shape). - if (!((!start_attr || start_attr->i() == 0) && - (!end_attr || end_attr->i() == std::numeric_limits::max()))) { - 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;