Method profileDataInParallel

compiler/record-minmax/src/RecordMinMax.cpp:495–579

}

void RecordMinMax::profileDataInParallel(const std::string &input_data_path)
{
  LOGGER(l);

  assert(_interpreters.size() == _threads_size);
  assert(_observers.size() == _threads_size);

  const long h5_file_size = getH5FileSize(input_data_path);

  if (h5_file_size > h5_max_size_bytes)
    throw std::runtime_error("H5 file size is too large for parallel recording");

  WholeOutput whole_output;
  try
  {
    whole_output = importH5Data(input_data_path);
  }
  catch (const std::bad_alloc &e)
  {
    throw std::runtime_error("Out of memory during h5 data load.");
  }

  const auto num_records = whole_output.size();
  const auto input_nodes = loco::input_nodes(_module->graph());

  // Start parallel part
  INFO(l) << _threads_size << " concurrent threads are supported." << std::endl;

  const auto run_threads = num_records < _threads_size ? num_records : _threads_size;

  const auto records_batch = static_cast<uint32_t>(num_records / run_threads);

  auto interpret_batch = [&whole_output, &input_nodes](int first_record, int last_record,
                                                       luci_interpreter::Interpreter *interpreter) {
    for (int record_index = first_record; record_index < last_record; ++record_index)
    {
      for (uint32_t input_idx = 0; input_idx < input_nodes.size(); input_idx++)
      {
        const auto *input_node = loco::must_cast<const luci::CircleInput *>(input_nodes[input_idx]);

        const auto &cur_input_data = whole_output[record_index][input_idx];
        interpreter->writeInputTensor(input_node, cur_input_data.data(), cur_input_data.size());
      }
      interpreter->interpret();
    }
  };

  std::vector<std::thread> threads;
  for (uint32_t t = 0; t < run_threads; ++t)
  {
    if (t < run_threads - 1)
    {
      threads.emplace_back(interpret_batch, records_batch * t, records_batch * (t + 1),
                           _interpreters[t].get());
    }
    else
    {
      threads.emplace_back(interpret_batch, records_batch * t, num_records, _interpreters[t].get());