diff --git a/onnxruntime/core/optimizer/gather_to_split_fusion.cc b/onnxruntime/core/optimizer/gather_to_split_fusion.cc new file mode 100644 index 0000000000..e0e5b1b2ef --- /dev/null +++ b/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()); + 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(shape->dim_size()); + + bool can_fuse = true; + bool first_edge = true; + int64_t split_axis = 0; + InlinedVector gather_outputs(output_count, nullptr); + InlinedVector> 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(axis)); + if (!utils::HasDimValue(dim) || dim.dim_value() != static_cast(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(output_count); + if (index < 0 || index >= static_cast(output_count) || gather_outputs[static_cast(index)]) { + can_fuse = false; + break; + } + Node& gather_node = *graph.GetNode(it->Index()); + nodes_to_fuse.emplace_back(gather_node); + gather_outputs[static_cast(index)] = gather_node.MutableOutputDefs()[0]; + } + + if (!can_fuse) continue; + + ONNX_NAMESPACE::TypeProto split_output_type; + const ONNX_NAMESPACE::TensorProto_DataType element_type = static_cast( + 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(i)); + } + } + + InlinedVector 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{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(1)); + axes_initializer_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); + InlinedVector 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 diff --git a/onnxruntime/core/optimizer/gather_to_split_fusion.h b/onnxruntime/core/optimizer/gather_to_split_fusion.h new file mode 100644 index 0000000000..4860393b47 --- /dev/null +++ b/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& 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 diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index b6cdacd994..45dc9cdcb9 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/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> GenerateTransformers( transformers.emplace_back(std::make_unique(cpu_cuda_rocm_eps)); transformers.emplace_back(std::make_unique(cpu_cuda_rocm_eps)); transformers.emplace_back(std::make_unique(cpu_cuda_rocm_eps)); + transformers.emplace_back(std::make_unique(cpu_cuda_rocm_eps)); transformers.emplace_back(std::make_unique(cpu_cuda_rocm_eps)); transformers.emplace_back(std::make_unique(cpu_cuda_rocm_eps)); diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 730dcf16c2..5ab7284113 100755 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/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({{54}}); + auto* shape_arg = builder.MakeInput({{1}}); + auto* reshape_out = builder.MakeIntermediate({{2, 3, 3, 3}}); + auto* gather_index_1 = builder.MakeInitializer({}, {static_cast(0)}); + auto* gather_index_2 = builder.MakeInitializer({}, {static_cast(1)}); + auto* gather_index_3 = builder.MakeInitializer({}, {static_cast(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(2)); + builder.AddNode("Gather", {reshape_out, gather_index_2}, {gahter_out_2}) + .AddAttribute("axis", static_cast(-2)); + builder.AddNode("Gather", {reshape_out, gather_index_3}, {gahter_out_3}) + .AddAttribute("axis", static_cast(2)); + builder.AddNode("Transpose", {gahter_out_1}, {transpose_out_1}).AddAttribute("perm", std::vector{0, 2, 1}); + builder.AddNode("Transpose", {gahter_out_2}, {transpose_out_2}).AddAttribute("perm", std::vector{0, 2, 1}); + builder.AddNode("Transpose", {gahter_out_3}, {transpose_out_3}).AddAttribute("perm", std::vector{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(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(attrs.at("axes").ints().at(0))); + } + } + }; + + std::unique_ptr transformer = std::make_unique(); + 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(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(*(init_const.data()))); + } + } + }; + + std::unique_ptr transformer = std::make_unique(); + 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({{72}}); + auto* shape_arg = builder.MakeInput({{1}}); + auto* reshape_out = builder.MakeIntermediate({{2, 3, 4, 3}}); + auto* gather_index_1 = builder.MakeInitializer({}, {static_cast(0)}); + auto* gather_index_2 = builder.MakeInitializer({}, {static_cast(1)}); + auto* gather_index_3 = builder.MakeInitializer({}, {static_cast(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(2)); + builder.AddNode("Gather", {reshape_out, gather_index_2}, {gahter_out_2}) + .AddAttribute("axis", static_cast(2)); + builder.AddNode("Gather", {reshape_out, gather_index_3}, {gahter_out_3}) + .AddAttribute("axis", static_cast(2)); + builder.AddNode("Transpose", {gahter_out_1}, {transpose_out_1}) + .AddAttribute("perm", std::vector{0, 2, 1}); + builder.AddNode("Transpose", {gahter_out_2}, {transpose_out_2}) + .AddAttribute("perm", std::vector{0, 2, 1}); + builder.AddNode("Transpose", {gahter_out_3}, {transpose_out_3}) + .AddAttribute("perm", std::vector{0, 2, 1}); + }; + + std::unique_ptr transformer = std::make_unique(); + 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({{54}}); + auto* shape_arg = builder.MakeInput({{1}}); + auto* reshape_out = builder.MakeIntermediate({{2, 3, 3, 3}}); + auto* gather_index_1 = builder.MakeInitializer({}, {static_cast(0)}); + auto* gather_index_2 = builder.MakeInitializer({}, {static_cast(1)}); + auto* gather_index_3 = builder.MakeInitializer({}, {static_cast(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(2)); + builder.AddNode("Gather", {reshape_out, gather_index_2}, {gahter_out_2}) + .AddAttribute("axis", static_cast(2)); + builder.AddNode("Gather", {reshape_out, gather_index_3}, {gahter_out_3}) + .AddAttribute("axis", static_cast(2)); + builder.AddNode("Transpose", {gahter_out_1}, {transpose_out_1}) + .AddAttribute("perm", std::vector{0, 2, 1}); + builder.AddNode("Transpose", {gahter_out_2}, {transpose_out_2}) + .AddAttribute("perm", std::vector{0, 2, 1}); + builder.AddNode("Transpose", {gahter_out_3}, {transpose_out_3}) + .AddAttribute("perm", std::vector{0, 2, 1}); + }; + + std::unique_ptr transformer = std::make_unique(); + 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({{54}}); + auto* shape_arg = builder.MakeInput({{1}}); + auto* reshape_out = builder.MakeIntermediate({{2, 3, 3, 3}}); + auto* gather_index_1 = builder.MakeInitializer({}, {static_cast(0)}); + auto* gather_index_2 = builder.MakeInitializer({}, {static_cast(1)}); + auto* gather_index_3 = builder.MakeInitializer({}, {static_cast(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(1)); + builder.AddNode("Gather", {reshape_out, gather_index_2}, {gahter_out_2}) + .AddAttribute("axis", static_cast(2)); + builder.AddNode("Gather", {reshape_out, gather_index_3}, {gahter_out_3}) + .AddAttribute("axis", static_cast(3)); + builder.AddNode("Transpose", {gahter_out_1}, {transpose_out_1}) + .AddAttribute("perm", std::vector{0, 2, 1}); + builder.AddNode("Transpose", {gahter_out_2}, {transpose_out_2}) + .AddAttribute("perm", std::vector{0, 2, 1}); + builder.AddNode("Transpose", {gahter_out_3}, {transpose_out_3}) + .AddAttribute("perm", std::vector{0, 2, 1}); + }; + + std::unique_ptr transformer = std::make_unique(); + TestGraphTransformer(build_test_case, 14, *logger_, std::move(transformer), TransformerLevel::Level1, 1, + pre_graph_checker, post_graph_checker); + } +} + } // namespace test } // namespace onnxruntime diff --git a/orttraining/orttraining/core/optimizer/graph_transformer_utils.cc b/orttraining/orttraining/core/optimizer/graph_transformer_utils.cc index 2012979462..3149600ee8 100644 --- a/orttraining/orttraining/core/optimizer/graph_transformer_utils.cc +++ b/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> GeneratePreTrainingTransformers( transformers.emplace_back(std::make_unique(compatible_eps)); transformers.emplace_back(std::make_unique(compatible_eps)); transformers.emplace_back(std::make_unique(compatible_eps)); + transformers.emplace_back(std::make_unique(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