Method LoadIntoMemory

paddle/fluid/framework/data_set.cc:489–597

// which will later be fed into readers' channel
template <typename T>
void DatasetImpl<T>::LoadIntoMemory() {
  VLOG(3) << "DatasetImpl<T>::LoadIntoMemory() begin";
  platform::Timer timeline;
  timeline.Start();
  if (gpu_graph_mode_) {
    VLOG(1) << "in gpu_graph_mode";
#if defined(PADDLE_WITH_PSCORE) && defined(PADDLE_WITH_HETERPS)
    std::vector<std::future<void>> wait_futures;
    auto pool = GetReadThreadPool(thread_num_);
    for (size_t i = 0; i < readers_.size(); i++) {
      readers_[i]->SetGpuGraphMode(gpu_graph_mode_);
    }

    if (STAT_GET(STAT_epoch_finish) == 1) {
      VLOG(0) << "get epoch finish true";
      STAT_RESET(STAT_epoch_finish, 0);
      for (size_t i = 0; i < readers_.size(); i++) {
        readers_[i]->ResetPathNum();
        readers_[i]->ResetEpochFinish();
      }
    }

    for (int64_t i = 0; i < thread_num_; ++i) {
      wait_futures.emplace_back(
          pool[i]->Run([this, i]() { readers_[i]->DoWalkandSage(); }));
    }
    for (auto& th : wait_futures) {
      th.get();
    }
    wait_futures.clear();

    uint64_t node_num = 0;
    std::vector<uint64_t> offsets;
    offsets.resize(thread_num_);

    for (int i = 0; i < thread_num_; i++) {
      auto& host_vec = (*readers_[i]->GetHostVec());
      offsets[i] = node_num;
      node_num += host_vec.size();
    }
    gpu_graph_total_keys_.resize(node_num + 1);
    for (int i = 0; i < thread_num_; i++) {
      uint64_t off = offsets[i];
      wait_futures.emplace_back(pool[i]->Run([this, i, off]() {
        auto& host_vec = (*readers_[i]->GetHostVec());
        for (size_t j = 0; j < host_vec.size(); j++) {
          gpu_graph_total_keys_[off + j] = host_vec[j];
        }
        if (FLAGS_gpugraph_storage_mode != GpuGraphStorageMode::WHOLE_HBM) {
          readers_[i]->clear_gpu_mem();
        }
      }));
    }
    for (auto& th : wait_futures) {
      th.get();
    }
    wait_futures.clear();