diff --git a/onnxruntime/core/optimizer/gemm_transpose_fusion.cc b/onnxruntime/core/optimizer/gemm_transpose_fusion.cc index a66ad987cfaef..103b3a616ecc7 100644 --- a/onnxruntime/core/optimizer/gemm_transpose_fusion.cc +++ b/onnxruntime/core/optimizer/gemm_transpose_fusion.cc @@ -11,6 +11,22 @@ using namespace ONNX_NAMESPACE; using namespace onnxruntime::common; namespace onnxruntime { +namespace { + +bool IsMatrixTranspose(const Node& node) { + const auto& attributes = node.GetAttributes(); + const auto perm_it = attributes.find("perm"); + if (perm_it == attributes.end()) { + // Gemm inputs and outputs are rank 2, so the default reverse-axis permutation is [1, 0]. + return true; + } + + const auto& perm = perm_it->second.ints(); + return perm.size() == 2 && perm[0] == 1 && perm[1] == 0; +} + +} // namespace + 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); @@ -24,8 +40,8 @@ Status GemmTransposeFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& m std::vector> nodes_to_remove; 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") { + // check if input A is a matrix 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()) { @@ -43,8 +59,8 @@ Status GemmTransposeFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& m new_gemm_input_defs[0] = A_node.MutableInputDefs()[0]; } } - // check if input B is a Transpose - if (B_node_ptr != nullptr && B_node_ptr->OpType() == "Transpose") { + // check if input B is a matrix 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()); @@ -61,10 +77,11 @@ Status GemmTransposeFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& m nodes_to_remove.push_back(gemm_node); - // check if output node is Transpose + // check if output node is a matrix 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()); @@ -103,9 +120,10 @@ bool GemmTransposeFusion::SatisfyCondition(const Graph& graph, const Node& node, return false; } - // Fusion can be applied if there is a transpose at either of the inputs + // Fusion can be applied if there is a matrix 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()) { @@ -117,7 +135,7 @@ bool GemmTransposeFusion::SatisfyCondition(const Graph& graph, const Node& node, } } - // Fusion can be applied if there is a Transpose at the output of Gemm + // Fusion can be applied if there is a matrix Transpose at the output of Gemm // by the rule (AB)' = B'A' provided that C is missing // Supported for Opset >=11 as earlier opsets have C as a required input if (!graph_utils::IsSupportedOptypeVersionAndDomain(node, "Gemm", {11, 13}) || @@ -130,6 +148,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/gemm_transpose_fusion_test.cc b/onnxruntime/test/optimizer/gemm_transpose_fusion_test.cc new file mode 100644 index 0000000000000..4547d4dbe1f43 --- /dev/null +++ b/onnxruntime/test/optimizer/gemm_transpose_fusion_test.cc @@ -0,0 +1,136 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "gtest/gtest.h" + +#include "core/graph/model.h" +#include "core/optimizer/gemm_transpose_fusion.h" +#include "core/optimizer/rule_based_graph_transformer.h" +#include "test/test_environment.h" +#include "test/unittest_util/framework_test_utils.h" +#include "test/unittest_util/graph_transform_test_builder.h" +#include "test/util/include/asserts.h" + +namespace onnxruntime { +namespace test { + +namespace { + +void RunInputTransposeTest(const std::vector& perm, + const std::vector& input_shape, + bool expect_fusion, + int64_t expected_trans_a) { + auto& logger = DefaultLoggingManager().DefaultLogger(); + Model model("GemmTransposeFusionPermutation", false, ModelMetaData(), PathString(), + IOnnxRuntimeOpSchemaRegistryList(), {{kOnnxDomain, 13}}, {}, logger); + Graph& graph = model.MainGraph(); + ModelTestBuilder builder(graph); + + auto* input_a = builder.MakeInput(input_shape, -1.0f, 1.0f); + auto* transpose_output = builder.MakeIntermediate(); + Node& transpose = builder.AddNode("Transpose", {input_a}, {transpose_output}); + transpose.AddAttribute("perm", perm); + + auto* input_b = builder.MakeInput({3, 4}, -1.0f, 1.0f); + auto* output = builder.MakeOutput(); + Node& gemm = builder.AddNode("Gemm", {transpose_output, input_b}, {output}); + gemm.AddAttribute("transA", int64_t{0}); + gemm.AddAttribute("transB", int64_t{0}); + gemm.AddAttribute("alpha", 1.0f); + gemm.AddAttribute("beta", 1.0f); + + ASSERT_STATUS_OK(graph.Resolve()); + + RuleBasedGraphTransformer transformer("GemmTransposeFusionTest"); + ASSERT_STATUS_OK(transformer.Register(std::make_unique())); + + bool modified = false; + ASSERT_STATUS_OK(transformer.Apply(graph, modified, logger)); + ASSERT_STATUS_OK(graph.Resolve()); + + EXPECT_EQ(modified, expect_fusion); + auto op_counts = CountOpsInGraph(graph); + EXPECT_EQ(op_counts["Transpose"], expect_fusion ? 0 : 1); + EXPECT_EQ(op_counts["Gemm"], 1); + + const Node* resulting_gemm = nullptr; + for (const auto& graph_node : graph.Nodes()) { + if (graph_node.OpType() == "Gemm") { + resulting_gemm = &graph_node; + break; + } + } + + ASSERT_NE(resulting_gemm, nullptr); + EXPECT_EQ(resulting_gemm->GetAttributes().at("transA").i(), expected_trans_a); +} + +void RunWeightTransposeTest(const std::vector& perm, + const std::vector& weight_shape, + bool expect_fusion, + int64_t expected_trans_b) { + auto& logger = DefaultLoggingManager().DefaultLogger(); + Model model("GemmWeightTransposeFusionPermutation", false, ModelMetaData(), PathString(), + IOnnxRuntimeOpSchemaRegistryList(), {{kOnnxDomain, 13}}, {}, logger); + Graph& graph = model.MainGraph(); + ModelTestBuilder builder(graph); + + auto* input_a = builder.MakeInput({3, 4}, -1.0f, 1.0f); + auto* weight = builder.MakeInput(weight_shape, -1.0f, 1.0f); + auto* transpose_output = builder.MakeIntermediate(); + Node& transpose = builder.AddNode("Transpose", {weight}, {transpose_output}); + transpose.AddAttribute("perm", perm); + + auto* output = builder.MakeOutput(); + Node& gemm = builder.AddNode("Gemm", {input_a, transpose_output}, {output}); + gemm.AddAttribute("transA", int64_t{0}); + gemm.AddAttribute("transB", int64_t{0}); + gemm.AddAttribute("alpha", 1.0f); + gemm.AddAttribute("beta", 1.0f); + + ASSERT_STATUS_OK(graph.Resolve()); + + RuleBasedGraphTransformer transformer("GemmTransposeFusionTest"); + ASSERT_STATUS_OK(transformer.Register(std::make_unique())); + + bool modified = false; + ASSERT_STATUS_OK(transformer.Apply(graph, modified, logger)); + ASSERT_STATUS_OK(graph.Resolve()); + + EXPECT_EQ(modified, expect_fusion); + auto op_counts = CountOpsInGraph(graph); + EXPECT_EQ(op_counts["Transpose"], expect_fusion ? 0 : 1); + EXPECT_EQ(op_counts["Gemm"], 1); + + const Node* resulting_gemm = nullptr; + for (const auto& graph_node : graph.Nodes()) { + if (graph_node.OpType() == "Gemm") { + resulting_gemm = &graph_node; + break; + } + } + + ASSERT_NE(resulting_gemm, nullptr); + EXPECT_EQ(resulting_gemm->GetAttributes().at("transB").i(), expected_trans_b); +} + +} // namespace + +TEST(GemmTransposeFusionTest, IdentityInputTransposeIsNotFolded) { + RunInputTransposeTest({0, 1}, {2, 3}, false, 0); +} + +TEST(GemmTransposeFusionTest, MatrixInputTransposeIsFolded) { + RunInputTransposeTest({1, 0}, {3, 2}, true, 1); +} + +TEST(GemmTransposeFusionTest, IdentityWeightTransposeIsNotFolded) { + RunWeightTransposeTest({0, 1}, {4, 5}, false, 0); +} + +TEST(GemmTransposeFusionTest, MatrixWeightTransposeIsFolded) { + RunWeightTransposeTest({1, 0}, {5, 4}, true, 1); +} + +} // namespace test +} // namespace onnxruntime