From 7cfb2d142d18875d08a938cf2e7aecca7ed2723b Mon Sep 17 00:00:00 2001 From: luozijian Date: Fri, 4 Sep 2026 14:25:42 +0800 Subject: [PATCH 1/4] fix: avoid fusing identity transpose into Gemm Signed-off-by: luozijian --- .../core/optimizer/gemm_transpose_fusion.cc | 17 ++++++++++-- .../test/optimizer/graph_transform_test.cc | 27 +++++++++++++++++++ 2 files changed, 42 insertions(+), 2 deletions(-) diff --git a/onnxruntime/core/optimizer/gemm_transpose_fusion.cc b/onnxruntime/core/optimizer/gemm_transpose_fusion.cc index a66ad987cfaef..9a4be5ce4c5a4 100644 --- a/onnxruntime/core/optimizer/gemm_transpose_fusion.cc +++ b/onnxruntime/core/optimizer/gemm_transpose_fusion.cc @@ -11,6 +11,17 @@ using namespace ONNX_NAMESPACE; using namespace onnxruntime::common; namespace onnxruntime { +static bool IsMatrixTranspose(const Node& transpose_node) { + const auto perm_attr = transpose_node.GetAttributes().find("perm"); + if (perm_attr != transpose_node.GetAttributes().end()) { + const auto perms = RetrieveValues(perm_attr->second); + return perms.size() == 2 && perms[0] == 1 && perms[1] == 0; + } + + const auto* shape = transpose_node.InputDefs()[0]->Shape(); + return shape != nullptr && shape->dim_size() == 2; +} + Status GemmTransposeFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& modified, const logging::Logger&) const { auto& gemm_node = node; const Node* A_node_ptr = graph_utils::GetInputNode(gemm_node, 0); @@ -25,7 +36,7 @@ Status GemmTransposeFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& m auto new_gemm_input_defs = gemm_node.MutableInputDefs(); // check if input A is a Transpose - if (A_node_ptr != nullptr && A_node_ptr->OpType() == "Transpose") { + if (A_node_ptr != nullptr && A_node_ptr->OpType() == "Transpose" && IsMatrixTranspose(*A_node_ptr)) { // make sure all consumers are gemm nodes to avoid possible double transpose std::vector gemm_nodes = graph_utils::FindChildrenByType(*A_node_ptr, "Gemm"); if (gemm_nodes.size() == A_node_ptr->GetOutputEdgesCount()) { @@ -44,7 +55,7 @@ Status GemmTransposeFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& m } } // check if input B is a Transpose - if (B_node_ptr != nullptr && B_node_ptr->OpType() == "Transpose") { + if (B_node_ptr != nullptr && B_node_ptr->OpType() == "Transpose" && IsMatrixTranspose(*B_node_ptr)) { std::vector gemm_nodes = graph_utils::FindChildrenByType(*B_node_ptr, "Gemm"); if (gemm_nodes.size() == B_node_ptr->GetOutputEdgesCount()) { Node& B_node = *graph.GetNode(B_node_ptr->Index()); @@ -106,6 +117,7 @@ bool GemmTransposeFusion::SatisfyCondition(const Graph& graph, const Node& node, // Fusion can be applied if there is a transpose at either of the inputs for (auto node_it = node.InputNodesBegin(); node_it != node.InputNodesEnd(); ++node_it) { if (graph_utils::IsSupportedOptypeVersionAndDomain(*node_it, "Transpose", {1, 13, 21, 23, 24, 25}) && + IsMatrixTranspose(*node_it) && !graph.NodeProducesGraphOutput(*node_it) && // Make sure the two nodes do not span execution providers. node_it->GetExecutionProviderType() == node.GetExecutionProviderType()) { @@ -130,6 +142,7 @@ bool GemmTransposeFusion::SatisfyCondition(const Graph& graph, const Node& node, const auto next_node_it = node.OutputNodesBegin(); if (next_node_it != node.OutputNodesEnd() && graph_utils::IsSupportedOptypeVersionAndDomain(*next_node_it, "Transpose", {1, 13, 21, 23, 24, 25}) && + IsMatrixTranspose(*next_node_it) && next_node_it->GetInputEdgesCount() == 1 && // Make sure the two nodes do not span execution providers. next_node_it->GetExecutionProviderType() == node.GetExecutionProviderType()) { diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 074d5d5f6f32e..40965089b6703 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -5412,6 +5412,33 @@ TEST_F(GraphTransformationTests, GemmTransposeFusion2Inputs) { ASSERT_TRUE(new_input_defs[1]->Name() == "B"); } +TEST_F(GraphTransformationTests, GemmTransposeFusionDoesNotFuseIdentityTranspose) { + auto build_test_case = [](ModelTestBuilder& builder) { + auto* input = builder.MakeInput({3, 4}); + auto* weight = builder.MakeInput({4, 5}); + auto* transposed_weight = builder.MakeIntermediate(std::vector{4, 5}); + auto* output = builder.MakeOutput(std::vector{3, 5}); + + builder.AddNode("Transpose", {weight}, {transposed_weight}).AddAttribute("perm", std::vector{0, 1}); + auto& gemm = builder.AddNode("Gemm", {input, transposed_weight}, {output}); + gemm.AddAttribute("transA", int64_t{0}); + gemm.AddAttribute("transB", int64_t{0}); + gemm.AddAttribute("alpha", 1.0f); + gemm.AddAttribute("beta", 1.0f); + }; + + auto check_graph = [](Graph& graph) { + TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Transpose"] == 1); + TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Gemm"] == 1); + return Status::OK(); + }; + + auto rule_transformer = std::make_unique("RuleTransformer"); + ASSERT_STATUS_OK(rule_transformer->Register(std::make_unique())); + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 13, *logger_, std::move(rule_transformer), TransformerLevel::Level1, + 1, check_graph, check_graph)); +} + // (A')'B' = AB' where transpose has multiple consumers TEST_F(GraphTransformationTests, GemmTransposeFusion2OutputsFromTranspose) { constexpr const ORTCHAR_T* model_uri = MODEL_FOLDER "fusion/gemm_transpose_2outputs_from_transpose.onnx"; From 8b7e6ff1740288223e50e8cbaf230e54147d4ad9 Mon Sep 17 00:00:00 2001 From: luozijian Date: Fri, 4 Sep 2026 14:39:13 +0800 Subject: [PATCH 2/4] test: cover identity output transpose in Gemm fusion Signed-off-by: luozijian --- .../test/optimizer/graph_transform_test.cc | 27 +++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 40965089b6703..d85bcbbd90cf6 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -5439,6 +5439,33 @@ TEST_F(GraphTransformationTests, GemmTransposeFusionDoesNotFuseIdentityTranspose 1, check_graph, check_graph)); } +TEST_F(GraphTransformationTests, GemmTransposeFusionDoesNotFuseIdentityTransposeAtOutput) { + auto build_test_case = [](ModelTestBuilder& builder) { + auto* input = builder.MakeInput({3, 4}); + auto* weight = builder.MakeInput({4, 5}); + auto* gemm_output = builder.MakeIntermediate(std::vector{3, 5}); + auto* output = builder.MakeOutput(std::vector{3, 5}); + + auto& gemm = builder.AddNode("Gemm", {input, weight}, {gemm_output}); + gemm.AddAttribute("transA", int64_t{0}); + gemm.AddAttribute("transB", int64_t{0}); + gemm.AddAttribute("alpha", 1.0f); + gemm.AddAttribute("beta", 1.0f); + builder.AddNode("Transpose", {gemm_output}, {output}).AddAttribute("perm", std::vector{0, 1}); + }; + + auto check_graph = [](Graph& graph) { + TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Transpose"] == 1); + TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Gemm"] == 1); + return Status::OK(); + }; + + auto rule_transformer = std::make_unique("RuleTransformer"); + ASSERT_STATUS_OK(rule_transformer->Register(std::make_unique())); + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 13, *logger_, std::move(rule_transformer), TransformerLevel::Level1, + 1, check_graph, check_graph)); +} + // (A')'B' = AB' where transpose has multiple consumers TEST_F(GraphTransformationTests, GemmTransposeFusion2OutputsFromTranspose) { constexpr const ORTCHAR_T* model_uri = MODEL_FOLDER "fusion/gemm_transpose_2outputs_from_transpose.onnx"; From 031cc6710c72ac408a85ca2f6fd03fe49c6446bf Mon Sep 17 00:00:00 2001 From: luozijian Date: Fri, 4 Sep 2026 20:30:24 +0800 Subject: [PATCH 3/4] test: fix Gemm transpose input shape setup --- onnxruntime/test/optimizer/graph_transform_test.cc | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index d85bcbbd90cf6..909c66dfae044 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -5414,8 +5414,8 @@ TEST_F(GraphTransformationTests, GemmTransposeFusion2Inputs) { TEST_F(GraphTransformationTests, GemmTransposeFusionDoesNotFuseIdentityTranspose) { auto build_test_case = [](ModelTestBuilder& builder) { - auto* input = builder.MakeInput({3, 4}); - auto* weight = builder.MakeInput({4, 5}); + auto* input = builder.MakeInput({{3, 4}}); + auto* weight = builder.MakeInput({{4, 5}}); auto* transposed_weight = builder.MakeIntermediate(std::vector{4, 5}); auto* output = builder.MakeOutput(std::vector{3, 5}); @@ -5441,8 +5441,8 @@ TEST_F(GraphTransformationTests, GemmTransposeFusionDoesNotFuseIdentityTranspose TEST_F(GraphTransformationTests, GemmTransposeFusionDoesNotFuseIdentityTransposeAtOutput) { auto build_test_case = [](ModelTestBuilder& builder) { - auto* input = builder.MakeInput({3, 4}); - auto* weight = builder.MakeInput({4, 5}); + auto* input = builder.MakeInput({{3, 4}}); + auto* weight = builder.MakeInput({{4, 5}}); auto* gemm_output = builder.MakeIntermediate(std::vector{3, 5}); auto* output = builder.MakeOutput(std::vector{3, 5}); From ffa99759ba765c9d00c18fb47241d3e35556f135 Mon Sep 17 00:00:00 2001 From: luozijian Date: Sat, 5 Sep 2026 10:30:44 +0800 Subject: [PATCH 4/4] fix: preserve identity output transposes during Gemm input fusion --- .../core/optimizer/gemm_transpose_fusion.cc | 2 +- .../test/optimizer/graph_transform_test.cc | 72 +++++++++++++++++-- 2 files changed, 68 insertions(+), 6 deletions(-) diff --git a/onnxruntime/core/optimizer/gemm_transpose_fusion.cc b/onnxruntime/core/optimizer/gemm_transpose_fusion.cc index 9a4be5ce4c5a4..ed7675739b175 100644 --- a/onnxruntime/core/optimizer/gemm_transpose_fusion.cc +++ b/onnxruntime/core/optimizer/gemm_transpose_fusion.cc @@ -75,7 +75,7 @@ Status GemmTransposeFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& m // check if output node is Transpose if (output_node_ptr != gemm_node.OutputNodesEnd() && gemm_node.InputDefs().size() <= 2 && // C is missing - output_node_ptr->OpType() == "Transpose") { + output_node_ptr->OpType() == "Transpose" && IsMatrixTranspose(*output_node_ptr)) { Node& output_node = *graph.GetNode(output_node_ptr->Index()); // (AB)' = B'A' : reverse the inputs std::reverse(new_gemm_input_defs.begin(), new_gemm_input_defs.end()); diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 909c66dfae044..ba84f0416ce37 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -5441,22 +5441,32 @@ TEST_F(GraphTransformationTests, GemmTransposeFusionDoesNotFuseIdentityTranspose TEST_F(GraphTransformationTests, GemmTransposeFusionDoesNotFuseIdentityTransposeAtOutput) { auto build_test_case = [](ModelTestBuilder& builder) { - auto* input = builder.MakeInput({{3, 4}}); - auto* weight = builder.MakeInput({{4, 5}}); + auto* input = builder.MakeInput({{4, 3}}, "A"); + auto* weight = builder.MakeInput({{4, 5}}, "B"); auto* gemm_output = builder.MakeIntermediate(std::vector{3, 5}); auto* output = builder.MakeOutput(std::vector{3, 5}); auto& gemm = builder.AddNode("Gemm", {input, weight}, {gemm_output}); - gemm.AddAttribute("transA", int64_t{0}); + gemm.AddAttribute("transA", int64_t{1}); gemm.AddAttribute("transB", int64_t{0}); - gemm.AddAttribute("alpha", 1.0f); - gemm.AddAttribute("beta", 1.0f); + gemm.AddAttribute("alpha", 2.0f); + gemm.AddAttribute("beta", 3.0f); builder.AddNode("Transpose", {gemm_output}, {output}).AddAttribute("perm", std::vector{0, 1}); }; auto check_graph = [](Graph& graph) { TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Transpose"] == 1); TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Gemm"] == 1); + for (const auto& node : graph.Nodes()) { + if (node.OpType() == "Gemm") { + TEST_RETURN_IF_NOT(node.GetAttributes().at("transA").i() == 1); + TEST_RETURN_IF_NOT(node.GetAttributes().at("transB").i() == 0); + TEST_RETURN_IF_NOT(node.GetAttributes().at("alpha").f() == 2.0f); + TEST_RETURN_IF_NOT(node.GetAttributes().at("beta").f() == 3.0f); + TEST_RETURN_IF_NOT(node.InputDefs()[0]->Name() == "A"); + TEST_RETURN_IF_NOT(node.InputDefs()[1]->Name() == "B"); + } + } return Status::OK(); }; @@ -5466,6 +5476,58 @@ TEST_F(GraphTransformationTests, GemmTransposeFusionDoesNotFuseIdentityTranspose 1, check_graph, check_graph)); } +TEST_F(GraphTransformationTests, GemmTransposeFusionPreservesIdentityOutputWhenFusingInput) { + for (bool transpose_input_b : {false, true}) { + SCOPED_TRACE(transpose_input_b); + auto build_test_case = [transpose_input_b](ModelTestBuilder& builder) { + auto* input = builder.MakeInput(transpose_input_b ? std::vector{3, 4} + : std::vector{4, 3}, + "A"); + auto* weight = builder.MakeInput(transpose_input_b ? std::vector{5, 4} + : std::vector{4, 5}, + "B"); + auto* transposed_input = builder.MakeIntermediate(); + auto* gemm_output = builder.MakeIntermediate(std::vector{3, 5}); + auto* output = builder.MakeOutput(std::vector{3, 5}); + + builder.AddNode("Transpose", {transpose_input_b ? weight : input}, {transposed_input}) + .AddAttribute("perm", std::vector{1, 0}); + auto& gemm = builder.AddNode("Gemm", {transpose_input_b ? input : transposed_input, transpose_input_b ? transposed_input : weight}, + {gemm_output}); + gemm.AddAttribute("transA", int64_t{0}); + gemm.AddAttribute("transB", int64_t{0}); + gemm.AddAttribute("alpha", 2.0f); + gemm.AddAttribute("beta", 3.0f); + builder.AddNode("Transpose", {gemm_output}, {output}).AddAttribute("perm", std::vector{0, 1}); + }; + + auto check_graph = [transpose_input_b](Graph& graph) { + TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Transpose"] == 1); + TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Gemm"] == 1); + for (const auto& node : graph.Nodes()) { + if (node.OpType() == "Gemm") { + TEST_RETURN_IF_NOT(node.GetAttributes().at("transA").i() == (transpose_input_b ? 0 : 1)); + TEST_RETURN_IF_NOT(node.GetAttributes().at("transB").i() == (transpose_input_b ? 1 : 0)); + TEST_RETURN_IF_NOT(node.GetAttributes().at("alpha").f() == 2.0f); + TEST_RETURN_IF_NOT(node.GetAttributes().at("beta").f() == 3.0f); + TEST_RETURN_IF_NOT(node.InputDefs()[0]->Name() == "A"); + TEST_RETURN_IF_NOT(node.InputDefs()[1]->Name() == "B"); + } else if (node.OpType() == "Transpose") { + const auto perm = RetrieveValues(node.GetAttributes().at("perm")); + TEST_RETURN_IF_NOT(perm == std::vector({0, 1})); + TEST_RETURN_IF_NOT(graph.NodeProducesGraphOutput(node)); + } + } + return Status::OK(); + }; + + auto rule_transformer = std::make_unique("RuleTransformer"); + ASSERT_STATUS_OK(rule_transformer->Register(std::make_unique())); + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 13, *logger_, std::move(rule_transformer), + TransformerLevel::Level1, 1, nullptr, check_graph)); + } +} + // (A')'B' = AB' where transpose has multiple consumers TEST_F(GraphTransformationTests, GemmTransposeFusion2OutputsFromTranspose) { constexpr const ORTCHAR_T* model_uri = MODEL_FOLDER "fusion/gemm_transpose_2outputs_from_transpose.onnx";