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

@@ -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<bool>(gemm_node.GetAttributes().at("transA").i());
+  const bool transB = static_cast<bool>(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<NodeArg*> 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<NodeArg*> 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<int64_t>(transA));
+  new_gemm_node.AddAttribute("transB", static_cast<int64_t>(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

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<std::string> 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

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<std::unique_ptr<RewriteRule>> GenerateRewriteRules(
       rules.push_back(std::make_unique<NoopElimination>());
       rules.push_back(std::make_unique<DivMulFusion>());
       rules.push_back(std::make_unique<FuseReluClip>());
+      rules.push_back(std::make_unique<GemmSumFusion>());
       rules.push_back(std::make_unique<GemmTransposeFusion>());
       rules.push_back(std::make_unique<NotWhereFusion>());
       rules.push_back(std::make_unique<ConvAddFusion>());

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<Model> p_model;
+  ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_));
+  Graph& graph = p_model->MainGraph();
+  std::map<std::string, int> 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<RuleBasedGraphTransformer>("RuleTransformer1");
+  ASSERT_STATUS_OK(rule_transformer_L1->Register(std::make_unique<GemmSumFusion>()));
+  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<bool>(node.GetAttributes().at("transA").i()));
+  ASSERT_FALSE(static_cast<bool>(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<Model> p_model;
+  ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_));
+  Graph& graph = p_model->MainGraph();
+  std::map<std::string, int> 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<RuleBasedGraphTransformer>("RuleTransformer1");
+  ASSERT_STATUS_OK(rule_transformer_L1->Register(std::make_unique<GemmSumFusion>()));
+  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<bool>(node.GetAttributes().at("transA").i()));
+  ASSERT_TRUE(static_cast<bool>(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<Model> p_model;
+  ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_));
+  Graph& graph = p_model->MainGraph();
+  std::map<std::string, int> 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<RuleBasedGraphTransformer>("RuleTransformer1");
+  ASSERT_STATUS_OK(rule_transformer_L1->Register(std::make_unique<GemmSumFusion>()));
+  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<bool>(node.GetAttributes().at("transA").i()));
+      ASSERT_FALSE(static_cast<bool>(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<Model> p_model;
+  ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_));
+  Graph& graph = p_model->MainGraph();
+  std::map<std::string, int> 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<RuleBasedGraphTransformer>("RuleTransformer1");
+  ASSERT_STATUS_OK(rule_transformer_L1->Register(std::make_unique<GemmSumFusion>()));
+  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<Model> p_model;
+  ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_));
+  Graph& graph = p_model->MainGraph();
+  std::map<std::string, int> 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<RuleBasedGraphTransformer>("RuleTransformer1");
+  ASSERT_STATUS_OK(rule_transformer_L1->Register(std::make_unique<GemmSumFusion>()));
+  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<Model> p_model;
+  ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_));
+  Graph& graph = p_model->MainGraph();
+  std::map<std::string, int> 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<RuleBasedGraphTransformer>("RuleTransformer1");
+  ASSERT_STATUS_OK(rule_transformer_L1->Register(std::make_unique<GemmSumFusion>()));
+  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<bool>(node.GetAttributes().at("transA").i()));
+  ASSERT_FALSE(static_cast<bool>(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<Model> p_model;
+  ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_));
+  Graph& graph = p_model->MainGraph();
+  std::map<std::string, int> 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<RuleBasedGraphTransformer>("RuleTransformer1");
+  ASSERT_STATUS_OK(rule_transformer_L1->Register(std::make_unique<GemmSumFusion>()));
+  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<Model> p_model;
+  ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_));
+  Graph& graph = p_model->MainGraph();
+  std::map<std::string, int> 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<RuleBasedGraphTransformer>("RuleTransformer1");
+  ASSERT_STATUS_OK(rule_transformer_L1->Register(std::make_unique<GemmSumFusion>()));
+  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";
 

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")

orttraining/orttraining/core/optimizer/graph_transformer_utils.cc

@@ -5,6 +5,7 @@
 #include "orttraining/core/optimizer/graph_transformer_utils.h"
 
 #include "core/mlas/inc/mlas.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"
@@ -25,9 +26,11 @@
 #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/graph_transformer_utils.h"
 #include "core/optimizer/identity_elimination.h"
+#include "core/optimizer/isinf_reducesum_fusion.h"
 #include "core/optimizer/layer_norm_fusion.h"
 #include "core/optimizer/matmul_add_fusion.h"
 #include "core/optimizer/matmul_scale_fusion.h"
@@ -35,23 +38,21 @@
 #include "core/optimizer/nchwc_transformer.h"
 #include "core/optimizer/noop_elimination.h"
 #include "core/optimizer/not_where_fusion.h"
+#include "core/optimizer/propagate_cast_ops.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/isinf_reducesum_fusion.h"
-#include "core/optimizer/propagate_cast_ops.h"
 #include "core/session/inference_session.h"
 #include "orttraining/core/framework/distributed_run_context.h"
-#include "core/optimizer/bias_dropout_fusion.h"
-#include "orttraining/core/optimizer/concat_replacement.h"
 #include "orttraining/core/optimizer/batchnorm_replacement.h"
+#include "orttraining/core/optimizer/concat_replacement.h"
 #include "orttraining/core/optimizer/insert_output_rewriter.h"
 #include "orttraining/core/optimizer/localized_recompute.h"
-#include "orttraining/core/optimizer/transformer_layer_recompute.h"
 #include "orttraining/core/optimizer/loss_rewriter.h"
+#include "orttraining/core/optimizer/transformer_layer_recompute.h"
 
 namespace onnxruntime {
 namespace training {
@@ -83,6 +84,7 @@ std::vector<std::unique_ptr<GraphTransformer>> GeneratePreTrainingTransformers(
       ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique<NoopElimination>()));
       ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique<DivMulFusion>()));
       ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique<EliminateDropout>()));
+      ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique<GemmSumFusion>()));
       ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique<GemmTransposeFusion>()));
       ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique<NotWhereFusion>()));
       ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique<InsertSoftmaxCrossEntropyLossOutput>()));