Optimize GatherND (#4097)

* Optimize GatherND
* Refine the code, Fix few comments

onnxruntime/core/optimizer/computation_reduction.cc

@@ -0,0 +1,304 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/common/safeint.h"
+#include "core/graph/graph_utils.h"
+#include "core/optimizer/initializer.h"
+#include "core/optimizer/utils.h"
+#include "core/optimizer/computation_reduction.h"
+
+using namespace ONNX_NAMESPACE;
+using namespace ::onnxruntime::common;
+namespace onnxruntime {
+typedef std::function<Status(Graph&, Node&, Node&)> Handler;
+
+constexpr int GATHERND_BATCH_DIM = 1;
+
+static bool IsLeadingDimsEqual(const TensorShapeProto* input_shape, const TensorShapeProto* output_shape,
+                               const int num_dim_to_check) {
+  ORT_ENFORCE(output_shape->dim_size() >= num_dim_to_check && input_shape->dim_size() >= num_dim_to_check);
+
+  for (int i = 0; i < num_dim_to_check; ++i) {
+    auto& output_dim = output_shape->dim(i);
+    auto& input_dim = input_shape->dim(i);
+    if (output_dim.has_dim_value() && input_dim.has_dim_value()) {
+      if (output_dim.dim_value() != input_dim.dim_value()) {
+        return false;
+      }
+    } else if (output_dim.has_dim_param() && input_dim.has_dim_param()) {
+      if (output_dim.dim_param() != input_dim.dim_param()) {
+        return false;
+      }
+    } else {
+      return false;
+    }
+  }
+
+  return true;
+}
+
+static int GetValidInputForGatherND(const Node& target_node) {
+  // target_node is the producer of GatherND's input.
+  // If target_node's some input tensors have exactly same shape with
+  // target_node output tensor shape, then it is safe to gather using
+  // original slice ranges.
+  int candidate_input_index = -1;
+  auto output_shape = target_node.OutputDefs()[0]->Shape();
+  const int output_rank = output_shape->dim_size();
+  for (size_t i = 0; i < target_node.InputDefs().size(); ++i) {
+    auto input_shape = target_node.InputDefs()[i]->Shape();
+    const int input_rank = input_shape->dim_size();
+    if (input_rank != output_rank) {
+      continue;
+    }
+
+    if (IsLeadingDimsEqual(input_shape, output_shape, GATHERND_BATCH_DIM + 1)) {
+      candidate_input_index = SafeInt<int>(i);
+      break;
+    }
+  }
+
+  return candidate_input_index;
+}
+
+static TensorShapeProto ReplaceSymbolicDimValue(const TensorShapeProto* shape, const int replacement_axis,
+                                                const std::string& replacement_dim_value) {
+  ORT_ENFORCE(replacement_axis >= 0 && replacement_axis < shape->dim_size());
+  TensorShapeProto output_shape;
+  for (int i = 0; i < shape->dim_size(); ++i) {
+    auto& dim = shape->dim(i);
+    if (i == replacement_axis) {
+      output_shape.add_dim()->set_dim_param(replacement_dim_value);
+      continue;
+    }
+
+    if (dim.has_dim_value()) {
+      output_shape.add_dim()->set_dim_value(dim.dim_value());
+    } else if (dim.has_dim_param()) {
+      output_shape.add_dim()->set_dim_param(dim.dim_param());
+    } else {
+      ORT_THROW("Invalid dim found in ReplaceSymbolicDimValue");
+    }
+  }
+
+  return output_shape;
+}
+
+static Status SwapGatherNDWithTargetNode(Graph& graph, Node& gathernd_node, Node& target_node,
+                                         const int target_node_input_index = 0) {
+  auto new_input_arg_for_gathernd = target_node.MutableInputDefs()[target_node_input_index];
+  auto target_node_out_arg = target_node.MutableOutputDefs()[0];
+  auto gathernd_out_arg = gathernd_node.MutableOutputDefs()[0];
+  auto gathernd_old_consumers = graph.GetConsumerNodes(gathernd_out_arg->Name());
+  const auto& graph_outputs = graph.GetOutputs();
+  bool need_update_graph_output = false;
+  if (std::find(graph_outputs.begin(), graph_outputs.end(), gathernd_out_arg) != graph_outputs.end()) {
+    need_update_graph_output = true;
+  }
+
+  const std::string& gathered_dim_param = gathernd_out_arg->Shape()->dim(GATHERND_BATCH_DIM).dim_param();
+  TensorShapeProto new_output_shape_for_gathernd =
+      ReplaceSymbolicDimValue(new_input_arg_for_gathernd->Shape(), GATHERND_BATCH_DIM, gathered_dim_param);
+
+  TensorShapeProto new_output_shape_for_target_node =
+      ReplaceSymbolicDimValue(target_node_out_arg->Shape(), GATHERND_BATCH_DIM, gathered_dim_param);
+
+  // update input/output definitions.
+  int output_index = optimizer_utils::IndexOfNodeOutput(target_node, *gathernd_node.MutableInputDefs()[0]);
+  graph.RemoveEdge(target_node.Index(), gathernd_node.Index(), output_index, 0);
+  const Node* target_node_input_node = graph.GetProducerNode(new_input_arg_for_gathernd->Name());
+  if (target_node_input_node != nullptr) {
+    output_index = optimizer_utils::IndexOfNodeOutput(*target_node_input_node, *new_input_arg_for_gathernd);
+    graph.AddEdge(target_node_input_node->Index(), gathernd_node.Index(), output_index, 0);
+  } else {
+    // new_input_arg_for_gathernd is graph input
+    graph_utils::ReplaceNodeInput(gathernd_node, 0, *new_input_arg_for_gathernd);
+  }
+
+  graph_utils::ReplaceDownstreamNodeInput(graph, gathernd_node, 0 /*output_idx*/,
+                                          target_node, 0 /*replacement_output_idx*/);
+
+  if (target_node_input_node != nullptr) {
+    graph.RemoveEdge(target_node_input_node->Index(), target_node.Index(), output_index, target_node_input_index);
+  }
+  graph.AddEdge(gathernd_node.Index(), target_node.Index(), 0, target_node_input_index);
+
+  // update consumer relation ship
+  if (!gathernd_old_consumers.empty()) {
+    graph.UpdateConsumerNodes(target_node_out_arg->Name(), {const_cast<Node*>(gathernd_old_consumers[0])});
+  }
+  graph.UpdateConsumerNodes(gathernd_out_arg->Name(), {&target_node});
+
+  // update shapes
+  gathernd_out_arg->SetShape(new_output_shape_for_gathernd);
+  target_node_out_arg->SetShape(new_output_shape_for_target_node);
+
+  if (need_update_graph_output) {
+    std::vector<const NodeArg*> graph_new_outputs;
+    for (auto out_arg : graph_outputs) {
+      if (out_arg->Name().compare(gathernd_out_arg->Name()) == 0) {
+        graph_new_outputs.push_back(target_node_out_arg);
+      } else {
+        graph_new_outputs.push_back(out_arg);
+      }
+    }
+    graph.SetOutputs(graph_new_outputs);
+    graph.SetGraphResolveNeeded();
+    graph.SetGraphProtoSyncNeeded();
+  }
+
+  return Status::OK();
+}
+
+static Status SimpleHandler(Graph& graph, Node& gathernd_node, Node& target_node) {
+  return SwapGatherNDWithTargetNode(graph, gathernd_node, target_node, 0);
+}
+
+/*
+  This handler change the graphs this way:
+    Before:
+      input_1[b,s,h]   weight_2[h]
+                  \         /
+                    Add[b,s,h]    indices[b,p_s,1]
+                      |            /
+                  GatherND[b,p_s,h]
+                      |
+                <subsquent graphs>
+    After :
+               input_1[b,s,h]      indices[b,p_s,1]
+                       |            /
+                    GatherND[b,p_s,h]    weight_2[h]
+                              \              /
+                              Add[b,p_s,h]
+                                    |
+                            <subsquent graphs>
+  Note: b: batch, s: sequence_length, h: hidden_size, p_s: dynamic_prediction_count
+*/
+static Status BinaryElementwiseHandler(Graph& graph, Node& gathernd_node, Node& target_node) {
+  int target_node_input_index = GetValidInputForGatherND(target_node);
+  ORT_RETURN_IF_NOT(target_node_input_index != -1);
+  return SwapGatherNDWithTargetNode(graph, gathernd_node, target_node, target_node_input_index);
+}
+
+/*
+  This handler change the graphs this way:
+    Before:
+      input_1[b,s,h]    weight_2[h, 2h]
+                  \         /
+                  MatMul[b,s,2h]    indices[b,p_s,1]
+                      |            /
+                  GatherND[b,p_s,2h]
+                      |
+                <subsquent graphs>
+    After :
+               input_1[b,s,h]      indices[b,p_s,1]
+                       |            /
+                    GatherND[b,p_s,h]    weight_2[h,2h]
+                              \              /
+                              MatMul[b,p_s,2h]
+                                    |
+                            <subsquent graphs>
+  Note: b: batch, s: sequence_length, h: hidden_size, p_s: dynamic_prediction_count
+*/
+static Status MatMulHandler(Graph& graph, Node& gathernd_node, Node& target_node) {
+  int target_node_input_index = GetValidInputForGatherND(target_node);
+  ORT_RETURN_IF_NOT(target_node_input_index == 0);
+  return SwapGatherNDWithTargetNode(graph, gathernd_node, target_node, target_node_input_index);
+}
+
+static std::unordered_map<std::string, Handler> handlers = {
+    {"Add", BinaryElementwiseHandler},
+    {"Div", BinaryElementwiseHandler},
+    {"Gelu", SimpleHandler},
+    {"LayerNormalization", SimpleHandler},
+    {"MatMul", MatMulHandler}};
+
+static Status Delegate(Graph& graph, Node& gathernd_node, Node& target_node) {
+  const std::string& op_type = target_node.OpType();
+  if (handlers.count(op_type)) {
+    return handlers[op_type](graph, gathernd_node, target_node);
+  } else {
+    return common::Status(common::ONNXRUNTIME, common::NOT_IMPLEMENTED, op_type + " handler is not implemented");
+  }
+}
+
+Status ComputationReductionTransformer::ApplyImpl(Graph& graph, bool& modified, int graph_level,
+                                                  const logging::Logger& logger) const {
+  GraphViewer graph_viewer(graph);
+  const auto& order = graph_viewer.GetNodesInTopologicalOrder();
+
+  for (auto index : order) {
+    auto* node_ptr = graph.GetNode(index);
+    if (!node_ptr)
+      // node was removed, this should not happen since we are not removing nodes.
+      continue;
+
+    auto& node = *node_ptr;
+    ORT_RETURN_IF_ERROR(Recurse(node, modified, graph_level, logger));
+
+    // Same ideas might apply for Gather, GatherElements, Slice, Split, etc.
+    // Todo: let's review the real cases to make the logic more generic.
+    if (!graph_utils::IsSupportedOptypeVersionAndDomain(node, "GatherND", {1, 12, 13}, kOnnxDomain) ||
+        !graph_utils::IsSupportedProvider(node, GetCompatibleExecutionProviders()) ||
+        node.GetOutputEdgesCount() > 1) {  // allow GatherND have no out edges in case it is graph output.
+      continue;
+    }
+
+    auto batch_dims = static_cast<int64_t>(node.GetAttributes().at("batch_dims").i());
+    if (batch_dims != GATHERND_BATCH_DIM) {
+      continue;
+    }
+
+    auto indices_shape = node.MutableInputDefs()[1]->Shape();
+    if (indices_shape == nullptr) {
+      continue;
+    }
+
+    const auto indices_rank = indices_shape->dim_size();
+    auto& indices_last_dim = indices_shape->dim(indices_rank - 1);
+    // Since GatherND is assumed to have batch_dims=1, if the input data's shape is [batch, sequence, ..., ... ],
+    // limiting indices_rank=3 will make sure produced output is in shape [batch, sliced_sequence, ..., ...]
+    // and the rank did not change.
+    if (!(indices_last_dim.has_dim_value() && indices_last_dim.dim_value() == 1 && indices_rank == 3)) {
+      continue;
+    }
+
+    // Todo: check whether we want to move GatherND up, for example, if GatherND's outputs are larger
+    // than inputs, we should NOT probably bring it ahead.
+    bool stop = false;
+    while (!stop) {
+      const Node* gathernd_data_producer = graph.GetProducerNode(node.MutableInputDefs()[0]->Name());
+      if (gathernd_data_producer == nullptr) {
+        break;
+      }
+      Node* input_node = const_cast<Node*>(gathernd_data_producer);
+      if (graph.GetConsumerNodes(input_node->MutableOutputDefs()[0]->Name()).size() > 1) {
+        LOGS_DEFAULT(WARNING) << "node " << node.Name() << " stopped at node "
+                              << input_node->Name();
+        break;
+      }
+
+      auto ret = Delegate(graph, node, *input_node);
+      if (ret.IsOK()) {
+        LOGS_DEFAULT(WARNING) << "node " << node.Name() << " up across node "
+                              << input_node->Name() << std::endl;
+        modified = true;
+      } else if (ret.Code() == common::NOT_IMPLEMENTED) {
+        LOGS_DEFAULT(WARNING) << "node " << node.Name() << " stopped at node "
+                              << input_node->Name();
+        break;
+      } else {
+        LOGS_DEFAULT(WARNING) << " terminate due to unexpected error, node names:" << node.Name()
+                              << ", " << input_node->Name() << ", error " << ret.ErrorMessage() << std::endl;
+        stop = true;
+      }
+    }
+  }
+
+  if (modified) {
+    graph.SetGraphResolveNeeded();
+  }
+  return Status::OK();
+}
+
+}  // namespace onnxruntime

onnxruntime/core/optimizer/computation_reduction.h

@@ -0,0 +1,18 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "core/optimizer/graph_transformer.h"
+
+namespace onnxruntime {
+
+class ComputationReductionTransformer : public GraphTransformer {
+ public:
+  ComputationReductionTransformer(const std::unordered_set<std::string>& compatible_execution_providers = {}) noexcept
+      : GraphTransformer("ComputationReductionTransformer", compatible_execution_providers) {}
+
+  Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override;
+};
+
+}  // namespace onnxruntime

onnxruntime/core/optimizer/utils.cc

@@ -217,6 +217,18 @@ int32_t IndexOfNodeInput(const Node& node, const NodeArg& node_arg) {
   return -1;
 }
 
+int32_t IndexOfNodeOutput(const Node& node, const NodeArg& node_arg) {
+  int32_t index = 0;
+  for (auto& output_arg : node.OutputDefs()) {
+    if (output_arg->Name().compare(node_arg.Name()) == 0) {
+      return index;
+    }
+    index++;
+  }
+
+  return -1;
+}
+
 bool IsSupportedDataType(const Node& node, const std::vector<std::string>& supported_data_types) {
   for (const auto& input_arg : node.InputDefs()) {
     if (std::find(supported_data_types.begin(), supported_data_types.end(),

onnxruntime/core/optimizer/utils.h

@@ -57,10 +57,15 @@ bool ValidateShape(const NodeArg& node_arg, const std::initializer_list<int64_t>
 bool IsShapeKnownOnAllDims(const NodeArg& node_arg, int expected_dim_size);
 
 /** Get the index of node_arg among the node's all inputs.
-@remarks -1 when node_arg is not in node's inputs..
+@remarks -1 when node_arg is not in node's inputs.
 */
 int32_t IndexOfNodeInput(const Node& node, const NodeArg& node_arg);
 
+/** Get the index of node_arg among the node's all outputs.
+@remarks -1 when node_arg is not in node's outputs.
+*/
+int32_t IndexOfNodeOutput(const Node& node, const NodeArg& node_arg);
+
 /** Check whether node's input data types are in supported data type list.
 @param supported_data_types specify the supported data types.
 */

onnxruntime/test/optimizer/graph_transform_test.cc

@@ -5,6 +5,7 @@
 #pragma warning(disable : 4244)
 #endif
 
+#include <random>
 #include "core/graph/onnx_protobuf.h"
 
 #include "asserts.h"
@@ -15,6 +16,7 @@
 #include "core/graph/model.h"
 #include "core/optimizer/attention_fusion.h"
 #include "core/optimizer/bias_gelu_fusion.h"
+#include "core/optimizer/computation_reduction.h"
 #include "core/optimizer/cast_elimination.h"
 #include "core/optimizer/constant_folding.h"
 #include "core/optimizer/conv_activation_fusion.h"
@@ -55,6 +57,7 @@
 #include "test/optimizer/graph_transform_test_fixture.h"
 #include "test/providers/provider_test_utils.h"
 #include "test/test_environment.h"
+#include "test/util/include/default_providers.h"
 
 using namespace std;
 using namespace ONNX_NAMESPACE;
@@ -107,7 +110,7 @@ TEST_F(GraphTransformationTests, DropoutElimination) {
 }
 
 TEST_F(GraphTransformationTests, SliceElimination) {
-  std::vector<std::basic_string<ORTCHAR_T> > model_names = {ORT_TSTR("slice-v1-elim.onnx"), ORT_TSTR("slice-v11-elim.onnx")};
+  std::vector<std::basic_string<ORTCHAR_T>> model_names = {ORT_TSTR("slice-v1-elim.onnx"), ORT_TSTR("slice-v11-elim.onnx")};
   for (const auto& model_name : model_names) {
     auto model_uri = MODEL_FOLDER + model_name;
     std::shared_ptr<Model> model;
@@ -398,9 +401,9 @@ TEST_F(GraphTransformationTests, DontFuseConvWithBNWithOptionalOutputs) {
 }
 
 TEST_F(GraphTransformationTests, FuseConvBNMulAddUnsqueeze) {
-  std::vector<std::basic_string<ORTCHAR_T> > test_models = {ORT_TSTR("fusion/fuse-conv-bn-mul-add-unsqueeze.onnx"),
-                                                            ORT_TSTR("fusion/fuse-conv-bn-mul-add-unsqueeze.negative_axes.onnx"),
-                                                            ORT_TSTR("fusion/fuse-conv-bn-mul-add-unsqueeze-no-bias.onnx")};
+  std::vector<std::basic_string<ORTCHAR_T>> test_models = {ORT_TSTR("fusion/fuse-conv-bn-mul-add-unsqueeze.onnx"),
+                                                           ORT_TSTR("fusion/fuse-conv-bn-mul-add-unsqueeze.negative_axes.onnx"),
+                                                           ORT_TSTR("fusion/fuse-conv-bn-mul-add-unsqueeze-no-bias.onnx")};
   for (const auto& model : test_models) {
     auto model_uri = MODEL_FOLDER + model;
 
@@ -2552,5 +2555,194 @@ TEST_F(GraphTransformationTests, MatMulIntegerToFloatTest) {
 
 #endif
 
+// LayerNormalization implementation is in contrib namespace (OnnxDomain 1), so
+// Without contib_ops enabled, we cannot parse the graph correctly.
+#ifndef DISABLE_CONTRIB_OPS
+// We used Opset 12 for testing to make sure we are not using GatherND OnnxDomain Opset 1.
+static void GatherNDComputationReductionTest(const std::string op_type, logging::Logger& logger) {
+  std::string op_type_lower = op_type;
+  std::transform(op_type_lower.begin(), op_type_lower.end(), op_type_lower.begin(), [](unsigned char c) { return std::tolower(c); });
+  std::string file_path = std::string("testdata/transform/computation_reduction/gathernd_") + op_type_lower + std::string(".onnx");
+  std::shared_ptr<Model> model;
+  ASSERT_STATUS_OK(Model::Load(ToPathString(file_path), model, nullptr, logger));
+  Graph& graph = model->MainGraph();
+  std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
+
+  onnxruntime::GraphTransformerManager graph_transformation_mgr{1};
+  graph_transformation_mgr.Register(onnxruntime::make_unique<ComputationReductionTransformer>(), TransformerLevel::Level1);
+  ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, logger));
+
+  GraphViewer graph_viewer(graph);
+  const auto& node_topology_list = graph_viewer.GetNodesInTopologicalOrder();
+
+  Node* gathernd_node = nullptr;
+  for (auto node_index : node_topology_list) {
+    Node* p_node = graph.GetNode(node_index);
+    ASSERT_FALSE(p_node == nullptr);
+    if (p_node->OpType().compare("GatherND") == 0) {
+      gathernd_node = p_node;
+      EXPECT_EQ(gathernd_node->MutableInputDefs()[0]->Name(), "input");
+      const auto& consumers = graph.GetConsumerNodes(gathernd_node->MutableOutputDefs()[0]->Name());
+      EXPECT_EQ(consumers[0]->OpType(), op_type);
+    }
+  }
+
+  ASSERT_FALSE(gathernd_node == nullptr);
+}
+
+TEST_F(GraphTransformationTests, ComputationReductionTransformer_GatherND_Gelu) {
+  GatherNDComputationReductionTest("Gelu", *logger_);
+}
+
+TEST_F(GraphTransformationTests, ComputationReductionTransformer_GatherND_Add) {
+  GatherNDComputationReductionTest("Add", *logger_);
+}
+
+TEST_F(GraphTransformationTests, ComputationReductionTransformer_GatherND_LayerNormalization) {
+  GatherNDComputationReductionTest("LayerNormalization", *logger_);
+}
+
+TEST_F(GraphTransformationTests, ComputationReductionTransformer_GatherND_MatMul) {
+  GatherNDComputationReductionTest("MatMul", *logger_);
+}
+
+static void RunGatherNDE2EGraph(std::vector<OrtValue>& run_results, const PathString& model_uri,
+                                const std::string session_log_id, const std::string& provider_type,
+                                const std::vector<int64_t>& dims_input,
+                                const std::vector<float>& input_values,
+                                const std::vector<int64_t>& dims_unsqueezed_masked_lm_positions,
+                                const std::vector<int64_t>& values_unsqueezed_masked_lm_positions) {
+  SessionOptions so;
+  // we don't want any transformation here.
+  so.graph_optimization_level = TransformerLevel::Default;
+  so.session_logid = session_log_id;
+
+  InferenceSession session_object{so, GetEnvironment()};
+  std::unique_ptr<IExecutionProvider> execution_provider;
+  if (provider_type == onnxruntime::kCpuExecutionProvider)
+    execution_provider = DefaultCpuExecutionProvider();
+  else if (provider_type == onnxruntime::kCudaExecutionProvider)
+    execution_provider = DefaultCudaExecutionProvider();
+  EXPECT_TRUE(session_object.RegisterExecutionProvider(std::move(execution_provider)).IsOK());
+
+  Status st;
+  ASSERT_TRUE((st = session_object.Load(model_uri)).IsOK()) << st;
+  ASSERT_TRUE((st = session_object.Initialize()).IsOK()) << st;
+
+  OrtValue input1;
+  CreateMLValue<float>(TestCPUExecutionProvider()->GetAllocator(0, OrtMemTypeDefault), dims_input, input_values, &input1);
+  OrtValue input2;
+  CreateMLValue<int64_t>(TestCPUExecutionProvider()->GetAllocator(0, OrtMemTypeDefault), dims_unsqueezed_masked_lm_positions,
+                         values_unsqueezed_masked_lm_positions, &input2);
+
+  NameMLValMap feeds;
+  feeds.insert(std::make_pair("input", input1));
+  feeds.insert(std::make_pair("unsqueezed_masked_lm_positions", input2));
+
+  // prepare outputs
+  std::vector<std::string> output_names;
+  output_names.push_back("output");
+  output_names.push_back("gather_output");
+
+  // Now run
+  RunOptions run_options;
+  st = session_object.Run(run_options, feeds, output_names, &run_results);
+
+  EXPECT_TRUE(st.IsOK());
+}
+
+TEST_F(GraphTransformationTests, ComputationReductionTransformer_GatherND_E2E) {
+  auto model_uri = MODEL_FOLDER "computation_reduction/e2e.onnx";
+  std::shared_ptr<Model> model;
+  ASSERT_STATUS_OK(Model::Load(model_uri, model, nullptr, *logger_));
+  Graph& graph = model->MainGraph();
+  std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
+
+  onnxruntime::GraphTransformerManager graph_transformation_mgr{5};
+  graph_transformation_mgr.Register(onnxruntime::make_unique<ComputationReductionTransformer>(), TransformerLevel::Level1);
+  ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_));
+
+  // check the expected node orders.
+  {
+    GraphViewer graph_viewer(graph);
+    const auto& node_topology_list = graph_viewer.GetNodesInTopologicalOrder();
+
+    Node* gathernd_node = nullptr;
+    for (auto node_index : node_topology_list) {
+      Node* p_node = graph.GetNode(node_index);
+      ASSERT_FALSE(p_node == nullptr);
+      if (p_node->OpType().compare("GatherND") == 0) {
+        gathernd_node = p_node;
+        const Node* layer_norm_node = graph.GetProducerNode(gathernd_node->MutableInputDefs()[0]->Name());
+        EXPECT_EQ(layer_norm_node->OpType(), "LayerNormalization");
+        EXPECT_EQ(layer_norm_node->Name(), "layer_norm_1");
+        const auto& consumers = graph.GetConsumerNodes(gathernd_node->MutableOutputDefs()[0]->Name());
+        EXPECT_EQ(consumers[0]->OpType(), "MatMul");
+        EXPECT_EQ(consumers[0]->Name(), "matmul_1");
+        break;
+      }
+    }
+
+    ASSERT_FALSE(gathernd_node == nullptr);
+  }
+
+  // check result diff after the re-order
+  auto new_model_uri = "computation_reduction_transformer_after.onnx";
+  Model::Save(*model, new_model_uri);
+
+  float scale = 1.f;
+  float mean = 0.f;
+  float seed = 123.f;
+  std::default_random_engine generator_float{gsl::narrow_cast<uint32_t>(seed)};
+  std::normal_distribution<float> distribution_float{mean, scale};
+
+  int batch_size = 8;
+  int sequence = 128;
+  int hidden_size = 128;
+  int dynamic_predict_count = 20;
+  const std::vector<int64_t> dims_input = {batch_size, sequence, hidden_size};
+  std::vector<float> input_values(TensorShape(dims_input).Size());
+  std::for_each(input_values.begin(), input_values.end(),
+                [&generator_float, &distribution_float](float& value) { value = distribution_float(generator_float); });
+
+  const std::vector<int64_t> dims_unsqueezed_masked_lm_positions = {batch_size, dynamic_predict_count, 1};
+  std::vector<int64_t> values_unsqueezed_masked_lm_positions(TensorShape(dims_unsqueezed_masked_lm_positions).Size());
+
+  std::random_device rd;                                   // obtain a random number from hardware
+  std::mt19937 eng(rd());                                  // seed the generator
+  std::uniform_int_distribution<> distr(0, sequence - 1);  // define the range
+  std::for_each(values_unsqueezed_masked_lm_positions.begin(), values_unsqueezed_masked_lm_positions.end(),
+                [&distr, &eng](int64_t& value) { value = distr(eng); });
+
+  static const std::string all_provider_types[] = {
+      onnxruntime::kCpuExecutionProvider,
+#ifdef USE_CUDA
+      onnxruntime::kCudaExecutionProvider,
+#endif
+  };
+
+  for (auto& provider_type : all_provider_types) {
+    std::vector<OrtValue> expected_ort_values;
+    RunGatherNDE2EGraph(expected_ort_values, model_uri, std::string("RawGraphRun"), provider_type,
+                        dims_input, input_values, dims_unsqueezed_masked_lm_positions,
+                        values_unsqueezed_masked_lm_positions);
+
+    std::vector<OrtValue> actual_ort_values;
+    RunGatherNDE2EGraph(actual_ort_values, ToPathString(new_model_uri), std::string("OptimizedGraphRun"), provider_type,
+                        dims_input, input_values, dims_unsqueezed_masked_lm_positions,
+                        values_unsqueezed_masked_lm_positions);
+
+    ASSERT_TRUE(expected_ort_values.size() == actual_ort_values.size());
+    const double per_sample_tolerance = 1e-4;
+    const double relative_per_sample_tolerance = 1e-4;
+    for (size_t i = 0; i < expected_ort_values.size(); i++) {
+      auto ret = CompareOrtValue(actual_ort_values[i], expected_ort_values[i],
+                                 per_sample_tolerance, relative_per_sample_tolerance, false);
+      EXPECT_EQ(ret.first, COMPARE_RESULT::SUCCESS) << ret.second;
+    }
+  }
+}
+#endif
+
 }  // namespace test
 }  // namespace onnxruntime

onnxruntime/test/testdata/transform/computation_reduction.py

@@ -0,0 +1,101 @@
+import onnx
+from onnx import helper
+from onnx import AttributeProto, TensorProto, GraphProto, OperatorSetIdProto
+from onnx import numpy_helper
+import numpy as np
+
+vocab_size=256 #30258
+
+X = helper.make_tensor_value_info('input', TensorProto.FLOAT, ["batch", "seqlen", 128])
+unsqueezed_masked_lm_positions = helper.make_tensor_value_info('unsqueezed_masked_lm_positions', 
+                                                               TensorProto.INT64, ["batch", "dynamic_prediction_count", 1])
+Y = helper.make_tensor_value_info('output', TensorProto.FLOAT, ["batch", "dynamic_prediction_count", vocab_size])
+Gather_Y = helper.make_tensor_value_info('gather_output', TensorProto.FLOAT, ["batch", 128])
+
+layer_norm1_weight_np_vals = np.random.uniform(0.0, 1.0, (128)).astype(np.float32).reshape((128))
+layer_norm1_weight_initializer = numpy_helper.from_array(layer_norm1_weight_np_vals, "bert.encoder.layer.2.output.LayerNorm.weight")
+
+layer_norm1_bias_np_vals = np.random.uniform(0.0, 1.0, (128)).astype(np.float32).reshape((128))
+layer_norm1_bias_initializer = numpy_helper.from_array(layer_norm1_bias_np_vals, "bert.encoder.layer.2.output.LayerNorm.bias")
+
+matmul1_np_vals = np.random.uniform(0.0, 1.0, (128, 128)).astype(np.float32).reshape((128, 128))
+matmul1_initializer = numpy_helper.from_array(matmul1_np_vals, "matmul1_initializer")
+
+add1_np_vals = np.random.uniform(0.0, 1.0, (128)).astype(np.float32).reshape((128))
+add1_initializer = numpy_helper.from_array(add1_np_vals, "add1_initializerr")
+
+layer_norm2_weight_np_vals = np.random.uniform(0.0, 1.0, (128)).astype(np.float32).reshape((128))
+layer_norm2_weight_initializer = numpy_helper.from_array(layer_norm2_weight_np_vals, "cls.predictions.transform.LayerNorm.weight")
+
+layer_norm2_bias_np_vals = np.random.uniform(0.0, 1.0, (128)).astype(np.float32).reshape((128))
+layer_norm2_bias_initializer = numpy_helper.from_array(layer_norm2_bias_np_vals, "cls.predictions.transform.LayerNorm.bias")
+
+matmul2_np_vals = np.random.uniform(0.0, 1.0, (128, vocab_size)).astype(np.float32).reshape((128, vocab_size))
+matmul2_initializer = numpy_helper.from_array(matmul2_np_vals, "bert.embeddings.word_embeddings.weight_transposed")
+
+add2_np_vals = np.random.uniform(0.0, 1.0, (vocab_size)).astype(np.float32).reshape((vocab_size))
+add2_initializer = numpy_helper.from_array(add2_np_vals, "cls.predictions.bias")
+
+gather_indice_np_vals = np.asarray([0]).astype(np.int64).reshape(())
+gather_indice_initializer = numpy_helper.from_array(gather_indice_np_vals, "gather_indice_initializer")
+
+nodes=[]
+layer_norm1 = helper.make_node('LayerNormalization', 
+                               ['input', layer_norm1_weight_initializer.name, layer_norm1_bias_initializer.name],
+                               ['layer_norm1', 'saved_mean1', 'saved_inv_std_var1'],
+                               name='layer_norm_1', epsilon=9.999999960041972e-13, axis=-1)
+nodes.append(layer_norm1)
+
+gather1 = helper.make_node('Gather', ['layer_norm1', gather_indice_initializer.name], ['gather_output'], name="gather_output", axis=1)
+nodes.append(gather1)
+
+matmul1 = helper.make_node('MatMul', ['layer_norm1', matmul1_initializer.name], ['matmul1'], name="matmul_1")
+nodes.append(matmul1)
+
+add1 = helper.make_node('Add', [add1_initializer.name, 'matmul1'], ['add1'], name="add_1")
+nodes.append(add1)
+
+gelu1 = helper.make_node('Gelu', ['add1'], ['gelu1'], name='gelu_1', domain='com.microsoft')
+nodes.append(gelu1)
+
+layer_norm2 = helper.make_node('LayerNormalization',
+                               ['gelu1', layer_norm2_weight_initializer.name, layer_norm2_bias_initializer.name],
+                               ['layer_norm2', 'saved_mean2', 'saved_inv_std_var2'],
+                               name='layer_norm_2', epsilon=9.999999960041972e-13, axis=-1)
+nodes.append(layer_norm2)
+
+matmul2 = helper.make_node('MatMul', ['layer_norm2', matmul2_initializer.name], ['matmul2'], name="matmul_2")
+nodes.append(matmul2)
+
+add2 = helper.make_node('Add', ['matmul2', add2_initializer.name], ['add2'], name="add_2")
+nodes.append(add2)
+
+gathernd1 = helper.make_node('GatherND', ['add2', 'unsqueezed_masked_lm_positions'], ['gathernd1'], name="gathernd_1", batch_dims=1)
+nodes.append(gathernd1)
+
+identity1 = helper.make_node('Identity', ['gathernd1'], ['output'], name="output")
+nodes.append(identity1)
+
+initializers=[layer_norm1_weight_initializer, layer_norm1_bias_initializer, matmul1_initializer, add1_initializer,
+              layer_norm2_weight_initializer, layer_norm2_bias_initializer, matmul2_initializer, add2_initializer,
+              gather_indice_initializer]
+# Create the graph (GraphProto)
+graph_def = helper.make_graph(nodes, 'test-model', [X, unsqueezed_masked_lm_positions], [Y, Gather_Y], initializers)
+
+opsets = []
+onnxdomain = OperatorSetIdProto()
+onnxdomain.version = 12
+onnxdomain.domain = ""  # The empty string ("") or absence of this field implies the operator set that is defined as part of the ONNX specification.
+opsets.append(onnxdomain)
+
+msdomain = OperatorSetIdProto()
+msdomain.version = 1
+msdomain.domain = "com.microsoft"
+
+opsets.append(msdomain)
+kwargs = {}
+kwargs["opset_imports"] = opsets
+
+model_def = helper.make_model(graph_def, producer_name='onnx-example', **kwargs)
+
+onnx.save(model_def, "computation_reduction_transformer.onnx")

onnxruntime/test/testdata/transform/computation_reduction/e2e.py

@@ -0,0 +1,101 @@
+import onnx
+from onnx import helper
+from onnx import AttributeProto, TensorProto, GraphProto, OperatorSetIdProto
+from onnx import numpy_helper
+import numpy as np
+
+vocab_size=256 #30258
+
+X = helper.make_tensor_value_info('input', TensorProto.FLOAT, ["batch", "seqlen", 128])
+unsqueezed_masked_lm_positions = helper.make_tensor_value_info('unsqueezed_masked_lm_positions', 
+                                                               TensorProto.INT64, ["batch", "dynamic_prediction_count", 1])
+Y = helper.make_tensor_value_info('output', TensorProto.FLOAT, ["batch", "dynamic_prediction_count", vocab_size])
+Gather_Y = helper.make_tensor_value_info('gather_output', TensorProto.FLOAT, ["batch", 128])
+
+layer_norm1_weight_np_vals = np.random.uniform(0.0, 1.0, (128)).astype(np.float32).reshape((128))
+layer_norm1_weight_initializer = numpy_helper.from_array(layer_norm1_weight_np_vals, "bert.encoder.layer.2.output.LayerNorm.weight")
+
+layer_norm1_bias_np_vals = np.random.uniform(0.0, 1.0, (128)).astype(np.float32).reshape((128))
+layer_norm1_bias_initializer = numpy_helper.from_array(layer_norm1_bias_np_vals, "bert.encoder.layer.2.output.LayerNorm.bias")
+
+matmul1_np_vals = np.random.uniform(0.0, 1.0, (128, 128)).astype(np.float32).reshape((128, 128))
+matmul1_initializer = numpy_helper.from_array(matmul1_np_vals, "matmul1_initializer")
+
+add1_np_vals = np.random.uniform(0.0, 1.0, (128)).astype(np.float32).reshape((128))
+add1_initializer = numpy_helper.from_array(add1_np_vals, "add1_initializerr")
+
+layer_norm2_weight_np_vals = np.random.uniform(0.0, 1.0, (128)).astype(np.float32).reshape((128))
+layer_norm2_weight_initializer = numpy_helper.from_array(layer_norm2_weight_np_vals, "cls.predictions.transform.LayerNorm.weight")
+
+layer_norm2_bias_np_vals = np.random.uniform(0.0, 1.0, (128)).astype(np.float32).reshape((128))
+layer_norm2_bias_initializer = numpy_helper.from_array(layer_norm2_bias_np_vals, "cls.predictions.transform.LayerNorm.bias")
+
+matmul2_np_vals = np.random.uniform(0.0, 1.0, (128, vocab_size)).astype(np.float32).reshape((128, vocab_size))
+matmul2_initializer = numpy_helper.from_array(matmul2_np_vals, "bert.embeddings.word_embeddings.weight_transposed")
+
+add2_np_vals = np.random.uniform(0.0, 1.0, (vocab_size)).astype(np.float32).reshape((vocab_size))
+add2_initializer = numpy_helper.from_array(add2_np_vals, "cls.predictions.bias")
+
+gather_indice_np_vals = np.asarray([0]).astype(np.int64).reshape(())
+gather_indice_initializer = numpy_helper.from_array(gather_indice_np_vals, "gather_indice_initializer")
+
+nodes=[]
+layer_norm1 = helper.make_node('LayerNormalization', 
+                               ['input', layer_norm1_weight_initializer.name, layer_norm1_bias_initializer.name],
+                               ['layer_norm1', 'saved_mean1', 'saved_inv_std_var1'],
+                               name='layer_norm_1', epsilon=9.999999960041972e-13, axis=-1)
+nodes.append(layer_norm1)
+
+gather1 = helper.make_node('Gather', ['layer_norm1', gather_indice_initializer.name], ['gather_output'], name="gather_output", axis=1)
+nodes.append(gather1)
+
+matmul1 = helper.make_node('MatMul', ['layer_norm1', matmul1_initializer.name], ['matmul1'], name="matmul_1")
+nodes.append(matmul1)
+
+add1 = helper.make_node('Add', [add1_initializer.name, 'matmul1'], ['add1'], name="add_1")
+nodes.append(add1)
+
+gelu1 = helper.make_node('Gelu', ['add1'], ['gelu1'], name='gelu_1', domain='com.microsoft')
+nodes.append(gelu1)
+
+layer_norm2 = helper.make_node('LayerNormalization',
+                               ['gelu1', layer_norm2_weight_initializer.name, layer_norm2_bias_initializer.name],
+                               ['layer_norm2', 'saved_mean2', 'saved_inv_std_var2'],
+                               name='layer_norm_2', epsilon=9.999999960041972e-13, axis=-1)
+nodes.append(layer_norm2)
+
+matmul2 = helper.make_node('MatMul', ['layer_norm2', matmul2_initializer.name], ['matmul2'], name="matmul_2")
+nodes.append(matmul2)
+
+add2 = helper.make_node('Add', ['matmul2', add2_initializer.name], ['add2'], name="add_2")
+nodes.append(add2)
+
+gathernd1 = helper.make_node('GatherND', ['add2', 'unsqueezed_masked_lm_positions'], ['gathernd1'], name="gathernd_1", batch_dims=1)
+nodes.append(gathernd1)
+
+identity1 = helper.make_node('Identity', ['gathernd1'], ['output'], name="output")
+nodes.append(identity1)
+
+initializers=[layer_norm1_weight_initializer, layer_norm1_bias_initializer, matmul1_initializer, add1_initializer,
+              layer_norm2_weight_initializer, layer_norm2_bias_initializer, matmul2_initializer, add2_initializer,
+              gather_indice_initializer]
+# Create the graph (GraphProto)
+graph_def = helper.make_graph(nodes, 'test-model', [X, unsqueezed_masked_lm_positions], [Y, Gather_Y], initializers)
+
+opsets = []
+onnxdomain = OperatorSetIdProto()
+onnxdomain.version = 12
+onnxdomain.domain = ""  # The empty string ("") or absence of this field implies the operator set that is defined as part of the ONNX specification.
+opsets.append(onnxdomain)
+
+msdomain = OperatorSetIdProto()
+msdomain.version = 1
+msdomain.domain = "com.microsoft"
+
+opsets.append(msdomain)
+kwargs = {}
+kwargs["opset_imports"] = opsets
+
+model_def = helper.make_model(graph_def, producer_name='onnx-example', **kwargs)
+
+onnx.save(model_def, "e2e.onnx")

onnxruntime/test/testdata/transform/computation_reduction/gathernd_add.py

@@ -0,0 +1,50 @@
+import onnx
+from onnx import helper
+from onnx import AttributeProto, TensorProto, GraphProto, OperatorSetIdProto
+from onnx import numpy_helper
+import numpy as np
+
+X = helper.make_tensor_value_info('input', TensorProto.FLOAT, ["batch", "seqlen", 128])
+unsqueezed_masked_lm_positions = helper.make_tensor_value_info('unsqueezed_masked_lm_positions', 
+                                                               TensorProto.INT64, ["batch", "dynamic_prediction_count", 1])
+Y = helper.make_tensor_value_info('output', TensorProto.FLOAT, ["batch", "dynamic_prediction_count", 128])
+Y2 = helper.make_tensor_value_info('output2', TensorProto.FLOAT, ["batch", "dynamic_prediction_count", 128])
+nodes = []
+
+# case 1
+bias_np_val = np.random.uniform(0.0, 1.0, (128)).astype(np.float32).reshape((128))
+bias_initializer = numpy_helper.from_array(bias_np_val, "bias")
+add1 = helper.make_node('Add', ['input', 'bias'], ['add_1'], name="add_1")
+nodes.append(add1)
+
+gathernd1 = helper.make_node('GatherND', ['add_1', 'unsqueezed_masked_lm_positions'], ['output'], name="gathernd_1", batch_dims=1)
+nodes.append(gathernd1)
+
+# case 2
+bias2_np_val = np.random.uniform(0.0, 1.0, (128)).astype(np.float32).reshape((128))
+bias2_initializer = numpy_helper.from_array(bias2_np_val, "bias2")
+add2 = helper.make_node('Add', ['bias2', 'input'], ['add_2'], name="add_2")
+nodes.append(add2)
+
+gathernd2 = helper.make_node('GatherND', ['add_2', 'unsqueezed_masked_lm_positions'], ['output2'], name="gathernd_2", batch_dims=1)
+nodes.append(gathernd2)
+
+graph_def = helper.make_graph(nodes, 'test-model', [X, unsqueezed_masked_lm_positions], [Y, Y2], [bias_initializer, bias2_initializer])
+
+opsets = []
+onnxdomain = OperatorSetIdProto()
+onnxdomain.version = 12
+onnxdomain.domain = ""  # The empty string ("") or absence of this field implies the operator set that is defined as part of the ONNX specification.
+opsets.append(onnxdomain)
+
+msdomain = OperatorSetIdProto()
+msdomain.version = 1
+msdomain.domain = "com.microsoft"
+
+opsets.append(msdomain)
+kwargs = {}
+kwargs["opset_imports"] = opsets
+
+model_def = helper.make_model(graph_def, producer_name='onnx-example', **kwargs)
+
+onnx.save(model_def, "gathernd_add.onnx")

onnxruntime/test/testdata/transform/computation_reduction/gathernd_div.py

@@ -0,0 +1,50 @@
+import onnx
+from onnx import helper
+from onnx import AttributeProto, TensorProto, GraphProto, OperatorSetIdProto
+from onnx import numpy_helper
+import numpy as np
+
+X = helper.make_tensor_value_info('input', TensorProto.FLOAT, ["batch", "seqlen", 128])
+unsqueezed_masked_lm_positions = helper.make_tensor_value_info('unsqueezed_masked_lm_positions', 
+                                                               TensorProto.INT64, ["batch", "dynamic_prediction_count", 1])
+Y = helper.make_tensor_value_info('output', TensorProto.FLOAT, ["batch", "dynamic_prediction_count", 128])
+Y2 = helper.make_tensor_value_info('output2', TensorProto.FLOAT, ["batch", "dynamic_prediction_count", 128])
+nodes = []
+
+# case 1
+divisor_np_val = np.random.uniform(0.0, 1.0, (128)).astype(np.float32).reshape((128))
+divisor_initializer = numpy_helper.from_array(divisor_np_val, "divisor")
+div1 = helper.make_node('Div', ['input', 'divisor'], ['div_1'], name="div_1")
+nodes.append(div1)
+
+gathernd1 = helper.make_node('GatherND', ['div_1', 'unsqueezed_masked_lm_positions'], ['output'], name="gathernd_1", batch_dims=1)
+nodes.append(gathernd1)
+
+# case 2
+divisor2_np_val = np.random.uniform(0.0, 1.0, (128)).astype(np.float32).reshape((128))
+divisor2_initializer = numpy_helper.from_array(divisor2_np_val, "divisor2")
+div2 = helper.make_node('Div', ['divisor2', 'input'], ['div_2'], name="div_2")
+nodes.append(div2)
+
+gathernd2 = helper.make_node('GatherND', ['div_2', 'unsqueezed_masked_lm_positions'], ['output2'], name="gathernd_2", batch_dims=1)
+nodes.append(gathernd2)
+
+graph_def = helper.make_graph(nodes, 'test-model', [X, unsqueezed_masked_lm_positions], [Y, Y2], [divisor_initializer, divisor2_initializer])
+
+opsets = []
+onnxdomain = OperatorSetIdProto()
+onnxdomain.version = 12
+onnxdomain.domain = ""  # The empty string ("") or absence of this field implies the operator set that is defined as part of the ONNX specification.
+opsets.append(onnxdomain)
+
+msdomain = OperatorSetIdProto()
+msdomain.version = 1
+msdomain.domain = "com.microsoft"
+
+opsets.append(msdomain)
+kwargs = {}
+kwargs["opset_imports"] = opsets
+
+model_def = helper.make_model(graph_def, producer_name='onnx-example', **kwargs)
+
+onnx.save(model_def, "gathernd_div.onnx")

onnxruntime/test/testdata/transform/computation_reduction/gathernd_gelu.py

@@ -0,0 +1,38 @@
+import onnx
+from onnx import helper
+from onnx import AttributeProto, TensorProto, GraphProto, OperatorSetIdProto
+from onnx import numpy_helper
+import numpy as np
+
+X = helper.make_tensor_value_info('input', TensorProto.FLOAT, ["batch", "seqlen", 128])
+unsqueezed_masked_lm_positions = helper.make_tensor_value_info('unsqueezed_masked_lm_positions', 
+                                                               TensorProto.INT64, ["batch", "dynamic_prediction_count", 1])
+Y = helper.make_tensor_value_info('output', TensorProto.FLOAT, ["batch", "dynamic_prediction_count", 128])
+nodes = []
+
+# case 1
+gelu1 = helper.make_node('Gelu', ['input'], ['gelu_1'], name="gelu_1", domain='com.microsoft')
+nodes.append(gelu1)
+
+gathernd1 = helper.make_node('GatherND', ['gelu_1', 'unsqueezed_masked_lm_positions'], ['output'], name="gathernd_1", batch_dims=1)
+nodes.append(gathernd1)
+
+graph_def = helper.make_graph(nodes, 'test-model', [X, unsqueezed_masked_lm_positions], [Y])
+
+opsets = []
+onnxdomain = OperatorSetIdProto()
+onnxdomain.version = 12
+onnxdomain.domain = ""  # The empty string ("") or absence of this field implies the operator set that is defined as part of the ONNX specification.
+opsets.append(onnxdomain)
+
+msdomain = OperatorSetIdProto()
+msdomain.version = 1
+msdomain.domain = "com.microsoft"
+
+opsets.append(msdomain)
+kwargs = {}
+kwargs["opset_imports"] = opsets
+
+model_def = helper.make_model(graph_def, producer_name='onnx-example', **kwargs)
+
+onnx.save(model_def, "gathernd_gelu.onnx")

onnxruntime/test/testdata/transform/computation_reduction/gathernd_layernormalization.py

@@ -0,0 +1,46 @@
+import onnx
+from onnx import helper
+from onnx import AttributeProto, TensorProto, GraphProto, OperatorSetIdProto
+from onnx import numpy_helper
+import numpy as np
+
+X = helper.make_tensor_value_info('input', TensorProto.FLOAT, ["batch", "seqlen", 128])
+unsqueezed_masked_lm_positions = helper.make_tensor_value_info('unsqueezed_masked_lm_positions', 
+                                                               TensorProto.INT64, ["batch", "dynamic_prediction_count", 1])
+Y = helper.make_tensor_value_info('output', TensorProto.FLOAT, ["batch", "dynamic_prediction_count", 128])
+
+layer_norm1_weight_np_vals = np.random.uniform(0.0, 1.0, (128)).astype(np.float32).reshape((128))
+layer_norm1_weight_initializer = numpy_helper.from_array(layer_norm1_weight_np_vals, "bert.encoder.layer.2.output.LayerNorm.weight")
+
+layer_norm1_bias_np_vals = np.random.uniform(0.0, 1.0, (128)).astype(np.float32).reshape((128))
+layer_norm1_bias_initializer = numpy_helper.from_array(layer_norm1_bias_np_vals, "bert.encoder.layer.2.output.LayerNorm.bias")
+
+nodes=[]
+layer_norm1 = helper.make_node('LayerNormalization', 
+                               ['input', layer_norm1_weight_initializer.name, layer_norm1_bias_initializer.name],
+                               ['layer_norm1', 'saved_mean1', 'saved_inv_std_var1'],
+                               name='layer_norm_1', epsilon=9.999999960041972e-13, axis=-1)
+nodes.append(layer_norm1)
+
+gathernd1 = helper.make_node('GatherND', ['layer_norm1', 'unsqueezed_masked_lm_positions'], ['output'], name="gathernd_1", batch_dims=1)
+nodes.append(gathernd1)
+
+graph_def = helper.make_graph(nodes, 'test-model', [X, unsqueezed_masked_lm_positions], [Y], [layer_norm1_weight_initializer, layer_norm1_bias_initializer])
+
+opsets = []
+onnxdomain = OperatorSetIdProto()
+onnxdomain.version = 12
+onnxdomain.domain = ""  # The empty string ("") or absence of this field implies the operator set that is defined as part of the ONNX specification.
+opsets.append(onnxdomain)
+
+msdomain = OperatorSetIdProto()
+msdomain.version = 1
+msdomain.domain = "com.microsoft"
+
+opsets.append(msdomain)
+kwargs = {}
+kwargs["opset_imports"] = opsets
+
+model_def = helper.make_model(graph_def, producer_name='onnx-example', **kwargs)
+
+onnx.save(model_def, "gathernd_layernormalization.onnx")

onnxruntime/test/testdata/transform/computation_reduction/gathernd_matmul.py

@@ -0,0 +1,42 @@
+import onnx
+from onnx import helper
+from onnx import AttributeProto, TensorProto, GraphProto, OperatorSetIdProto
+from onnx import numpy_helper
+import numpy as np
+
+X = helper.make_tensor_value_info('input', TensorProto.FLOAT, ["batch", "seqlen", 128])
+unsqueezed_masked_lm_positions = helper.make_tensor_value_info('unsqueezed_masked_lm_positions', 
+                                                               TensorProto.INT64, ["batch", "dynamic_prediction_count", 1])
+Y = helper.make_tensor_value_info('output', TensorProto.FLOAT, ["batch", "dynamic_prediction_count", 128])
+
+matmul1_np_vals = np.random.uniform(0.0, 1.0, (128, 128)).astype(np.float32).reshape((128, 128))
+matmul1_initializer = numpy_helper.from_array(matmul1_np_vals, "matmul1_initializer")
+
+nodes=[]
+matmul1 = helper.make_node('MatMul', ['input', matmul1_initializer.name], ['matmul1'], name="matmul_1")
+nodes.append(matmul1)
+
+gathernd1 = helper.make_node('GatherND', ['matmul1', 'unsqueezed_masked_lm_positions'], ['output'], name="gathernd_1", batch_dims=1)
+nodes.append(gathernd1)
+
+initializers=[matmul1_initializer]
+
+graph_def = helper.make_graph(nodes, 'test-model', [X, unsqueezed_masked_lm_positions], [Y], initializers)
+
+opsets = []
+onnxdomain = OperatorSetIdProto()
+onnxdomain.version = 12
+onnxdomain.domain = ""  # The empty string ("") or absence of this field implies the operator set that is defined as part of the ONNX specification.
+opsets.append(onnxdomain)
+
+msdomain = OperatorSetIdProto()
+msdomain.version = 1
+msdomain.domain = "com.microsoft"
+
+opsets.append(msdomain)
+kwargs = {}
+kwargs["opset_imports"] = opsets
+
+model_def = helper.make_model(graph_def, producer_name='onnx-example', **kwargs)
+
+onnx.save(model_def, "gathernd_matmul.onnx")

orttraining/orttraining/core/optimizer/graph_transformer_utils.cc

@@ -36,6 +36,7 @@
 #include "core/optimizer/fast_gelu_fusion.h"
 #include "core/optimizer/gelu_approximation.h"
 #include "core/optimizer/graph_transformer_utils.h"
+#include "core/optimizer/computation_reduction.h"
 #include "core/mlas/inc/mlas.h"
 #include "core/session/inference_session.h"
 
@@ -85,7 +86,7 @@ std::vector<std::unique_ptr<GraphTransformer>> GeneratePreTrainingTransformers(T
             training::DistributedRunContext::RankInGroup(training::WorkerGroupType::HorizontalParallel),
             horizontal_parallel_size, compatible_eps));
       }
-
+      transformers.emplace_back(onnxruntime::make_unique<ComputationReductionTransformer>(compatible_eps));
     } break;
 
     case TransformerLevel::Level2: {

orttraining/orttraining/core/session/training_session.cc

@@ -120,7 +120,7 @@ Status TrainingSession::ConfigureForTraining(
   TrainingConfigurationResult config_result{};
 
   ORT_ENFORCE(config.distributed_config.pipeline_parallel_size > 0,
-    "This parameter should be 1 if there is no pipelie parallelism. Otherwise, it's the number of pipeline stages.");
+              "This parameter should be 1 if there is no pipelie parallelism. Otherwise, it's the number of pipeline stages.");
 
   DistributedRunContext::CreateInstance({config.distributed_config.world_rank,
                                          config.distributed_config.world_size,
@@ -150,7 +150,7 @@ Status TrainingSession::ConfigureForTraining(
                            (!config.pipeline_config.has_value() ||
                             (config.distributed_config.world_rank + 1 == config.distributed_config.world_size));
   optional<std::string> loss_scale_input_name =
-        enable_loss_scale ? optional<std::string>{""} : optional<std::string>{};
+      enable_loss_scale ? optional<std::string>{""} : optional<std::string>{};
   if (config.pipeline_config.has_value()) {
     // if use pipeline, first check if model contains send op. If it does, set the
     // send node's output as the start tensor to build gradient graph
@@ -189,15 +189,14 @@ Status TrainingSession::ConfigureForTraining(
       !config.weight_names_to_train.empty()
           ? config.weight_names_to_train
           : GetTrainableModelInitializers(config.immutable_weights, loss_name);
-  if (config.weight_names_to_not_train.size() > 0)
-  {
+  if (config.weight_names_to_not_train.size() > 0) {
     LOGS(*session_logger_, INFO) << "Excluding following weights from trainable list as specified in configuration:";
     for (const auto& weight_name_to_not_train : config.weight_names_to_not_train) {
       trainable_initializers.erase(weight_name_to_not_train);
       LOGS(*session_logger_, INFO) << weight_name_to_not_train;
     }
   }
-  
+
   ORT_RETURN_IF_ERROR(ApplyTransformationsToMainGraph(trainable_initializers, config.enable_gelu_approximation));
 
   // derive actual set of weights to train
@@ -256,7 +255,7 @@ Status TrainingSession::ConfigureForTraining(
       pipeline_result.fetch_names.push_back(name);
     }
     pipeline_result.pipeline_stage_id = config.distributed_config.world_rank /
-      (config.distributed_config.data_parallel_size * config.distributed_config.horizontal_parallel_size);
+                                        (config.distributed_config.data_parallel_size * config.distributed_config.horizontal_parallel_size);
     config_result.pipeline_config_result = pipeline_result;
   }
 
@@ -264,15 +263,14 @@ Status TrainingSession::ConfigureForTraining(
   // TODO: this is a temp workaround for removing rank tensor before adding optimizer.
   // Re-visit after we port logic for model splitting and hence know the rank tensor name.
   for (auto it = weights_to_train_.begin(); it != weights_to_train_.end();) {
-      const auto* node_arg = model_->MainGraph().GetNodeArg(*it);
-      ORT_RETURN_IF_NOT(node_arg, "Failed to get NodeArg with name ", *it);
-      if (node_arg->TypeAsProto()->tensor_type().elem_type() != ONNX_NAMESPACE::TensorProto_DataType_FLOAT &&
-          node_arg->TypeAsProto()->tensor_type().elem_type() != ONNX_NAMESPACE::TensorProto_DataType_FLOAT16) {
-        it = weights_to_train_.erase(it);
-      }
-      else{
-          ++it;
-      }
+    const auto* node_arg = model_->MainGraph().GetNodeArg(*it);
+    ORT_RETURN_IF_NOT(node_arg, "Failed to get NodeArg with name ", *it);
+    if (node_arg->TypeAsProto()->tensor_type().elem_type() != ONNX_NAMESPACE::TensorProto_DataType_FLOAT &&
+        node_arg->TypeAsProto()->tensor_type().elem_type() != ONNX_NAMESPACE::TensorProto_DataType_FLOAT16) {
+      it = weights_to_train_.erase(it);
+    } else {
+      ++it;
+    }
   }
 
   // add optimizer or gradient accumulation
@@ -341,8 +339,7 @@ Status TrainingSession::ConfigureForTraining(
   // Note: in the pipeline case, different ranks may resident in the same node. This could lead to a potential write
   // conflict. It is user's responsibility to make sure different rank is passed in with different
   // model_with_training_graph_path value.
-  if ((IsRootNode(config) || config.pipeline_config.has_value())
-    && config.model_with_training_graph_path.has_value()) {
+  if ((IsRootNode(config) || config.pipeline_config.has_value()) && config.model_with_training_graph_path.has_value()) {
     ORT_IGNORE_RETURN_VALUE(Save(
         config.model_with_training_graph_path.value(), SaveOption::NO_RELOAD));
   }
@@ -561,34 +558,34 @@ Status TrainingSession::AddTensorboard(const std::string& summary_name,
 }
 
 Status TrainingSession::InsertPipelineOps(
-  const std::unordered_set<std::string>& initializer_names_to_preserve,
-  std::string& forward_waited_event_name,
-  std::string& forward_recorded_event_name,
-  std::string& backward_waited_event_name,
-  std::string& backward_recorded_event_name,
-  std::string& forward_wait_output_name,
-  std::string& forward_record_output_name,
-  std::string& backward_wait_output_name,
-  std::string& backward_record_output_name,
-  std::string& forward_waited_event_after_recv_name,
-  std::string& forward_recorded_event_before_send_name,
-  std::string& backward_waited_event_after_recv_name,
-  std::string& backward_recorded_event_before_send_name) {
+    const std::unordered_set<std::string>& initializer_names_to_preserve,
+    std::string& forward_waited_event_name,
+    std::string& forward_recorded_event_name,
+    std::string& backward_waited_event_name,
+    std::string& backward_recorded_event_name,
+    std::string& forward_wait_output_name,
+    std::string& forward_record_output_name,
+    std::string& backward_wait_output_name,
+    std::string& backward_record_output_name,
+    std::string& forward_waited_event_after_recv_name,
+    std::string& forward_recorded_event_before_send_name,
+    std::string& backward_waited_event_after_recv_name,
+    std::string& backward_recorded_event_before_send_name) {
   ORT_RETURN_IF_ERROR(TransformGraphForPipeline(
-    model_->MainGraph(),
-    initializer_names_to_preserve,
-    forward_waited_event_name,
-    forward_recorded_event_name,
-    backward_waited_event_name,
-    backward_recorded_event_name,
-    forward_wait_output_name,
-    forward_record_output_name,
-    backward_wait_output_name,
-    backward_record_output_name,
-    forward_waited_event_after_recv_name,
-    forward_recorded_event_before_send_name,
-    backward_waited_event_after_recv_name,
-    backward_recorded_event_before_send_name));
+      model_->MainGraph(),
+      initializer_names_to_preserve,
+      forward_waited_event_name,
+      forward_recorded_event_name,
+      backward_waited_event_name,
+      backward_recorded_event_name,
+      forward_wait_output_name,
+      forward_record_output_name,
+      backward_wait_output_name,
+      backward_record_output_name,
+      forward_waited_event_after_recv_name,
+      forward_recorded_event_before_send_name,
+      backward_waited_event_after_recv_name,
+      backward_recorded_event_before_send_name));
   return DoPostLoadProcessing(*model_);
 }
 
@@ -821,7 +818,7 @@ common::Status TrainingSession::Run(const RunOptions& run_options, IOBinding& io
       const auto& session_state = GetSessionState();
       auto default_cpu_alloc_info = session_state.GetExecutionProviders().GetDefaultCpuMemoryInfo();
       auto cpu_allocator = session_state.GetAllocator(default_cpu_alloc_info);
-      
+
       feed_value = onnxruntime::MakeScalarMLValue<float>(cpu_allocator, 0.f, true /*is_1d*/);
       // Bind new feed to graph input.
       ORT_RETURN_IF_ERROR(io_binding.BindInput(drop_ratio, feed_value));
@@ -851,11 +848,11 @@ Status TrainingSession::SetDropoutEvalFeedNames() {
 
   for (const auto& node : graph.Nodes()) {
     auto it = Dropout_Nodes.find(node.OpType());
-    if(it != Dropout_Nodes.cend()) {
+    if (it != Dropout_Nodes.cend()) {
       auto& ratio_name = node.InputDefs()[1]->Name();
       dropout_eval_feeds_.insert(ratio_name);
       ORT_ENFORCE(model_->MainGraph().GetProducerNode(ratio_name) == nullptr,
-      "Input: " + ratio_name + " should not have any producer node.");
+                  "Input: " + ratio_name + " should not have any producer node.");
       defs.AddGraphInputs({ratio_name});
     }
   }
@@ -971,41 +968,39 @@ bool TrainingSession::IsImmutableWeight(const ImmutableWeights& immutable_weight
 
 std::unordered_set<std::string> TrainingSession::GetTrainableModelInitializers(
     const ImmutableWeights& immutable_weights, const std::string& loss_name) const {
-
   const Graph& graph = model_->MainGraph();
   const auto& initialized_tensors = graph.GetAllInitializedTensors();
   std::unordered_set<std::string> trainable_initializers;
 
-  auto add_trainable_initializers = [&](const Node* node) { 
-      for (auto input : node->InputDefs()) {
-        std::string initializer_name = input->Name();
-        if (initialized_tensors.count(initializer_name) == 0)
-          continue;
+  auto add_trainable_initializers = [&](const Node* node) {
+    for (auto input : node->InputDefs()) {
+      std::string initializer_name = input->Name();
+      if (initialized_tensors.count(initializer_name) == 0)
+        continue;
 
-        if (IsUntrainable(node, initializer_name, session_logger_) ||
-            IsImmutableWeight(immutable_weights, node, initialized_tensors.at(initializer_name), session_logger_))
-          continue;
+      if (IsUntrainable(node, initializer_name, session_logger_) ||
+          IsImmutableWeight(immutable_weights, node, initialized_tensors.at(initializer_name), session_logger_))
+        continue;
 
-        trainable_initializers.insert(initializer_name);
-      }
+      trainable_initializers.insert(initializer_name);
+    }
   };
 
   auto stop_at_untrainable = [&](const Node* from, const Node* to) {
+    auto is_trainable_from_to_link = [&](Node::EdgeEnd e) {
+      if (&e.GetNode() != to)
+        return false;
 
-      auto is_trainable_from_to_link = [&](Node::EdgeEnd e) { 
-        if (&e.GetNode() != to) 
-          return false;
+      std::string input_name = from->InputDefs()[e.GetDstArgIndex()]->Name();
+      return !IsUntrainable(from, input_name, session_logger_);
+    };
 
-        std::string input_name = from->InputDefs()[e.GetDstArgIndex()]->Name();
-        return !IsUntrainable(from, input_name, session_logger_);
-      };
+    bool proceed = std::any_of(from->InputEdgesBegin(), from->InputEdgesEnd(), is_trainable_from_to_link);
+    if (!proceed && session_logger_) {
+      VLOGS(*session_logger_, 1) << "Stopping training parameters discovery traversal from " << from->Name() << " to " << to->Name() << std::endl;
+    }
 
-      bool proceed = std::any_of(from->InputEdgesBegin(), from->InputEdgesEnd(), is_trainable_from_to_link);
-      if (!proceed && session_logger_) {
-        VLOGS(*session_logger_, 1) << "Stopping training parameters discovery traversal from " << from->Name() << " to " << to->Name() << std::endl;
-      }
-
-      return !proceed;
+    return !proceed;
   };
 
   // perform reverse dfs from output node to discover trainable parameters

orttraining/orttraining/core/session/training_session.h

@@ -302,7 +302,7 @@ class TrainingSession : public InferenceSession {
    */
   std::unordered_set<std::string> GetDropoutEvalFeeds() const { return dropout_eval_feeds_; }
   /** Override Run function in InferenceSession to inject some training-specific logics **/
-  using InferenceSession::Run; // For overload resolution.
+  using InferenceSession::Run;  // For overload resolution.
   common::Status Run(const RunOptions& run_options, IOBinding& io_binding) override;
 
   common::Status Run(IOBinding& io_binding) override;
@@ -441,7 +441,7 @@ class TrainingSession : public InferenceSession {
   @param immutable_weights do not include initializers matching an (op_type, input_index, value) entry from this table
   @param backprop_source_name reverse DFS back propagation source name (i.e. loss name or pipeline send output name)
   */
-  std::unordered_set<std::string> GetTrainableModelInitializers(const ImmutableWeights& immutable_weights, 
+  std::unordered_set<std::string> GetTrainableModelInitializers(const ImmutableWeights& immutable_weights,
                                                                 const std::string& backprop_source_name) const;
 
   std::unordered_set<std::string> GetStateTensorNames() const;

orttraining/orttraining/models/bert/main.cc

@@ -331,6 +331,7 @@ Status ParseArguments(int argc, char* argv[], BertParameters& params, OrtParamet
 
     params.deepspeed_zero = ZeROConfig(flags["deepspeed_zero_stage"].as<int>());
     params.enable_grad_norm_clip = flags["enable_grad_norm_clip"].as<bool>();
+
     float alpha = flags["alpha"].as<float>();
     float beta = flags["beta"].as<float>();
     float lambda = flags["lambda"].as<float>();
@@ -507,8 +508,8 @@ float GetLossValue(const Tensor& loss_tensor) {
 // batch is not part of the initial tensor shape vector till later
 // see GetTensorDimensionsFromInputs() in training_util.h and training_runner.cc for more details
 const std::map<std::string, std::pair<std::string, size_t>> input_to_dimension_mapping = {
-  {"input1", {"SeqLen", 0}},                   // int64[batch,sequence]    "sequence" -> "SeqLen", 0
-  {"masked_lm_ids", {"PredictionsPerSeq", 0}}  // int64[batch,dynamic_prediction_count]
+    {"input1", {"SeqLen", 0}},                   // int64[batch,sequence]    "sequence" -> "SeqLen", 0
+    {"masked_lm_ids", {"PredictionsPerSeq", 0}}  // int64[batch,dynamic_prediction_count]
 };
 
 // generic properties for storing perf metrics

orttraining/orttraining/models/runner/training_runner.cc

@@ -652,23 +652,24 @@ void TrainingRunner::RunWithUpdate(VectorString& feed_names,
 
   Status status = Status::OK();
   pipeline_worker_pool_.workers[worker_id] = std::thread([&](
-      const size_t worker_id, const size_t step) {
+                                                             const size_t worker_id, const size_t step) {
 #ifdef ENABLE_NVTX_PROFILE
     // Store the tag for the thread which runs session_.Run(...).
     // It will be used to name range in Nvidia's visual profiler.
     auto& profile_context = profile::Context::GetInstance();
     profile_context.SetThreadTag(
-      std::this_thread::get_id(), std::to_string(step));
+        std::this_thread::get_id(), std::to_string(step));
 #else
     ORT_UNUSED_PARAMETER(step);
 #endif
     status = session_.Run(
-      RunOptions(),
-      pipeline_worker_pool_.worker_states[worker_id].feed_names,
-      pipeline_worker_pool_.worker_states[worker_id].feeds,
-      pipeline_worker_pool_.worker_states[worker_id].fetch_names,
-      &(pipeline_worker_pool_.worker_states[worker_id].fetches));
-  }, worker_id, step_);
+        RunOptions(),
+        pipeline_worker_pool_.worker_states[worker_id].feed_names,
+        pipeline_worker_pool_.worker_states[worker_id].feeds,
+        pipeline_worker_pool_.worker_states[worker_id].fetch_names,
+        &(pipeline_worker_pool_.worker_states[worker_id].fetches));
+  },
+                                                         worker_id, step_);
 
   // Wait all workers to finish this round of pipeline parallelism.
   // The last batch in a pipeline collects gradient and update the model.
@@ -738,26 +739,27 @@ void TrainingRunner::RunWithoutUpdate(VectorString& feed_names,
 
   // Async launch of a session.
   pipeline_worker_pool_.workers[worker_id] = std::thread([&](
-      const size_t worker_id, const size_t step) {
+                                                             const size_t worker_id, const size_t step) {
 #ifdef ENABLE_NVTX_PROFILE
     // Store the tag for the thread which runs session_.Run(...).
     // It will be used to name range in Nvidia's visual profiler.
     auto& profile_context = profile::Context::GetInstance();
     profile_context.SetThreadTag(
-      std::this_thread::get_id(), std::to_string(step));
+        std::this_thread::get_id(), std::to_string(step));
 #else
     ORT_UNUSED_PARAMETER(step);
 #endif
     RunOptions run_options;
     run_options.only_execute_path_to_fetches = true;
     auto status = session_.Run(
-      run_options,
-      pipeline_worker_pool_.worker_states[worker_id].feed_names,
-      pipeline_worker_pool_.worker_states[worker_id].feeds,
-      pipeline_worker_pool_.worker_states[worker_id].fetch_names,
-      &(pipeline_worker_pool_.worker_states[worker_id].fetches));
+        run_options,
+        pipeline_worker_pool_.worker_states[worker_id].feed_names,
+        pipeline_worker_pool_.worker_states[worker_id].feeds,
+        pipeline_worker_pool_.worker_states[worker_id].fetch_names,
+        &(pipeline_worker_pool_.worker_states[worker_id].fetches));
     ORT_THROW_IF_ERROR(status);
-  }, worker_id, step_);
+  },
+                                                         worker_id, step_);
 
   // Add one after process one batch.
   ++step_;
@@ -836,31 +838,31 @@ Status TrainingRunner::TrainingLoop(IDataLoader& training_data_loader, IDataLoad
 
         if (is_weight_update_step) {
           ORT_RETURN_IF_ERROR(PrepareFeedNamesAndFeeds(ModelUpdateStep,
-                                  training_data_loader,
-                                  *training_data,
-                                  lr_scheduler.get(),
-                                  batch,
-                                  feed_names,
-                                  feeds));
+                                                       training_data_loader,
+                                                       *training_data,
+                                                       lr_scheduler.get(),
+                                                       batch,
+                                                       feed_names,
+                                                       feeds));
           ORT_RETURN_IF_ERROR(
-            PrepareFetchNamesAndFetches(ModelUpdateStep,
-                                        fetch_names,
-                                        fetches));
+              PrepareFetchNamesAndFetches(ModelUpdateStep,
+                                          fetch_names,
+                                          fetches));
           RunWithUpdate(feed_names, fetch_names, feeds, fetches);
         } else {
           ORT_RETURN_IF_ERROR(PrepareFeedNamesAndFeeds(GradientAccumulateStep,
-                                  training_data_loader,
-                                  *training_data,
-                                  lr_scheduler.get(),
-                                  batch,
-                                  feed_names,
-                                  feeds));
+                                                       training_data_loader,
+                                                       *training_data,
+                                                       lr_scheduler.get(),
+                                                       batch,
+                                                       feed_names,
+                                                       feeds));
           ORT_RETURN_IF_ERROR(
-            PrepareFetchNamesAndFetches(GradientAccumulateStep,
-                                        fetch_names,
-                                        fetches));
+              PrepareFetchNamesAndFetches(GradientAccumulateStep,
+                                          fetch_names,
+                                          fetches));
           RunWithoutUpdate(feed_names, fetch_names, feeds,
-                            gradient_accumulation_step_count);
+                           gradient_accumulation_step_count);
         }
 
         // at this point, step_ already be increased by 1.
@@ -966,7 +968,7 @@ Status TrainingRunner::TrainingLoop(IDataLoader& training_data_loader, IDataLoad
             << "Throughput: " << throughput << " Examples / Second\n"
             << "Stabilized Throughput: " << stabilized_throughput << " Examples / Second\n"
             << "EndToEnd Throughput: " << e2e_throughput << " Examples / Second\n"
-            << "Average Step Time: " << all_steps_duration_seconds.count() / (step_ - step_start)<< " Second\n"
+            << "Average Step Time: " << all_steps_duration_seconds.count() / (step_ - step_start) << " Second\n"
             << "Average Step Throughput: " << params_.batch_size * (step_ - step_start) / (all_steps_duration_seconds.count()) << " Examples / Second\n";
 
   return Status::OK();
@@ -1158,11 +1160,11 @@ Status TrainingRunner::Evaluate(InferenceSession& session, IDataLoader& data_loa
         RunOptions run_options;
         run_options.only_execute_path_to_fetches = true;
         status = session.Run(
-          run_options,
-          feed_names,
-          feeds,
-          fetch_names,
-          &fetches);
+            run_options,
+            feed_names,
+            feeds,
+            fetch_names,
+            &fetches);
       });
       // Wait Run(...) to finish.
       pipeline_worker_pool_.Join(worker_id);
@@ -1172,10 +1174,10 @@ Status TrainingRunner::Evaluate(InferenceSession& session, IDataLoader& data_loa
       // Pipeline cannot reuse training threads to do evaluation.
       // Otherwise, deadlock may happens.
       ORT_RETURN_IF_ERROR(session.Run(run_options,
-        feed_names,
-        feeds,
-        fetch_names,
-        &fetches));
+                                      feed_names,
+                                      feeds,
+                                      fetch_names,
+                                      &fetches));
     }
 
     // Assume that user-specified fetches are avaliable only on the last pipeline stage.

orttraining/orttraining/models/runner/training_runner.h

@@ -30,9 +30,9 @@ class TrainingRunner {
 
     PathString train_data_dir;
     PathString test_data_dir;
-    PathString output_dir;  // Output of training, e.g., trained model files.
-    PathString perf_output_dir; // training perf metrics
-    std::string model_type; // bert/gpt2/...
+    PathString output_dir;       // Output of training, e.g., trained model files.
+    PathString perf_output_dir;  // training perf metrics
+    std::string model_type;      // bert/gpt2/...
 
     LossFunctionInfo loss_func_info;
 
@@ -180,7 +180,7 @@ class TrainingRunner {
   common::Status Initialize();
 
   common::Status Run(IDataLoader* training_data_loader, IDataLoader* test_data_loader,
-    const MapStringToString& mapped_dimensions = {});
+                     const MapStringToString& mapped_dimensions = {});
 
   common::Status EndTraining(IDataLoader* data_loader);
 
@@ -192,7 +192,9 @@ class TrainingRunner {
   TrainingSession& GetSession() { return session_; }
 
  private:
-  enum SessionMode: int {ModelUpdateStep, GradientAccumulateStep, EvaluateStep};
+  enum SessionMode : int { ModelUpdateStep,
+                           GradientAccumulateStep,
+                           EvaluateStep };
   Status PrepareFeedNamesAndFeeds(const SessionMode mode,
                                   IDataLoader& training_data_loader,
                                   DataSet& training_data,
@@ -212,7 +214,7 @@ class TrainingRunner {
                         std::vector<MLValue>& feeds,
                         size_t& gradient_accumulation_step_count);
   Status TrainingLoop(IDataLoader& training_data_loader, IDataLoader* test_data_loader,
-    const MapStringToString& mapped_dimensions);
+                      const MapStringToString& mapped_dimensions);
   Status Evaluate(InferenceSession& session, IDataLoader& data_loader);
 
   Status SaveCheckpoint(const PathString& checkpoint_path);