Fix GemmTransposeFusion handling of identity Transpose

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<std::reference_wrapper<Node>> 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<const Node*> 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<const Node*> 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()) {

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<int64_t>& perm,
+                           const std::vector<int64_t>& 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<float>(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<float>({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<GemmTransposeFusion>()));
+
+  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<int64_t>& perm,
+                            const std::vector<int64_t>& 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<float>({3, 4}, -1.0f, 1.0f);
+  auto* weight = builder.MakeInput<float>(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<GemmTransposeFusion>()));
+
+  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