From f6edf135132bbcbeb8e2294078d305c4353422cc Mon Sep 17 00:00:00 2001 From: Aidan Beggs Date: Thu, 11 Nov 2021 18:33:13 -0800 Subject: [PATCH] Implement a Gemm/Sum fusion pattern (#9699) When the pattern Sum(Gemm(A, B), C) exists, we can convert it to Gemm(A, B, C), assuming that C the output of the original Gemm is not used elsewhere, and this change does not break broadcasting. --- onnxruntime/core/optimizer/gemm_sum_fusion.cc | 167 +++++++++ onnxruntime/core/optimizer/gemm_sum_fusion.h | 52 +++ .../core/optimizer/graph_transformer_utils.cc | 17 +- .../test/optimizer/graph_transform_test.cc | 326 ++++++++++++++++++ .../transform/fusion/gemm_sum_attributes.onnx | Bin 0 -> 263 bytes .../transform/fusion/gemm_sum_basic.onnx | Bin 0 -> 200 bytes .../fusion/gemm_sum_fusion_broadcast.onnx | Bin 0 -> 199 bytes .../testdata/transform/fusion/gemm_sum_gen.py | 187 ++++++++++ .../fusion/gemm_sum_internal_nodes.onnx | Bin 0 -> 288 bytes .../gemm_sum_no_fusion_broadcast_failure.onnx | Bin 0 -> 208 bytes .../fusion/gemm_sum_no_fusion_c_used.onnx | Bin 0 -> 226 bytes ...m_no_fusion_original_gemm_output_used.onnx | Bin 0 -> 225 bytes ...emm_sum_no_fusion_sum_multiple_inputs.onnx | Bin 0 -> 226 bytes .../core/optimizer/graph_transformer_utils.cc | 12 +- 14 files changed, 749 insertions(+), 12 deletions(-) create mode 100644 onnxruntime/core/optimizer/gemm_sum_fusion.cc create mode 100644 onnxruntime/core/optimizer/gemm_sum_fusion.h create mode 100644 onnxruntime/test/testdata/transform/fusion/gemm_sum_attributes.onnx create mode 100644 onnxruntime/test/testdata/transform/fusion/gemm_sum_basic.onnx create mode 100644 onnxruntime/test/testdata/transform/fusion/gemm_sum_fusion_broadcast.onnx create mode 100644 onnxruntime/test/testdata/transform/fusion/gemm_sum_gen.py create mode 100644 onnxruntime/test/testdata/transform/fusion/gemm_sum_internal_nodes.onnx create mode 100644 onnxruntime/test/testdata/transform/fusion/gemm_sum_no_fusion_broadcast_failure.onnx create mode 100644 onnxruntime/test/testdata/transform/fusion/gemm_sum_no_fusion_c_used.onnx create mode 100644 onnxruntime/test/testdata/transform/fusion/gemm_sum_no_fusion_original_gemm_output_used.onnx create mode 100644 onnxruntime/test/testdata/transform/fusion/gemm_sum_no_fusion_sum_multiple_inputs.onnx diff --git a/onnxruntime/core/optimizer/gemm_sum_fusion.cc b/onnxruntime/core/optimizer/gemm_sum_fusion.cc new file mode 100644 index 0000000000..c3010b5d54 --- /dev/null +++ b/onnxruntime/core/optimizer/gemm_sum_fusion.cc @@ -0,0 +1,167 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/optimizer/gemm_sum_fusion.h" + +#include "core/graph/graph_utils.h" +#include "core/optimizer/initializer.h" +#include "core/optimizer/utils.h" + +using namespace ONNX_NAMESPACE; +using namespace onnxruntime::common; +namespace onnxruntime { + +Status GemmSumFusion::Apply(Graph& graph, Node& gemm_node, RewriteRuleEffect& modified, const logging::Logger&) const { + // Get currently set attributes of Gemm. Beta will become 1.0. + const bool transA = static_cast(gemm_node.GetAttributes().at("transA").i()); + const bool transB = static_cast(gemm_node.GetAttributes().at("transB").i()); + const float alpha = gemm_node.GetAttributes().at("alpha").f(); + const float beta = 1.0f; + + Node& sum_node = *graph.GetNode(gemm_node.OutputEdgesBegin()->GetNode().Index()); + + // The first two input defs for our new gemm (aka tensor A and tensor B in GemmSumFusion's + // documentation) are exactly the same as the old gemm's input defs. + std::vector new_gemm_input_defs = gemm_node.MutableInputDefs(); + + // The other new gemm's input def is the old sum's other input def (aka tensor C in + // GemmSumFusion's documentation). + if (sum_node.InputDefs()[0]->Name() == gemm_node.OutputDefs()[0]->Name()) { + new_gemm_input_defs.push_back(sum_node.MutableInputDefs()[1]); + } else { + new_gemm_input_defs.push_back(sum_node.MutableInputDefs()[0]); + } + ORT_ENFORCE(new_gemm_input_defs.size() == 3); + + std::vector new_gemm_output_defs = sum_node.MutableOutputDefs(); + ORT_ENFORCE(new_gemm_output_defs.size() == 1); + + Node& new_gemm_node = graph.AddNode(graph.GenerateNodeName(gemm_node.Name() + "_sum_transformed"), + gemm_node.OpType(), + "Fused Gemm with Sum", + new_gemm_input_defs, + new_gemm_output_defs, + {}, + gemm_node.Domain()); + new_gemm_node.AddAttribute("transA", static_cast(transA)); + new_gemm_node.AddAttribute("transB", static_cast(transB)); + new_gemm_node.AddAttribute("alpha", alpha); + new_gemm_node.AddAttribute("beta", beta); + + // Move both input edges from original gemm to new gemm. + for (auto gemm_input_edge : graph_utils::GraphEdge::GetNodeInputEdges(gemm_node)) { + ORT_ENFORCE(gemm_input_edge.src_arg_index < 2); + graph.AddEdge(gemm_input_edge.src_node, new_gemm_node.Index(), gemm_input_edge.src_arg_index, gemm_input_edge.dst_arg_index); + graph.RemoveEdge(gemm_input_edge.src_node, gemm_input_edge.dst_node, gemm_input_edge.src_arg_index, gemm_input_edge.dst_arg_index); + } + + // Move all output edges from sum to new gemm. + for (auto sum_output_edge : graph_utils::GraphEdge::GetNodeOutputEdges(sum_node)) { + ORT_ENFORCE(sum_output_edge.src_arg_index == 0); + graph.AddEdge(new_gemm_node.Index(), sum_output_edge.dst_node, sum_output_edge.src_arg_index, sum_output_edge.dst_arg_index); + graph.RemoveEdge(sum_output_edge.src_node, sum_output_edge.dst_node, sum_output_edge.src_arg_index, sum_output_edge.dst_arg_index); + } + + // Finally, move the other sum input edge to "C" for the new gemm node. + // If The other sum input def is a a graph input, there is no edge to move. + bool sum_input_moved = false; + for (auto sum_input_edge : graph_utils::GraphEdge::GetNodeInputEdges(sum_node)) { + // The C tensor in GemmSumFusion's documentation is the sum output which does not come from + // the fused gemm. The following condition will be true when seeing the input of C to sum. + if (sum_input_edge.src_node != gemm_node.Index()) { + ORT_ENFORCE(!sum_input_moved); + graph.AddEdge(sum_input_edge.src_node, new_gemm_node.Index(), sum_input_edge.src_arg_index, 2); + graph.RemoveEdge(sum_input_edge.src_node, sum_input_edge.dst_node, sum_input_edge.src_arg_index, sum_input_edge.dst_arg_index); + sum_input_moved = true; + } + } + // Old gemm node output is no longer needed. It was previously fed into the + // sum node which is now also handled by the new gemm. Remove this output edge + // to allow the node to be removed from the graph. + graph_utils::RemoveNodeOutputEdges(graph, gemm_node); + ORT_ENFORCE(graph.RemoveNode(gemm_node.Index())); + + // All output edges of the sum node should have been removed already, so we can + // remove it from the graph. + ORT_ENFORCE(sum_node.GetOutputEdgesCount() == 0); + ORT_ENFORCE(graph.RemoveNode(sum_node.Index())); + + modified = RewriteRuleEffect::kRemovedCurrentNode; + + return Status::OK(); +} + +bool GemmSumFusion::SatisfyCondition(const Graph& graph, const Node& node, const logging::Logger&) const { + // Perform a series of checks. If any fail, fusion may not be performed. + + // Original gemm's C must be missing for this fusion pattern to be valid. + // Supported for Opset >=11 as earlier opsets have C as a required input + if (!graph_utils::IsSupportedOptypeVersionAndDomain(node, "Gemm", {11, 13}) || + graph.NodeProducesGraphOutput(node) || + // 2 inputs means that the gemm has A and B present, but not C. + node.InputDefs().size() != 2) { + return false; + } + + // This gemm node must have exactly one output for this fusion pattern to be valid. We have already + // verified that this node does not produce any graph outputs, so we can check output edges. + if (node.GetOutputEdgesCount() != 1) { + return false; + } + + const NodeArg* node_output = node.OutputDefs()[0]; + const Node& output_node = node.OutputEdgesBegin()->GetNode(); + + // Fusion can be applied if the only output node is a Sum with exactly two inputs. + if ( + !graph_utils::IsSupportedOptypeVersionAndDomain(output_node, "Sum", {1, 6, 8, 13}) || + output_node.InputDefs().size() != 2 || + // Make sure the two nodes do not span execution providers. + output_node.GetExecutionProviderType() != node.GetExecutionProviderType()) { + return false; + } + + // Sum must have the same input types, data types do not need to be checked. + + // Get the other sum input. + const NodeArg* other_sum_input = nullptr; + if (output_node.InputDefs()[0]->Name() == node_output->Name()) { + other_sum_input = output_node.InputDefs()[1]; + } else { + other_sum_input = output_node.InputDefs()[0]; + } + ORT_ENFORCE(other_sum_input != nullptr); + + // valid bias_shapes are (N) or (1, N) or (M, 1) or (M, N) as + // GEMM only supports unidirectional broadcast on the bias input C + // + // TODO: verify if scalar shape works here together with matmul_add_fusion. + if (!other_sum_input->Shape()) { + return false; + } + + if (!node_output->Shape() || node_output->Shape()->dim_size() != 2) { + return false; + } + + const auto& bias_shape = *other_sum_input->Shape(); + const auto& matmul_output_shape = *node_output->Shape(); + const auto& M = matmul_output_shape.dim()[0]; + const auto& N = matmul_output_shape.dim()[1]; + auto dim_has_value_1 = [](const TensorShapeProto_Dimension& dim) { + return dim.has_dim_value() && dim.dim_value() == 1; + }; + + const bool valid = ((bias_shape.dim_size() == 1 && bias_shape.dim()[0] == N) || + (bias_shape.dim_size() == 2 && dim_has_value_1(bias_shape.dim()[0]) && bias_shape.dim()[1] == N) || + (bias_shape.dim_size() == 2 && bias_shape.dim()[0] == M && + (dim_has_value_1(bias_shape.dim()[1]) || bias_shape.dim()[1] == N))); + if (!valid) { + return false; + } + + // If none of the above checks specify render this fusion invalid, fusion is valid. + return true; +} + +} // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/gemm_sum_fusion.h b/onnxruntime/core/optimizer/gemm_sum_fusion.h new file mode 100644 index 0000000000..5ed8db159b --- /dev/null +++ b/onnxruntime/core/optimizer/gemm_sum_fusion.h @@ -0,0 +1,52 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/optimizer/rewrite_rule.h" + +namespace onnxruntime { + +/** +@Class GemmSumFusion + +Rewrite rule that fuses Gemm and Sum nodes to a single Gemm node. +This fusion can be applied in the following scenario: +1) Sum at output of Gemm: when the output of a Gemm is immedietly summed with + exactly one other element, we can fuse this Sum with Gemm by using the other + Sum input as C, provided that the C input to the Gemm is missing. + This is supported for opset >= 11, as this is when Gemm input C became optional. + +TODO: Support the Add use case: Sum(x, y) ~= Add. + +This patterm is attempted to be triggered only on nodes with op type "Gemm". + +A --> Gemm --> D --> Sum --> E + ^ ^ + | | +B -----+ C + +is equivalent to + +A --> Gemm --> E + ^ ^ + | | +B ----+ C + +Where each letter represents a tensor. +*/ +class GemmSumFusion : public RewriteRule { + public: + GemmSumFusion() noexcept : RewriteRule("GemmSumFusion") {} + + std::vector TargetOpTypes() const noexcept override { + return {"Gemm"}; + } + + private: + bool SatisfyCondition(const Graph& graph, const Node& node, const logging::Logger& logger) const override; + + Status Apply(Graph& graph, Node& node, RewriteRuleEffect& rule_effect, const logging::Logger& logger) const override; +}; + +} // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index 84a9819d55..8c63b595c2 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/onnxruntime/core/optimizer/graph_transformer_utils.cc @@ -5,6 +5,7 @@ #include "core/mlas/inc/mlas.h" #include "core/optimizer/attention_fusion.h" +#include "core/optimizer/bias_dropout_fusion.h" #include "core/optimizer/bias_gelu_fusion.h" #include "core/optimizer/bias_softmax_fusion.h" #include "core/optimizer/cast_elimination.h" @@ -24,30 +25,31 @@ #include "core/optimizer/gelu_approximation.h" #include "core/optimizer/gelu_fusion.h" #include "core/optimizer/gemm_activation_fusion.h" +#include "core/optimizer/gemm_sum_fusion.h" #include "core/optimizer/gemm_transpose_fusion.h" #include "core/optimizer/identity_elimination.h" #include "core/optimizer/layer_norm_fusion.h" #include "core/optimizer/matmul_add_fusion.h" #include "core/optimizer/matmul_integer_to_float.h" #include "core/optimizer/matmul_scale_fusion.h" +#include "core/optimizer/matmul_transpose_fusion.h" #include "core/optimizer/nchwc_transformer.h" #include "core/optimizer/nhwc_transformer.h" #include "core/optimizer/noop_elimination.h" #include "core/optimizer/not_where_fusion.h" +#include "core/optimizer/qdq_transformer/qdq_propagation.h" +#include "core/optimizer/qdq_transformer/qdq_s8_to_u8.h" +#include "core/optimizer/qdq_transformer/relu_quantizelinear.h" +#include "core/optimizer/qdq_transformer/selectors_actions/qdq_selector_action_transformer.h" #include "core/optimizer/relu_clip_fusion.h" #include "core/optimizer/reshape_fusion.h" #include "core/optimizer/rule_based_graph_transformer.h" #include "core/optimizer/skip_layer_norm_fusion.h" #include "core/optimizer/slice_elimination.h" -#include "core/optimizer/unsqueeze_elimination.h" -#include "core/optimizer/qdq_transformer/qdq_propagation.h" -#include "core/optimizer/qdq_transformer/qdq_s8_to_u8.h" -#include "core/optimizer/qdq_transformer/relu_quantizelinear.h" -#include "core/optimizer/qdq_transformer/selectors_actions/qdq_selector_action_transformer.h" #include "core/optimizer/transpose_optimizer/ort_transpose_optimizer.h" +#include "core/optimizer/unsqueeze_elimination.h" #include "core/session/onnxruntime_session_options_config_keys.h" -#include "core/optimizer/matmul_transpose_fusion.h" -#include "core/optimizer/bias_dropout_fusion.h" + namespace onnxruntime { class IExecutionProvider; @@ -73,6 +75,7 @@ std::vector> GenerateRewriteRules( rules.push_back(std::make_unique()); rules.push_back(std::make_unique()); rules.push_back(std::make_unique()); + rules.push_back(std::make_unique()); rules.push_back(std::make_unique()); rules.push_back(std::make_unique()); rules.push_back(std::make_unique()); diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 2fda29c58f..0d63ba7040 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -35,6 +35,7 @@ #include "core/optimizer/fast_gelu_fusion.h" #include "core/optimizer/gelu_approximation.h" #include "core/optimizer/gelu_fusion.h" +#include "core/optimizer/gemm_sum_fusion.h" #include "core/optimizer/gemm_activation_fusion.h" #include "core/optimizer/gemm_transpose_fusion.h" #include "core/optimizer/graph_transformer.h" @@ -1367,6 +1368,331 @@ TEST_F(GraphTransformationTests, GemmTransposeFusionInputOutput) { ASSERT_TRUE(new_input_defs[1]->Name() == "A"); } +// Sum(Gemm(A, B, _), C) -> Gemm(A, B, C) +TEST_F(GraphTransformationTests, GemmSumFusionBasic) { + auto model_uri = MODEL_FOLDER "fusion/gemm_sum_basic.onnx"; + std::shared_ptr p_model; + ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_)); + Graph& graph = p_model->MainGraph(); + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_EQ(op_to_count["Sum"], 1); + ASSERT_EQ(op_to_count["Gemm"], 1); + ASSERT_EQ(graph.NumberOfNodes(), 2); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + auto rule_transformer_L1 = std::make_unique("RuleTransformer1"); + ASSERT_STATUS_OK(rule_transformer_L1->Register(std::make_unique())); + ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::move(rule_transformer_L1), TransformerLevel::Level1)); + ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_)); + + op_to_count = CountOpsInGraph(graph); + ASSERT_EQ(op_to_count["Sum"], 0); + ASSERT_EQ(op_to_count["Gemm"], 1); + ASSERT_EQ(graph.NumberOfNodes(), 1); + + auto& node = *graph.Nodes().begin(); + ASSERT_TRUE(node.OpType() == "Gemm"); + ASSERT_FALSE(static_cast(node.GetAttributes().at("transA").i())); + ASSERT_FALSE(static_cast(node.GetAttributes().at("transB").i())); + ASSERT_EQ(node.GetAttributes().at("alpha").f(), 1.0); + ASSERT_EQ(node.GetAttributes().at("beta").f(), 1.0); + auto new_input_defs = node.InputDefs(); + ASSERT_EQ(new_input_defs.size(), 3u); + ASSERT_TRUE(new_input_defs[0]->Name() == "A"); + ASSERT_TRUE(new_input_defs[1]->Name() == "B"); + ASSERT_TRUE(new_input_defs[2]->Name() == "C"); + auto new_output_defs = node.OutputDefs(); + ASSERT_EQ(new_output_defs.size(), 1u); + ASSERT_TRUE(new_output_defs[0]->Name() == "output"); +} + +// Sum(Gemm(A, B, _), C) -> Gemm(A, B, C), with attributes +TEST_F(GraphTransformationTests, GemmSumFusionAttributes) { + auto model_uri = MODEL_FOLDER "fusion/gemm_sum_attributes.onnx"; + std::shared_ptr p_model; + ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_)); + Graph& graph = p_model->MainGraph(); + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_EQ(op_to_count["Sum"], 1); + ASSERT_EQ(op_to_count["Gemm"], 1); + ASSERT_EQ(graph.NumberOfNodes(), 2); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + auto rule_transformer_L1 = std::make_unique("RuleTransformer1"); + ASSERT_STATUS_OK(rule_transformer_L1->Register(std::make_unique())); + ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::move(rule_transformer_L1), TransformerLevel::Level1)); + ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_)); + + op_to_count = CountOpsInGraph(graph); + ASSERT_EQ(op_to_count["Sum"], 0); + ASSERT_EQ(op_to_count["Gemm"], 1); + ASSERT_EQ(graph.NumberOfNodes(), 1); + + auto& node = *graph.Nodes().begin(); + ASSERT_TRUE(node.OpType() == "Gemm"); + ASSERT_FALSE(static_cast(node.GetAttributes().at("transA").i())); + ASSERT_TRUE(static_cast(node.GetAttributes().at("transB").i())); + ASSERT_EQ(node.GetAttributes().at("alpha").f(), 3.5); + ASSERT_EQ(node.GetAttributes().at("beta").f(), 1.0); + auto new_input_defs = node.InputDefs(); + ASSERT_EQ(new_input_defs.size(), 3u); + ASSERT_TRUE(new_input_defs[0]->Name() == "A"); + ASSERT_TRUE(new_input_defs[1]->Name() == "B"); + ASSERT_TRUE(new_input_defs[2]->Name() == "C"); + auto new_output_defs = node.OutputDefs(); + ASSERT_EQ(new_output_defs.size(), 1u); + ASSERT_TRUE(new_output_defs[0]->Name() == "output"); +} + +// Identity(Sum(Gemm(Identity(A), Identity(B), _), Identity(C)) should still fuse Gemm/Sum internally. +TEST_F(GraphTransformationTests, GemmSumFusionInternalNodes) { + auto model_uri = MODEL_FOLDER "fusion/gemm_sum_internal_nodes.onnx"; + std::shared_ptr p_model; + ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_)); + Graph& graph = p_model->MainGraph(); + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_EQ(op_to_count["Sum"], 1); + ASSERT_EQ(op_to_count["Gemm"], 1); + ASSERT_EQ(op_to_count["Identity"], 4); + ASSERT_EQ(graph.NumberOfNodes(), 6); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + auto rule_transformer_L1 = std::make_unique("RuleTransformer1"); + ASSERT_STATUS_OK(rule_transformer_L1->Register(std::make_unique())); + ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::move(rule_transformer_L1), TransformerLevel::Level1)); + ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_)); + + op_to_count = CountOpsInGraph(graph); + ASSERT_EQ(op_to_count["Sum"], 0); + ASSERT_EQ(op_to_count["Gemm"], 1); + ASSERT_EQ(op_to_count["Identity"], 4); + ASSERT_EQ(graph.NumberOfNodes(), 5); + + for(Node &node : graph.Nodes()) { + if(node.OpType() == "Gemm") { + ASSERT_FALSE(static_cast(node.GetAttributes().at("transA").i())); + ASSERT_FALSE(static_cast(node.GetAttributes().at("transB").i())); + ASSERT_EQ(node.GetAttributes().at("alpha").f(), 1.0); + ASSERT_EQ(node.GetAttributes().at("beta").f(), 1.0); + + auto new_input_defs = node.InputDefs(); + ASSERT_EQ(new_input_defs.size(), 3u); + ASSERT_TRUE(new_input_defs[0]->Name() == "tp0"); + ASSERT_TRUE(new_input_defs[1]->Name() == "tp1"); + ASSERT_TRUE(new_input_defs[2]->Name() == "tp3"); + auto new_output_defs = node.OutputDefs(); + ASSERT_EQ(new_output_defs.size(), 1u); + ASSERT_TRUE(new_output_defs[0]->Name() == "tp4"); + } + } +} + +// Sum(Gemm(A, B, C), D) does not perform transform. +TEST_F(GraphTransformationTests, GemmSumFusionNoFusionCUsed) { + auto model_uri = MODEL_FOLDER "fusion/gemm_sum_no_fusion_c_used.onnx"; + std::shared_ptr p_model; + ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_)); + Graph& graph = p_model->MainGraph(); + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_EQ(op_to_count["Sum"], 1); + ASSERT_EQ(op_to_count["Gemm"], 1); + ASSERT_EQ(graph.NumberOfNodes(), 2); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + auto rule_transformer_L1 = std::make_unique("RuleTransformer1"); + ASSERT_STATUS_OK(rule_transformer_L1->Register(std::make_unique())); + ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::move(rule_transformer_L1), TransformerLevel::Level1)); + ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_)); + + op_to_count = CountOpsInGraph(graph); + // Assert that the Sum still exists. Fusion should not occur with this pattern. + ASSERT_EQ(op_to_count["Sum"], 1); + ASSERT_EQ(op_to_count["Gemm"], 1); + ASSERT_EQ(graph.NumberOfNodes(), 2); + + for (Node& node : graph.Nodes()) { + if (node.OpType() == "Gemm") { + auto new_input_defs = node.InputDefs(); + ASSERT_EQ(new_input_defs.size(), 3u); + ASSERT_TRUE(new_input_defs[0]->Name() == "A"); + ASSERT_TRUE(new_input_defs[1]->Name() == "B"); + ASSERT_TRUE(new_input_defs[2]->Name() == "C"); + } else if (node.OpType() == "Sum") { + auto new_input_defs = node.InputDefs(); + ASSERT_EQ(new_input_defs.size(), 2u); + ASSERT_TRUE(new_input_defs[1]->Name() == "D"); + auto new_output_defs = node.OutputDefs(); + ASSERT_EQ(new_output_defs.size(), 1u); + ASSERT_TRUE(new_output_defs[0]->Name() == "output"); + } else { + FAIL(); + } + } +} + +// Sum(Gemm(A, B), C, D) does not perform transform. +TEST_F(GraphTransformationTests, GemmSumFusionNoFusionSumMultipleInputs) { + auto model_uri = MODEL_FOLDER "fusion/gemm_sum_no_fusion_sum_multiple_inputs.onnx"; + std::shared_ptr p_model; + ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_)); + Graph& graph = p_model->MainGraph(); + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_EQ(op_to_count["Sum"], 1); + ASSERT_EQ(op_to_count["Gemm"], 1); + ASSERT_EQ(graph.NumberOfNodes(), 2); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + auto rule_transformer_L1 = std::make_unique("RuleTransformer1"); + ASSERT_STATUS_OK(rule_transformer_L1->Register(std::make_unique())); + ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::move(rule_transformer_L1), TransformerLevel::Level1)); + ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_)); + + op_to_count = CountOpsInGraph(graph); + // Assert that the Sum still exists. Fusion should not occur with this pattern. + ASSERT_EQ(op_to_count["Sum"], 1); + ASSERT_EQ(op_to_count["Gemm"], 1); + ASSERT_EQ(graph.NumberOfNodes(), 2); + + for (Node& node : graph.Nodes()) { + if (node.OpType() == "Gemm") { + auto new_input_defs = node.InputDefs(); + ASSERT_EQ(new_input_defs.size(), 2u); + ASSERT_TRUE(new_input_defs[0]->Name() == "A"); + ASSERT_TRUE(new_input_defs[1]->Name() == "B"); + } else if (node.OpType() == "Sum") { + auto new_input_defs = node.InputDefs(); + ASSERT_EQ(new_input_defs.size(), 3u); + ASSERT_TRUE(new_input_defs[1]->Name() == "C"); + ASSERT_TRUE(new_input_defs[2]->Name() == "D"); + auto new_output_defs = node.OutputDefs(); + ASSERT_EQ(new_output_defs.size(), 1u); + ASSERT_TRUE(new_output_defs[0]->Name() == "output"); + } else { + FAIL(); + } + } +} + +// Sum(Gemm(A, B, _), C) -> Gemm(A, B, C), with broadcast. +TEST_F(GraphTransformationTests, GemmSumFusionBroadcast) { + auto model_uri = MODEL_FOLDER "fusion/gemm_sum_fusion_broadcast.onnx"; + std::shared_ptr p_model; + ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_)); + Graph& graph = p_model->MainGraph(); + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_EQ(op_to_count["Sum"], 1); + ASSERT_EQ(op_to_count["Gemm"], 1); + ASSERT_EQ(graph.NumberOfNodes(), 2); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + auto rule_transformer_L1 = std::make_unique("RuleTransformer1"); + ASSERT_STATUS_OK(rule_transformer_L1->Register(std::make_unique())); + ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::move(rule_transformer_L1), TransformerLevel::Level1)); + ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_)); + + op_to_count = CountOpsInGraph(graph); + ASSERT_EQ(op_to_count["Sum"], 0); + ASSERT_EQ(op_to_count["Gemm"], 1); + ASSERT_EQ(graph.NumberOfNodes(), 1); + + auto& node = *graph.Nodes().begin(); + ASSERT_TRUE(node.OpType() == "Gemm"); + ASSERT_FALSE(static_cast(node.GetAttributes().at("transA").i())); + ASSERT_FALSE(static_cast(node.GetAttributes().at("transB").i())); + ASSERT_EQ(node.GetAttributes().at("alpha").f(), 1.0); + ASSERT_EQ(node.GetAttributes().at("beta").f(), 1.0); + auto new_input_defs = node.InputDefs(); + ASSERT_EQ(new_input_defs.size(), 3u); + ASSERT_TRUE(new_input_defs[0]->Name() == "A"); + ASSERT_TRUE(new_input_defs[1]->Name() == "B"); + ASSERT_TRUE(new_input_defs[2]->Name() == "C"); + auto new_output_defs = node.OutputDefs(); + ASSERT_EQ(new_output_defs.size(), 1u); + ASSERT_TRUE(new_output_defs[0]->Name() == "output"); +} + +// Sum(Gemm(A, B, _), C), with invalid broadcasting (no fusion performed). +TEST_F(GraphTransformationTests, GemmSumFusionNoFusionBroadcastFailure) { + auto model_uri = MODEL_FOLDER "fusion/gemm_sum_no_fusion_broadcast_failure.onnx"; + std::shared_ptr p_model; + ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_)); + Graph& graph = p_model->MainGraph(); + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_EQ(op_to_count["Sum"], 1); + ASSERT_EQ(op_to_count["Gemm"], 1); + ASSERT_EQ(graph.NumberOfNodes(), 2); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + auto rule_transformer_L1 = std::make_unique("RuleTransformer1"); + ASSERT_STATUS_OK(rule_transformer_L1->Register(std::make_unique())); + ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::move(rule_transformer_L1), TransformerLevel::Level1)); + ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_)); + + op_to_count = CountOpsInGraph(graph); + ASSERT_EQ(op_to_count["Sum"], 1); + ASSERT_EQ(op_to_count["Gemm"], 1); + ASSERT_EQ(graph.NumberOfNodes(), 2); + + for (Node& node : graph.Nodes()) { + if (node.OpType() == "Gemm") { + auto new_input_defs = node.InputDefs(); + ASSERT_EQ(new_input_defs.size(), 2u); + ASSERT_TRUE(new_input_defs[0]->Name() == "A"); + ASSERT_TRUE(new_input_defs[1]->Name() == "B"); + } else if (node.OpType() == "Sum") { + auto new_input_defs = node.InputDefs(); + ASSERT_EQ(new_input_defs.size(), 2u); + ASSERT_TRUE(new_input_defs[1]->Name() == "C"); + auto new_output_defs = node.OutputDefs(); + ASSERT_EQ(new_output_defs.size(), 1u); + ASSERT_TRUE(new_output_defs[0]->Name() == "output"); + } else { + FAIL(); + } + } +} + +// Sum(Gemm(A, B, _), C) where intermediate Gemm output is used, so fusion cannot be performed. +TEST_F(GraphTransformationTests, GemmSumFusionNoFusionOriginalGemmOutputUsed) { + auto model_uri = MODEL_FOLDER "fusion/gemm_sum_no_fusion_original_gemm_output_used.onnx"; + std::shared_ptr p_model; + ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_)); + Graph& graph = p_model->MainGraph(); + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_EQ(op_to_count["Sum"], 1); + ASSERT_EQ(op_to_count["Gemm"], 1); + ASSERT_EQ(graph.NumberOfNodes(), 2); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + auto rule_transformer_L1 = std::make_unique("RuleTransformer1"); + ASSERT_STATUS_OK(rule_transformer_L1->Register(std::make_unique())); + ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::move(rule_transformer_L1), TransformerLevel::Level1)); + ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_)); + + op_to_count = CountOpsInGraph(graph); + ASSERT_EQ(op_to_count["Sum"], 1); + ASSERT_EQ(op_to_count["Gemm"], 1); + ASSERT_EQ(graph.NumberOfNodes(), 2); + + for (Node& node : graph.Nodes()) { + if (node.OpType() == "Gemm") { + auto new_input_defs = node.InputDefs(); + ASSERT_EQ(new_input_defs.size(), 2u); + ASSERT_TRUE(new_input_defs[0]->Name() == "A"); + ASSERT_TRUE(new_input_defs[1]->Name() == "B"); + } else if (node.OpType() == "Sum") { + auto new_input_defs = node.InputDefs(); + ASSERT_EQ(new_input_defs.size(), 2u); + ASSERT_TRUE(new_input_defs[1]->Name() == "C"); + auto new_output_defs = node.OutputDefs(); + ASSERT_EQ(new_output_defs.size(), 1u); + ASSERT_TRUE(new_output_defs[0]->Name() == "output"); + } else { + FAIL(); + } + } +} + TEST_F(GraphTransformationTests, FuseConvBnAddMulFloat16) { auto model_uri = MODEL_FOLDER "fusion/fuse-conv-bn-add-mul-float16.onnx"; diff --git a/onnxruntime/test/testdata/transform/fusion/gemm_sum_attributes.onnx b/onnxruntime/test/testdata/transform/fusion/gemm_sum_attributes.onnx new file mode 100644 index 0000000000000000000000000000000000000000..6ee43562c9bd54e5291b664ae898ee40dc965f56 GIT binary patch literal 263 zcmdW+oTq7w4yy2rvQwa3m$_ literal 0 HcmV?d00001 diff --git a/onnxruntime/test/testdata/transform/fusion/gemm_sum_fusion_broadcast.onnx b/onnxruntime/test/testdata/transform/fusion/gemm_sum_fusion_broadcast.onnx new file mode 100644 index 0000000000000000000000000000000000000000..e510f19f1ac6a68c0780c9a53434d3f546afc89d GIT binary patch literal 199 zcmYj}u?~VT5QaTVOg%-*MHWYPjkB9|a3Hei#3;dpp_ITi@lAYPTjIcSx%>Y6KPgGs zw*B<3<;K(pt_MDU0TYA^f+oS;DT{8kF_>aF5U#YajyuPhcsN6=5c=imn`P>2q;MmF z5Cy?Kj67gH#d*L*xeKx_#{PSWa|t7T?eXuQHbzmQBb=$VPB5;lIhv;G?O@-m;0Hf{ BC5-?8 literal 0 HcmV?d00001 diff --git a/onnxruntime/test/testdata/transform/fusion/gemm_sum_gen.py b/onnxruntime/test/testdata/transform/fusion/gemm_sum_gen.py new file mode 100644 index 0000000000..7ae435e99a --- /dev/null +++ b/onnxruntime/test/testdata/transform/fusion/gemm_sum_gen.py @@ -0,0 +1,187 @@ +import onnx +from onnx import helper +from onnx import TensorProto +from onnx import OperatorSetIdProto + +onnxdomain = OperatorSetIdProto() +onnxdomain.version = 12 +# The empty string ("") or absence of this field implies the operator set that is defined as part of the ONNX specification. +onnxdomain.domain = "" +msdomain = OperatorSetIdProto() +msdomain.version = 1 +msdomain.domain = "com.microsoft" +opsets = [onnxdomain, msdomain] + + +def save(model_path, nodes, inputs, outputs, initializers): + graph = helper.make_graph( + nodes, + "GemmSumTest", + inputs, outputs, initializers) + + model = helper.make_model( + graph, opset_imports=opsets, producer_name="onnxruntime-test") + + print(model_path) + onnx.save(model, model_path) + +def gen_gemm_sum_basic(model_path): + nodes = [ + helper.make_node(op_type="Gemm", inputs=["A", "B"], outputs=["tp0"]), + helper.make_node(op_type="Sum", inputs=["tp0", "C"], outputs=["output"]), + ] + + inputs = [ + helper.make_tensor_value_info("A", TensorProto.FLOAT, ['M', 'K']), + helper.make_tensor_value_info("B", TensorProto.FLOAT, ['K', 'N']), + helper.make_tensor_value_info("C", TensorProto.FLOAT, ['M', 'N']) + ] + + outputs = [ + helper.make_tensor_value_info("output", TensorProto.FLOAT, ['M', 'N']) + ] + + save(model_path, nodes, inputs, outputs, initializers=[]) + +def gen_gemm_sum_attributes(model_path): + nodes = [ + helper.make_node(op_type="Gemm", inputs=["A", "B"], outputs=["tp0"], alpha=3.5, beta=6.25, transA=False, transB=True), + helper.make_node(op_type="Sum", inputs=["tp0", "C"], outputs=["output"]), + ] + + inputs = [ + helper.make_tensor_value_info("A", TensorProto.FLOAT, ['M', 'K']), + helper.make_tensor_value_info("B", TensorProto.FLOAT, ['N', 'K']), + helper.make_tensor_value_info("C", TensorProto.FLOAT, ['M', 'N']) + ] + + outputs = [ + helper.make_tensor_value_info("output", TensorProto.FLOAT, ['M', 'N']) + ] + + save(model_path, nodes, inputs, outputs, initializers=[]) + +def gen_gemm_sum_internal_nodes(model_path): + nodes = [ + helper.make_node(op_type="Identity", inputs=["A"], outputs=["tp0"]), + helper.make_node(op_type="Identity", inputs=["B"], outputs=["tp1"]), + helper.make_node(op_type="Gemm", inputs=["tp0", "tp1"], outputs=["tp2"]), + helper.make_node(op_type="Identity", inputs=["C"], outputs=["tp3"]), + helper.make_node(op_type="Sum", inputs=["tp2", "tp3"], outputs=["tp4"]), + helper.make_node(op_type="Identity", inputs=["tp4"], outputs=["output"]), + ] + + inputs = [ + helper.make_tensor_value_info("A", TensorProto.FLOAT, ['M', 'K']), + helper.make_tensor_value_info("B", TensorProto.FLOAT, ['K', 'N']), + helper.make_tensor_value_info("C", TensorProto.FLOAT, ['M', 'N']) + ] + + outputs = [ + helper.make_tensor_value_info("output", TensorProto.FLOAT, ['M', 'N']) + ] + + save(model_path, nodes, inputs, outputs, initializers=[]) + +def gen_gemm_sum_no_fusion_c_used(model_path): + nodes = [ + helper.make_node(op_type="Gemm", inputs=["A", "B", "C"], outputs=["tp0"]), + helper.make_node(op_type="Sum", inputs=["tp0", "D"], outputs=["output"]), + ] + + inputs = [ + helper.make_tensor_value_info("A", TensorProto.FLOAT, ['M', 'K']), + helper.make_tensor_value_info("B", TensorProto.FLOAT, ['K', 'N']), + helper.make_tensor_value_info("C", TensorProto.FLOAT, ['M', 'N']), + helper.make_tensor_value_info("D", TensorProto.FLOAT, ['M', 'N']), + ] + + outputs = [ + helper.make_tensor_value_info("output", TensorProto.FLOAT, ['M', 'N']) + ] + + save(model_path, nodes, inputs, outputs, initializers=[]) + +def gen_gemm_sum_no_fusion_sum_multiple_inputs(model_path): + nodes = [ + helper.make_node(op_type="Gemm", inputs=["A", "B"], outputs=["tp0"]), + helper.make_node(op_type="Sum", inputs=["tp0", "C", "D"], outputs=["output"]), + ] + + inputs = [ + helper.make_tensor_value_info("A", TensorProto.FLOAT, ['M', 'K']), + helper.make_tensor_value_info("B", TensorProto.FLOAT, ['K', 'N']), + helper.make_tensor_value_info("C", TensorProto.FLOAT, ['M', 'N']), + helper.make_tensor_value_info("D", TensorProto.FLOAT, ['M', 'N']), + ] + + outputs = [ + helper.make_tensor_value_info("output", TensorProto.FLOAT, ['M', 'N']) + ] + + save(model_path, nodes, inputs, outputs, initializers=[]) + +def gen_gemm_sum_fusion_broadcast(model_path): + nodes = [ + helper.make_node(op_type="Gemm", inputs=["A", "B"], outputs=["tp0"]), + helper.make_node(op_type="Sum", inputs=["tp0", "C"], outputs=["output"]), + ] + + inputs = [ + helper.make_tensor_value_info("A", TensorProto.FLOAT, ['M', 'K']), + helper.make_tensor_value_info("B", TensorProto.FLOAT, ['K', 'N']), + helper.make_tensor_value_info("C", TensorProto.FLOAT, [1, 'N']), + ] + + outputs = [ + helper.make_tensor_value_info("output", TensorProto.FLOAT, ['M', 'N']) + ] + + save(model_path, nodes, inputs, outputs, initializers=[]) + +def gen_gemm_sum_no_fusion_broadcast_failure(model_path): + nodes = [ + helper.make_node(op_type="Gemm", inputs=["A", "B"], outputs=["tp0"]), + helper.make_node(op_type="Sum", inputs=["tp0", "C"], outputs=["output"]), + ] + + inputs = [ + helper.make_tensor_value_info("A", TensorProto.FLOAT, ['M', 'K']), + helper.make_tensor_value_info("B", TensorProto.FLOAT, ['K', 'N']), + # should work with multidirectional broadcast as second argument, but not unidirectional. + helper.make_tensor_value_info("C", TensorProto.FLOAT, [1, 'M', 'N']), + ] + + outputs = [ + helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, 'M', 'N']) + ] + + save(model_path, nodes, inputs, outputs, initializers=[]) + +def gen_gemm_sum_no_fusion_original_gemm_output_used(model_path): + nodes = [ + helper.make_node(op_type="Gemm", inputs=["A", "B"], outputs=["tp0"]), + helper.make_node(op_type="Sum", inputs=["tp0", "C"], outputs=["output"]), + ] + + inputs = [ + helper.make_tensor_value_info("A", TensorProto.FLOAT, ['M', 'K']), + helper.make_tensor_value_info("B", TensorProto.FLOAT, ['K', 'N']), + helper.make_tensor_value_info("C", TensorProto.FLOAT, ['M', 'N']), + ] + + outputs = [ + helper.make_tensor_value_info("tp0", TensorProto.FLOAT, ['M', 'N']), + helper.make_tensor_value_info("output", TensorProto.FLOAT, ['M', 'N']) + ] + + save(model_path, nodes, inputs, outputs, initializers=[]) + +gen_gemm_sum_basic("gemm_sum_basic.onnx") +gen_gemm_sum_attributes("gemm_sum_attributes.onnx") +gen_gemm_sum_internal_nodes("gemm_sum_internal_nodes.onnx") +gen_gemm_sum_no_fusion_c_used("gemm_sum_no_fusion_c_used.onnx") +gen_gemm_sum_no_fusion_sum_multiple_inputs("gemm_sum_no_fusion_sum_multiple_inputs.onnx") +gen_gemm_sum_fusion_broadcast("gemm_sum_fusion_broadcast.onnx") +gen_gemm_sum_no_fusion_broadcast_failure("gemm_sum_no_fusion_broadcast_failure.onnx") +gen_gemm_sum_no_fusion_original_gemm_output_used("gemm_sum_no_fusion_original_gemm_output_used.onnx") \ No newline at end of file diff --git a/onnxruntime/test/testdata/transform/fusion/gemm_sum_internal_nodes.onnx b/onnxruntime/test/testdata/transform/fusion/gemm_sum_internal_nodes.onnx new file mode 100644 index 0000000000000000000000000000000000000000..d3cb7d377daeac368605fa08d51c3b112e82600c GIT binary patch literal 288 zcmZXOy9&ZU5Jg=QWEmgC#UNOwSIncbGbu!ZX|%F5Mo6*QkZeHz+D|Z71Q%>B%(;h| z;aoVXsvb>O=~5OeU9@`pVkC&|gr)0^#}Bun@AMNa2J;Q(5+sC=D)8)mAtlrVb7-h@ zSOFqL6?xXFle7Sc#87p*?(|O=BP{BTbN|uB#XAm7ITJ*dU`Ip>9Gf@=&VUmWf9&kK br0D>CCUZj za51_FvE`SR6qJ@IF$b6C3UPxZfs7EK<|t91P9Xs!bwtGPcj49{D!o?uK<0Qz%o1CAkmz$YflwX{mRwBR%0Pra&*8l(j literal 0 HcmV?d00001 diff --git a/onnxruntime/test/testdata/transform/fusion/gemm_sum_no_fusion_sum_multiple_inputs.onnx b/onnxruntime/test/testdata/transform/fusion/gemm_sum_no_fusion_sum_multiple_inputs.onnx new file mode 100644 index 0000000000000000000000000000000000000000..faef5e4d8f17fb10e18c4558dba7e4fcc41766dc GIT binary patch literal 226 zcmd!bwtGP> GeneratePreTrainingTransformers( ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique())); ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique())); ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique())); + ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique())); ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique())); ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique())); ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique()));