// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <thread>

#include "gtest/gtest.h"
#include "orttraining/core/optimizer/gist_encode_decode.h"
#include "test/providers/provider_test_utils.h"
#include "test/framework/test_utils.h"
#include "test/util/include/default_providers.h"
#include "core/common/path_utils.h"
#include "core/providers/cpu/cpu_execution_provider.h"
#include "core/session/environment.h"
#include "orttraining/core/framework/distributed_run_context.h"
#include "orttraining/models/runner/training_runner.h"
#include "orttraining/test/session/training_session_test_utils.h"

#include "orttraining/training_ops/cpu/controlflow/event_pool.h"  // TODO: move with PipelineBatchPlanner

#if defined(USE_CUDA) || defined(USE_ROCM)
#include "bert_toy_fetches.h"
#endif

using namespace onnxruntime::logging;
using namespace onnxruntime::training;
using namespace onnxruntime::path_utils;
using namespace onnxruntime::test::training_session_test_utils;

namespace onnxruntime {
namespace test {

namespace {
constexpr auto CONCAT_MODEL_PATH = ORT_TSTR("testdata/transform/concat_trainable.onnx");
}  // namespace

static Status BuildBackPropGraph(
    const PathString& forward_model_file,
    const TrainingSession::TrainingConfiguration& config,
    PathString& backward_model_file) {
  std::unique_ptr<Environment> env;
  ORT_RETURN_IF_ERROR(Environment::Create(nullptr, env));

  SessionOptions so{};
  TrainingSession training_session{so, *env};

  std::cout << "Loading source model file = " << ToUTF8String(forward_model_file) << "\n";

  ORT_RETURN_IF_ERROR(training_session.Load(forward_model_file));

  TrainingSession::TrainingConfigurationResult config_result{};
  ORT_RETURN_IF_ERROR(training_session.ConfigureForTraining(config, config_result));

  backward_model_file = config.model_with_training_graph_path.value();

  return Status::OK();
}

/**
 * Run a training session for this model for 1 epoch, using batch size of 1 and synthetic input data.
 * @param so - SessionOptions for this run.
 * @param backprop_model_file - Model file to be run. This should already contain loss function and backward prop nodes.
 * @return TrainingSession for this run.
 */
static std::unique_ptr<TrainingSession> RunTrainingSessionWithChecks(
    const SessionOptions& so, const PathString& backprop_model_file) {
  std::unique_ptr<Environment> env;
  ORT_THROW_IF_ERROR(Environment::Create(nullptr, env));

  std::unique_ptr<TrainingSession> training_session = std::make_unique<TrainingSession>(so, *env);

  ORT_THROW_IF_ERROR(training_session->Load(backprop_model_file));

  std::pair<common::Status, const ModelMetadata*> res = training_session->GetModelMetadata();
  ORT_THROW_IF_ERROR(res.first);
  ORT_ENFORCE(res.second != nullptr);
  auto model_metadata = res.second;
  std::cout << "Loaded " << model_metadata->graph_name << '\n';

  ORT_THROW_IF_ERROR(training_session->Initialize());

  std::vector<OrtValue> gradient_fetches;
  RunOptions run_options;
  run_options.run_log_verbosity_level = so.session_log_verbosity_level;
  run_options.run_tag = so.session_logid;
  run_options.training_mode = true;

  // Create dummy feeds
  std::vector<int64_t> image_dims = {1, 784};
  std::vector<int64_t> label_dims = {1, 10};
  std::vector<float> image_value(784, 1);
  std::vector<float> label_value(10, 1);

  OrtValue imageMLValue;
  TrainingUtil::CreateCpuMLValue(image_dims, image_value, &imageMLValue);
  OrtValue labelMLValue;
  TrainingUtil::CreateCpuMLValue(label_dims, label_value, &labelMLValue);

  auto fw_feeds = std::make_pair<std::vector<std::string>, std::vector<OrtValue>>({"X", "labels"}, {imageMLValue, labelMLValue});

  auto output_names_include_gradients = GetModelOutputNames(*training_session);
  std::vector<std::string> training_output_names(output_names_include_gradients.begin(), output_names_include_gradients.end());

  auto start_time = std::chrono::high_resolution_clock::now();

  ORT_THROW_IF_ERROR(training_session->Run(run_options, fw_feeds.first, fw_feeds.second, training_output_names, &gradient_fetches));

  auto end_time = std::chrono::high_resolution_clock::now();
  auto elapsed = TimeDiffMicroSeconds(start_time, end_time);
  std::cout << "Training session run completed in " << elapsed << " microseconds.\n";

  return training_session;
}

TEST(GradientGraphBuilderTest, BuildGradientGraphTest) {
  const auto config = MakeBasicTrainingConfig();
  PathString backprop_model_file;
  ASSERT_STATUS_OK(BuildBackPropGraph(ORIGINAL_MODEL_PATH, config, backprop_model_file));

  std::shared_ptr<Model> pModel;
  ASSERT_STATUS_OK(Model::Load(backprop_model_file, pModel, nullptr, DefaultLoggingManager().DefaultLogger()));

  Graph& graph = pModel->MainGraph();
  EXPECT_FALSE(graph.GraphResolveNeeded());
  EXPECT_TRUE(graph.NumberOfNodes() > 0);
  EXPECT_TRUE(graph.MaxNodeIndex() > 0);

  std::cout << "Graph input names = [\n";
  for (const NodeArg* p_node_arg : graph.GetInputs()) {
    std::cout << '\t' << p_node_arg->Name() << '\n';
  }
  std::cout << "]\n";

  std::cout << "Graph output names = [\n";
  for (const NodeArg* p_node_arg : graph.GetOutputs()) {
    std::cout << '\t' << p_node_arg->Name() << '\n';
  }
  std::cout << "]\n";

  for (Node& node : graph.Nodes()) {
    const NodeIndex node_index = node.Index();
    const std::string& node_name = node.Name();
    const std::string& op_type = node.OpType();

    std::cout << "Operation node:"
              << " Index=" << node_index
              << (node.NodeType() == Node::Type::Fused ? "-(FUSED)" : "")
              << " OpType=" << op_type
              << " Name=" << node_name
              << '\n';
  }
}

TEST(GradientGraphBuilderTest, BuildConcatGradientGraphTest) {
  const auto config = MakeBasicTrainingConfig();
  PathString backprop_model_file;
  ASSERT_STATUS_OK(BuildBackPropGraph(CONCAT_MODEL_PATH, config, backprop_model_file));

  std::shared_ptr<Model> pModel;
  ASSERT_STATUS_OK(Model::Load(backprop_model_file, pModel, nullptr, DefaultLoggingManager().DefaultLogger()));

  Graph& graph = pModel->MainGraph();
  EXPECT_FALSE(graph.GraphResolveNeeded());
  EXPECT_TRUE(graph.NumberOfNodes() > 0);
  EXPECT_TRUE(graph.MaxNodeIndex() > 0);

  std::map<std::string, int> op_to_count = CountOpsInGraph(graph);

  ASSERT_EQ(op_to_count["Concat"], 0);
  ASSERT_EQ(op_to_count["Split"], 0);
  ASSERT_EQ(op_to_count["com.microsoft.ConcatTraining"], 1);
  ASSERT_EQ(op_to_count["com.microsoft.SplitTraining"], 1);
}

TEST(GradientGraphBuilderTest, TrainingSession_Basic) {
  const auto config = MakeBasicTrainingConfig();
  PathString backprop_model_file;
  ASSERT_STATUS_OK(BuildBackPropGraph(ORIGINAL_MODEL_PATH, config, backprop_model_file));

  SessionOptions so{};
  RunTrainingSessionWithChecks(so, backprop_model_file);
}

TEST(GradientGraphBuilderTest, GraphTransformation_WithGist) {
  // Setup training session configuration
  auto config = MakeBasicTrainingConfig();
  constexpr int op_type_max = 9;
  const vector<std::string> compr_type_vec = {"GistBinarize", "GistPack8", "GistPack16", "GistPackMsfp15"};

  PathString backprop_model_file;
  for (auto& compr_type : compr_type_vec) {
    // Add GIST config to training session (op_type_max ensures GIST is applied to all applicable ops)
    TrainingSession::TrainingConfiguration::GistConfiguration gist{};
    gist.op_type = op_type_max;
    gist.compr_type = compr_type;
    config.gist_config = gist;

    // Create backward graph with gist transformations
    ASSERT_STATUS_OK(BuildBackPropGraph(ORIGINAL_MODEL_PATH, config, backprop_model_file));

    // Check correctness of GIST transformation
    backprop_model_file = config.model_with_training_graph_path.value();
    std::shared_ptr<Model> pModel;
    ASSERT_STATUS_OK(Model::Load(backprop_model_file, pModel, nullptr, DefaultLoggingManager().DefaultLogger()));
    Graph& graph = pModel->MainGraph();

    std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
    std::cout << "Type: "
              << "com.microsoft." + gist.compr_type + "Encoder"
              << ", Count: " << op_to_count["com.microsoft." + gist.compr_type + "Encoder"] << std::endl;
    ASSERT_TRUE(op_to_count["com.microsoft.GistPack1Encoder"] == op_to_count["com.microsoft.GistPack1Decoder"]);
    ASSERT_TRUE(op_to_count["com.microsoft." + gist.compr_type + "Encoder"] == op_to_count["com.microsoft." + gist.compr_type + "Decoder"]);
    ASSERT_TRUE(op_to_count["com.microsoft." + gist.compr_type + "Encoder"] + op_to_count["com.microsoft.GistPack1Encoder"] > 0);
  }
}

#ifdef USE_CUDA
TEST(GradientGraphBuilderTest, TrainingSession_WithGist) {
  // Setup training session configuration including GIST config (op_flag 9 ensures GIST will be applied to all possible supported node types)
  auto config = MakeBasicTrainingConfig();
  TrainingSession::TrainingConfiguration::GistConfiguration gist{};
  gist.op_type = 9;  // Apply Gist to all applicable operator types
  gist.compr_type = "GistPack8";
  config.gist_config = gist;

  // Create backward graph with gist transformations
  const PathString& forward_model_file = ORIGINAL_MODEL_PATH;
  std::unique_ptr<Environment> env;
  ORT_THROW_IF_ERROR(Environment::Create(nullptr, env));

  SessionOptions so{};
  TrainingSession training_session{so, *env};

  std::cout << "Loading source model file = " << ToUTF8String(forward_model_file) << "\n";

  ORT_THROW_IF_ERROR(training_session.Load(forward_model_file));

  TrainingSession::TrainingConfigurationResult config_result{};
  ORT_THROW_IF_ERROR(training_session.ConfigureForTraining(config, config_result));

  // Check correctness of GIST transformation
  PathString backprop_model_file = config.model_with_training_graph_path.value();
  std::shared_ptr<Model> pModel;
  ASSERT_STATUS_OK(Model::Load(backprop_model_file, pModel, nullptr, DefaultLoggingManager().DefaultLogger()));
  Graph& graph = pModel->MainGraph();

  std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
  std::cout << "Type: "
            << "com.microsoft." + gist.compr_type + "Encoder"
            << ", Count: " << op_to_count["com.microsoft." + gist.compr_type + "Encoder"] << std::endl;
  ASSERT_TRUE(op_to_count["com.microsoft.GistPack1Encoder"] == op_to_count["com.microsoft.GistPack1Decoder"]);
  ASSERT_TRUE(op_to_count["com.microsoft." + gist.compr_type + "Encoder"] == op_to_count["com.microsoft." + gist.compr_type + "Decoder"]);
  ASSERT_TRUE(op_to_count["com.microsoft." + gist.compr_type + "Encoder"] + op_to_count["com.microsoft.GistPack1Encoder"] > 0);

  // Run training session with checks
  std::pair<common::Status, const ModelMetadata*> res = training_session.GetModelMetadata();
  ORT_THROW_IF_ERROR(res.first);
  ORT_ENFORCE(res.second != nullptr);
  auto model_metadata = res.second;
  std::cout << "Loaded " << model_metadata->graph_name << '\n';

  // Add cuda execution provider for gist encode/decode nodes
  ASSERT_STATUS_OK(training_session.RegisterExecutionProvider(DefaultCudaExecutionProvider()));

  ORT_THROW_IF_ERROR(training_session.Initialize());

  std::vector<OrtValue> gradient_fetches;
  RunOptions run_options;
  run_options.run_log_verbosity_level = so.session_log_verbosity_level;
  run_options.run_tag = so.session_logid;
  run_options.training_mode = true;

  // Create dummy feeds
  std::vector<int64_t> image_dims = {1, 784};
  std::vector<int64_t> label_dims = {1, 10};
  std::vector<float> image_value(784, 1);
  std::vector<float> label_value(10, 1);

  OrtValue imageMLValue;
  TrainingUtil::CreateCpuMLValue(image_dims, image_value, &imageMLValue);
  OrtValue labelMLValue;
  TrainingUtil::CreateCpuMLValue(label_dims, label_value, &labelMLValue);

  auto fw_feeds = std::make_pair<std::vector<std::string>, std::vector<OrtValue>>({"X", "labels"}, {imageMLValue, labelMLValue});

  auto output_names_include_gradients = GetModelOutputNames(training_session);
  std::vector<std::string> training_output_names(output_names_include_gradients.begin(), output_names_include_gradients.end());

  auto start_time = std::chrono::high_resolution_clock::now();

  ORT_THROW_IF_ERROR(training_session.Run(run_options, fw_feeds.first, fw_feeds.second, training_output_names, &gradient_fetches));

  auto end_time = std::chrono::high_resolution_clock::now();
  auto elapsed = TimeDiffMicroSeconds(start_time, end_time);
  std::cout << "Training session run completed in " << elapsed << " microseconds.\n";
}
#endif

TEST(GradientGraphBuilderTest, TrainingSession_WithLogging) {
  const auto& log_manager = DefaultLoggingManager();
  const auto& default_logger = log_manager.DefaultLogger();
  log_manager.SetDefaultLoggerSeverity(Severity::kINFO);

  EXPECT_TRUE(default_logger.OutputIsEnabled(Severity::kERROR, ::onnxruntime::logging::DataType::USER)) << "ERROR level logging enabled.";
  EXPECT_TRUE(default_logger.OutputIsEnabled(Severity::kWARNING, ::onnxruntime::logging::DataType::USER)) << "WARNING level logging enabled.";
  EXPECT_TRUE(default_logger.OutputIsEnabled(Severity::kINFO, ::onnxruntime::logging::DataType::USER)) << "INFO level logging enabled.";

  const auto config = MakeBasicTrainingConfig();
  PathString backprop_model_file;
  ASSERT_STATUS_OK(BuildBackPropGraph(ORIGINAL_MODEL_PATH, config, backprop_model_file));

  SessionOptions so;
  so.session_logid = "training_session_with_logging";
  so.session_log_verbosity_level = 1;  // 1 == detailed logging

  std::unique_ptr<TrainingSession> training_session = RunTrainingSessionWithChecks(so, backprop_model_file);

  EXPECT_TRUE(default_logger.OutputIsEnabled(Severity::kERROR, ::onnxruntime::logging::DataType::USER)) << "ERROR level logging still enabled.";
  EXPECT_TRUE(default_logger.OutputIsEnabled(Severity::kWARNING, ::onnxruntime::logging::DataType::USER)) << "WARNING level logging still enabled.";
  EXPECT_TRUE(default_logger.OutputIsEnabled(Severity::kINFO, ::onnxruntime::logging::DataType::USER)) << "INFO level logging still enabled.";

  std::string profile_file = training_session->EndProfiling();

  log_manager.SetDefaultLoggerSeverity(Severity::kWARNING);

  EXPECT_EQ(profile_file, std::string()) << "There should be no profile output file.";
}

TEST(GradientGraphBuilderTest, TrainingSession_WithProfiler) {
  const auto config = MakeBasicTrainingConfig();
  PathString backprop_model_file;
  ASSERT_STATUS_OK(BuildBackPropGraph(ORIGINAL_MODEL_PATH, config, backprop_model_file));

  SessionOptions so;
  so.enable_profiling = true;
  so.profile_file_prefix = ORT_TSTR("onnx_training_profiler_test");

  std::unique_ptr<TrainingSession> training_session = RunTrainingSessionWithChecks(so, backprop_model_file);

  std::string profile_file = training_session->EndProfiling();

  std::cout << "Profile output file = " << profile_file << '\n';

  std::ifstream profile(profile_file);
  ASSERT_TRUE(profile);

  std::vector<std::string> core_trace_fields = {"pid", "dur", "ts", "ph", "X", "name", "args"};
  std::vector<std::string> fiddle_profile_data_fields = {"dur", "activation_size", "parameter_size", "output_size"};

  int count = 0;
  std::string line;
  while (std::getline(profile, line)) {
    if (count == 0) {
      ASSERT_TRUE(line.find('[') != std::string::npos)
          << "Missing opening array marker in first trace record: " << line;
      // Opening array marker found.
    } else if (line.find(']') != std::string::npos) {
      // Closing array marker found.
      break;
    } else if (count >= 1) {
      if (count == 1) {
        auto s = "model_loading_uri";
        ASSERT_TRUE(line.find(s) != std::string::npos)
            << "Missing field '" << s << "' in trace record: " << line;
      }

      // Check we have the core fields in each trace record.
      for (auto& s : core_trace_fields) {
        ASSERT_TRUE(line.find(s) != std::string::npos)
            << "Missing core trace field '" << s << "' in trace record: " << line;
      }

      // Check we have the data profile fields output for each kernel operation.
      if (line.find("_kernel_time") != std::string::npos) {
        for (auto& s : fiddle_profile_data_fields) {
          ASSERT_TRUE(line.find(s) != std::string::npos)
              << "Missing data profile field '" << s << "' in trace record: " << line;
        }
      }
    }

    count++;
  }
  ASSERT_TRUE(count > 1);
}

#if defined(USE_CUDA) || defined(USE_ROCM)
static void RunBertTrainingWithChecks(
    const SessionOptions& so,
    const PathString& backprop_model_file) {
  std::unique_ptr<Environment> env;
  ASSERT_STATUS_OK(Environment::Create(nullptr, env));

  std::unique_ptr<TrainingSession> training_session = std::make_unique<TrainingSession>(so, *env);

  ASSERT_STATUS_OK(training_session->Load(backprop_model_file));

  std::pair<common::Status, const ModelMetadata*> res = training_session->GetModelMetadata();
  ASSERT_STATUS_OK(res.first);
  ASSERT_TRUE(res.second != nullptr);
  auto model_metadata = res.second;
  std::cout << "Loaded " << model_metadata->graph_name << '\n';

#ifdef USE_CUDA
  ASSERT_STATUS_OK(training_session->RegisterExecutionProvider(DefaultCudaExecutionProvider()));
#elif USE_ROCM
  ASSERT_STATUS_OK(training_session->RegisterExecutionProvider(DefaultRocmExecutionProvider()));
#endif
  ASSERT_STATUS_OK(training_session->Initialize());

  RunOptions run_options;
  run_options.run_log_verbosity_level = so.session_log_verbosity_level;
  run_options.run_tag = so.session_logid;
  run_options.training_mode = true;

  // Creating feeds
  int batch_size = 13;
  int max_seq_len_in_batch = 7;
  std::vector<std::string> feed_names = {
      "input_ids",
      "token_type_ids",
      "input_mask",
      "masked_lm_ids",
      "next_sentence_labels",
      "masked_lm_positions",
  };
  std::vector<TensorShape> tensor_shapes = {
      {batch_size, max_seq_len_in_batch},
      {batch_size, max_seq_len_in_batch},
      {batch_size, max_seq_len_in_batch},
      {batch_size, max_seq_len_in_batch},
      {batch_size},
      {batch_size, max_seq_len_in_batch},
      {batch_size, max_seq_len_in_batch}};

  std::vector<std::vector<int64_t>> tensor_values = {
      /*input_ids*/
      {49, 97, 53, 5, 33, 65, 62,
       51, 38, 61, 45, 74, 27, 64,
       17, 36, 17, 96, 12, 79, 32,
       68, 90, 77, 18, 39, 12, 93,
       9, 87, 42, 60, 71, 12, 45,
       55, 40, 78, 81, 26, 70, 61,
       56, 66, 33, 7, 70, 1, 11,
       92, 51, 90, 85, 80, 0, 78,
       63, 42, 31, 93, 41, 90, 8,
       24, 72, 28, 30, 18, 69, 57,
       11, 10, 40, 65, 62, 13, 38,
       70, 37, 90, 15, 70, 42, 69,
       26, 77, 70, 75, 36, 56, 11},
      /*token_type_ids*/
      {12, 13, 1, 8, 15, 12, 9,
       15, 11, 6, 4, 9, 4, 3,
       8, 4, 9, 3, 2, 10, 15,
       3, 11, 13, 10, 6, 15, 14,
       8, 1, 0, 2, 12, 0, 15,
       10, 7, 10, 2, 6, 7, 7,
       4, 14, 2, 2, 10, 15, 3,
       9, 9, 3, 10, 6, 9, 14,
       2, 12, 10, 7, 9, 5, 6,
       5, 1, 8, 15, 2, 2, 4,
       4, 1, 2, 12, 8, 7, 6,
       13, 8, 14, 15, 11, 2, 10,
       3, 15, 10, 6, 7, 0, 8},
      /*input_mask*/
      {1, 1, 0, 1, 1, 1, 1,
       1, 1, 0, 0, 1, 0, 0,
       1, 0, 1, 0, 0, 1, 1,
       0, 1, 1, 1, 0, 1, 1,
       1, 0, 0, 0, 1, 0, 1,
       1, 0, 1, 0, 0, 0, 0,
       0, 1, 0, 0, 1, 1, 0,
       1, 1, 0, 1, 0, 1, 1,
       0, 1, 1, 0, 1, 0, 0,
       0, 0, 1, 1, 0, 0, 0,
       0, 0, 0, 1, 1, 0, 0,
       1, 1, 1, 1, 1, 0, 1,
       0, 1, 1, 0, 0, 0, 1},
      /*masked_lm_ids*/
      {1, 1, 0, 1, 2, 1, 1,
       1, 1, 1, 2, 0, 2, 0,
       1, 0, 0, 2, 1, 2, 2,
       2, 0, 1, 0, 2, 0, 2,
       1, 1, 2, 0, 1, 1, 1,
       2, 2, 0, 2, 1, 1, 2,
       1, 0, 2, 0, 0, 2, 1,
       2, 2, 2, 0, 2, 1, 1,
       0, 2, 1, 2, 0, 0, 2,
       0, 0, 0, 2, 1, 0, 0,
       1, 2, 1, 0, 1, 2, 1,
       2, 0, 2, 1, 2, 0, 2,
       2, 2, 1, 1, 0, 2, 1},
      /*next_sentence_labels*/
      {1, 1, 0, 1, 1, 1, 1, 1, 1, 0, 0, 1, 0},
      /*masked_lm_positions*/
      {0, 1, 2, 3, 4, 5, 6,
       0, 1, 2, 3, 4, 5, 6,
       0, 1, 2, 3, 4, 5, 6,
       0, 1, 2, 3, 4, 5, 6,
       0, 1, 2, 3, 4, 5, 6,
       0, 1, 2, 3, 4, 5, 6,
       0, 1, 2, 3, 4, 5, 6,
       0, 1, 2, 3, 4, 5, 6,
       0, 1, 2, 3, 4, 5, 6,
       0, 1, 2, 3, 4, 5, 6,
       0, 1, 2, 3, 4, 5, 6,
       0, 1, 2, 3, 4, 5, 6,
       0, 1, 2, 3, 4, 5, 6}};

  std::vector<OrtValue> feeds(feed_names.size());
  for (size_t i = 0; i < 6; ++i) {
    TrainingUtil::CreateCpuMLValue(tensor_shapes[i].GetDims(), tensor_values[i], &feeds[i]);
  }

  auto output_names_include_gradients = GetModelOutputNames(*training_session);
  std::vector<std::string> fetch_names(output_names_include_gradients.begin(), output_names_include_gradients.end());

  std::vector<OrtValue> fetches;

  ASSERT_STATUS_OK(training_session->Run(run_options, feed_names, feeds, fetch_names, &fetches));

  for (size_t i = 0; i < fetch_names.size(); ++i) {
    if (!fetches[i].IsAllocated() || !!fetches[i].IsTensor())
      continue;

    const Tensor& tensor = fetches[i].Get<Tensor>();
    if (DataTypeImpl::GetType<float>() != tensor.DataType()) {
      continue;
    }

    const std::string& name = fetch_names[i];
    if (BERT_TOY_FETCHES.find(name) == BERT_TOY_FETCHES.end()) {
      continue;
    }

    auto gradient_ref = BERT_TOY_FETCHES.at(name);
    EXPECT_TRUE(static_cast<size_t>(tensor.Shape().Size()) == gradient_ref.size());

    float max_diff = 0;
    float max_percent_diff = 0;
    const float* data = tensor.template Data<float>();
    for (size_t idx = 0; idx < gradient_ref.size(); ++idx) {
      float diff = std::fabs(static_cast<float>(gradient_ref[idx]) - data[idx]);
      max_diff = std::fmax(max_diff, diff);
      max_percent_diff = std::fmax(max_percent_diff, diff / data[idx]);
    }
    EXPECT_TRUE(max_diff < 1e-5) << name << " is incorrect: max_diff is " << max_diff;
    if (max_diff > 1e-10) {
      EXPECT_TRUE(max_percent_diff < 0.01f) << name << " is incorrect: max_percent_diff is "
                                            << max_percent_diff;
    }
  }
}
#endif
TEST(GradientGraphBuilderTest, TrainingSession_BertToy) {
  const auto model_path = ORT_TSTR("testdata/bert_toy_optimized.onnx");

  TrainingSession::TrainingConfiguration config{};
  config.model_with_training_graph_path = ORT_TSTR("testdata/bert_toy_optimized_bw.onnx");
  config.loss_function_config = TrainingSession::TrainingConfiguration::LossFunctionConfiguration{};
  config.loss_function_config.value().loss_function_info =
      LossFunctionInfo(OpDef("BertLoss", kOnnxDomain),
                       "total_loss",
                       {/*prediction_masked_lm*/ "prediction_scores",
                        /*prediction_next_sentence*/ "seq_relationship_score",
                        /*masked_lm_positions*/ "masked_lm_positions",
                        /*masked_lm_ids*/ "masked_lm_ids",
                        /*next_sentence_labels*/ "next_sentence_labels",
                        /*mlm_loss*/ "mlm_loss",
                        /*nsp_loss*/ "nsp_loss"});
  config.weight_names_to_not_train = {
      "position_01",            // Slice's dat input
      "op_min_ends_expand_10",  // op_min_ends_expand_10
  };
  config.immutable_weights = {
      {"Div", {{1, 8.0f}, {1, 1.4142135381698608f}}},
      {"Add", {{1, 1.0f}, {1, 9.999999960041972e-13f}}},
      {"Mul", {{1, 0.5f}, {1, -10000.0f}}},
      {"Sub", {{0, 1.0f}}}};

  PathString backprop_model_file;
  ASSERT_STATUS_OK(BuildBackPropGraph(model_path, config, backprop_model_file));

#if defined(USE_CUDA) || defined(USE_ROCM)
  SessionOptions so;
  RunBertTrainingWithChecks(so, backprop_model_file);
#endif
}

class PipelineSplitter {
 public:
  struct UnidirectionCutInfo {
    // nodes are identified by its output[0]
    std::vector<std::string> nodes;
    // inputs for sync between sub models
    std::vector<std::string> sync_inputs;
    // outputs for sync between sub models
    // note there might be some graph ouputs do not need to sync
    std::vector<std::string> sync_outputs;
    // dependencies for maintaining topological order
    std::vector<std::string> wait_depends;
    std::vector<std::string> record_depends;
  };

  struct CutInfo {
    UnidirectionCutInfo fw;
    UnidirectionCutInfo bw;
  };

  PipelineSplitter() = default;

  void Split(
      PathString backprop_model_file,
      const std::vector<PathString>& sub_model_files,
      const std::vector<CutInfo>& cuts) {
    const auto num_subs = cuts.size();

    std::shared_ptr<Model> model;
    ASSERT_STATUS_OK(Model::Load(
        backprop_model_file, model, nullptr, DefaultLoggingManager().DefaultLogger()));

    const auto& main_graph = model->MainGraph();
    const auto mp = model->ToProto();
    const auto& main_gp = mp.graph();

    auto lookup_main_graph_node_arg_proto =
        [&main_graph](const std::string& node_arg_name) -> const ONNX_NAMESPACE::ValueInfoProto* {
      const auto* node_arg = main_graph.GetNodeArg(node_arg_name);
      if (!node_arg) return nullptr;
      return &node_arg->ToProto();
    };

    std::vector<ONNX_NAMESPACE::ModelProto> sub_mps(num_subs, mp);
    for (size_t i = 0; i < num_subs; ++i) {
      auto& sub = sub_mps[i];
      sub.clear_graph();
      FillInputWait(sub.mutable_graph(),
                    main_gp,
                    lookup_main_graph_node_arg_proto,
                    cuts[i].fw.sync_inputs,
                    cuts[i].fw.wait_depends,
                    i,
                    /*bw*/ false);
    }

    for (const auto& n : main_gp.node()) {
      // check which sub_model the node should be in
      const size_t sub_id = [&]() {
        for (size_t i = 0; i < num_subs; ++i) {
          const auto& cut = cuts[i];
          if (std::count(cut.fw.nodes.cbegin(), cut.fw.nodes.cend(), n.output()[0])) {
            return i;
          }
          if (std::count(cut.bw.nodes.cbegin(), cut.bw.nodes.cend(), n.output()[0])) {
            return i;
          }
        }
        ORT_THROW("Failed to find sub-model containing node: ", n.name());
      }();

      auto* sub_gp = sub_mps[sub_id].mutable_graph();
      const auto& cut = cuts[sub_id];

      // add WaitEvent node at the beginning of bw
      if (!cut.bw.nodes.empty() && n.output()[0] == cut.bw.nodes.front()) {
        FillInputWait(sub_gp,
                      main_gp,
                      lookup_main_graph_node_arg_proto,
                      cut.bw.sync_inputs,
                      cut.bw.wait_depends,
                      sub_id,
                      /*bw*/ true);
      }

      // copy node to sub model
      sub_gp->mutable_node()->Add()->CopyFrom(n);
      for (auto i = n.input().cbegin(); i != n.input().cend(); ++i) {
        AddItemByName(sub_gp->mutable_initializer(), main_gp.initializer(), *i, *i);
        if (0 == std::count(cut.fw.sync_inputs.cbegin(), cut.fw.sync_inputs.cend(), *i) &&
            0 == std::count(cut.bw.sync_inputs.cbegin(), cut.bw.sync_inputs.cend(), *i)) {
          // carry over original graph's input, if not in sync_inputs
          AddItemByName(sub_gp->mutable_input(), main_gp.input(), *i, *i);
        }
      }
      for (auto i = n.output().cbegin(); i != n.output().cend(); ++i) {
        if (std::count(cut.fw.sync_outputs.cbegin(), cut.fw.sync_outputs.cend(), *i) ||
            std::count(cut.bw.sync_outputs.cbegin(), cut.bw.sync_outputs.cend(), *i))
          continue;  // sync_ouputs already handled, skip

        // add graph output
        if (!AddItemByName(sub_gp->mutable_output(), main_gp.output(), *i, *i)) {
          // for non-output, add shape info
          AddValueInfoByNameFromLookup(*sub_gp->mutable_value_info(),
                                       lookup_main_graph_node_arg_proto,
                                       *i,
                                       *i);
        }
      }

      // add RecordEvent node at the end of fw and bw
      if ((!cut.fw.nodes.empty() && n.output()[0] == cut.fw.nodes.back()) ||
          (!cut.bw.nodes.empty() && n.output()[0] == cut.bw.nodes.back())) {
        bool bw = (n.output()[0] == cut.bw.nodes.back());
        const auto& sync_outputs = (bw ? cut.bw.sync_outputs : cut.fw.sync_outputs);
        const auto& dependencies = (bw ? cut.bw.record_depends : cut.fw.record_depends);

        FillOutputRecord(sub_gp,
                         main_gp,
                         lookup_main_graph_node_arg_proto,
                         sync_outputs,
                         dependencies,
                         sub_id,
                         bw);
      }
    }

    // save sub models
    for (size_t sub_id = 0; sub_id < num_subs; ++sub_id) {
      std::ofstream ofs(sub_model_files[sub_id], std::ofstream::binary);
      sub_mps[sub_id].SerializeToOstream(&ofs);
      ofs.close();
    }
  }

 private:
  // add RepeatedField item by name from another RepeatedFields
  // return true if the name exists in dst
  template <typename TD, typename TS>
  bool AddItemByName(TD* dst, const TS& src, const std::string& name, const std::string& new_name) {
    for (auto iter = dst->cbegin(); iter != dst->cend(); ++iter) {
      if (iter->name() == new_name) {
        return true;
      }
    }
    for (auto iter = src.cbegin(); iter != src.cend(); ++iter) {
      if (iter->name() == name) {
        auto* p = dst->Add();
        p->CopyFrom(*iter);
        *p->mutable_name() = new_name;
        return true;
      }
    }
    return false;
  }

  // expected signature of TValueInfoLookupFn:
  //   const ValueInfoProto* TValueInfoLookupFn(const std::string& name)
  template <typename TValueInfoLookupFn>
  bool AddValueInfoByNameFromLookup(
      google::protobuf::RepeatedPtrField<ONNX_NAMESPACE::ValueInfoProto>& dst, TValueInfoLookupFn lookup,
      const std::string& name, const std::string& new_name) {
    for (auto iter = dst.cbegin(); iter != dst.cend(); ++iter) {
      if (iter->name() == new_name) {
        return true;
      }
    }

    const ONNX_NAMESPACE::ValueInfoProto* value_info = lookup(name);
    if (value_info) {
      auto* p = dst.Add();
      p->CopyFrom(*value_info);
      *p->mutable_name() = new_name;
      return true;
    }

    return false;
  }

  template <typename TValueInfoLookupFn>
  void FillInputWait(
      ONNX_NAMESPACE::GraphProto* sub_gp,
      const ONNX_NAMESPACE::GraphProto& main_gp,
      TValueInfoLookupFn main_graph_lookup,
      const std::vector<std::string>& sync_inputs,
      const std::vector<std::string>& dependencies,
      size_t sub_id,
      bool bw) {
    // input/output with Wait/RecordEvent
    // Note data is gated by Wait/RecordEvent, so name with postfix "_sync"
    // In distributed training, the pattern is:
    //   wait_data -> recv -> wait_pipeline -> fw/bw -> record_pipeline -> send -> record_data
    // Here wait_data/record_data is to ensure execution order due to data dependency (same batch across pipelines),
    // while wait_pipeline/recorde_pipeline is to ensure pipeline execution order.
    // This test simplifies the graph to omit send/recv,
    // but we still need to have double wait and record to sync data and pipeline separately
    ONNX_NAMESPACE::NodeProto* wait_data_np = nullptr;
    ONNX_NAMESPACE::NodeProto* wait_pipeline_np = nullptr;
    std::string wait_data_id = "wait_data_" + std::to_string(sub_id) + (bw ? "_bw" : "_fw");
    std::string wait_pipeline_id = "wait_pipeline_" + std::to_string(sub_id) + (bw ? "_bw" : "_fw");
    bool is_first = (sub_id == 0 && !bw);
    if (sync_inputs.size() + dependencies.size() > 0) {
      if (!is_first) {
        wait_data_np = sub_gp->add_node();
        *wait_data_np->mutable_op_type() = "WaitEvent";
        *wait_data_np->mutable_domain() = kMSDomain;
        *wait_data_np->mutable_input()->Add() = wait_data_id;
      }
      wait_pipeline_np = sub_gp->add_node();
      *wait_pipeline_np->mutable_op_type() = "WaitEvent";
      *wait_pipeline_np->mutable_domain() = kMSDomain;
      *wait_pipeline_np->mutable_input()->Add() = wait_pipeline_id;
    }
    for (const auto& name : sync_inputs) {
      std::string input_name = name + "_sync";
      std::string recv_name = name + "_recv";
      if (wait_data_np) {
        *wait_data_np->mutable_input()->Add() = input_name;
        *wait_data_np->mutable_output()->Add() = recv_name;
        *wait_pipeline_np->mutable_input()->Add() = recv_name;
      } else {
        *wait_pipeline_np->mutable_input()->Add() = input_name;
      }
      *wait_pipeline_np->mutable_output()->Add() = name;
      // some input comes graph input
      if (AddItemByName(sub_gp->mutable_input(),
                        main_gp.input(),
                        name,
                        input_name)) {
        ASSERT_TRUE(is_first);
        // add shape info
        EXPECT_TRUE(AddValueInfoByNameFromLookup(*sub_gp->mutable_value_info(),
                                                 main_graph_lookup,
                                                 name,
                                                 name));
      } else {
        // some input comes from the middle of the graph
        AddValueInfoByNameFromLookup(*sub_gp->mutable_input(),
                                     main_graph_lookup,
                                     name,
                                     input_name);
        // add shape info
        AddValueInfoByNameFromLookup(*sub_gp->mutable_value_info(),
                                     main_graph_lookup,
                                     name,
                                     recv_name);
        AddValueInfoByNameFromLookup(*sub_gp->mutable_value_info(),
                                     main_graph_lookup,
                                     name,
                                     name);
      }
    }

    if (wait_pipeline_np) {
      // add dependencies on the first wait
      auto* wait_np = wait_data_np ? wait_data_np : wait_pipeline_np;
      for (const auto& dep : dependencies) {
        *wait_np->mutable_input()->Add() = dep;
      }

      // add input for event ids
      if (wait_data_np) {
        auto* p = sub_gp->mutable_input()->Add();
        *p->mutable_name() = wait_data_id;
        p->mutable_type()->mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_INT64);
      }
      auto* p = sub_gp->mutable_input()->Add();
      *p->mutable_name() = wait_pipeline_id;
      p->mutable_type()->mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_INT64);
    }
  }

  template <typename TValueInfoLookupFn>
  void FillOutputRecord(
      ONNX_NAMESPACE::GraphProto* sub_gp,
      const ONNX_NAMESPACE::GraphProto& /*main_gp*/,
      TValueInfoLookupFn main_graph_lookup,
      const std::vector<std::string>& sync_outputs,
      const std::vector<std::string>& dependencies,
      size_t sub_id,
      bool bw) {
    ONNX_NAMESPACE::NodeProto* record_pipeline_np = nullptr;
    ONNX_NAMESPACE::NodeProto* record_data_np = nullptr;
    std::string record_pipeline_id = "record_pipeline_" + std::to_string(sub_id) + (bw ? "_bw" : "_fw");
    std::string record_data_id = "record_data_" + std::to_string(sub_id) + (bw ? "_bw" : "_fw");
    bool is_last = (sub_id == 0 && bw);
    if (sync_outputs.size() + dependencies.size() > 0) {
      record_pipeline_np = sub_gp->add_node();
      *record_pipeline_np->mutable_op_type() = "RecordEvent";
      *record_pipeline_np->mutable_domain() = kMSDomain;
      *record_pipeline_np->mutable_input()->Add() = record_pipeline_id;

      if (!is_last) {
        record_data_np = sub_gp->add_node();
        *record_data_np->mutable_op_type() = "RecordEvent";
        *record_data_np->mutable_domain() = kMSDomain;
        *record_data_np->mutable_input()->Add() = record_data_id;
      }
    }

    if (sync_outputs.size() > 0) {
      for (const auto& name : sync_outputs) {
        *record_pipeline_np->mutable_input()->Add() = name;
        if (record_data_np) {
          *record_pipeline_np->mutable_output()->Add() = name + "_send";
          *record_data_np->mutable_input()->Add() = name + "_send";
          *record_data_np->mutable_output()->Add() = name + "_sync";
        } else {
          *record_pipeline_np->mutable_output()->Add() = name + "_sync";
        }
      }
    }

    if (record_pipeline_np) {
      for (const auto& name : dependencies)
        *record_pipeline_np->mutable_input()->Add() = name;

      // add input for event_id
      auto* p = sub_gp->mutable_input()->Add();
      *p->mutable_name() = record_pipeline_id;
      p->mutable_type()->mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_INT64);

      if (record_data_np) {
        p = sub_gp->mutable_input()->Add();
        *p->mutable_name() = record_data_id;
        p->mutable_type()->mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_INT64);
      }
    }

    // add graph output and shape info
    for (const auto& name : sync_outputs) {
      AddValueInfoByNameFromLookup(*sub_gp->mutable_output(),
                                   main_graph_lookup,
                                   name,
                                   name + "_sync");
      if (!is_last) {
        AddValueInfoByNameFromLookup(*sub_gp->mutable_value_info(),
                                     main_graph_lookup,
                                     name,
                                     name + "_send");
      }
      AddValueInfoByNameFromLookup(*sub_gp->mutable_value_info(),
                                   main_graph_lookup,
                                   name,
                                   name);
    }
  }
};

// TODO: move to a proper location for pipeline training

// pipeline plan for each batch
struct PipelineBatchInfo {
  // Event pairs for each pipeline slot to WaitEvent when start, and RecordEvent when end
  std::vector<std::pair<int64_t, int64_t>> events;
  // indices of retired batches, so their data could be reused
  // a batch can only be retired after finished backward in stage 0
  // this can be used to join worker threads or reuse buffers
  // for example, in a node with N GPUs and B batches to run in pipeline (with one stage for each GPU)
  // there will be (N * B) threads created, and by being able to retire,
  // only at most (N * (2 * N - 1)) concurrent threads are needed
  // for small number of B, there's no retired threads so total count would be the same.
  // for big number of B, this would be helpful
  std::vector<int64_t> retired_batches;
};

class PipelineTimeline {
 public:
  struct Slot {
    enum class Type {
      Unused,
      Forward,
      Backward
    };
    Type type;
    size_t batch_id;

    Slot() : type(Type::Unused) {}
  };

  PipelineTimeline() = default;

  void Initialize(size_t num_stages, size_t num_slots) {
    slots_.resize(num_stages);
    for (size_t s = 0; s < num_stages; ++s) {
      slots_[s].resize(num_slots);
    }
  }

  bool IsOccupied(size_t s, size_t t) const {
    return slots_[s][t].type != Slot::Type::Unused;
  }

  const Slot& Get(size_t s, size_t t) const {
    return slots_[s][t];
  }

  size_t GetNumSlots() const {
    return slots_[0].size();
  }

  void Occupy(size_t s, size_t t, size_t batch_id, Slot::Type st) {
    Slot& slot = slots_[s][t];
    ORT_ENFORCE(slot.type == Slot::Type::Unused);
    slot.type = st;
    slot.batch_id = batch_id;
  }

 private:
  std::vector<std::vector<Slot>> slots_;
};

// pipeline planner for batches
class PipelineBatchPlanner {
 private:
  int64_t max_id_;
  PipelineTimeline timeline_;

 public:
  PipelineBatchPlanner()
      : max_id_(::onnxruntime::contrib::OrtEventPool::GetPoolSize() - 1) {
  }

  // Generate timeline for one-forward-one-backward scheduling,
  // which schedules execution in batch order to minimize latency for onging batches
  // each stage requires 2 pair of wait/record events for FW/BW
  void GenerateOneFWOneBWTimeline(size_t num_stages, size_t num_batches) {
    // The first batch has 2 * (num_stages - 1) gaps between FW and BW
    // then 2 slots for FW/BW in each batch
    size_t num_slots = 2 * (num_stages - 1) + num_batches * 2;
    timeline_.Initialize(num_stages, num_slots);

    // fw time slot to start the search for empty ones in each stage
    std::vector<size_t> t_fw(num_stages, 0);
    // bw time slot to start the search for empty ones in each stage
    std::vector<size_t> t_bw(num_stages, 0);

    // generate timeline in batch order to minimize latency for ongoing batches
    for (size_t batch_id = 0; batch_id < num_batches; ++batch_id) {
      // plan for FW
      for (size_t s = 0; s < num_stages; ++s) {
        while (timeline_.IsOccupied(s, t_fw[s])) {
          ++t_fw[s];
        }
        // after find a slot, update t[s+1..] if needed
        for (size_t ss = s + 1; ss < num_stages; ++ss) {
          t_fw[ss] = std::max(t_fw[ss], t_fw[s] + (ss - s));
        }
        // occupy slot in timeline
        timeline_.Occupy(s, t_fw[s]++, batch_id, PipelineTimeline::Slot::Type::Forward);
      }
      // plan for BW
      for (int s = gsl::narrow<int>(num_stages - 1); s >= 0; --s) {
        t_bw[s] = std::max(t_fw[s], t_bw[s]);
        while (timeline_.IsOccupied(s, t_bw[s])) {
          t_bw[s]++;
        }
        // after find a slot, update t_bw[s-1..]
        for (int ss = s - 1; ss >= 0; --ss) {
          t_bw[ss] = std::max(t_bw[ss], t_bw[s] + (s - ss));
        }
        // occupy slot in timeline
        timeline_.Occupy(s, t_bw[s], batch_id, PipelineTimeline::Slot::Type::Backward);
      }
    }
  }

  // create pipeline plans according to generated timeline
  // with start_event_id = s, the output of each stage is [-1, s], [s, s+1], [s+1, s+2]... for each occupied slot
  // and will be assigned to each batch's PipelineBatchInfo
  // returns the first unused event_id after creating
  int64_t CreatePlan(int64_t start_event_id, const size_t stage, std::vector<PipelineBatchInfo>& plan) {
    // fill in plan
    int64_t prev_event_id = -1;
    int64_t event_id = start_event_id;
    std::vector<int64_t> retired_batches;
    for (size_t t = 0; t < timeline_.GetNumSlots(); ++t) {
      if (!timeline_.IsOccupied(stage, t))
        continue;

      const auto& slot = timeline_.Get(stage, t);
      ORT_ENFORCE(event_id < max_id_);
      if (stage == 0) {
        if (slot.type == PipelineTimeline::Slot::Type::Forward) {
          // set retired batches when starting a new batch (s == 0 && !bw)
          plan[slot.batch_id].retired_batches = retired_batches;
          retired_batches.clear();
        } else if (slot.type == PipelineTimeline::Slot::Type::Backward) {
          // add to retired batches after backward of stage 0
          retired_batches.push_back(gsl::narrow<int64_t>(slot.batch_id));
        }
      }
      // add a pair of wait/record event ids to given batch_id
      plan[slot.batch_id].events.emplace_back(prev_event_id, event_id);
      prev_event_id = event_id;
      ++event_id;
    }
    return event_id;
  }
};

void RetrieveEventOperators(
    Graph& graph,
    const int stage_index,
    const int num_stages,
    Node** forward_recv_wait,
    Node** forward_recv_record,
    Node** forward_compute_wait,
    Node** forward_compute_record,
    Node** forward_send_wait,
    Node** forward_send_record,
    Node** backward_recv_wait,
    Node** backward_recv_record,
    Node** backward_compute_wait,
    Node** backward_compute_record,
    Node** backward_send_wait,
    Node** backward_send_record) {
  // Initialize retrieved nodes.
  // Non-existing nodes may hold NULL forever.
  // Existing nodes may get valid pointers below.
  *forward_recv_wait = nullptr;
  *forward_recv_record = nullptr;
  *forward_compute_wait = nullptr;
  *forward_compute_record = nullptr;
  *forward_send_wait = nullptr;
  *forward_send_record = nullptr;
  *backward_recv_wait = nullptr;
  *backward_recv_record = nullptr;
  *backward_compute_wait = nullptr;
  *backward_compute_record = nullptr;
  *backward_send_wait = nullptr;
  *backward_send_record = nullptr;

  // Declare container for WaitEvent's in topological order.
  std::vector<Node*> waits;
  // Declare container for RecordEvent's in topological order.
  std::vector<Node*> records;

  // Find out WaitEvent's and RecordEvent's.
  GraphViewer graph_viewer(graph);
  for (auto& node_idx : graph_viewer.GetNodesInTopologicalOrder()) {
    Node* node = graph.GetNode(node_idx);
    if (node->OpType() == "WaitEvent") {
      waits.push_back(node);
    } else if (node->OpType() == "RecordEvent") {
      records.push_back(node);
    }
  }

  // For different stages, assign nodes based on different rules.
  if (stage_index != 0 && stage_index != num_stages - 1) {
    // Wait/Record patterns at middle stages:
    //   Wait -> Recv -> Record -> Wait -> FW -> Record -> Wait -> Send -> Record ->
    //   Wait -> Recv -> Record -> Wait -> BW -> Record -> Wait -> Send -> Record

    ORT_ENFORCE(waits.size() == 6, " size is ", waits.size(), " at stage ", stage_index);
    *forward_recv_wait = waits[0];
    *forward_compute_wait = waits[1];
    *forward_send_wait = waits[2];
    *backward_recv_wait = waits[3];
    *backward_compute_wait = waits[4];
    *backward_send_wait = waits[5];

    ORT_ENFORCE(records.size() == 6, " size is ", waits.size(), " at stage ", stage_index);
    *forward_recv_record = records[0];
    *forward_compute_record = records[1];
    *forward_send_record = records[2];
    *backward_recv_record = records[3];
    *backward_compute_record = records[4];
    *backward_send_record = records[5];
  } else if (stage_index == 0) {
    // Wait/Record patterns at the 1st stages:
    //                             Wait -> FW -> Record -> Wait -> Send -> Record ->
    //   Wait -> Recv -> Record -> Wait -> BW -> Record

    ORT_ENFORCE(waits.size() == 4, " size is ", waits.size(), " at stage ", stage_index);
    *forward_compute_wait = waits[0];
    *forward_send_wait = waits[1];
    *backward_recv_wait = waits[2];
    *backward_compute_wait = waits[3];

    ORT_ENFORCE(records.size() == 4, " size is ", waits.size(), " at stage ", stage_index);
    *forward_compute_record = records[0];
    *forward_send_record = records[1];
    *backward_recv_record = records[2];
    *backward_compute_record = records[3];
  } else if (stage_index == num_stages - 1) {
    // Wait/Record patterns at the last stages:
    //   Wait -> Recv -> Record -> Wait -> FW ->
    //                                     BW -> Record -> Wait -> Send -> Record

    ORT_ENFORCE(waits.size() == 3, " size is ", waits.size(), " at stage ", stage_index);
    *forward_recv_wait = waits[0];
    *forward_compute_wait = waits[1];
    *backward_send_wait = waits[2];

    ORT_ENFORCE(records.size() == 3, " size is ", waits.size(), " at stage ", stage_index);
    *forward_recv_record = records[0];
    *backward_compute_record = records[1];
    *backward_send_record = records[2];
  } else {
    ORT_THROW("Wrong number of WaitEvent operators: ",
              waits.size(), " allowed value range is [0, ", num_stages - 1, ").");
  }
}

void RetrieveSendRecvOperators(
    Graph& graph,
    Node** forward_recv,
    Node** forward_send,
    Node** backward_recv,
    Node** backward_send) {
  // Initialize retrieved nodes.
  // Non-existing nodes may hold NULL forever.
  // Existing nodes may get valid pointers below.
  *forward_recv = nullptr;
  *forward_send = nullptr;
  *backward_recv = nullptr;
  *backward_send = nullptr;

  auto is_backward = [](Node& node) {
    return (node.Description() == "Backward pass");
  };

  // Search for Send's and Recv's by assuming that
  // there are only one Send and one Recv in forward/backward.
  for (auto& node : graph.Nodes()) {
    if (node.OpType() == "Send") {
      if (is_backward(node)) {
        // backward_send can only be assigned one valid pointer.
        // If it is assigned more than once, it means we have multiple
        // Send in backward pass and therefore our assumption doesn't hold.
        // This check ensure that only we only update *backward_send when
        // its value is NULL and guards our one-Recv assumption.
        ASSERT_TRUE(!(*backward_send));
        *backward_send = &node;
      } else {
        // Guard the uniqueness of Send in the forward pass by throwing
        // when *forward_send already carries a valid pointer.
        ASSERT_TRUE(!(*forward_send));
        *forward_send = &node;
      }
    } else if (node.OpType() == "Recv") {
      if (is_backward(node)) {
        // Guard the uniqueness of Recv in the backward pass by throwing
        // when *backward_recv already carries a valid pointer.
        ASSERT_TRUE(!(*backward_recv));
        *backward_recv = &node;
      } else {
        // Guard the uniqueness of Recv in the forwaard pass by throwing
        // when *forward_recv already carries a valid pointer.
        *forward_recv = &node;
      }
    }
  }
}

PathString GenerateFileNameWithIndex(const std::string& base_str, int index, const std::string& file_suffix) {
  return path_utils::MakePathString(base_str, index, file_suffix);
}

// DistributedRunTestContext provides a method to override existing DistributedRunTestContext instance.
// This is for test purpose only. Please don't use it for other scenarios.
class DistributedRunTestContext : public DistributedRunContext {
 public:
  DistributedRunTestContext(const TrainingSession::TrainingConfiguration& config)
      : DistributedRunContext(config.distributed_config.world_rank,
                              config.distributed_config.world_size,
                              config.distributed_config.local_rank,
                              config.distributed_config.local_size,
                              config.distributed_config.data_parallel_size,
                              config.distributed_config.horizontal_parallel_size,
                              config.distributed_config.pipeline_parallel_size) {
  }

  // Reset the static DistributedRunContext object with new value.
  void ResetDistributedRunContext() {
    DistributedRunContext::GetRunConfig() = params_;
    auto& dp_group = DistributedRunContext::GetWorkerGroup(WorkerGroupType::DataParallel);
    dp_group = groups_[WorkerGroupType::DataParallel];

    auto& hp_group = DistributedRunContext::GetWorkerGroup(WorkerGroupType::HorizontalParallel);
    hp_group = groups_[WorkerGroupType::HorizontalParallel];

    auto& mp_group = DistributedRunContext::GetInstance().GetWorkerGroup(WorkerGroupType::PipelineParallel);
    mp_group = groups_[WorkerGroupType::PipelineParallel];
  }
};

void OverwritePipelineRank(const TrainingSession::TrainingConfiguration& config, const int pipeline_rank) {
  // DistributedRunContext is a static global. Create one if it hasn't been created yet.
  DistributedRunContext::CreateInstance({config.distributed_config.world_rank,
                                         config.distributed_config.world_size,
                                         config.distributed_config.local_rank,
                                         config.distributed_config.local_size,
                                         config.distributed_config.data_parallel_size,
                                         config.distributed_config.horizontal_parallel_size,
                                         config.distributed_config.pipeline_parallel_size});

  // If DistributedRunContext has already been created prior to this test, the CreateInstance() call above won't
  // create a new DistributedRunContext instance, as it is statically cached in the process.
  // In this case, we create a DistributedRunTestContext object and assign its value to the static object's field.
  DistributedRunTestContext ctx(config);
  ctx.ResetDistributedRunContext();

  // Overwrite the pipeline rank in case the static DistributedRunContext has been created and is stale and not up-to-date.
  auto& mp_group = DistributedRunContext::GetInstance().GetWorkerGroup(WorkerGroupType::PipelineParallel);
  mp_group.rank_in_group = pipeline_rank;
}

TEST(GradientGraphBuilderTest, PipelineOnlinePartition_bert_tiny) {
  const auto model_path = ORT_TSTR("testdata/bert_toy_optimized.onnx");

  constexpr size_t total_partition_count = 3;
  TrainingSession::TrainingConfiguration::PipelineConfiguration pipe{};
  pipe.do_partition = true;

  // cut model in 3 partitions
  TrainingSession::TrainingConfiguration::CutInfo cut0 = {
      onnxruntime::training::TrainingSession::TrainingConfiguration::CutEdge("326"),
      onnxruntime::training::TrainingSession::TrainingConfiguration::CutEdge("103", {"413", "529"})};

  TrainingSession::TrainingConfiguration::CutInfo cut1 = {
      onnxruntime::training::TrainingSession::TrainingConfiguration::CutEdge("558"),
      onnxruntime::training::TrainingSession::TrainingConfiguration::CutEdge("103", {"645"})};

  pipe.cut_list.emplace_back(cut0);
  pipe.cut_list.emplace_back(cut1);

  TrainingSession::TrainingConfiguration::MixedPrecisionConfiguration mixed_precision_config{};
  mixed_precision_config.use_mixed_precision_initializers = true;

  // 2 test variations - full precision and mixed precision
  const std::vector<bool> test_with_fp32{true, false};
  for (auto is_fp32 : test_with_fp32) {
    // graph is partitioned into 3 parts.
    for (int i = 0; i < static_cast<int>(total_partition_count); ++i) {
      PathString partition_file = GenerateFileNameWithIndex("pipeline_partition_", i, ".onnx");
      PathString output_file = GenerateFileNameWithIndex("pipeline_partition_", i, "_back.onnx");
      auto config = MakeBasicTrainingConfig();

      if (i == static_cast<int>(total_partition_count - 1)) {
        config.loss_function_config = TrainingSession::TrainingConfiguration::LossFunctionConfiguration{};
        config.loss_function_config.value().loss_function_info =
            LossFunctionInfo(OpDef("BertLoss", kOnnxDomain),
                             "total_loss",
                             {/*prediction_masked_lm*/ "prediction_scores",
                              /*prediction_next_sentence*/ "seq_relationship_score",
                              /*masked_lm_positions*/ "masked_lm_positions",
                              /*masked_lm_ids*/ "masked_lm_ids",
                              /*next_sentence_labels*/ "next_sentence_labels",
                              /*mlm_loss*/ "mlm_loss",
                              /*nsp_loss*/ "nsp_loss"});
      }

      // Add weight_names_to_not_train to avoid generating backward graph on those tensor
      config.weight_names_to_not_train = {
          "position_01",            // Slice's dat input
          "op_min_ends_expand_10",  // op_min_ends_expand_10
      };
      pipe.partitioned_model_path = partition_file;
      config.pipeline_config = pipe;
      config.distributed_config.world_rank = i;
      config.distributed_config.world_size = total_partition_count;
      config.distributed_config.local_rank = i;
      config.distributed_config.local_size = total_partition_count;
      config.distributed_config.data_parallel_size = 1;
      config.distributed_config.horizontal_parallel_size = 1;
      config.distributed_config.pipeline_parallel_size = total_partition_count;
      config.model_with_training_graph_path = output_file;

      OverwritePipelineRank(config, i);

      if (!is_fp32) {
        config.mixed_precision_config = mixed_precision_config;
      }

      PathString backprop_model_file;
      Status status = BuildBackPropGraph(model_path, config, backprop_model_file);
      ASSERT_TRUE(status.IsOK()) << status << " (is_fp32 = " << is_fp32 << ", stage = " << i << ").\n";

      // Skip the re-load for mixed-precision case. This model contains grad op that has function body,
      // which takes a const tensor input. Const cast for input in function body won't be saved in the output
      // model so reload will run into error.
      // For the purpose of testing mixed-precision, BuildBackPropGraph above will be sufficient to verify the
      // partition logic and validate the graph.
      if (is_fp32) {
        std::shared_ptr<Model> model;
        // Ensure the partitioned model load.
        status = Model::Load(backprop_model_file, model, nullptr, DefaultLoggingManager().DefaultLogger());
        ASSERT_TRUE(status.IsOK()) << status << " (is_fp32 = " << is_fp32 << ", stage = " << i << ").\n";

        // verify the first stage contains word embedding as input and the last stage doesn't
        auto model_proto = model->ToProto();
        const auto& graph_proto = model_proto.graph();

        bool found_word_embedding = false;
        for (auto& tensor : graph_proto.initializer()) {
          if (tensor.name() == "bert.embeddings.word_embeddings.weight") {
            found_word_embedding = true;
          }
        }
        if (i == 0) {
          ASSERT_TRUE(found_word_embedding) << " (is_fp32 = " << is_fp32 << ", stage = " << i << ").\n";
        } else {
          ASSERT_FALSE(found_word_embedding) << " (is_fp32 = " << is_fp32 << ", stage = " << i << ").\n";
        }
      }
    }
  }
}

TEST(GradientGraphBuilderTest, PipelineOnlinePartition_MLP) {
  auto model_uri = ORIGINAL_MODEL_PATH;

  TrainingSession::TrainingConfiguration::PipelineConfiguration pipe{};
  pipe.do_partition = true;

  // evenly cut the MLP model in 3 partitions
  TrainingSession::TrainingConfiguration::CutInfo cut0 = {TrainingSession::TrainingConfiguration::CutEdge("T3")};
  TrainingSession::TrainingConfiguration::CutInfo cut1 = {TrainingSession::TrainingConfiguration::CutEdge("T6")};
  pipe.cut_list.emplace_back(cut0);
  pipe.cut_list.emplace_back(cut1);

  TrainingSession::TrainingConfiguration::MixedPrecisionConfiguration mixed_precision_config{};
  mixed_precision_config.use_mixed_precision_initializers = true;

  // 2 test variations - full precision and mixed precision
  const std::vector<bool> test_with_fp32{true, false};
  for (auto is_fp32 : test_with_fp32) {
    // graph is partitioned into 3 parts.
    for (int i = 0; i < 3; ++i) {
      PathString output_file = GenerateFileNameWithIndex("pipeline_partition_", i, "_back.onnx");

      auto config = MakeBasicTrainingConfig();

      config.pipeline_config = pipe;
      config.distributed_config.world_rank = i;
      config.distributed_config.world_size = 3;
      config.distributed_config.local_rank = i;
      config.distributed_config.local_size = 3;
      config.distributed_config.data_parallel_size = 1;
      config.distributed_config.horizontal_parallel_size = 1;
      config.distributed_config.pipeline_parallel_size = 3;
      config.model_with_training_graph_path = output_file;

      OverwritePipelineRank(config, i);

      if (!is_fp32) {
        config.mixed_precision_config = mixed_precision_config;
      }

      PathString backprop_model_file;
      Status status = BuildBackPropGraph(model_uri, config, backprop_model_file);
      ASSERT_TRUE(status.IsOK()) << status << " (is_fp32 = " << is_fp32 << ", stage = " << i << ").\n";

      // Skip the re-load for mixed-precision case. This model contains grad op that has function body,
      // which takes a const tensor input. Const cast for input in function body won't be saved in the output
      // model so reload will run into error.
      // For the purpose of testing mixed-precision, BuildBackPropGraph above will be sufficient to verify the
      // partition logic and validate the graph.
      if (is_fp32) {
        std::shared_ptr<Model> model;
        // Ensure the partitioned model load.
        status = Model::Load(backprop_model_file, model, nullptr, DefaultLoggingManager().DefaultLogger());
        ASSERT_TRUE(status.IsOK()) << status << " (is_fp32 = " << is_fp32 << ", stage = " << i << ").\n";
      }
    }
  }
}

Status RunOnlinePartition(const std::vector<TrainingSession::TrainingConfiguration::CutInfo>& cut_list,
                          int pipeline_stage_size) {
  auto model_uri = ORIGINAL_MODEL_PATH;

  TrainingSession::TrainingConfiguration::PipelineConfiguration pipe{};
  pipe.do_partition = true;
  pipe.cut_list = cut_list;

  for (int i = 0; i < pipeline_stage_size; ++i) {
    PathString output_file = GenerateFileNameWithIndex("pipeline_partition_", i, "_back.onnx");

    auto config = MakeBasicTrainingConfig();
    config.pipeline_config = pipe;

    config.distributed_config.world_rank = i;
    config.distributed_config.world_size = pipeline_stage_size;
    config.distributed_config.local_rank = i;
    config.distributed_config.local_size = pipeline_stage_size;
    config.distributed_config.data_parallel_size = 1;
    config.distributed_config.horizontal_parallel_size = 1;
    config.distributed_config.pipeline_parallel_size = pipeline_stage_size;

    OverwritePipelineRank(config, i);

    config.model_with_training_graph_path = output_file;

    PathString backprop_model_file;
    auto status = BuildBackPropGraph(model_uri, config, backprop_model_file);
    EXPECT_FALSE(status.IsOK());
  }
  return Status::OK();
}

TEST(GradientGraphBuilderTest, PipelineOnlinePartition_Invalid_Input) {
  using CutEdge = TrainingSession::TrainingConfiguration::CutEdge;
  using CutInfo = TrainingSession::TrainingConfiguration::CutInfo;

  // Test with invalid cut edge
  CutInfo invalid_cut_edge = {CutEdge("3")};
  ASSERT_STATUS_OK(RunOnlinePartition(std::vector<CutInfo>{invalid_cut_edge}, 2 /* pipeline_stage_size */));

  // Test mis-matched cut list with stage size
  CutInfo cut_edge = {CutEdge("T3")};
  ASSERT_STATUS_OK(RunOnlinePartition(std::vector<CutInfo>{cut_edge}, 3 /* pipeline_stage_size */));

  // Test unordered cut_info list
  CutInfo cut0 = {CutEdge("T3")};
  CutInfo cut1 = {CutEdge("T6")};
  ASSERT_STATUS_OK(RunOnlinePartition(std::vector<CutInfo>{cut1, cut0}, 3 /* pipeline_stage_size */));
}

// verify pipeline config can load and gradient graph can construct.
TEST(GradientGraphBuilderTest, TrainingSession_PipelineTransform_base) {
  std::string filename_base = "testdata/test_training_model_";

  auto load_and_check_gradient_graph = [](int stageIdx, PathString& input_file, PathString& output_file) {
    auto config = MakeBasicTrainingConfig();

    TrainingSession::TrainingConfiguration::PipelineConfiguration pipe_config{};
    config.pipeline_config = pipe_config;

    PathString backprop_model_file;
    ASSERT_STATUS_OK(BuildBackPropGraph(input_file, config, backprop_model_file));

    std::shared_ptr<Model> model;
    ASSERT_STATUS_OK(Model::Load(backprop_model_file, model, nullptr, DefaultLoggingManager().DefaultLogger()));

    Graph& graph = model->MainGraph();

    // Declare forward event nodes.
    // The nodes are declared according to their topological order.
    Node* forward_recv_wait{nullptr};
    Node* forward_recv_record{nullptr};
    Node* forward_compute_wait{nullptr};
    Node* forward_compute_record{nullptr};
    Node* forward_send_wait{nullptr};
    Node* forward_send_record{nullptr};

    // Declare backward event nodes.
    // The nodes are declared according to their topological order.
    Node* backward_recv_wait{nullptr};
    Node* backward_recv_record{nullptr};
    Node* backward_compute_wait{nullptr};
    Node* backward_compute_record{nullptr};
    Node* backward_send_wait{nullptr};
    Node* backward_send_record{nullptr};

    // Find event nodes.
    RetrieveEventOperators(
        graph,
        stageIdx,
        3,
        &forward_recv_wait,
        &forward_recv_record,
        &forward_compute_wait,
        &forward_compute_record,
        &forward_send_wait,
        &forward_send_record,
        &backward_recv_wait,
        &backward_recv_record,
        &backward_compute_wait,
        &backward_compute_record,
        &backward_send_wait,
        &backward_send_record);

    // Check event nodes.
    if (stageIdx == 2) {
      // Last stage's event pattern:
      //   Wait -> Recv -> Record -> Wait -> FW ->
      //                                     BW -> Record -> Wait -> Send -> Record
      ASSERT_TRUE(forward_recv_wait);
      ASSERT_TRUE(forward_recv_record);
      ASSERT_TRUE(forward_compute_wait);
      ASSERT_TRUE(!forward_compute_record);
      ASSERT_TRUE(!forward_send_wait);
      ASSERT_TRUE(!forward_send_record);

      ASSERT_TRUE(!backward_recv_wait);
      ASSERT_TRUE(!backward_recv_record);
      ASSERT_TRUE(!backward_compute_wait);
      ASSERT_TRUE(backward_compute_record);
      ASSERT_TRUE(backward_send_wait);
      ASSERT_TRUE(backward_send_record);
    } else if (stageIdx == 1) {
      // Middle stage's event pattern:
      //   Wait -> Recv -> Record -> Wait -> FW -> Record -> Wait -> Send -> Record ->
      //   Wait -> Recv -> Record -> Wait -> BW -> Record -> Wait -> Send -> Record
      ASSERT_TRUE(forward_recv_wait);
      ASSERT_TRUE(forward_recv_record);
      ASSERT_TRUE(forward_compute_wait);
      ASSERT_TRUE(forward_compute_record);
      ASSERT_TRUE(forward_send_wait);
      ASSERT_TRUE(forward_send_record);

      ASSERT_TRUE(backward_recv_wait);
      ASSERT_TRUE(backward_recv_record);
      ASSERT_TRUE(backward_compute_wait);
      ASSERT_TRUE(backward_compute_record);
      ASSERT_TRUE(backward_send_wait);
      ASSERT_TRUE(backward_send_record);
    } else {
      // First stage's event pattern:
      //                             Wait -> FW -> Record -> Wait -> Send -> Record ->
      //   Wait -> Recv -> Record -> Wait -> BW -> Record
      ASSERT_TRUE(!forward_recv_wait);
      ASSERT_TRUE(!forward_recv_record);
      ASSERT_TRUE(forward_compute_wait);
      ASSERT_TRUE(forward_compute_record);
      ASSERT_TRUE(forward_send_wait);
      ASSERT_TRUE(forward_send_record);

      ASSERT_TRUE(backward_recv_wait);
      ASSERT_TRUE(backward_recv_record);
      ASSERT_TRUE(backward_compute_wait);
      ASSERT_TRUE(backward_compute_record);
      ASSERT_TRUE(!backward_send_wait);
      ASSERT_TRUE(!backward_send_record);
    }

    Node* forward_send{nullptr};
    Node* forward_recv{nullptr};
    Node* backward_recv{nullptr};
    Node* backward_send{nullptr};

    RetrieveSendRecvOperators(
        graph,
        &forward_recv,
        &forward_send,
        &backward_recv,
        &backward_send);

    // Except the last partion, each partition should have send forward and recv backward.
    if (stageIdx == 0 || stageIdx == 1) {
      ASSERT_TRUE(forward_send && backward_recv);
    } else {
      ASSERT_TRUE(!forward_send && !backward_recv);
    }
    // Except the first partion, each partition should have recv forward and send backward.
    if (stageIdx == 1 || stageIdx == 2) {
      ASSERT_TRUE(forward_recv && backward_send);
    } else {
      ASSERT_TRUE(!forward_recv && !backward_send);
    }

    auto mp = model->ToProto();
    std::ofstream ofs(output_file, std::ofstream::binary);
    mp.SerializeToOstream(&ofs);
    ofs.close();
  };

  for (int i = 0; i < 3; ++i) {
    PathString input_file = GenerateFileNameWithIndex(filename_base, i, ".onnx");
    PathString output_file = GenerateFileNameWithIndex(filename_base, i, "_back.onnx");

    load_and_check_gradient_graph(i, input_file, output_file);
  }
}

TEST(GradientGraphBuilderTest, TrainingSession_WithPipeline) {
  auto config = MakeBasicTrainingConfig();
  // config.set_gradients_as_graph_outputs = true;
  PathString backprop_model_file;
  ASSERT_STATUS_OK(BuildBackPropGraph(ORIGINAL_MODEL_PATH, config, backprop_model_file));

  // cut the model using outputs
  const std::vector<PipelineSplitter::CutInfo> cuts = {
      // sub model 0
      {{{"T1", "T2", "T3"},
        {"X"},
        {"T3"},
        {},
        {}},
       {{"T2_grad", "T1_grad", "B1_grad", "W1_grad"},
        {"T3_grad"},
        {},
        {"T3_sync"},
        {"B1_grad", "W1_grad"}}},
      // sub model 1
      {{{"T4", "T5", "T6"},
        {"T3"},
        {"T6"},
        {},
        {}},
       {{"T5_grad", "T4_grad", "T3_grad", "B2_grad", "W2_grad"},
        {"T6_grad"},
        {"T3_grad"},
        {"T6_sync"},
        {"B2_grad", "W2_grad"}}},
      // sub model 2
      {{{"T7", "MeanSquaredError_diff", "MeanSquaredError_diff_square", "loss", "predictions"},
        {"T6"},
        {},
        {},
        {}},
       {{"MeanSquaredError_reduce_mean_Grad/Sized_X",
         "MeanSquaredError_reduce_mean_Grad/Sized_Grad",
         "MeanSquaredError_reduce_mean_Grad/Scale",
         "MeanSquaredError_reduce_mean_Grad/Scaled_Grad",
         "MeanSquaredError_reduce_mean_Grad/Shaped_X",
         "MeanSquaredError_diff_square_grad",
         "MeanSquaredError_pow_Grad/Sub_I1",
         "MeanSquaredError_pow_Grad/Pow_I0",
         "MeanSquaredError_pow_Grad/Mul_Pow_I0_I1",
         "MeanSquaredError_diff_grad",
         "predictions_grad",
         "B3_grad",
         "T7_grad",
         "W3_grad",
         "T6_grad"},
        {},
        {"T6_grad"},
        {},
        {"loss", "predictions", "B3_grad", "W3_grad"}}}};

  const auto num_subs = cuts.size();

  std::vector<PathString> sub_model_files(num_subs);
  for (size_t sub_id = 0; sub_id < num_subs; ++sub_id) {
    sub_model_files[sub_id] = GenerateFileNameWithIndex("sub_", static_cast<int>(sub_id), ".onnx");
  }

  PipelineSplitter splitter;
  splitter.Split(backprop_model_file, sub_model_files, cuts);

  // create training sessions
  std::unique_ptr<Environment> env;
  ASSERT_STATUS_OK(Environment::Create(nullptr, env));

  struct SubSession {
    std::unique_ptr<TrainingSession> sess;
    SessionOptions so;
    RunOptions run_options;
  };

  std::vector<SubSession> subs(num_subs);
  for (size_t sub_id = 0; sub_id < num_subs; ++sub_id) {
    auto& sub_sess = subs[sub_id];
    sub_sess.so.enable_profiling = true;
    sub_sess.so.profile_file_prefix = GenerateFileNameWithIndex("pipeline", static_cast<int>(sub_id), "");

    sub_sess.run_options.run_log_verbosity_level = sub_sess.so.session_log_verbosity_level;
    sub_sess.run_options.run_tag = sub_sess.so.session_logid;
    sub_sess.run_options.training_mode = true;

    sub_sess.sess = std::make_unique<TrainingSession>(sub_sess.so, *env);
    ASSERT_STATUS_OK(sub_sess.sess->Load(sub_model_files[sub_id]));
    ASSERT_STATUS_OK(sub_sess.sess->Initialize());
  }

  // pipeline inputs for each batch
  struct PipelineFeed {
    OrtValue x_value;
    OrtValue label_value;
    std::vector<OrtValue> record_data_values;
    std::vector<std::pair<OrtValue, OrtValue>> wait_record_pipeline_values;

    void SetInputs(const std::vector<float>& x, const std::vector<float>& label) {
      // dummy data for model inputs
      std::vector<int64_t> x_dims = {1, 784};
      std::vector<int64_t> label_dims = {1, 10};
      TrainingUtil::CreateCpuMLValue<float>(x_dims, x, &x_value);
      TrainingUtil::CreateCpuMLValue<float>(label_dims, label, &label_value);
    }

    void SetEvents(const std::vector<int64_t>& record_data,
                   const std::vector<std::pair<int64_t, int64_t>>& wait_record_pipeline) {
      record_data_values.resize(record_data.size());
      for (size_t i = 0; i < record_data.size(); ++i) {
        TrainingUtil::CreateCpuMLValue<int64_t>({}, {record_data[i]}, &record_data_values[i]);
      }
      wait_record_pipeline_values.resize(wait_record_pipeline.size());
      for (size_t i = 0; i < wait_record_pipeline.size(); ++i) {
        TrainingUtil::CreateCpuMLValue<int64_t>(
            {}, {wait_record_pipeline[i].first},
            &wait_record_pipeline_values[i].first);
        TrainingUtil::CreateCpuMLValue<int64_t>(
            {}, {wait_record_pipeline[i].second},
            &wait_record_pipeline_values[i].second);
      }
    }
  };

  // pipeline data for each batch
  struct PipelineData : public PipelineFeed {
    OrtValue t3_value;
    OrtValue t3_grad_value;
    OrtValue t6_value;
    OrtValue t6_grad_value;

    PipelineData() {
      std::vector<int64_t> t3_dims = {1, 128};
      std::vector<int64_t> t6_dims = {1, 32};
      std::vector<float> t3_data(128);
      std::vector<float> t6_data(32);
      TrainingUtil::CreateCpuMLValue<float>(t3_dims, t3_data, &t3_value);
      TrainingUtil::CreateCpuMLValue<float>(t3_dims, t3_data, &t3_grad_value);
      TrainingUtil::CreateCpuMLValue<float>(t6_dims, t6_data, &t6_value);
      TrainingUtil::CreateCpuMLValue<float>(t6_dims, t6_data, &t6_grad_value);
    };
  };

  auto worker = [&subs](size_t sub_id, PipelineData& data) {
    std::vector<std::string> input_names;
    std::vector<OrtValue> input_values;
    std::vector<std::string> output_names;
    std::vector<OrtValue> output_values;
    switch (sub_id) {
      case 0:
        input_names = {
            "X_sync", "T3_grad_sync",
            "wait_pipeline_0_fw",
            "record_pipeline_0_fw", "record_data_0_fw",
            "wait_data_0_bw", "wait_pipeline_0_bw",
            "record_pipeline_0_bw"};
        input_values = {
            data.x_value, data.t3_grad_value,
            data.wait_record_pipeline_values[0].first,
            data.wait_record_pipeline_values[0].second,
            data.record_data_values[0],
            data.record_data_values[3],
            data.wait_record_pipeline_values[1].first,
            data.wait_record_pipeline_values[1].second};
        output_names = {"T3_sync"};
        output_values = {data.t3_value};
        break;
      case 1:
        input_names = {
            "T3_sync", "T6_grad_sync",
            "wait_data_1_fw", "wait_pipeline_1_fw",
            "record_pipeline_1_fw", "record_data_1_fw",
            "wait_data_1_bw", "wait_pipeline_1_bw",
            "record_pipeline_1_bw", "record_data_1_bw"};
        input_values = {
            data.t3_value, data.t6_grad_value,
            data.record_data_values[0],
            data.wait_record_pipeline_values[2].first,
            data.wait_record_pipeline_values[2].second,
            data.record_data_values[1],
            data.record_data_values[2],
            data.wait_record_pipeline_values[3].first,
            data.wait_record_pipeline_values[3].second,
            data.record_data_values[3]};
        output_names = {"T6_sync", "T3_grad_sync"};
        output_values = {data.t6_value, data.t3_grad_value};
        break;
      case 2:
        // note that last stage only need to wait on FW and record and BW
        // there's no wait/record in between
        input_names = {
            "T6_sync", "labels",
            "wait_data_2_fw", "wait_pipeline_2_fw",
            "record_pipeline_2_bw", "record_data_2_bw"};
        input_values = {
            data.t6_value, data.label_value,
            data.record_data_values[1],
            data.wait_record_pipeline_values[4].first,   // wait on FW
            data.wait_record_pipeline_values[5].second,  // record on BW
            data.record_data_values[2]};
        output_names = {"T6_grad_sync"};
        output_values = {data.t6_grad_value};
        break;
      default:
        ASSERT_TRUE(false);
    }
    EXPECT_STATUS_OK(subs[sub_id].sess->Run(subs[sub_id].run_options, input_names, input_values, output_names, &output_values));
  };

  const std::vector<int64_t> start_ids = {100, 200, 300};
  const std::vector<int64_t> expected_end_ids = {112, 212, 312};
  const size_t num_stages = start_ids.size();
  constexpr int num_batches = 6;
  std::vector<PipelineBatchInfo> plan(num_batches);
  PipelineBatchPlanner planner;
  planner.GenerateOneFWOneBWTimeline(num_stages, num_batches);

  // create plan for all stages for testing purpose
  // in actual execution, only one stage would be needed for each rank
  for (size_t stage = 0; stage < num_stages; ++stage) {
    int64_t end_id = planner.CreatePlan(start_ids[stage], stage, plan);
    EXPECT_TRUE(end_id == expected_end_ids[stage]);
  }

  // Timeline view of ground truth for plan
  // sub 0: F0 F1 F2 F3 F4 B0 F5 B1    B2    B3    B4    B5
  // sub 1:    F0 F1 F2 B0 F3 B1 F4 B2 F5 B3    B4    B5
  // sub 2:       F0 B0 F1 B1 F2 B2 F3 B3 F4 B4 F5 B5
  // Note that in distributed training, event id would be local to each pipeline
  // We use different ranges for event ids for pipelines here:
  // 0 -> 99: data dependencies for record_data
  // 100 -> 199: sub 0
  // 200 -> 299: sub 1
  // 300 -> 399: sub 2
  // so for sub 0, the schedule is:
  //   F0(-1, 100), F1(100,101), F2(101,102)...
  // for sub 1, the schedule is:
  //   F0(-1, 200), F1(200,201), F2(201,202)...
  // for sub 2, the schedule is:
  //   F0 (-1, 300), B0 (300, 301), F1 (301, 302), B1(302, 303)...
  // Note the chart above is timeline view, execution plan needs to change it to batch view
  const std::vector<PipelineBatchInfo> expected_plan = {
      // each batch event pairs are in order of:
      // batch 0 events on {sub0_fw, sub0_bw, sub1_fw, sub1_bw, sub2_fwbw}
      {{{-1, 100}, {104, 105}, {-1, 200}, {202, 203}, {-1, 300}, {300, 301}}, {}},
      // batch 1
      {{{100, 101}, {106, 107}, {200, 201}, {204, 205}, {301, 302}, {302, 303}}, {}},
      // batch 2
      {{{101, 102}, {107, 108}, {201, 202}, {206, 207}, {303, 304}, {304, 305}}, {}},
      // batch 3
      {{{102, 103}, {108, 109}, {203, 204}, {208, 209}, {305, 306}, {306, 307}}, {}},
      // batch 4
      {{{103, 104}, {109, 110}, {205, 206}, {209, 210}, {307, 308}, {308, 309}}, {}},
      // batch 5
      {{{105, 106}, {110, 111}, {207, 208}, {210, 211}, {309, 310}, {310, 311}}, {0}},
  };
  for (int batch = 0; batch < num_batches; ++batch) {
    EXPECT_TRUE(expected_plan[batch].retired_batches == plan[batch].retired_batches);
    EXPECT_TRUE(expected_plan[batch].events.size() == plan[batch].events.size());
    for (size_t evt_id = 0; evt_id < expected_plan[batch].events.size(); ++evt_id) {
      EXPECT_TRUE(expected_plan[batch].events[evt_id] == plan[batch].events[evt_id]);
    }
  }

  struct BatchContext {
    PipelineData data;
    std::vector<std::thread> workers;
  };
  std::unordered_map<int64_t, std::shared_ptr<BatchContext>> batch_ctx_pool;
  for (int batch_id = 0; batch_id < num_batches; ++batch_id) {
    std::vector<float> x(784);
    std::vector<float> label(10);
    std::shared_ptr<BatchContext> batch_ctx;
    bool reuse_batch_ctx = false;
    if (!plan[batch_id].retired_batches.empty()) {
      auto iter = batch_ctx_pool.find(plan[batch_id].retired_batches[0]);
      if (iter != batch_ctx_pool.end()) {
        batch_ctx = iter->second;
        // clean up retired batch, and reclaim data for reuse
        for (auto& w : batch_ctx->workers) {
          w.join();
        }
        batch_ctx->workers.resize(0);
        batch_ctx_pool.erase(plan[batch_id].retired_batches[0]);
        reuse_batch_ctx = true;
      }
    }
    if (!reuse_batch_ctx) {
      batch_ctx = std::make_shared<BatchContext>();
    }

    // set inputs
    batch_ctx->data.SetInputs(x, label);
    batch_ctx->data.SetEvents({batch_id * 4, batch_id * 4 + 1, batch_id * 4 + 2, batch_id * 4 + 3}, plan[batch_id].events);

    // create one worker thread for each batch and each pipeline stage
    for (size_t sub_id = 0; sub_id < num_subs; ++sub_id) {
      auto* pd = &(batch_ctx->data);
      batch_ctx->workers.emplace_back([&worker, pd, sub_id]() {
        worker(sub_id, *pd);
      });
    }
    batch_ctx_pool.emplace(batch_id, batch_ctx);
  }

  // wait until all workers done
  for (auto& pair : batch_ctx_pool) {
    for (auto& w : pair.second->workers) {
      w.join();
    }
  }

  // finish profiler
  for (size_t sub_id = 0; sub_id < num_subs; ++sub_id) {
    subs[sub_id].sess->EndProfiling();
  }
}

}  // namespace test
}  // namespace onnxruntime
