data_feed.h 61.8 KB
Newer Older
W
Wang Guibao 已提交
1 2 3 4 5 6
/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

L
lxsbupt 已提交
7
http://www.apache.org/licenses/LICENSE-2.0
W
Wang Guibao 已提交
8 9 10 11 12 13 14 15 16

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */

#pragma once

J
jiaqi 已提交
17 18 19 20 21
#if defined _WIN32 || defined __APPLE__
#else
#define _LINUX
#endif

W
Wang Guibao 已提交
22
#include <fstream>
23
#include <future>  // NOLINT
W
Wang Guibao 已提交
24 25
#include <memory>
#include <mutex>  // NOLINT
D
danleifeng 已提交
26
#include <random>
27
#include <sstream>
W
Wang Guibao 已提交
28 29
#include <string>
#include <thread>  // NOLINT
30
#include <unordered_map>
31
#include <unordered_set>
32
#include <utility>
33
#include <vector>
W
Wang Guibao 已提交
34

J
jiaqi 已提交
35
#include "paddle/fluid/framework/archive.h"
36
#include "paddle/fluid/framework/blocking_queue.h"
J
jiaqi 已提交
37
#include "paddle/fluid/framework/channel.h"
W
Wang Guibao 已提交
38
#include "paddle/fluid/framework/data_feed.pb.h"
39
#include "paddle/fluid/framework/fleet/fleet_wrapper.h"
W
Wang Guibao 已提交
40 41 42
#include "paddle/fluid/framework/lod_tensor.h"
#include "paddle/fluid/framework/reader.h"
#include "paddle/fluid/framework/variable.h"
Y
yaoxuefeng 已提交
43
#include "paddle/fluid/platform/timer.h"
44
#include "paddle/fluid/string/string_helper.h"
45
#if defined(PADDLE_WITH_CUDA)
D
danleifeng 已提交
46
#include "paddle/fluid/framework/fleet/heter_ps/gpu_graph_utils.h"
47 48
#include "paddle/fluid/platform/cuda_device_guard.h"
#include "paddle/fluid/platform/device/gpu/gpu_info.h"
P
pangengzheng 已提交
49
#include "paddle/phi/core/cuda_stream.h"
50
#endif
W
Wang Guibao 已提交
51

Y
yaoxuefeng 已提交
52 53 54 55 56
DECLARE_int32(record_pool_max_size);
DECLARE_int32(slotpool_thread_num);
DECLARE_bool(enable_slotpool_wait_release);
DECLARE_bool(enable_slotrecord_reset_shrink);

W
wanghuancoder 已提交
57 58 59 60 61
namespace paddle {
namespace framework {
class DataFeedDesc;
class Scope;
class Variable;
D
danleifeng 已提交
62 63
class NeighborSampleResult;
class NodeQueryResult;
L
lxsbupt 已提交
64 65
template <typename KeyType, typename ValType>
class HashTable;
W
wanghuancoder 已提交
66 67 68
}  // namespace framework
}  // namespace paddle

69
namespace phi {
70
class DenseTensor;
71
}  // namespace phi
72

W
Wang Guibao 已提交
73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91
namespace paddle {
namespace framework {

// DataFeed is the base virtual class for all ohther DataFeeds.
// It is used to read files and parse the data for subsequent trainer.
// Example:
//   DataFeed* reader =
//   paddle::framework::DataFeedFactory::CreateDataFeed(data_feed_name);
//   reader->Init(data_feed_desc); // data_feed_desc is a protobuf object
//   reader->SetFileList(filelist);
//   const std::vector<std::string> & use_slot_alias =
//   reader->GetUseSlotAlias();
//   for (auto name: use_slot_alias){ // for binding memory
//     reader->AddFeedVar(scope->Var(name), name);
//   }
//   reader->Start();
//   while (reader->Next()) {
//      // trainer do something
//   }
Y
yaoxuefeng 已提交
92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135

template <typename T>
struct SlotValues {
  std::vector<T> slot_values;
  std::vector<uint32_t> slot_offsets;

  void add_values(const T* values, uint32_t num) {
    if (slot_offsets.empty()) {
      slot_offsets.push_back(0);
    }
    if (num > 0) {
      slot_values.insert(slot_values.end(), values, values + num);
    }
    slot_offsets.push_back(static_cast<uint32_t>(slot_values.size()));
  }
  T* get_values(int idx, size_t* size) {
    uint32_t& offset = slot_offsets[idx];
    (*size) = slot_offsets[idx + 1] - offset;
    return &slot_values[offset];
  }
  void add_slot_feasigns(const std::vector<std::vector<T>>& slot_feasigns,
                         uint32_t fea_num) {
    slot_values.reserve(fea_num);
    int slot_num = static_cast<int>(slot_feasigns.size());
    slot_offsets.resize(slot_num + 1);
    for (int i = 0; i < slot_num; ++i) {
      auto& slot_val = slot_feasigns[i];
      slot_offsets[i] = static_cast<uint32_t>(slot_values.size());
      uint32_t num = static_cast<uint32_t>(slot_val.size());
      if (num > 0) {
        slot_values.insert(slot_values.end(), slot_val.begin(), slot_val.end());
      }
    }
    slot_offsets[slot_num] = slot_values.size();
  }
  void clear(bool shrink) {
    slot_offsets.clear();
    slot_values.clear();
    if (shrink) {
      slot_values.shrink_to_fit();
      slot_offsets.shrink_to_fit();
    }
  }
};
T
Thunderbrook 已提交
136
union FeatureFeasign {
137 138 139 140 141 142
  uint64_t uint64_feasign_;
  float float_feasign_;
};

struct FeatureItem {
  FeatureItem() {}
T
Thunderbrook 已提交
143
  FeatureItem(FeatureFeasign sign, uint16_t slot) {
144 145 146
    this->sign() = sign;
    this->slot() = slot;
  }
T
Thunderbrook 已提交
147 148 149 150 151 152
  FeatureFeasign& sign() {
    return *(reinterpret_cast<FeatureFeasign*>(sign_buffer()));
  }
  const FeatureFeasign& sign() const {
    const FeatureFeasign* ret =
        reinterpret_cast<FeatureFeasign*>(sign_buffer());
153 154 155 156 157 158 159
    return *ret;
  }
  uint16_t& slot() { return slot_; }
  const uint16_t& slot() const { return slot_; }

 private:
  char* sign_buffer() const { return const_cast<char*>(sign_); }
T
Thunderbrook 已提交
160
  char sign_[sizeof(FeatureFeasign)];
161 162 163
  uint16_t slot_;
};

Y
yaoxuefeng 已提交
164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195
struct AllSlotInfo {
  std::string slot;
  std::string type;
  int used_idx;
  int slot_value_idx;
};
struct UsedSlotInfo {
  int idx;
  int slot_value_idx;
  std::string slot;
  std::string type;
  bool dense;
  std::vector<int> local_shape;
  int total_dims_without_inductive;
  int inductive_shape_index;
};
struct SlotRecordObject {
  uint64_t search_id;
  uint32_t rank;
  uint32_t cmatch;
  std::string ins_id_;
  SlotValues<uint64_t> slot_uint64_feasigns_;
  SlotValues<float> slot_float_feasigns_;

  ~SlotRecordObject() { clear(true); }
  void reset(void) { clear(FLAGS_enable_slotrecord_reset_shrink); }
  void clear(bool shrink) {
    slot_uint64_feasigns_.clear(shrink);
    slot_float_feasigns_.clear(shrink);
  }
};
using SlotRecord = SlotRecordObject*;
196 197 198 199 200 201 202 203 204
// sizeof Record is much less than std::vector<MultiSlotType>
struct Record {
  std::vector<FeatureItem> uint64_feasigns_;
  std::vector<FeatureItem> float_feasigns_;
  std::string ins_id_;
  std::string content_;
  uint64_t search_id;
  uint32_t rank;
  uint32_t cmatch;
205
  std::string uid_;
206 207
};

Y
yaoxuefeng 已提交
208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380
inline SlotRecord make_slotrecord() {
  static const size_t slot_record_byte_size = sizeof(SlotRecordObject);
  void* p = malloc(slot_record_byte_size);
  new (p) SlotRecordObject;
  return reinterpret_cast<SlotRecordObject*>(p);
}

inline void free_slotrecord(SlotRecordObject* p) {
  p->~SlotRecordObject();
  free(p);
}

template <class T>
class SlotObjAllocator {
 public:
  explicit SlotObjAllocator(std::function<void(T*)> deleter)
      : free_nodes_(NULL), capacity_(0), deleter_(deleter) {}
  ~SlotObjAllocator() { clear(); }

  void clear() {
    T* tmp = NULL;
    while (free_nodes_ != NULL) {
      tmp = reinterpret_cast<T*>(reinterpret_cast<void*>(free_nodes_));
      free_nodes_ = free_nodes_->next;
      deleter_(tmp);
      --capacity_;
    }
    CHECK_EQ(capacity_, static_cast<size_t>(0));
  }
  T* acquire(void) {
    T* x = NULL;
    x = reinterpret_cast<T*>(reinterpret_cast<void*>(free_nodes_));
    free_nodes_ = free_nodes_->next;
    --capacity_;
    return x;
  }
  void release(T* x) {
    Node* node = reinterpret_cast<Node*>(reinterpret_cast<void*>(x));
    node->next = free_nodes_;
    free_nodes_ = node;
    ++capacity_;
  }
  size_t capacity(void) { return capacity_; }

 private:
  struct alignas(T) Node {
    union {
      Node* next;
      char data[sizeof(T)];
    };
  };
  Node* free_nodes_;  // a list
  size_t capacity_;
  std::function<void(T*)> deleter_ = nullptr;
};
static const int OBJPOOL_BLOCK_SIZE = 10000;
class SlotObjPool {
 public:
  SlotObjPool()
      : max_capacity_(FLAGS_record_pool_max_size), alloc_(free_slotrecord) {
    ins_chan_ = MakeChannel<SlotRecord>();
    ins_chan_->SetBlockSize(OBJPOOL_BLOCK_SIZE);
    for (int i = 0; i < FLAGS_slotpool_thread_num; ++i) {
      threads_.push_back(std::thread([this]() { run(); }));
    }
    disable_pool_ = false;
    count_ = 0;
  }
  ~SlotObjPool() {
    ins_chan_->Close();
    for (auto& t : threads_) {
      t.join();
    }
  }
  void disable_pool(bool disable) { disable_pool_ = disable; }
  void set_max_capacity(size_t max_capacity) { max_capacity_ = max_capacity; }
  void get(std::vector<SlotRecord>* output, int n) {
    output->resize(n);
    return get(&(*output)[0], n);
  }
  void get(SlotRecord* output, int n) {
    int size = 0;
    mutex_.lock();
    int left = static_cast<int>(alloc_.capacity());
    if (left > 0) {
      size = (left >= n) ? n : left;
      for (int i = 0; i < size; ++i) {
        output[i] = alloc_.acquire();
      }
    }
    mutex_.unlock();
    count_ += n;
    if (size == n) {
      return;
    }
    for (int i = size; i < n; ++i) {
      output[i] = make_slotrecord();
    }
  }
  void put(std::vector<SlotRecord>* input) {
    size_t size = input->size();
    if (size == 0) {
      return;
    }
    put(&(*input)[0], size);
    input->clear();
  }
  void put(SlotRecord* input, size_t size) {
    CHECK(ins_chan_->WriteMove(size, input) == size);
  }
  void run(void) {
    std::vector<SlotRecord> input;
    while (ins_chan_->ReadOnce(input, OBJPOOL_BLOCK_SIZE)) {
      if (input.empty()) {
        continue;
      }
      // over max capacity
      size_t n = input.size();
      count_ -= n;
      if (disable_pool_ || n + capacity() > max_capacity_) {
        for (auto& t : input) {
          free_slotrecord(t);
        }
      } else {
        for (auto& t : input) {
          t->reset();
        }
        mutex_.lock();
        for (auto& t : input) {
          alloc_.release(t);
        }
        mutex_.unlock();
      }
      input.clear();
    }
  }
  void clear(void) {
    platform::Timer timeline;
    timeline.Start();
    mutex_.lock();
    alloc_.clear();
    mutex_.unlock();
    // wait release channel data
    if (FLAGS_enable_slotpool_wait_release) {
      while (!ins_chan_->Empty()) {
        sleep(1);
      }
    }
    timeline.Pause();
    VLOG(3) << "clear slot pool data size=" << count_.load()
            << ", span=" << timeline.ElapsedSec();
  }
  size_t capacity(void) {
    mutex_.lock();
    size_t total = alloc_.capacity();
    mutex_.unlock();
    return total;
  }

 private:
  size_t max_capacity_;
  Channel<SlotRecord> ins_chan_;
  std::vector<std::thread> threads_;
  std::mutex mutex_;
  SlotObjAllocator<SlotRecordObject> alloc_;
  bool disable_pool_;
  std::atomic<long> count_;  // NOLINT
};

inline SlotObjPool& SlotRecordPool() {
  static SlotObjPool pool;
  return pool;
}
381 382 383 384 385 386 387 388 389
struct PvInstanceObject {
  std::vector<Record*> ads;
  void merge_instance(Record* ins) { ads.push_back(ins); }
};

using PvInstance = PvInstanceObject*;

inline PvInstance make_pv_instance() { return new PvInstanceObject(); }

T
Thunderbrook 已提交
390 391 392 393 394 395 396 397 398 399 400 401
struct SlotConf {
  std::string name;
  std::string type;
  int use_slots_index;
  int use_slots_is_dense;
};

class CustomParser {
 public:
  CustomParser() {}
  virtual ~CustomParser() {}
  virtual void Init(const std::vector<SlotConf>& slots) = 0;
T
Thunderbrook 已提交
402
  virtual bool Init(const std::vector<AllSlotInfo>& slots) = 0;
T
Thunderbrook 已提交
403
  virtual void ParseOneInstance(const char* str, Record* instance) = 0;
404 405
  virtual int ParseInstance(int len,
                            const char* str,
T
Thunderbrook 已提交
406 407
                            std::vector<Record>* instances) {
    return 0;
408
  }
Y
yaoxuefeng 已提交
409 410 411 412 413 414 415 416 417 418 419 420 421
  virtual bool ParseOneInstance(
      const std::string& line,
      std::function<void(std::vector<SlotRecord>&, int)>
          GetInsFunc) {  // NOLINT
    return true;
  }
  virtual bool ParseFileInstance(
      std::function<int(char* buf, int len)> ReadBuffFunc,
      std::function<void(std::vector<SlotRecord>&, int, int)>
          PullRecordsFunc,  // NOLINT
      int& lines) {         // NOLINT
    return false;
  }
T
Thunderbrook 已提交
422 423
};

424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485
struct UsedSlotGpuType {
  int is_uint64_value;
  int slot_value_idx;
};

#if defined(PADDLE_WITH_CUDA) && defined(PADDLE_WITH_HETERPS)
template <typename T>
struct CudaBuffer {
  T* cu_buffer;
  uint64_t buf_size;

  CudaBuffer<T>() {
    cu_buffer = NULL;
    buf_size = 0;
  }
  ~CudaBuffer<T>() { free(); }
  T* data() { return cu_buffer; }
  uint64_t size() { return buf_size; }
  void malloc(uint64_t size) {
    buf_size = size;
    CUDA_CHECK(
        cudaMalloc(reinterpret_cast<void**>(&cu_buffer), size * sizeof(T)));
  }
  void free() {
    if (cu_buffer != NULL) {
      CUDA_CHECK(cudaFree(cu_buffer));
      cu_buffer = NULL;
    }
    buf_size = 0;
  }
  void resize(uint64_t size) {
    if (size <= buf_size) {
      return;
    }
    free();
    malloc(size);
  }
};
template <typename T>
struct HostBuffer {
  T* host_buffer;
  size_t buf_size;
  size_t data_len;

  HostBuffer<T>() {
    host_buffer = NULL;
    buf_size = 0;
    data_len = 0;
  }
  ~HostBuffer<T>() { free(); }

  T* data() { return host_buffer; }
  const T* data() const { return host_buffer; }
  size_t size() const { return data_len; }
  void clear() { free(); }
  T& back() { return host_buffer[data_len - 1]; }

  T& operator[](size_t i) { return host_buffer[i]; }
  const T& operator[](size_t i) const { return host_buffer[i]; }
  void malloc(size_t len) {
    buf_size = len;
    CUDA_CHECK(cudaHostAlloc(reinterpret_cast<void**>(&host_buffer),
486 487
                             buf_size * sizeof(T),
                             cudaHostAllocDefault));
488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538
    CHECK(host_buffer != NULL);
  }
  void free() {
    if (host_buffer != NULL) {
      CUDA_CHECK(cudaFreeHost(host_buffer));
      host_buffer = NULL;
    }
    buf_size = 0;
  }
  void resize(size_t size) {
    if (size <= buf_size) {
      data_len = size;
      return;
    }
    data_len = size;
    free();
    malloc(size);
  }
};

struct BatchCPUValue {
  HostBuffer<int> h_uint64_lens;
  HostBuffer<uint64_t> h_uint64_keys;
  HostBuffer<int> h_uint64_offset;

  HostBuffer<int> h_float_lens;
  HostBuffer<float> h_float_keys;
  HostBuffer<int> h_float_offset;

  HostBuffer<int> h_rank;
  HostBuffer<int> h_cmatch;
  HostBuffer<int> h_ad_offset;
};

struct BatchGPUValue {
  CudaBuffer<int> d_uint64_lens;
  CudaBuffer<uint64_t> d_uint64_keys;
  CudaBuffer<int> d_uint64_offset;

  CudaBuffer<int> d_float_lens;
  CudaBuffer<float> d_float_keys;
  CudaBuffer<int> d_float_offset;

  CudaBuffer<int> d_rank;
  CudaBuffer<int> d_cmatch;
  CudaBuffer<int> d_ad_offset;
};

class MiniBatchGpuPack {
 public:
  MiniBatchGpuPack(const paddle::platform::Place& place,
P
pangengzheng 已提交
539 540
                   const std::vector<UsedSlotInfo>& infos,
                   phi::StreamId stream_id);
541
  ~MiniBatchGpuPack();
P
pangengzheng 已提交
542 543
  bool is_use() { return is_using_; }
  void set_use_flag(bool is_use) { is_using_ = is_use; }
544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570
  void reset(const paddle::platform::Place& place);
  void pack_instance(const SlotRecord* ins_vec, int num);
  int ins_num() { return ins_num_; }
  int pv_num() { return pv_num_; }
  BatchGPUValue& value() { return value_; }
  BatchCPUValue& cpu_value() { return buf_; }
  UsedSlotGpuType* get_gpu_slots(void) {
    return reinterpret_cast<UsedSlotGpuType*>(gpu_slots_.data());
  }
  SlotRecord* get_records(void) { return &ins_vec_[0]; }

  // tensor gpu memory reused
  void resize_tensor(void) {
    if (used_float_num_ > 0) {
      int float_total_len = buf_.h_float_lens.back();
      if (float_total_len > 0) {
        float_tensor_.mutable_data<float>({float_total_len, 1}, this->place_);
      }
    }
    if (used_uint64_num_ > 0) {
      int uint64_total_len = buf_.h_uint64_lens.back();
      if (uint64_total_len > 0) {
        uint64_tensor_.mutable_data<int64_t>({uint64_total_len, 1},
                                             this->place_);
      }
    }
  }
571 572
  phi::DenseTensor& float_tensor(void) { return float_tensor_; }
  phi::DenseTensor& uint64_tensor(void) { return uint64_tensor_; }
P
pangengzheng 已提交
573 574 575 576 577 578
  std::vector<phi::DenseTensor>& float_tensor_vec(void) {
    return float_tensor_vec_;
  }
  std::vector<phi::DenseTensor>& uint64_tensor_vec(void) {
    return uint64_tensor_vec_;
  }
579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602

  HostBuffer<size_t>& offsets(void) { return offsets_; }
  HostBuffer<void*>& h_tensor_ptrs(void) { return h_tensor_ptrs_; }

  void* gpu_slot_offsets(void) { return gpu_slot_offsets_->ptr(); }

  void* slot_buf_ptr(void) { return slot_buf_ptr_->ptr(); }

  void resize_gpu_slot_offsets(const size_t slot_total_bytes) {
    if (gpu_slot_offsets_ == nullptr) {
      gpu_slot_offsets_ = memory::AllocShared(place_, slot_total_bytes);
    } else if (gpu_slot_offsets_->size() < slot_total_bytes) {
      auto buf = memory::AllocShared(place_, slot_total_bytes);
      gpu_slot_offsets_.swap(buf);
      buf = nullptr;
    }
  }
  const std::string& get_lineid(int idx) {
    if (enable_pv_) {
      return ins_vec_[idx]->ins_id_;
    }
    return batch_ins_[idx]->ins_id_;
  }

P
pangengzheng 已提交
603 604
  cudaStream_t get_stream() { return stream_; }

605 606 607 608 609 610 611 612 613 614 615 616 617
 private:
  void transfer_to_gpu(void);
  void pack_all_data(const SlotRecord* ins_vec, int num);
  void pack_uint64_data(const SlotRecord* ins_vec, int num);
  void pack_float_data(const SlotRecord* ins_vec, int num);

 public:
  template <typename T>
  void copy_host2device(CudaBuffer<T>* buf, const T* val, size_t size) {
    if (size == 0) {
      return;
    }
    buf->resize(size);
618 619
    CUDA_CHECK(cudaMemcpyAsync(
        buf->data(), val, size * sizeof(T), cudaMemcpyHostToDevice, stream_));
620 621 622 623 624 625 626
  }
  template <typename T>
  void copy_host2device(CudaBuffer<T>* buf, const HostBuffer<T>& val) {
    copy_host2device(buf, val.data(), val.size());
  }

 private:
P
pangengzheng 已提交
627
  bool is_using_ = false;
628
  paddle::platform::Place place_;
P
pangengzheng 已提交
629
  std::unique_ptr<phi::CUDAStream> stream_holder_;
630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646
  cudaStream_t stream_;
  BatchGPUValue value_;
  BatchCPUValue buf_;
  int ins_num_ = 0;
  int pv_num_ = 0;

  bool enable_pv_ = false;
  int used_float_num_ = 0;
  int used_uint64_num_ = 0;
  int used_slot_size_ = 0;

  CudaBuffer<UsedSlotGpuType> gpu_slots_;
  std::vector<UsedSlotGpuType> gpu_used_slots_;
  std::vector<SlotRecord> ins_vec_;
  const SlotRecord* batch_ins_ = nullptr;

  // uint64 tensor
647
  phi::DenseTensor uint64_tensor_;
P
pangengzheng 已提交
648
  std::vector<phi::DenseTensor> uint64_tensor_vec_;
649
  // float tensor
650
  phi::DenseTensor float_tensor_;
P
pangengzheng 已提交
651
  std::vector<phi::DenseTensor> float_tensor_vec_;
652 653 654 655 656 657 658 659 660 661 662 663
  // batch
  HostBuffer<size_t> offsets_;
  HostBuffer<void*> h_tensor_ptrs_;

  std::shared_ptr<phi::Allocation> gpu_slot_offsets_ = nullptr;
  std::shared_ptr<phi::Allocation> slot_buf_ptr_ = nullptr;
};
class MiniBatchGpuPackMgr {
  static const int MAX_DEIVCE_NUM = 16;

 public:
  MiniBatchGpuPackMgr() {
P
pangengzheng 已提交
664
    pack_list_.resize(MAX_DEIVCE_NUM);
665
    for (int i = 0; i < MAX_DEIVCE_NUM; ++i) {
P
pangengzheng 已提交
666
      pack_list_[i].clear();
667 668 669 670
    }
  }
  ~MiniBatchGpuPackMgr() {
    for (int i = 0; i < MAX_DEIVCE_NUM; ++i) {
P
pangengzheng 已提交
671 672 673 674 675 676
      for (size_t j = 0; j < pack_list_[i].size(); j++) {
        if (pack_list_[i][j] == nullptr) {
          continue;
        }
        delete pack_list_[i][j];
        pack_list_[i][j] = nullptr;
677 678 679
      }
    }
  }
P
pangengzheng 已提交
680 681

  // thread unsafe
682 683 684
  MiniBatchGpuPack* get(const paddle::platform::Place& place,
                        const std::vector<UsedSlotInfo>& infos) {
    int device_id = place.GetDeviceId();
P
pangengzheng 已提交
685 686 687 688 689 690
    for (size_t i = 0; i < pack_list_[device_id].size(); i++) {
      if (!pack_list_[device_id][i]->is_use()) {
        pack_list_[device_id][i]->set_use_flag(true);
        pack_list_[device_id][i]->reset(place);
        return pack_list_[device_id][i];
      }
691
    }
P
pangengzheng 已提交
692 693 694 695 696 697 698 699 700 701 702 703
    {
      std::lock_guard<std::mutex> lock(mutex_);
      if (!alloc_stream_map_.count(device_id)) {
        alloc_stream_map_.emplace(device_id, new phi::CUDAStream(place));
      }
    }
    phi::StreamId alloc_stream_id = reinterpret_cast<phi::StreamId>(
        alloc_stream_map_[device_id]->raw_stream());
    auto* pack = new MiniBatchGpuPack(place, infos, alloc_stream_id);
    pack->set_use_flag(true);
    pack_list_[device_id].push_back(pack);
    return pack;
704 705 706
  }

 private:
P
pangengzheng 已提交
707 708 709
  std::vector<std::vector<MiniBatchGpuPack*>> pack_list_;
  std::unordered_map<int, std::unique_ptr<phi::CUDAStream>> alloc_stream_map_;
  std::mutex mutex_;
710 711 712 713 714 715 716 717
};
// global mgr
inline MiniBatchGpuPackMgr& BatchGpuPackMgr() {
  static MiniBatchGpuPackMgr mgr;
  return mgr;
}
#endif

T
Thunderbrook 已提交
718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752
typedef paddle::framework::CustomParser* (*CreateParserObjectFunc)();

class DLManager {
  struct DLHandle {
    void* module;
    paddle::framework::CustomParser* parser;
  };

 public:
  DLManager() {}

  ~DLManager() {
#ifdef _LINUX
    std::lock_guard<std::mutex> lock(mutex_);
    for (auto it = handle_map_.begin(); it != handle_map_.end(); ++it) {
      delete it->second.parser;
      dlclose(it->second.module);
    }
#endif
  }

  bool Close(const std::string& name) {
#ifdef _LINUX
    auto it = handle_map_.find(name);
    if (it == handle_map_.end()) {
      return true;
    }
    delete it->second.parser;
    dlclose(it->second.module);
#endif
    VLOG(0) << "Not implement in windows";
    return false;
  }

  paddle::framework::CustomParser* Load(const std::string& name,
Y
yaoxuefeng 已提交
753
                                        const std::vector<SlotConf>& conf) {
T
Thunderbrook 已提交
754 755 756 757 758 759 760 761 762 763
#ifdef _LINUX
    std::lock_guard<std::mutex> lock(mutex_);
    DLHandle handle;
    std::map<std::string, DLHandle>::iterator it = handle_map_.find(name);
    if (it != handle_map_.end()) {
      return it->second.parser;
    }

    handle.module = dlopen(name.c_str(), RTLD_NOW);
    if (handle.module == nullptr) {
T
Thunderbrook 已提交
764
      VLOG(0) << "Create so of " << name << " fail, " << dlerror();
T
Thunderbrook 已提交
765 766 767 768 769 770 771 772 773 774 775 776 777 778 779
      return nullptr;
    }

    CreateParserObjectFunc create_parser_func =
        (CreateParserObjectFunc)dlsym(handle.module, "CreateParserObject");
    handle.parser = create_parser_func();
    handle.parser->Init(conf);
    handle_map_.insert({name, handle});

    return handle.parser;
#endif
    VLOG(0) << "Not implement in windows";
    return nullptr;
  }

Y
yaoxuefeng 已提交
780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807
  paddle::framework::CustomParser* Load(const std::string& name,
                                        const std::vector<AllSlotInfo>& conf) {
#ifdef _LINUX
    std::lock_guard<std::mutex> lock(mutex_);
    DLHandle handle;
    std::map<std::string, DLHandle>::iterator it = handle_map_.find(name);
    if (it != handle_map_.end()) {
      return it->second.parser;
    }
    handle.module = dlopen(name.c_str(), RTLD_NOW);
    if (handle.module == nullptr) {
      VLOG(0) << "Create so of " << name << " fail";
      exit(-1);
      return nullptr;
    }

    CreateParserObjectFunc create_parser_func =
        (CreateParserObjectFunc)dlsym(handle.module, "CreateParserObject");
    handle.parser = create_parser_func();
    handle.parser->Init(conf);
    handle_map_.insert({name, handle});

    return handle.parser;
#endif
    VLOG(0) << "Not implement in windows";
    return nullptr;
  }

T
Thunderbrook 已提交
808
  paddle::framework::CustomParser* ReLoad(const std::string& name,
Y
yaoxuefeng 已提交
809
                                          const std::vector<SlotConf>& conf) {
T
Thunderbrook 已提交
810 811 812 813 814 815 816 817 818
    Close(name);
    return Load(name, conf);
  }

 private:
  std::mutex mutex_;
  std::map<std::string, DLHandle> handle_map_;
};

D
danleifeng 已提交
819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917
struct engine_wrapper_t {
  std::default_random_engine engine;
#if !defined(_WIN32)
  engine_wrapper_t() {
    struct timespec tp;
    clock_gettime(CLOCK_REALTIME, &tp);
    double cur_time = tp.tv_sec + tp.tv_nsec * 1e-9;
    static std::atomic<uint64_t> x(0);
    std::seed_seq sseq = {x++, x++, x++, (uint64_t)(cur_time * 1000)};
    engine.seed(sseq);
  }
#endif
};

struct BufState {
  int left;
  int right;
  int central_word;
  int step;
  engine_wrapper_t random_engine_;

  int len;
  int cursor;
  int row_num;

  int batch_size;
  int walk_len;
  std::vector<int>* window;

  BufState() {}
  ~BufState() {}

  void Init(int graph_batch_size,
            int graph_walk_len,
            std::vector<int>* graph_window) {
    batch_size = graph_batch_size;
    walk_len = graph_walk_len;
    window = graph_window;

    left = 0;
    right = window->size() - 1;
    central_word = -1;
    step = -1;

    len = 0;
    cursor = 0;
    row_num = 0;
    for (size_t i = 0; i < graph_window->size(); i++) {
      VLOG(2) << "graph_window[" << i << "] = " << (*graph_window)[i];
    }
  }

  void Reset(int total_rows) {
    cursor = 0;
    row_num = total_rows;
    int tmp_len = cursor + batch_size > row_num ? row_num - cursor : batch_size;
    len = tmp_len;
    central_word = -1;
    step = -1;
    GetNextCentrolWord();
  }

  int GetNextStep() {
    step++;
    if (step <= right && central_word + (*window)[step] < walk_len) {
      return 1;
    }
    return 0;
  }

  void Debug() {
    VLOG(2) << "left: " << left << " right: " << right
            << " central_word: " << central_word << " step: " << step
            << " cursor: " << cursor << " len: " << len
            << " row_num: " << row_num;
  }

  int GetNextCentrolWord() {
    if (++central_word >= walk_len) {
      return 0;
    }
    int window_size = window->size() / 2;
    int random_window = random_engine_.engine() % window_size + 1;
    left = window_size - random_window;
    right = window_size + random_window - 1;
    VLOG(2) << "random window: " << random_window << " window[" << left
            << "] = " << (*window)[left] << " window[" << right
            << "] = " << (*window)[right];

    for (step = left; step <= right; step++) {
      if (central_word + (*window)[step] >= 0) {
        return 1;
      }
    }
    return 0;
  }

  int GetNextBatch() {
    cursor += len;
L
lxsbupt 已提交
918 919 920
    if (row_num - cursor < 0) {
      return 0;
    }
D
danleifeng 已提交
921 922 923 924 925 926 927 928 929 930 931 932 933 934
    int tmp_len = cursor + batch_size > row_num ? row_num - cursor : batch_size;
    if (tmp_len == 0) {
      return 0;
    }
    len = tmp_len;
    central_word = -1;
    step = -1;
    GetNextCentrolWord();
    return tmp_len != 0;
  }
};

class GraphDataGenerator {
 public:
935 936
  GraphDataGenerator() {}
  virtual ~GraphDataGenerator() {}
D
danleifeng 已提交
937
  void SetConfig(const paddle::framework::DataFeedDesc& data_feed_desc);
L
lxsbupt 已提交
938 939 940
  void AllocResource(int thread_id, std::vector<phi::DenseTensor*> feed_vec);
  void AllocTrainResource(int thread_id);
  void SetFeedVec(std::vector<phi::DenseTensor*> feed_vec);
D
danleifeng 已提交
941 942
  int AcquireInstance(BufState* state);
  int GenerateBatch();
L
lxsbupt 已提交
943 944 945 946 947
  int FillWalkBuf();
  int FillWalkBufMultiPath();
  int FillInferBuf();
  void DoWalkandSage();
  int FillSlotFeature(uint64_t* d_walk);
D
danleifeng 已提交
948 949 950 951
  int FillFeatureBuf(uint64_t* d_walk, uint64_t* d_feature, size_t key_num);
  int FillFeatureBuf(std::shared_ptr<phi::Allocation> d_walk,
                     std::shared_ptr<phi::Allocation> d_feature);
  void FillOneStep(uint64_t* start_ids,
952
                   int etype_id,
D
danleifeng 已提交
953
                   uint64_t* walk,
954
                   uint8_t* walk_ntype,
D
danleifeng 已提交
955
                   int len,
956
                   NeighborSampleResult& sample_res,  // NOLINT
D
danleifeng 已提交
957 958 959
                   int cur_degree,
                   int step,
                   int* len_per_row);
L
lxsbupt 已提交
960 961 962 963 964 965 966 967 968 969 970 971 972 973 974 975 976 977
  int FillInsBuf(cudaStream_t stream);
  int FillIdShowClkTensor(int total_instance,
                          bool gpu_graph_training,
                          size_t cursor = 0);
  int FillGraphIdShowClkTensor(int uniq_instance,
                               int total_instance,
                               int index);
  int FillGraphSlotFeature(
      int total_instance,
      bool gpu_graph_training,
      std::shared_ptr<phi::Allocation> final_sage_nodes = nullptr);
  int FillSlotFeature(uint64_t* d_walk, size_t key_num);
  int MakeInsPair(cudaStream_t stream);
  uint64_t CopyUniqueNodes();
  int GetPathNum() { return total_row_; }
  void ResetPathNum() { total_row_ = 0; }
  void ResetEpochFinish() { epoch_finish_ = false; }
  void ClearSampleState();
978
  void DumpWalkPath(std::string dump_path, size_t dump_rate);
D
danleifeng 已提交
979
  void SetDeviceKeys(std::vector<uint64_t>* device_keys, int type) {
L
lxsbupt 已提交
980 981 982 983 984 985 986 987 988 989 990 991 992 993 994 995 996 997 998 999 1000 1001 1002 1003 1004 1005 1006
    // type_to_index_[type] = h_device_keys_.size();
    // h_device_keys_.push_back(device_keys);
  }

  std::vector<std::shared_ptr<phi::Allocation>> SampleNeighbors(
      int64_t* uniq_nodes,
      int len,
      int sample_size,
      std::vector<int>& edges_split_num,  // NOLINT
      int64_t* neighbor_len);
  std::shared_ptr<phi::Allocation> FillReindexHashTable(int64_t* input,
                                                        int num_input,
                                                        int64_t len_hashtable,
                                                        int64_t* keys,
                                                        int* values,
                                                        int* key_index,
                                                        int* final_nodes_len);
  std::shared_ptr<phi::Allocation> GetReindexResult(int64_t* reindex_src_data,
                                                    int64_t* center_nodes,
                                                    int* final_nodes_len,
                                                    int node_len,
                                                    int64_t neighbor_len);
  std::shared_ptr<phi::Allocation> GenerateSampleGraph(
      uint64_t* node_ids,
      int len,
      int* uniq_len,
      std::shared_ptr<phi::Allocation>& inverse);  // NOLINT
1007
  std::shared_ptr<phi::Allocation> GetNodeDegree(uint64_t* node_ids, int len);
L
lxsbupt 已提交
1008 1009 1010 1011 1012 1013
  int InsertTable(const uint64_t* d_keys,
                  uint64_t len,
                  std::shared_ptr<phi::Allocation> d_uniq_node_num);
  std::vector<uint64_t>& GetHostVec() { return host_vec_; }
  bool get_epoch_finish() { return epoch_finish_; }
  void clear_gpu_mem();
D
danleifeng 已提交
1014 1015

 protected:
L
lxsbupt 已提交
1016
  HashTable<uint64_t, uint64_t>* table_;
D
danleifeng 已提交
1017 1018 1019 1020 1021 1022
  int walk_degree_;
  int walk_len_;
  int window_;
  int once_sample_startid_len_;
  int gpuid_;
  size_t cursor_;
L
lxsbupt 已提交
1023
  int thread_id_;
D
danleifeng 已提交
1024
  size_t jump_rows_;
L
lxsbupt 已提交
1025
  int edge_to_id_len_;
D
danleifeng 已提交
1026
  int64_t* id_tensor_ptr_;
L
lxsbupt 已提交
1027
  int* index_tensor_ptr_;
D
danleifeng 已提交
1028 1029
  int64_t* show_tensor_ptr_;
  int64_t* clk_tensor_ptr_;
1030
  int* degree_tensor_ptr_;
L
lxsbupt 已提交
1031 1032 1033

  cudaStream_t train_stream_;
  cudaStream_t sample_stream_;
D
danleifeng 已提交
1034
  paddle::platform::Place place_;
1035
  std::vector<phi::DenseTensor*> feed_vec_;
D
danleifeng 已提交
1036 1037 1038
  std::vector<size_t> offset_;
  std::shared_ptr<phi::Allocation> d_prefix_sum_;
  std::vector<std::shared_ptr<phi::Allocation>> d_device_keys_;
L
lxsbupt 已提交
1039
  std::shared_ptr<phi::Allocation> d_train_metapath_keys_;
D
danleifeng 已提交
1040 1041

  std::shared_ptr<phi::Allocation> d_walk_;
1042 1043
  std::shared_ptr<phi::Allocation> d_walk_ntype_;
  std::shared_ptr<phi::Allocation> d_excluded_train_pair_;
L
lxsbupt 已提交
1044
  std::shared_ptr<phi::Allocation> d_feature_list_;
D
danleifeng 已提交
1045 1046 1047
  std::shared_ptr<phi::Allocation> d_feature_;
  std::shared_ptr<phi::Allocation> d_len_per_row_;
  std::shared_ptr<phi::Allocation> d_random_row_;
L
lxsbupt 已提交
1048 1049 1050 1051 1052
  std::shared_ptr<phi::Allocation> d_uniq_node_num_;
  std::shared_ptr<phi::Allocation> d_slot_feature_num_map_;
  std::shared_ptr<phi::Allocation> d_actual_slot_id_map_;
  std::shared_ptr<phi::Allocation> d_fea_offset_map_;

D
danleifeng 已提交
1053 1054 1055 1056 1057 1058 1059
  std::vector<std::shared_ptr<phi::Allocation>> d_sampleidx2rows_;
  int cur_sampleidx2row_;
  // record the keys to call graph_neighbor_sample
  std::shared_ptr<phi::Allocation> d_sample_keys_;
  int sample_keys_len_;

  std::shared_ptr<phi::Allocation> d_ins_buf_;
L
lxsbupt 已提交
1060 1061
  std::shared_ptr<phi::Allocation> d_feature_size_list_buf_;
  std::shared_ptr<phi::Allocation> d_feature_size_prefixsum_buf_;
D
danleifeng 已提交
1062 1063 1064
  std::shared_ptr<phi::Allocation> d_pair_num_;
  std::shared_ptr<phi::Allocation> d_slot_tensor_ptr_;
  std::shared_ptr<phi::Allocation> d_slot_lod_tensor_ptr_;
L
lxsbupt 已提交
1065 1066 1067 1068 1069 1070 1071 1072 1073 1074 1075 1076 1077
  std::shared_ptr<phi::Allocation> d_reindex_table_key_;
  std::shared_ptr<phi::Allocation> d_reindex_table_value_;
  std::shared_ptr<phi::Allocation> d_reindex_table_index_;
  std::vector<std::shared_ptr<phi::Allocation>> edge_type_graph_;
  std::shared_ptr<phi::Allocation> d_sorted_keys_;
  std::shared_ptr<phi::Allocation> d_sorted_idx_;
  std::shared_ptr<phi::Allocation> d_offset_;
  std::shared_ptr<phi::Allocation> d_merged_cnts_;
  std::shared_ptr<phi::Allocation> d_buf_;

  // sage mode batch data
  std::vector<std::shared_ptr<phi::Allocation>> inverse_vec_;
  std::vector<std::shared_ptr<phi::Allocation>> final_sage_nodes_vec_;
1078
  std::vector<std::shared_ptr<phi::Allocation>> node_degree_vec_;
L
lxsbupt 已提交
1079 1080 1081 1082 1083
  std::vector<int> uniq_instance_vec_;
  std::vector<int> total_instance_vec_;
  std::vector<std::vector<std::shared_ptr<phi::Allocation>>> graph_edges_vec_;
  std::vector<std::vector<std::vector<int>>> edges_split_num_vec_;

1084
  int excluded_train_pair_len_;
L
lxsbupt 已提交
1085 1086 1087
  int64_t reindex_table_size_;
  int sage_batch_count_;
  int sage_batch_num_;
D
danleifeng 已提交
1088
  int ins_buf_pair_len_;
L
lxsbupt 已提交
1089

D
danleifeng 已提交
1090 1091 1092 1093 1094 1095 1096
  // size of a d_walk buf
  size_t buf_size_;
  int repeat_time_;
  std::vector<int> window_step_;
  BufState buf_state_;
  int batch_size_;
  int slot_num_;
L
lxsbupt 已提交
1097 1098
  std::vector<int> h_slot_feature_num_map_;
  int fea_num_per_node_;
D
danleifeng 已提交
1099 1100 1101
  int shuffle_seed_;
  int debug_mode_;
  bool gpu_graph_training_;
L
lxsbupt 已提交
1102 1103 1104 1105 1106 1107 1108 1109 1110 1111 1112 1113
  bool sage_mode_;
  std::vector<int> samples_;
  bool epoch_finish_;
  std::vector<uint64_t> host_vec_;
  std::vector<uint64_t> h_device_keys_len_;
  uint64_t h_train_metapath_keys_len_;
  uint64_t train_table_cap_;
  uint64_t infer_table_cap_;
  uint64_t copy_unique_len_;
  int total_row_;
  size_t infer_node_start_;
  size_t infer_node_end_;
1114 1115 1116
  std::set<int> infer_node_type_index_set_;
  std::string infer_node_type_;
  bool get_degree_;
D
danleifeng 已提交
1117 1118
};

W
Wang Guibao 已提交
1119 1120
class DataFeed {
 public:
1121 1122 1123
  DataFeed() {
    mutex_for_pick_file_ = nullptr;
    file_idx_ = nullptr;
H
hutuxian 已提交
1124 1125
    mutex_for_fea_num_ = nullptr;
    total_fea_num_ = nullptr;
1126
  }
W
Wang Guibao 已提交
1127
  virtual ~DataFeed() {}
H
hutuxian 已提交
1128
  virtual void Init(const DataFeedDesc& data_feed_desc) = 0;
W
Wang Guibao 已提交
1129
  virtual bool CheckFile(const char* filename) {
1130 1131
    PADDLE_THROW(platform::errors::Unimplemented(
        "This function(CheckFile) is not implemented."));
W
Wang Guibao 已提交
1132 1133 1134 1135 1136 1137
  }
  // Set filelist for DataFeed.
  // Pay attention that it must init all readers before call this function.
  // Otherwise, Init() function will init finish_set_filelist_ flag.
  virtual bool SetFileList(const std::vector<std::string>& files);
  virtual bool Start() = 0;
D
dongdaxiang 已提交
1138

W
Wang Guibao 已提交
1139 1140 1141 1142 1143 1144 1145 1146 1147 1148 1149 1150 1151 1152 1153
  // The trainer calls the Next() function, and the DataFeed will load a new
  // batch to the feed_vec. The return value of this function is the batch
  // size of the current batch.
  virtual int Next() = 0;
  // Get all slots' alias which defined in protofile
  virtual const std::vector<std::string>& GetAllSlotAlias() {
    return all_slots_;
  }
  // Get used slots' alias which defined in protofile
  virtual const std::vector<std::string>& GetUseSlotAlias() {
    return use_slots_;
  }
  // This function is used for binding feed_vec memory
  virtual void AddFeedVar(Variable* var, const std::string& name);

H
hutuxian 已提交
1154 1155 1156
  // This function is used for binding feed_vec memory in a given scope
  virtual void AssignFeedVar(const Scope& scope);

P
pangengzheng 已提交
1157 1158 1159 1160
  virtual std::vector<std::string> GetInputVarNames() {
    return std::vector<std::string>();
  }

1161 1162 1163 1164 1165 1166 1167
  // This function will do nothing at default
  virtual void SetInputPvChannel(void* channel) {}
  // This function will do nothing at default
  virtual void SetOutputPvChannel(void* channel) {}
  // This function will do nothing at default
  virtual void SetConsumePvChannel(void* channel) {}

1168
  // This function will do nothing at default
J
jiaqi 已提交
1169 1170 1171
  virtual void SetInputChannel(void* channel) {}
  // This function will do nothing at default
  virtual void SetOutputChannel(void* channel) {}
1172
  // This function will do nothing at default
J
jiaqi 已提交
1173
  virtual void SetConsumeChannel(void* channel) {}
1174
  // This function will do nothing at default
1175
  virtual void SetThreadId(int thread_id) {}
1176
  // This function will do nothing at default
1177
  virtual void SetThreadNum(int thread_num) {}
1178 1179
  // This function will do nothing at default
  virtual void SetParseInsId(bool parse_ins_id) {}
1180
  virtual void SetParseUid(bool parse_uid) {}
1181
  virtual void SetParseContent(bool parse_content) {}
1182 1183 1184
  virtual void SetParseLogKey(bool parse_logkey) {}
  virtual void SetEnablePvMerge(bool enable_pv_merge) {}
  virtual void SetCurrentPhase(int current_phase) {}
D
danleifeng 已提交
1185
#if defined(PADDLE_WITH_GPU_GRAPH) && defined(PADDLE_WITH_HETERPS)
L
lxsbupt 已提交
1186 1187 1188
  virtual void InitGraphResource() {}
  virtual void InitGraphTrainResource() {}
  virtual void SetDeviceKeys(std::vector<uint64_t>* device_keys, int type) {
D
danleifeng 已提交
1189 1190
    gpu_graph_data_generator_.SetDeviceKeys(device_keys, type);
  }
L
lxsbupt 已提交
1191 1192
#endif

D
danleifeng 已提交
1193 1194 1195
  virtual void SetGpuGraphMode(int gpu_graph_mode) {
    gpu_graph_mode_ = gpu_graph_mode;
  }
1196 1197 1198
  virtual void SetFileListMutex(std::mutex* mutex) {
    mutex_for_pick_file_ = mutex;
  }
H
hutuxian 已提交
1199
  virtual void SetFeaNumMutex(std::mutex* mutex) { mutex_for_fea_num_ = mutex; }
1200
  virtual void SetFileListIndex(size_t* file_index) { file_idx_ = file_index; }
H
hutuxian 已提交
1201
  virtual void SetFeaNum(uint64_t* fea_num) { total_fea_num_ = fea_num; }
1202 1203 1204 1205 1206 1207
  virtual const std::vector<std::string>& GetInsIdVec() const {
    return ins_id_vec_;
  }
  virtual const std::vector<std::string>& GetInsContentVec() const {
    return ins_content_vec_;
  }
P
pangengzheng 已提交
1208 1209 1210
  virtual void SetCurBatchSize(const int batch_size) {
    batch_size_ = batch_size;
  }
1211
  virtual int GetCurBatchSize() { return batch_size_; }
L
lxsbupt 已提交
1212 1213 1214 1215 1216 1217 1218 1219 1220 1221 1222 1223 1224 1225 1226 1227 1228 1229 1230 1231 1232 1233 1234 1235 1236 1237 1238 1239 1240 1241 1242 1243 1244 1245 1246 1247
  virtual int GetGraphPathNum() {
#if defined(PADDLE_WITH_GPU_GRAPH) && defined(PADDLE_WITH_HETERPS)
    return gpu_graph_data_generator_.GetPathNum();
#else
    return 0;
#endif
  }

#if defined(PADDLE_WITH_GPU_GRAPH) && defined(PADDLE_WITH_HETERPS)
  virtual const std::vector<uint64_t>* GetHostVec() {
    return &(gpu_graph_data_generator_.GetHostVec());
  }

  virtual void clear_gpu_mem() { gpu_graph_data_generator_.clear_gpu_mem(); }

  virtual bool get_epoch_finish() {
    return gpu_graph_data_generator_.get_epoch_finish();
  }

  virtual void ResetPathNum() { gpu_graph_data_generator_.ResetPathNum(); }

  virtual void ClearSampleState() {
    gpu_graph_data_generator_.ClearSampleState();
  }

  virtual void ResetEpochFinish() {
    gpu_graph_data_generator_.ResetEpochFinish();
  }

  virtual void DoWalkandSage() {
    PADDLE_THROW(platform::errors::Unimplemented(
        "This function(DoWalkandSage) is not implemented."));
  }
#endif

  virtual bool IsTrainMode() { return train_mode_; }
1248
  virtual void LoadIntoMemory() {
1249 1250
    PADDLE_THROW(platform::errors::Unimplemented(
        "This function(LoadIntoMemory) is not implemented."));
1251
  }
1252 1253 1254 1255
  virtual void SetPlace(const paddle::platform::Place& place) {
    place_ = place;
  }
  virtual const paddle::platform::Place& GetPlace() const { return place_; }
1256

P
pangengzheng 已提交
1257
#if defined(PADDLE_WITH_CUDA) && defined(PADDLE_WITH_HETERPS)
P
pangengzheng 已提交
1258 1259 1260 1261
  virtual MiniBatchGpuPack* get_pack(MiniBatchGpuPack* last_pack) {
    return nullptr;
  }

P
pangengzheng 已提交
1262 1263 1264 1265
  virtual void PackToScope(MiniBatchGpuPack* pack, const Scope* scope) {
    PADDLE_THROW(platform::errors::Unimplemented(
        "This function(PackToScope) is not implemented."));
  }
P
pangengzheng 已提交
1266
  virtual void SetInsIdVec(MiniBatchGpuPack* pack) {}
P
pangengzheng 已提交
1267 1268
#endif

1269 1270 1271 1272 1273
  virtual void DumpWalkPath(std::string dump_path, size_t dump_rate) {
    PADDLE_THROW(platform::errors::Unimplemented(
        "This function(DumpWalkPath) is not implemented."));
  }

W
Wang Guibao 已提交
1274 1275 1276 1277 1278 1279 1280 1281 1282 1283 1284 1285
 protected:
  // The following three functions are used to check if it is executed in this
  // order:
  //   Init() -> SetFileList() -> Start() -> Next()
  virtual void CheckInit();
  virtual void CheckSetFileList();
  virtual void CheckStart();
  virtual void SetBatchSize(
      int batch);  // batch size will be set in Init() function
  // This function is used to pick one file from the global filelist(thread
  // safe).
  virtual bool PickOneFile(std::string* filename);
1286
  virtual void CopyToFeedTensor(void* dst, const void* src, size_t size);
W
Wang Guibao 已提交
1287

1288 1289 1290
  std::vector<std::string> filelist_;
  size_t* file_idx_;
  std::mutex* mutex_for_pick_file_;
H
hutuxian 已提交
1291 1292 1293
  std::mutex* mutex_for_fea_num_ = nullptr;
  uint64_t* total_fea_num_ = nullptr;
  uint64_t fea_num_ = 0;
W
Wang Guibao 已提交
1294 1295 1296 1297 1298 1299 1300 1301 1302 1303

  // the alias of used slots, and its order is determined by
  // data_feed_desc(proto object)
  std::vector<std::string> use_slots_;
  std::vector<bool> use_slots_is_dense_;

  // the alias of all slots, and its order is determined by data_feed_desc(proto
  // object)
  std::vector<std::string> all_slots_;
  std::vector<std::string> all_slots_type_;
1304
  std::vector<std::vector<int>> use_slots_shape_;
1305 1306
  std::vector<int> inductive_shape_index_;
  std::vector<int> total_dims_without_inductive_;
H
hutuxian 已提交
1307 1308
  // For the inductive shape passed within data
  std::vector<std::vector<int>> multi_inductive_shape_index_;
W
Wang Guibao 已提交
1309 1310 1311 1312
  std::vector<int>
      use_slots_index_;  // -1: not used; >=0: the index of use_slots_

  // The data read by DataFeed will be stored here
1313
  std::vector<phi::DenseTensor*> feed_vec_;
W
Wang Guibao 已提交
1314

1315
  phi::DenseTensor* rank_offset_;
1316

W
Wang Guibao 已提交
1317 1318 1319 1320 1321 1322
  // the batch size defined by user
  int default_batch_size_;
  // current batch size
  int batch_size_;

  bool finish_init_;
1323
  bool finish_set_filelist_;
W
Wang Guibao 已提交
1324
  bool finish_start_;
1325
  std::string pipe_command_;
T
Thunderbrook 已提交
1326 1327
  std::string so_parser_name_;
  std::vector<SlotConf> slot_conf_;
1328 1329
  std::vector<std::string> ins_id_vec_;
  std::vector<std::string> ins_content_vec_;
1330
  platform::Place place_;
1331
  std::string uid_slot_;
1332 1333 1334

  // The input type of pipe reader, 0 for one sample, 1 for one batch
  int input_type_;
D
danleifeng 已提交
1335 1336 1337 1338
  int gpu_graph_mode_ = 0;
#if defined(PADDLE_WITH_GPU_GRAPH) && defined(PADDLE_WITH_HETERPS)
  GraphDataGenerator gpu_graph_data_generator_;
#endif
L
lxsbupt 已提交
1339
  bool train_mode_;
W
Wang Guibao 已提交
1340 1341 1342 1343 1344 1345 1346 1347 1348 1349 1350 1351 1352 1353 1354 1355 1356 1357 1358 1359 1360
};

// PrivateQueueDataFeed is the base virtual class for ohther DataFeeds.
// It use a read-thread to read file and parse data to a private-queue
// (thread level), and get data from this queue when trainer call Next().
template <typename T>
class PrivateQueueDataFeed : public DataFeed {
 public:
  PrivateQueueDataFeed() {}
  virtual ~PrivateQueueDataFeed() {}
  virtual bool Start();
  virtual int Next();

 protected:
  // The thread implementation function for reading file and parse.
  virtual void ReadThread();
  // This function is used to set private-queue size, and the most
  // efficient when the queue size is close to the batch size.
  virtual void SetQueueSize(int queue_size);
  // The reading and parsing method called in the ReadThread.
  virtual bool ParseOneInstance(T* instance) = 0;
D
dongdaxiang 已提交
1361
  virtual bool ParseOneInstanceFromPipe(T* instance) = 0;
W
Wang Guibao 已提交
1362
  // This function is used to put instance to vec_ins
1363 1364
  virtual void AddInstanceToInsVec(T* vec_ins,
                                   const T& instance,
W
Wang Guibao 已提交
1365 1366 1367 1368 1369 1370 1371 1372 1373 1374 1375 1376
                                   int index) = 0;
  // This function is used to put ins_vec to feed_vec
  virtual void PutToFeedVec(const T& ins_vec) = 0;

  // The thread for read files
  std::thread read_thread_;
  // using ifstream one line and one line parse is faster
  // than using fread one buffer and one buffer parse.
  //   for a 601M real data:
  //     ifstream one line and one line parse: 6034 ms
  //     fread one buffer and one buffer parse: 7097 ms
  std::ifstream file_;
D
dongdaxiang 已提交
1377
  std::shared_ptr<FILE> fp_;
W
Wang Guibao 已提交
1378
  size_t queue_size_;
1379
  string::LineFileReader reader_;
W
Wang Guibao 已提交
1380
  // The queue for store parsed data
1381
  std::shared_ptr<paddle::framework::ChannelObject<T>> queue_;
W
Wang Guibao 已提交
1382 1383
};

1384
template <typename T>
J
jiaqi 已提交
1385
class InMemoryDataFeed : public DataFeed {
1386 1387 1388
 public:
  InMemoryDataFeed();
  virtual ~InMemoryDataFeed() {}
H
hutuxian 已提交
1389
  virtual void Init(const DataFeedDesc& data_feed_desc) = 0;
1390 1391
  virtual bool Start();
  virtual int Next();
1392 1393 1394 1395
  virtual void SetInputPvChannel(void* channel);
  virtual void SetOutputPvChannel(void* channel);
  virtual void SetConsumePvChannel(void* channel);

J
jiaqi 已提交
1396 1397 1398
  virtual void SetInputChannel(void* channel);
  virtual void SetOutputChannel(void* channel);
  virtual void SetConsumeChannel(void* channel);
1399 1400
  virtual void SetThreadId(int thread_id);
  virtual void SetThreadNum(int thread_num);
1401
  virtual void SetParseInsId(bool parse_ins_id);
1402
  virtual void SetParseUid(bool parse_uid);
1403
  virtual void SetParseContent(bool parse_content);
1404 1405 1406
  virtual void SetParseLogKey(bool parse_logkey);
  virtual void SetEnablePvMerge(bool enable_pv_merge);
  virtual void SetCurrentPhase(int current_phase);
1407
  virtual void LoadIntoMemory();
T
Thunderbrook 已提交
1408
  virtual void LoadIntoMemoryFromSo();
Y
yaoxuefeng 已提交
1409 1410 1411 1412 1413
  virtual void SetRecord(T* records) { records_ = records; }
  int GetDefaultBatchSize() { return default_batch_size_; }
  void AddBatchOffset(const std::pair<int, int>& offset) {
    batch_offsets_.push_back(offset);
  }
X
xujiaqi01 已提交
1414

1415 1416 1417
 protected:
  virtual bool ParseOneInstance(T* instance) = 0;
  virtual bool ParseOneInstanceFromPipe(T* instance) = 0;
1418 1419
  virtual void ParseOneInstanceFromSo(const char* str,
                                      T* instance,
T
Thunderbrook 已提交
1420
                                      CustomParser* parser) {}
1421 1422
  virtual int ParseInstanceFromSo(int len,
                                  const char* str,
T
Thunderbrook 已提交
1423 1424 1425 1426
                                  std::vector<T>* instances,
                                  CustomParser* parser) {
    return 0;
  }
J
jiaqi 已提交
1427
  virtual void PutToFeedVec(const std::vector<T>& ins_vec) = 0;
Y
yaoxuefeng 已提交
1428
  virtual void PutToFeedVec(const T* ins_vec, int num) = 0;
1429

Y
yaoxuefeng 已提交
1430 1431 1432 1433 1434
  std::vector<std::vector<float>> batch_float_feasigns_;
  std::vector<std::vector<uint64_t>> batch_uint64_feasigns_;
  std::vector<std::vector<size_t>> offset_;
  std::vector<bool> visit_;

1435 1436
  int thread_id_;
  int thread_num_;
1437
  bool parse_ins_id_;
1438
  bool parse_uid_;
1439
  bool parse_content_;
1440 1441 1442
  bool parse_logkey_;
  bool enable_pv_merge_;
  int current_phase_{-1};  // only for untest
J
jiaqi 已提交
1443 1444 1445 1446 1447
  std::ifstream file_;
  std::shared_ptr<FILE> fp_;
  paddle::framework::ChannelObject<T>* input_channel_;
  paddle::framework::ChannelObject<T>* output_channel_;
  paddle::framework::ChannelObject<T>* consume_channel_;
1448 1449 1450 1451

  paddle::framework::ChannelObject<PvInstance>* input_pv_channel_;
  paddle::framework::ChannelObject<PvInstance>* output_pv_channel_;
  paddle::framework::ChannelObject<PvInstance>* consume_pv_channel_;
Y
yaoxuefeng 已提交
1452 1453 1454 1455 1456

  std::vector<std::pair<int, int>> batch_offsets_;
  uint64_t offset_index_ = 0;
  bool enable_heterps_ = false;
  T* records_ = nullptr;
1457 1458
};

W
Wang Guibao 已提交
1459 1460 1461 1462 1463
// This class define the data type of instance(ins_vec) in MultiSlotDataFeed
class MultiSlotType {
 public:
  MultiSlotType() {}
  ~MultiSlotType() {}
H
hutuxian 已提交
1464
  void Init(const std::string& type, size_t reserved_size = 0) {
W
Wang Guibao 已提交
1465 1466 1467
    CheckType(type);
    if (type_[0] == 'f') {
      float_feasign_.clear();
H
hutuxian 已提交
1468 1469 1470
      if (reserved_size) {
        float_feasign_.reserve(reserved_size);
      }
W
Wang Guibao 已提交
1471 1472
    } else if (type_[0] == 'u') {
      uint64_feasign_.clear();
H
hutuxian 已提交
1473 1474 1475
      if (reserved_size) {
        uint64_feasign_.reserve(reserved_size);
      }
W
Wang Guibao 已提交
1476 1477 1478
    }
    type_ = type;
  }
H
hutuxian 已提交
1479 1480 1481 1482
  void InitOffset(size_t max_batch_size = 0) {
    if (max_batch_size > 0) {
      offset_.reserve(max_batch_size + 1);
    }
W
Wang Guibao 已提交
1483 1484 1485 1486 1487 1488
    offset_.resize(1);
    // LoDTensor' lod is counted from 0, the size of lod
    // is one size larger than the size of data.
    offset_[0] = 0;
  }
  const std::vector<size_t>& GetOffset() const { return offset_; }
1489
  std::vector<size_t>& MutableOffset() { return offset_; }
W
Wang Guibao 已提交
1490 1491 1492 1493 1494 1495 1496 1497
  void AddValue(const float v) {
    CheckFloat();
    float_feasign_.push_back(v);
  }
  void AddValue(const uint64_t v) {
    CheckUint64();
    uint64_feasign_.push_back(v);
  }
H
hutuxian 已提交
1498 1499 1500 1501 1502 1503 1504 1505 1506 1507
  void CopyValues(const float* input, size_t size) {
    CheckFloat();
    float_feasign_.resize(size);
    memcpy(float_feasign_.data(), input, size * sizeof(float));
  }
  void CopyValues(const uint64_t* input, size_t size) {
    CheckUint64();
    uint64_feasign_.resize(size);
    memcpy(uint64_feasign_.data(), input, size * sizeof(uint64_t));
  }
W
Wang Guibao 已提交
1508 1509 1510 1511 1512 1513 1514 1515 1516 1517 1518 1519 1520
  void AddIns(const MultiSlotType& ins) {
    if (ins.GetType()[0] == 'f') {  // float
      CheckFloat();
      auto& vec = ins.GetFloatData();
      offset_.push_back(offset_.back() + vec.size());
      float_feasign_.insert(float_feasign_.end(), vec.begin(), vec.end());
    } else if (ins.GetType()[0] == 'u') {  // uint64
      CheckUint64();
      auto& vec = ins.GetUint64Data();
      offset_.push_back(offset_.back() + vec.size());
      uint64_feasign_.insert(uint64_feasign_.end(), vec.begin(), vec.end());
    }
  }
H
hutuxian 已提交
1521 1522 1523 1524 1525 1526 1527 1528
  void AppendValues(const uint64_t* input, size_t size) {
    CheckUint64();
    offset_.push_back(offset_.back() + size);
    uint64_feasign_.insert(uint64_feasign_.end(), input, input + size);
  }
  void AppendValues(const float* input, size_t size) {
    CheckFloat();
    offset_.push_back(offset_.back() + size);
1529

H
hutuxian 已提交
1530 1531
    float_feasign_.insert(float_feasign_.end(), input, input + size);
  }
W
Wang Guibao 已提交
1532
  const std::vector<float>& GetFloatData() const { return float_feasign_; }
1533
  std::vector<float>& MutableFloatData() { return float_feasign_; }
W
Wang Guibao 已提交
1534
  const std::vector<uint64_t>& GetUint64Data() const { return uint64_feasign_; }
1535
  std::vector<uint64_t>& MutableUint64Data() { return uint64_feasign_; }
W
Wang Guibao 已提交
1536
  const std::string& GetType() const { return type_; }
H
hutuxian 已提交
1537
  size_t GetBatchSize() { return offset_.size() - 1; }
1538
  std::string& MutableType() { return type_; }
W
Wang Guibao 已提交
1539

X
xujiaqi01 已提交
1540 1541
  std::string DebugString() {
    std::stringstream ss;
W
wanghuancoder 已提交
1542

1543 1544
    ss << "\ntype: " << type_ << "\n";
    ss << "offset: ";
X
xujiaqi01 已提交
1545 1546 1547 1548
    ss << "[";
    for (const size_t& i : offset_) {
      ss << offset_[i] << ",";
    }
1549
    ss << "]\ndata: [";
X
xujiaqi01 已提交
1550 1551 1552 1553 1554 1555 1556 1557 1558 1559 1560 1561 1562
    if (type_[0] == 'f') {
      for (const float& i : float_feasign_) {
        ss << i << ",";
      }
    } else {
      for (const uint64_t& i : uint64_feasign_) {
        ss << i << ",";
      }
    }
    ss << "]\n";
    return ss.str();
  }

W
Wang Guibao 已提交
1563 1564
 private:
  void CheckType(const std::string& type) const {
1565 1566
    PADDLE_ENFORCE_EQ((type == "uint64" || type == "float"),
                      true,
1567 1568 1569 1570
                      platform::errors::InvalidArgument(
                          "MultiSlotType error, expect type is uint64 or "
                          "float, but received type is %s.",
                          type));
W
Wang Guibao 已提交
1571 1572
  }
  void CheckFloat() const {
1573
    PADDLE_ENFORCE_EQ(
1574 1575
        type_[0],
        'f',
1576 1577
        platform::errors::InvalidArgument(
            "MultiSlotType error, add %s value to float slot.", type_));
W
Wang Guibao 已提交
1578 1579
  }
  void CheckUint64() const {
1580
    PADDLE_ENFORCE_EQ(
1581 1582
        type_[0],
        'u',
1583 1584
        platform::errors::InvalidArgument(
            "MultiSlotType error, add %s value to uint64 slot.", type_));
W
Wang Guibao 已提交
1585 1586 1587 1588 1589 1590 1591
  }
  std::vector<float> float_feasign_;
  std::vector<uint64_t> uint64_feasign_;
  std::string type_;
  std::vector<size_t> offset_;
};

J
jiaqi 已提交
1592 1593 1594 1595 1596 1597 1598 1599 1600 1601 1602 1603 1604 1605 1606 1607 1608 1609 1610 1611 1612 1613 1614 1615 1616 1617 1618 1619 1620 1621
template <class AR>
paddle::framework::Archive<AR>& operator<<(paddle::framework::Archive<AR>& ar,
                                           const MultiSlotType& ins) {
  ar << ins.GetType();
#ifdef _LINUX
  ar << ins.GetOffset();
#else
  const auto& offset = ins.GetOffset();
  ar << (uint64_t)offset.size();
  for (const size_t& x : offset) {
    ar << (const uint64_t)x;
  }
#endif
  ar << ins.GetFloatData();
  ar << ins.GetUint64Data();
  return ar;
}

template <class AR>
paddle::framework::Archive<AR>& operator>>(paddle::framework::Archive<AR>& ar,
                                           MultiSlotType& ins) {
  ar >> ins.MutableType();
#ifdef _LINUX
  ar >> ins.MutableOffset();
#else
  auto& offset = ins.MutableOffset();
  offset.resize(ar.template Get<uint64_t>());
  for (size_t& x : offset) {
    uint64_t t;
    ar >> t;
Y
yaoxuefeng 已提交
1622
    x = static_cast<size_t>(t);
J
jiaqi 已提交
1623 1624 1625 1626 1627 1628 1629
  }
#endif
  ar >> ins.MutableFloatData();
  ar >> ins.MutableUint64Data();
  return ar;
}

1630 1631
struct RecordCandidate {
  std::string ins_id_;
T
Thunderbrook 已提交
1632
  std::unordered_multimap<uint16_t, FeatureFeasign> feas_;
1633 1634 1635 1636 1637 1638 1639 1640 1641 1642 1643 1644
  size_t shadow_index_ = -1;  // Optimization for Reservoir Sample

  RecordCandidate() {}
  RecordCandidate(const Record& rec,
                  const std::unordered_set<uint16_t>& slot_index_to_replace) {
    for (const auto& fea : rec.uint64_feasigns_) {
      if (slot_index_to_replace.find(fea.slot()) !=
          slot_index_to_replace.end()) {
        feas_.insert({fea.slot(), fea.sign()});
      }
    }
  }
1645 1646

  RecordCandidate& operator=(const Record& rec) {
1647
    feas_.clear();
1648 1649
    ins_id_ = rec.ins_id_;
    for (auto& fea : rec.uint64_feasigns_) {
1650
      feas_.insert({fea.slot(), fea.sign()});
1651 1652 1653 1654 1655 1656 1657 1658
    }
    return *this;
  }
};

class RecordCandidateList {
 public:
  RecordCandidateList() = default;
1659
  RecordCandidateList(const RecordCandidateList&) {}
1660

1661
  size_t Size() { return cur_size_; }
1662 1663 1664
  void ReSize(size_t length);

  void ReInit();
1665 1666 1667 1668 1669 1670 1671 1672 1673 1674 1675 1676
  void ReInitPass() {
    for (size_t i = 0; i < cur_size_; ++i) {
      if (candidate_list_[i].shadow_index_ != i) {
        candidate_list_[i].ins_id_ =
            candidate_list_[candidate_list_[i].shadow_index_].ins_id_;
        candidate_list_[i].feas_.swap(
            candidate_list_[candidate_list_[i].shadow_index_].feas_);
        candidate_list_[i].shadow_index_ = i;
      }
    }
    candidate_list_.resize(cur_size_);
  }
1677 1678

  void AddAndGet(const Record& record, RecordCandidate* result);
1679 1680 1681 1682 1683 1684 1685 1686 1687 1688 1689 1690 1691 1692 1693 1694 1695 1696 1697 1698 1699 1700
  void AddAndGet(const Record& record, size_t& index_result) {  // NOLINT
    // std::unique_lock<std::mutex> lock(mutex_);
    size_t index = 0;
    ++total_size_;
    auto fleet_ptr = FleetWrapper::GetInstance();
    if (!full_) {
      candidate_list_.emplace_back(record, slot_index_to_replace_);
      candidate_list_.back().shadow_index_ = cur_size_;
      ++cur_size_;
      full_ = (cur_size_ == capacity_);
    } else {
      index = fleet_ptr->LocalRandomEngine()() % total_size_;
      if (index < capacity_) {
        candidate_list_.emplace_back(record, slot_index_to_replace_);
        candidate_list_[index].shadow_index_ = candidate_list_.size() - 1;
      }
    }
    index = fleet_ptr->LocalRandomEngine()() % cur_size_;
    index_result = candidate_list_[index].shadow_index_;
  }
  const RecordCandidate& Get(size_t index) const {
    PADDLE_ENFORCE_LT(
1701 1702
        index,
        candidate_list_.size(),
1703 1704
        platform::errors::OutOfRange("Your index [%lu] exceeds the number of "
                                     "elements in candidate_list[%lu].",
1705 1706
                                     index,
                                     candidate_list_.size()));
1707 1708 1709 1710 1711 1712
    return candidate_list_[index];
  }
  void SetSlotIndexToReplace(
      const std::unordered_set<uint16_t>& slot_index_to_replace) {
    slot_index_to_replace_ = slot_index_to_replace;
  }
1713 1714

 private:
1715 1716 1717 1718 1719 1720 1721
  size_t capacity_ = 0;
  std::mutex mutex_;
  bool full_ = false;
  size_t cur_size_ = 0;
  size_t total_size_ = 0;
  std::vector<RecordCandidate> candidate_list_;
  std::unordered_set<uint16_t> slot_index_to_replace_;
1722 1723
};

J
jiaqi 已提交
1724 1725
template <class AR>
paddle::framework::Archive<AR>& operator<<(paddle::framework::Archive<AR>& ar,
T
Thunderbrook 已提交
1726
                                           const FeatureFeasign& fk) {
J
jiaqi 已提交
1727 1728 1729 1730 1731 1732 1733
  ar << fk.uint64_feasign_;
  ar << fk.float_feasign_;
  return ar;
}

template <class AR>
paddle::framework::Archive<AR>& operator>>(paddle::framework::Archive<AR>& ar,
T
Thunderbrook 已提交
1734
                                           FeatureFeasign& fk) {
J
jiaqi 已提交
1735 1736 1737 1738 1739 1740 1741 1742 1743 1744 1745 1746 1747 1748 1749 1750 1751 1752 1753 1754 1755 1756 1757 1758 1759 1760 1761 1762 1763 1764 1765 1766 1767 1768 1769 1770 1771 1772 1773
  ar >> fk.uint64_feasign_;
  ar >> fk.float_feasign_;
  return ar;
}

template <class AR>
paddle::framework::Archive<AR>& operator<<(paddle::framework::Archive<AR>& ar,
                                           const FeatureItem& fi) {
  ar << fi.sign();
  ar << fi.slot();
  return ar;
}

template <class AR>
paddle::framework::Archive<AR>& operator>>(paddle::framework::Archive<AR>& ar,
                                           FeatureItem& fi) {
  ar >> fi.sign();
  ar >> fi.slot();
  return ar;
}

template <class AR>
paddle::framework::Archive<AR>& operator<<(paddle::framework::Archive<AR>& ar,
                                           const Record& r) {
  ar << r.uint64_feasigns_;
  ar << r.float_feasigns_;
  ar << r.ins_id_;
  return ar;
}

template <class AR>
paddle::framework::Archive<AR>& operator>>(paddle::framework::Archive<AR>& ar,
                                           Record& r) {
  ar >> r.uint64_feasigns_;
  ar >> r.float_feasigns_;
  ar >> r.ins_id_;
  return ar;
}

W
Wang Guibao 已提交
1774 1775 1776 1777 1778 1779 1780 1781
// This DataFeed is used to feed multi-slot type data.
// The format of multi-slot type data:
//   [n feasign_0 feasign_1 ... feasign_n]*
class MultiSlotDataFeed
    : public PrivateQueueDataFeed<std::vector<MultiSlotType>> {
 public:
  MultiSlotDataFeed() {}
  virtual ~MultiSlotDataFeed() {}
H
hutuxian 已提交
1782
  virtual void Init(const DataFeedDesc& data_feed_desc);
W
Wang Guibao 已提交
1783 1784 1785
  virtual bool CheckFile(const char* filename);

 protected:
D
dongdaxiang 已提交
1786
  virtual void ReadThread();
W
Wang Guibao 已提交
1787 1788 1789 1790
  virtual void AddInstanceToInsVec(std::vector<MultiSlotType>* vec_ins,
                                   const std::vector<MultiSlotType>& instance,
                                   int index);
  virtual bool ParseOneInstance(std::vector<MultiSlotType>* instance);
D
dongdaxiang 已提交
1791
  virtual bool ParseOneInstanceFromPipe(std::vector<MultiSlotType>* instance);
W
Wang Guibao 已提交
1792 1793
  virtual void PutToFeedVec(const std::vector<MultiSlotType>& ins_vec);
};
1794

J
jiaqi 已提交
1795
class MultiSlotInMemoryDataFeed : public InMemoryDataFeed<Record> {
1796 1797 1798
 public:
  MultiSlotInMemoryDataFeed() {}
  virtual ~MultiSlotInMemoryDataFeed() {}
H
hutuxian 已提交
1799
  virtual void Init(const DataFeedDesc& data_feed_desc);
Y
yaoxuefeng 已提交
1800
  // void SetRecord(Record* records) { records_ = records; }
1801

1802
 protected:
J
jiaqi 已提交
1803 1804
  virtual bool ParseOneInstance(Record* instance);
  virtual bool ParseOneInstanceFromPipe(Record* instance);
1805 1806
  virtual void ParseOneInstanceFromSo(const char* str,
                                      Record* instance,
1807
                                      CustomParser* parser) {}
1808 1809
  virtual int ParseInstanceFromSo(int len,
                                  const char* str,
T
Thunderbrook 已提交
1810 1811
                                  std::vector<Record>* instances,
                                  CustomParser* parser);
J
jiaqi 已提交
1812
  virtual void PutToFeedVec(const std::vector<Record>& ins_vec);
1813 1814 1815 1816
  virtual void GetMsgFromLogKey(const std::string& log_key,
                                uint64_t* search_id,
                                uint32_t* cmatch,
                                uint32_t* rank);
Y
yaoxuefeng 已提交
1817
  virtual void PutToFeedVec(const Record* ins_vec, int num);
1818 1819
};

Y
yaoxuefeng 已提交
1820 1821 1822
class SlotRecordInMemoryDataFeed : public InMemoryDataFeed<SlotRecord> {
 public:
  SlotRecordInMemoryDataFeed() {}
P
pangengzheng 已提交
1823
  virtual ~SlotRecordInMemoryDataFeed();
P
pangengzheng 已提交
1824 1825
  void Init(const DataFeedDesc& data_feed_desc) override;
  void LoadIntoMemory() override;
Y
yaoxuefeng 已提交
1826 1827 1828
  void ExpandSlotRecord(SlotRecord* ins);

 protected:
P
pangengzheng 已提交
1829 1830 1831 1832
  bool Start() override;
  int Next() override;
  bool ParseOneInstance(SlotRecord* instance) override { return false; }
  bool ParseOneInstanceFromPipe(SlotRecord* instance) override { return false; }
Y
yaoxuefeng 已提交
1833 1834
  // virtual void ParseOneInstanceFromSo(const char* str, T* instance,
  //                                    CustomParser* parser) {}
P
pangengzheng 已提交
1835
  void PutToFeedVec(const std::vector<SlotRecord>& ins_vec) override {}
Y
yaoxuefeng 已提交
1836 1837 1838 1839 1840

  virtual void LoadIntoMemoryByCommand(void);
  virtual void LoadIntoMemoryByLib(void);
  virtual void LoadIntoMemoryByLine(void);
  virtual void LoadIntoMemoryByFile(void);
P
pangengzheng 已提交
1841
  void SetInputChannel(void* channel) override {
Y
yaoxuefeng 已提交
1842 1843 1844
    input_channel_ = static_cast<ChannelObject<SlotRecord>*>(channel);
  }
  bool ParseOneInstance(const std::string& line, SlotRecord* rec);
P
pangengzheng 已提交
1845 1846 1847 1848 1849 1850 1851 1852 1853
  void PutToFeedVec(const SlotRecord* ins_vec, int num) override;
  void AssignFeedVar(const Scope& scope) override;
  std::vector<std::string> GetInputVarNames() override {
    std::vector<std::string> var_names;
    for (int i = 0; i < use_slot_size_; ++i) {
      var_names.push_back(used_slots_info_[i].slot);
    }
    return var_names;
  }
1854
#if defined(PADDLE_WITH_CUDA) && defined(PADDLE_WITH_HETERPS)
P
pangengzheng 已提交
1855 1856
  void BuildSlotBatchGPU(const int ins_num, MiniBatchGpuPack* pack);

P
pangengzheng 已提交
1857 1858
  virtual MiniBatchGpuPack* get_pack(MiniBatchGpuPack* last_pack);

P
pangengzheng 已提交
1859 1860 1861
  virtual void PackToScope(MiniBatchGpuPack* pack,
                           const Scope* scope = nullptr);

1862 1863
  void FillSlotValueOffset(const int ins_num,
                           const int used_slot_num,
1864 1865
                           size_t* slot_value_offsets,
                           const int* uint64_offsets,
1866 1867
                           const int uint64_slot_size,
                           const int* float_offsets,
1868
                           const int float_slot_size,
P
pangengzheng 已提交
1869 1870
                           const UsedSlotGpuType* used_slots,
                           cudaStream_t stream);
1871 1872 1873
  void CopyForTensor(const int ins_num,
                     const int used_slot_num,
                     void** dest,
1874
                     const size_t* slot_value_offsets,
1875 1876 1877 1878 1879 1880 1881 1882
                     const uint64_t* uint64_feas,
                     const int* uint64_offsets,
                     const int* uint64_ins_lens,
                     const int uint64_slot_size,
                     const float* float_feas,
                     const int* float_offsets,
                     const int* float_ins_lens,
                     const int float_slot_size,
P
pangengzheng 已提交
1883 1884
                     const UsedSlotGpuType* used_slots,
                     cudaStream_t stream);
1885
#endif
L
lxsbupt 已提交
1886 1887 1888 1889 1890

#if defined(PADDLE_WITH_GPU_GRAPH) && defined(PADDLE_WITH_HETERPS)
  virtual void InitGraphResource(void);
  virtual void InitGraphTrainResource(void);
  virtual void DoWalkandSage();
P
pangengzheng 已提交
1891 1892 1893 1894 1895 1896 1897 1898 1899 1900
  void SetInsIdVec(MiniBatchGpuPack* pack) override {
    if (parse_ins_id_) {
      size_t ins_num = pack->ins_num();
      ins_id_vec_.clear();
      ins_id_vec_.resize(ins_num);
      for (size_t i = 0; i < ins_num; i++) {
        ins_id_vec_[i] = pack->get_lineid(i);
      }
    }
  }
L
lxsbupt 已提交
1901
#endif
P
pangengzheng 已提交
1902
  void DumpWalkPath(std::string dump_path, size_t dump_rate) override;
L
lxsbupt 已提交
1903

Y
yaoxuefeng 已提交
1904 1905 1906 1907 1908 1909 1910 1911
  float sample_rate_ = 1.0f;
  int use_slot_size_ = 0;
  int float_use_slot_size_ = 0;
  int uint64_use_slot_size_ = 0;
  std::vector<AllSlotInfo> all_slots_info_;
  std::vector<UsedSlotInfo> used_slots_info_;
  size_t float_total_dims_size_ = 0;
  std::vector<int> float_total_dims_without_inductives_;
1912 1913

#if defined(PADDLE_WITH_CUDA) && defined(PADDLE_WITH_HETERPS)
P
pangengzheng 已提交
1914 1915 1916 1917 1918 1919 1920 1921 1922 1923 1924 1925 1926 1927
  int pack_thread_num_{5};
  std::vector<std::thread> pack_threads_;
  std::vector<MiniBatchGpuPack*> pack_vec_;
  BlockingQueue<MiniBatchGpuPack*> free_pack_queue_;
  BlockingQueue<MiniBatchGpuPack*> using_pack_queue_;
  std::atomic<bool> pack_is_end_{false};
  std::atomic<uint64_t> pack_offset_index_{0};
  MiniBatchGpuPack* last_pack_{nullptr};
  std::atomic<bool> stop_token_{false};
  std::atomic<int> thread_count_{0};
  std::mutex pack_mutex_;

  // async infershape
  std::map<const Scope*, std::vector<phi::DenseTensor*>> scpoe_feed_vec_;
1928
#endif
Y
yaoxuefeng 已提交
1929 1930
};

1931 1932 1933 1934 1935 1936 1937 1938 1939 1940 1941 1942 1943 1944 1945 1946 1947
class PaddleBoxDataFeed : public MultiSlotInMemoryDataFeed {
 public:
  PaddleBoxDataFeed() {}
  virtual ~PaddleBoxDataFeed() {}

 protected:
  virtual void Init(const DataFeedDesc& data_feed_desc);
  virtual bool Start();
  virtual int Next();
  virtual void AssignFeedVar(const Scope& scope);
  virtual void PutToFeedVec(const std::vector<PvInstance>& pv_vec);
  virtual void PutToFeedVec(const std::vector<Record*>& ins_vec);
  virtual int GetCurrentPhase();
  virtual void GetRankOffset(const std::vector<PvInstance>& pv_vec,
                             int ins_number);
  std::string rank_offset_name_;
  int pv_batch_size_;
1948 1949
};

1950
#if (defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)) && !defined(_WIN32)
H
hutuxian 已提交
1951 1952 1953 1954 1955 1956 1957 1958 1959 1960 1961 1962 1963 1964 1965 1966 1967 1968 1969 1970 1971 1972 1973 1974 1975 1976 1977 1978 1979 1980 1981 1982 1983 1984 1985 1986 1987 1988 1989 1990 1991 1992 1993 1994 1995 1996 1997 1998
template <typename T>
class PrivateInstantDataFeed : public DataFeed {
 public:
  PrivateInstantDataFeed() {}
  virtual ~PrivateInstantDataFeed() {}
  void Init(const DataFeedDesc& data_feed_desc) override;
  bool Start() override { return true; }
  int Next() override;

 protected:
  // The batched data buffer
  std::vector<MultiSlotType> ins_vec_;

  // This function is used to preprocess with a given filename, e.g. open it or
  // mmap
  virtual bool Preprocess(const std::string& filename) = 0;

  // This function is used to postprocess system resource such as closing file
  // NOTICE: Ensure that it is safe to call before Preprocess
  virtual bool Postprocess() = 0;

  // The reading and parsing method.
  virtual bool ParseOneMiniBatch() = 0;

  // This function is used to put ins_vec to feed_vec
  virtual void PutToFeedVec();
};

class MultiSlotFileInstantDataFeed
    : public PrivateInstantDataFeed<std::vector<MultiSlotType>> {
 public:
  MultiSlotFileInstantDataFeed() {}
  virtual ~MultiSlotFileInstantDataFeed() {}

 protected:
  int fd_{-1};
  char* buffer_{nullptr};
  size_t end_{0};
  size_t offset_{0};

  bool Preprocess(const std::string& filename) override;

  bool Postprocess() override;

  bool ParseOneMiniBatch() override;
};
#endif

W
Wang Guibao 已提交
1999 2000
}  // namespace framework
}  // namespace paddle