fix: avoid fusing identity transpose into Gemm

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<int64_t>(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<const Node*> 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<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());
@@ -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()) {

onnxruntime/test/optimizer/graph_transform_test.cc

@@ -5412,6 +5412,60 @@
   ASSERT_TRUE(new_input_defs[1]->Name() == "B");
 }
 
+TEST_F(GraphTransformationTests, GemmTransposeFusionDoesNotFuseIdentityTranspose) {
+  auto build_test_case = [](ModelTestBuilder& builder) {
+    auto* input = builder.MakeInput<float>({3, 4});
+    auto* weight = builder.MakeInput<float>({4, 5});
+    auto* transposed_weight = builder.MakeIntermediate<float>(std::vector<int64_t>{4, 5});
+    auto* output = builder.MakeOutput<float>(std::vector<int64_t>{3, 5});
+
+    builder.AddNode("Transpose", {weight}, {transposed_weight}).AddAttribute("perm", std::vector<int64_t>{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<RuleBasedGraphTransformer>("RuleTransformer");
+  ASSERT_STATUS_OK(rule_transformer->Register(std::make_unique<GemmTransposeFusion>()));
+  ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 13, *logger_, std::move(rule_transformer), TransformerLevel::Level1,
+                                        1, check_graph, check_graph));
+}
+
+TEST_F(GraphTransformationTests, GemmTransposeFusionDoesNotFuseIdentityTransposeAtOutput) {
+  auto build_test_case = [](ModelTestBuilder& builder) {
+    auto* input = builder.MakeInput<float>({3, 4});
+    auto* weight = builder.MakeInput<float>({4, 5});
+    auto* gemm_output = builder.MakeIntermediate<float>(std::vector<int64_t>{3, 5});
+    auto* output = builder.MakeOutput<float>(std::vector<int64_t>{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<int64_t>{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<RuleBasedGraphTransformer>("RuleTransformer");
+  ASSERT_STATUS_OK(rule_transformer->Register(std::make_unique<GemmTransposeFusion>()));
+  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";