Multiple Gather to Split Fusion (#13095)

For below code in some transformers models:
```
fused_qkv = fused_qkv.view(batch_size, seq_length, self.num_heads, 3, self.head_dim)
return fused_qkv[..., 0, :], fused_qkv[..., 1, :], fused_qkv[..., 2, :]
```

The exported graph will contains 3 Gather nodes, currently ORT's
GatherGrad CUDA implementation is slow. This pattern can be fused to use
one Split, so that we can launch less kernels for the compute, the perf
of Split/Concat (for grad) is also better than Gather/GatherGrad.

In a real example, one GatherGrad will take 15ms and there are 3 for
each layer in the graph, after the fusion, one Concat takes only 35us.
The total time of a step is improved from 1.5s to 0.4s.

onnxruntime/core/optimizer/gather_to_split_fusion.cc

@@ -0,0 +1,178 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/optimizer/gather_to_split_fusion.h"
+
+#include "core/graph/graph_utils.h"
+#include "core/optimizer/initializer.h"
+#include "core/optimizer/utils.h"
+
+namespace onnxruntime {
+
+bool GatherToSplitFusion::IsSupportedGather(const Graph& graph, const Node& node, int64_t& index, int64_t& axis) const {
+  if (!graph_utils::IsSupportedOptypeVersionAndDomain(node, "Gather", {1, 11, 13}) ||
+      !graph_utils::IsSupportedProvider(node, GetCompatibleExecutionProviders())) {
+    return false;
+  }
+
+  const NodeArg& input_arg = *(node.InputDefs()[1]);
+  if (!optimizer_utils::IsScalar(input_arg)) return false;
+  const ONNX_NAMESPACE::TensorProto* tensor_proto = graph_utils::GetConstantInitializer(graph, input_arg.Name());
+  if (!tensor_proto) return false;
+  Initializer init_const{*tensor_proto, graph.ModelPath()};
+  if (tensor_proto->data_type() != ONNX_NAMESPACE::TensorProto_DataType_INT64) return false;
+  index = *(init_const.data<int64_t>());
+  axis = 0;  // Default value.
+  auto& attrs = node.GetAttributes();
+  if (attrs.find("axis") != attrs.end()) {
+    auto& axis_attr = attrs.at("axis");
+    if (utils::HasInt(axis_attr)) axis = axis_attr.i();
+  }
+  return true;
+}
+
+/*
+GatherToSplitFusion is to fuse:
+Node -> Gather(index=0, axis=axis)
+    |-> Gather(index=1, axis=axis)
+    |-> Gather(index=2, axis=axis)
+    |...
+
+To
+
+Node -> Split -> Squeeze(axis=axis)
+             |-> Squeeze(axis=axis)
+             |-> Squeeze(axis=axis)
+             |...
+
+So that we can use one kernel to finish the job.
+*/
+Status GatherToSplitFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level,
+                                      const logging::Logger& logger) const {
+  GraphViewer graph_viewer(graph);
+  const auto& node_topology_list = graph_viewer.GetNodesInTopologicalOrder();
+
+  for (auto node_index : node_topology_list) {
+    auto* p_node = graph.GetNode(node_index);
+    if (p_node == nullptr) continue;  // we removed the node as part of an earlier fusion
+    Node& node = *p_node;
+
+    ORT_RETURN_IF_ERROR(Recurse(node, modified, graph_level, logger));
+
+    // Gather following Shape is a common case but not the target case to fuse here as its compute is normally very
+    // quick.
+    if (node.OpType() == "Shape") continue;
+
+    // Ideally it's possible that the node has multiple outputs and the Gather nodes can consume one of them.
+    // To make the fusion simple, here we require the node has only one output, if in the future we observe
+    // new pattern we can optimize the fusion here.
+    if (node.MutableOutputDefs().size() > 1) continue;
+
+    // No need to fuse if there is only one output or no output.
+    size_t output_count = node.GetOutputEdgesCount();
+    if (output_count <= 1) continue;
+
+    auto shape = node.MutableOutputDefs()[0]->Shape();
+    if (!shape) continue;
+    int64_t rank = static_cast<int64_t>(shape->dim_size());
+
+    bool can_fuse = true;
+    bool first_edge = true;
+    int64_t split_axis = 0;
+    InlinedVector<NodeArg*> gather_outputs(output_count, nullptr);
+    InlinedVector<std::reference_wrapper<Node>> nodes_to_fuse;
+    for (auto it = node.OutputNodesBegin(); it != node.OutputNodesEnd(); ++it) {
+      int64_t index, axis;
+      if (!IsSupportedGather(graph, *it, index, axis)) {
+        can_fuse = false;
+        break;
+      }
+      if (axis < 0) axis += rank;
+      if (first_edge) {
+        auto dim = shape->dim(static_cast<int>(axis));
+        if (!utils::HasDimValue(dim) || dim.dim_value() != static_cast<int64_t>(output_count)) {
+          can_fuse = false;
+          break;
+        }
+        split_axis = axis;
+        first_edge = false;
+      } else if (axis != split_axis) {
+        can_fuse = false;
+        break;
+      }
+      if (index < 0) index += static_cast<int64_t>(output_count);
+      if (index < 0 || index >= static_cast<int64_t>(output_count) || gather_outputs[static_cast<size_t>(index)]) {
+        can_fuse = false;
+        break;
+      }
+      Node& gather_node = *graph.GetNode(it->Index());
+      nodes_to_fuse.emplace_back(gather_node);
+      gather_outputs[static_cast<size_t>(index)] = gather_node.MutableOutputDefs()[0];
+    }
+
+    if (!can_fuse) continue;
+
+    ONNX_NAMESPACE::TypeProto split_output_type;
+    const ONNX_NAMESPACE::TensorProto_DataType element_type = static_cast<ONNX_NAMESPACE::TensorProto_DataType>(
+        node.MutableOutputDefs()[0]->TypeAsProto()->tensor_type().elem_type());
+    split_output_type.mutable_tensor_type()->set_elem_type(element_type);
+    for (int64_t i = 0; i < rank; ++i) {
+      if (i == split_axis) {
+        split_output_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1LL);
+      } else {
+        *(split_output_type.mutable_tensor_type()->mutable_shape()->add_dim()) = shape->dim(static_cast<int>(i));
+      }
+    }
+
+    InlinedVector<NodeArg*> split_outputs;
+    for (size_t i = 0; i < output_count; ++i) {
+      split_outputs.emplace_back(
+          &graph.GetOrCreateNodeArg(graph.GenerateNodeArgName("split" + std::to_string(i)), &split_output_type));
+    }
+
+    Node& split_node = graph.AddNode(graph.GenerateNodeName("Split"), "Split", "Split for Fused Gather nodes",
+                                     {node.MutableOutputDefs()[0]}, split_outputs);
+    split_node.AddAttribute("axis", split_axis);
+    split_node.SetExecutionProviderType(node.GetExecutionProviderType());
+
+    // Squeeze before and after OpSet-13 have different schemas.
+    int onnx_opset_version = -1;
+    if (graph.DomainToVersionMap().find(kOnnxDomain) != graph.DomainToVersionMap().end()) {
+      onnx_opset_version = graph.DomainToVersionMap().at(kOnnxDomain);
+    }
+
+    if (onnx_opset_version < 13) {
+      for (size_t i = 0; i < output_count; ++i) {
+        Node& squeeze_node = graph.AddNode(graph.GenerateNodeName("Squeeze" + std::to_string(i)), "Squeeze",
+                                           "Squeeze for Fused Gather nodes", {split_outputs[i]}, {gather_outputs[i]});
+        squeeze_node.AddAttribute("axes", std::vector<int64_t>{split_axis});
+        squeeze_node.SetExecutionProviderType(node.GetExecutionProviderType());
+      }
+    } else {
+      ONNX_NAMESPACE::TensorProto axes_initializer_proto;
+      axes_initializer_proto.set_name(graph.GenerateNodeName("SqueezeAxesInitializer"));
+      axes_initializer_proto.add_dims(static_cast<int64_t>(1));
+      axes_initializer_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64);
+      InlinedVector<int64_t> axes_value{split_axis};
+      axes_initializer_proto.set_raw_data(axes_value.data(), axes_value.size() * sizeof(int64_t));
+      NodeArg* axes_arg = &graph_utils::AddInitializer(graph, axes_initializer_proto);
+
+      for (size_t i = 0; i < output_count; ++i) {
+        Node& squeeze_node =
+            graph.AddNode(graph.GenerateNodeName("Squeeze" + std::to_string(i)), "Squeeze",
+                          "Squeeze for Fused Gather nodes", {split_outputs[i], axes_arg}, {gather_outputs[i]});
+        squeeze_node.SetExecutionProviderType(node.GetExecutionProviderType());
+      }
+    }
+
+    for (Node& n : nodes_to_fuse) {
+      graph_utils::RemoveNodeOutputEdges(graph, n);
+      graph.RemoveNode(n.Index());
+    }
+
+    modified = true;
+  }
+
+  return Status::OK();
+}
+}  // namespace onnxruntime

onnxruntime/core/optimizer/gather_to_split_fusion.h

@@ -0,0 +1,26 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "core/optimizer/graph_transformer.h"
+
+namespace onnxruntime {
+
+/**
+@Class GatherToSplitFusion
+
+Fuse multiple Gather nodes that comsuming one output to one Split node.
+*/
+class GatherToSplitFusion : public GraphTransformer {
+ public:
+  GatherToSplitFusion(const InlinedHashSet<std::string_view>& compatible_execution_providers = {}) noexcept
+      : GraphTransformer("GatherToSplitFusion", compatible_execution_providers) {}
+
+  Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override;
+
+ private:
+  bool IsSupportedGather(const Graph& graph, const Node& node, int64_t& index, int64_t& axis) const;
+};
+
+}  // namespace onnxruntime

onnxruntime/core/optimizer/graph_transformer_utils.cc

@@ -37,6 +37,7 @@
 #include "core/optimizer/expand_elimination.h"
 #include "core/optimizer/fast_gelu_fusion.h"
 #include "core/optimizer/free_dim_override_transformer.h"
+#include "core/optimizer/gather_to_split_fusion.h"
 #include "core/optimizer/gelu_approximation.h"
 #include "core/optimizer/gelu_fusion.h"
 #include "core/optimizer/gemm_activation_fusion.h"
@@ -249,6 +250,7 @@ InlinedVector<std::unique_ptr<GraphTransformer>> GenerateTransformers(
       transformers.emplace_back(std::make_unique<SimplifiedLayerNormFusion>(cpu_cuda_rocm_eps));
       transformers.emplace_back(std::make_unique<AttentionFusion>(cpu_cuda_rocm_eps));
       transformers.emplace_back(std::make_unique<EmbedLayerNormFusion>(cpu_cuda_rocm_eps));
+      transformers.emplace_back(std::make_unique<GatherToSplitFusion>(cpu_cuda_rocm_eps));
 
       transformers.emplace_back(std::make_unique<MatmulTransposeFusion>(cpu_cuda_rocm_eps));
       transformers.emplace_back(std::make_unique<BiasGeluFusion>(cpu_cuda_rocm_eps));

onnxruntime/test/optimizer/graph_transform_test.cc

@@ -38,6 +38,7 @@
 #include "core/optimizer/embed_layer_norm_fusion.h"
 #include "core/optimizer/expand_elimination.h"
 #include "core/optimizer/fast_gelu_fusion.h"
+#include "core/optimizer/gather_to_split_fusion.h"
 #include "core/optimizer/gelu_approximation.h"
 #include "core/optimizer/gelu_fusion.h"
 #include "core/optimizer/gemm_activation_fusion.h"
@@ -5191,5 +5192,204 @@ TEST_F(GraphTransformationTests, PropagateCastOpsTests_Gelu) {
 }
 #endif
 
+TEST_F(GraphTransformationTests, GatherToSplitFusion) {
+  auto build_test_case = [&](ModelTestBuilder& builder) {
+    auto* data_arg = builder.MakeInput<float>({{54}});
+    auto* shape_arg = builder.MakeInput<int64_t>({{1}});
+    auto* reshape_out = builder.MakeIntermediate<float>({{2, 3, 3, 3}});
+    auto* gather_index_1 = builder.MakeInitializer<int64_t>({}, {static_cast<int64_t>(0)});
+    auto* gather_index_2 = builder.MakeInitializer<int64_t>({}, {static_cast<int64_t>(1)});
+    auto* gather_index_3 = builder.MakeInitializer<int64_t>({}, {static_cast<int64_t>(2)});
+    auto* gahter_out_1 = builder.MakeIntermediate();
+    auto* gahter_out_2 = builder.MakeIntermediate();
+    auto* gahter_out_3 = builder.MakeIntermediate();
+    auto* transpose_out_1 = builder.MakeOutput();
+    auto* transpose_out_2 = builder.MakeOutput();
+    auto* transpose_out_3 = builder.MakeOutput();
+
+    builder.AddNode("Reshape", {data_arg, shape_arg}, {reshape_out});
+    builder.AddNode("Gather", {reshape_out, gather_index_1}, {gahter_out_1})
+        .AddAttribute("axis", static_cast<int64_t>(2));
+    builder.AddNode("Gather", {reshape_out, gather_index_2}, {gahter_out_2})
+        .AddAttribute("axis", static_cast<int64_t>(-2));
+    builder.AddNode("Gather", {reshape_out, gather_index_3}, {gahter_out_3})
+        .AddAttribute("axis", static_cast<int64_t>(2));
+    builder.AddNode("Transpose", {gahter_out_1}, {transpose_out_1}).AddAttribute("perm", std::vector<int64_t>{0, 2, 1});
+    builder.AddNode("Transpose", {gahter_out_2}, {transpose_out_2}).AddAttribute("perm", std::vector<int64_t>{0, 2, 1});
+    builder.AddNode("Transpose", {gahter_out_3}, {transpose_out_3}).AddAttribute("perm", std::vector<int64_t>{0, 2, 1});
+  };
+
+  auto pre_graph_checker = [&](Graph& graph) { ASSERT_EQ(CountOpsInGraph(graph)["Gather"], 3); };
+
+  // OpSet-12
+  {
+    auto post_graph_checker = [&](Graph& graph) {
+      ASSERT_EQ(CountOpsInGraph(graph)["Gather"], 0);
+      ASSERT_EQ(CountOpsInGraph(graph)["Split"], 1);
+      ASSERT_EQ(CountOpsInGraph(graph)["Squeeze"], 3);
+      for (auto& node : graph.Nodes()) {
+        if (node.OpType() == "Split") {
+          auto& attrs = node.GetAttributes();
+          ASSERT_TRUE(attrs.find("axis") != attrs.end());
+          ASSERT_EQ(2, static_cast<int>(attrs.at("axis").i()));
+        } else if (node.OpType() == "Squeeze") {
+          auto& attrs = node.GetAttributes();
+          ASSERT_TRUE(attrs.find("axes") != attrs.end());
+          ASSERT_EQ(2, static_cast<int>(attrs.at("axes").ints().at(0)));
+        }
+      }
+    };
+
+    std::unique_ptr<GraphTransformer> transformer = std::make_unique<GatherToSplitFusion>();
+    TestGraphTransformer(build_test_case, 12, *logger_, std::move(transformer), TransformerLevel::Level1, 1,
+                         pre_graph_checker, post_graph_checker);
+  }
+
+  // OpSet-14
+  {
+    auto post_graph_checker = [&](Graph& graph) {
+      ASSERT_EQ(CountOpsInGraph(graph)["Gather"], 0);
+      ASSERT_EQ(CountOpsInGraph(graph)["Split"], 1);
+      ASSERT_EQ(CountOpsInGraph(graph)["Squeeze"], 3);
+      for (auto& node : graph.Nodes()) {
+        if (node.OpType() == "Split") {
+          auto& attrs = node.GetAttributes();
+          ASSERT_TRUE(attrs.find("axis") != attrs.end());
+          ASSERT_EQ(2, static_cast<int>(attrs.at("axis").i()));
+        } else if (node.OpType() == "Squeeze") {
+          const NodeArg& input_arg = *(node.InputDefs()[1]);
+          const ONNX_NAMESPACE::TensorProto* tensor_proto =
+              graph_utils::GetConstantInitializer(graph, input_arg.Name());
+          ASSERT_TRUE(tensor_proto != nullptr);
+          Initializer init_const{*tensor_proto, graph.ModelPath()};
+          ASSERT_TRUE(tensor_proto->data_type() == ONNX_NAMESPACE::TensorProto_DataType_INT64);
+          ASSERT_EQ(2, static_cast<int>(*(init_const.data<int64_t>())));
+        }
+      }
+    };
+
+    std::unique_ptr<GraphTransformer> transformer = std::make_unique<GatherToSplitFusion>();
+    TestGraphTransformer(build_test_case, 14, *logger_, std::move(transformer), TransformerLevel::Level1, 1,
+                         pre_graph_checker, post_graph_checker);
+  }
+}
+
+TEST_F(GraphTransformationTests, GatherToSplitFusion_Invalid) {
+  auto pre_graph_checker = [&](Graph& graph) { ASSERT_EQ(CountOpsInGraph(graph)["Gather"], 3); };
+  auto post_graph_checker = [&](Graph& graph) {
+    ASSERT_EQ(CountOpsInGraph(graph)["Gather"], 3);
+    ASSERT_EQ(CountOpsInGraph(graph)["Split"], 0);
+    ASSERT_EQ(CountOpsInGraph(graph)["Squeeze"], 0);
+  };
+
+  // Invalid shape.
+  {
+    auto build_test_case = [&](ModelTestBuilder& builder) {
+      auto* data_arg = builder.MakeInput<float>({{72}});
+      auto* shape_arg = builder.MakeInput<int64_t>({{1}});
+      auto* reshape_out = builder.MakeIntermediate<float>({{2, 3, 4, 3}});
+      auto* gather_index_1 = builder.MakeInitializer<int64_t>({}, {static_cast<int64_t>(0)});
+      auto* gather_index_2 = builder.MakeInitializer<int64_t>({}, {static_cast<int64_t>(1)});
+      auto* gather_index_3 = builder.MakeInitializer<int64_t>({}, {static_cast<int64_t>(2)});
+      auto* gahter_out_1 = builder.MakeIntermediate();
+      auto* gahter_out_2 = builder.MakeIntermediate();
+      auto* gahter_out_3 = builder.MakeIntermediate();
+      auto* transpose_out_1 = builder.MakeOutput();
+      auto* transpose_out_2 = builder.MakeOutput();
+      auto* transpose_out_3 = builder.MakeOutput();
+
+      builder.AddNode("Reshape", {data_arg, shape_arg}, {reshape_out});
+      builder.AddNode("Gather", {reshape_out, gather_index_1}, {gahter_out_1})
+          .AddAttribute("axis", static_cast<int64_t>(2));
+      builder.AddNode("Gather", {reshape_out, gather_index_2}, {gahter_out_2})
+          .AddAttribute("axis", static_cast<int64_t>(2));
+      builder.AddNode("Gather", {reshape_out, gather_index_3}, {gahter_out_3})
+          .AddAttribute("axis", static_cast<int64_t>(2));
+      builder.AddNode("Transpose", {gahter_out_1}, {transpose_out_1})
+          .AddAttribute("perm", std::vector<int64_t>{0, 2, 1});
+      builder.AddNode("Transpose", {gahter_out_2}, {transpose_out_2})
+          .AddAttribute("perm", std::vector<int64_t>{0, 2, 1});
+      builder.AddNode("Transpose", {gahter_out_3}, {transpose_out_3})
+          .AddAttribute("perm", std::vector<int64_t>{0, 2, 1});
+    };
+
+    std::unique_ptr<GraphTransformer> transformer = std::make_unique<GatherToSplitFusion>();
+    TestGraphTransformer(build_test_case, 12, *logger_, std::move(transformer), TransformerLevel::Level1, 1,
+                         pre_graph_checker, post_graph_checker);
+  }
+
+  // Invalid Gather indices.
+  {
+    auto build_test_case = [&](ModelTestBuilder& builder) {
+      auto* data_arg = builder.MakeInput<float>({{54}});
+      auto* shape_arg = builder.MakeInput<int64_t>({{1}});
+      auto* reshape_out = builder.MakeIntermediate<float>({{2, 3, 3, 3}});
+      auto* gather_index_1 = builder.MakeInitializer<int64_t>({}, {static_cast<int64_t>(0)});
+      auto* gather_index_2 = builder.MakeInitializer<int64_t>({}, {static_cast<int64_t>(1)});
+      auto* gather_index_3 = builder.MakeInitializer<int64_t>({}, {static_cast<int64_t>(1)});
+      auto* gahter_out_1 = builder.MakeIntermediate();
+      auto* gahter_out_2 = builder.MakeIntermediate();
+      auto* gahter_out_3 = builder.MakeIntermediate();
+      auto* transpose_out_1 = builder.MakeOutput();
+      auto* transpose_out_2 = builder.MakeOutput();
+      auto* transpose_out_3 = builder.MakeOutput();
+
+      builder.AddNode("Reshape", {data_arg, shape_arg}, {reshape_out});
+      builder.AddNode("Gather", {reshape_out, gather_index_1}, {gahter_out_1})
+          .AddAttribute("axis", static_cast<int64_t>(2));
+      builder.AddNode("Gather", {reshape_out, gather_index_2}, {gahter_out_2})
+          .AddAttribute("axis", static_cast<int64_t>(2));
+      builder.AddNode("Gather", {reshape_out, gather_index_3}, {gahter_out_3})
+          .AddAttribute("axis", static_cast<int64_t>(2));
+      builder.AddNode("Transpose", {gahter_out_1}, {transpose_out_1})
+          .AddAttribute("perm", std::vector<int64_t>{0, 2, 1});
+      builder.AddNode("Transpose", {gahter_out_2}, {transpose_out_2})
+          .AddAttribute("perm", std::vector<int64_t>{0, 2, 1});
+      builder.AddNode("Transpose", {gahter_out_3}, {transpose_out_3})
+          .AddAttribute("perm", std::vector<int64_t>{0, 2, 1});
+    };
+
+    std::unique_ptr<GraphTransformer> transformer = std::make_unique<GatherToSplitFusion>();
+    TestGraphTransformer(build_test_case, 14, *logger_, std::move(transformer), TransformerLevel::Level1, 1,
+                         pre_graph_checker, post_graph_checker);
+  }
+
+  // Invalid Gather axis.
+  {
+    auto build_test_case = [&](ModelTestBuilder& builder) {
+      auto* data_arg = builder.MakeInput<float>({{54}});
+      auto* shape_arg = builder.MakeInput<int64_t>({{1}});
+      auto* reshape_out = builder.MakeIntermediate<float>({{2, 3, 3, 3}});
+      auto* gather_index_1 = builder.MakeInitializer<int64_t>({}, {static_cast<int64_t>(0)});
+      auto* gather_index_2 = builder.MakeInitializer<int64_t>({}, {static_cast<int64_t>(1)});
+      auto* gather_index_3 = builder.MakeInitializer<int64_t>({}, {static_cast<int64_t>(2)});
+      auto* gahter_out_1 = builder.MakeIntermediate();
+      auto* gahter_out_2 = builder.MakeIntermediate();
+      auto* gahter_out_3 = builder.MakeIntermediate();
+      auto* transpose_out_1 = builder.MakeOutput();
+      auto* transpose_out_2 = builder.MakeOutput();
+      auto* transpose_out_3 = builder.MakeOutput();
+
+      builder.AddNode("Reshape", {data_arg, shape_arg}, {reshape_out});
+      builder.AddNode("Gather", {reshape_out, gather_index_1}, {gahter_out_1})
+          .AddAttribute("axis", static_cast<int64_t>(1));
+      builder.AddNode("Gather", {reshape_out, gather_index_2}, {gahter_out_2})
+          .AddAttribute("axis", static_cast<int64_t>(2));
+      builder.AddNode("Gather", {reshape_out, gather_index_3}, {gahter_out_3})
+          .AddAttribute("axis", static_cast<int64_t>(3));
+      builder.AddNode("Transpose", {gahter_out_1}, {transpose_out_1})
+          .AddAttribute("perm", std::vector<int64_t>{0, 2, 1});
+      builder.AddNode("Transpose", {gahter_out_2}, {transpose_out_2})
+          .AddAttribute("perm", std::vector<int64_t>{0, 2, 1});
+      builder.AddNode("Transpose", {gahter_out_3}, {transpose_out_3})
+          .AddAttribute("perm", std::vector<int64_t>{0, 2, 1});
+    };
+
+    std::unique_ptr<GraphTransformer> transformer = std::make_unique<GatherToSplitFusion>();
+    TestGraphTransformer(build_test_case, 14, *logger_, std::move(transformer), TransformerLevel::Level1, 1,
+                         pre_graph_checker, post_graph_checker);
+  }
+}
+
 }  // namespace test
 }  // namespace onnxruntime

orttraining/orttraining/core/optimizer/graph_transformer_utils.cc

@@ -23,6 +23,7 @@
 #include "core/optimizer/expand_elimination.h"
 #include "core/optimizer/fast_gelu_fusion.h"
 #include "core/optimizer/free_dim_override_transformer.h"
+#include "core/optimizer/gather_to_split_fusion.h"
 #include "core/optimizer/gelu_approximation.h"
 #include "core/optimizer/gelu_fusion.h"
 #include "core/optimizer/gemm_activation_fusion.h"
@@ -99,6 +100,7 @@ std::vector<std::unique_ptr<GraphTransformer>> GeneratePreTrainingTransformers(
       transformers.emplace_back(std::make_unique<SimplifiedLayerNormFusion>(compatible_eps));
       transformers.emplace_back(std::make_unique<FastGeluFusion>(compatible_eps));
       transformers.emplace_back(std::make_unique<SoftmaxCrossEntropyLossInternalFusion>(compatible_eps));
+      transformers.emplace_back(std::make_unique<GatherToSplitFusion>(compatible_eps));
 
 #if defined(USE_CUDA) || defined(USE_ROCM)
       // We are supposed to use execution provider as indicator, but here we don't have access to the registered EP at this point