ps_gpu_wrapper.cc 31.4 KB
Newer Older
T
Thunderbrook 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28
// Copyright (c) 2019 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
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// 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.

/* 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

  http://www.apache.org/licenses/LICENSE-2.0

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. */

T
Thunderbrook 已提交
29
#ifdef PADDLE_WITH_HETERPS
Y
yaoxuefeng 已提交
30

T
Thunderbrook 已提交
31
#include <algorithm>
Y
yaoxuefeng 已提交
32 33
#include <deque>

T
Thunderbrook 已提交
34 35 36 37 38 39 40 41
#include "paddle/fluid/framework/fleet/ps_gpu_wrapper.h"
#include "paddle/fluid/platform/timer.h"

namespace paddle {
namespace framework {

std::shared_ptr<PSGPUWrapper> PSGPUWrapper::s_instance_ = NULL;
bool PSGPUWrapper::is_initialized_ = false;
T
Thunderbrook 已提交
42 43 44 45 46 47 48 49 50 51 52 53 54
#ifdef PADDLE_WITH_PSLIB
void PSGPUWrapper::InitAfsApi(const std::string& fs_name,
                              const std::string& fs_user,
                              const std::string& pass_wd,
                              const std::string& conf) {
  int ret = afs_handler_.init(fs_name.c_str(), fs_user.c_str(), pass_wd.c_str(),
                              conf.c_str());
  if (ret != 0) {
    LOG(ERROR) << "AFS Init Error";
  }
  use_afs_api_ = 1;
}
#endif
55
void PSGPUWrapper::PreBuildTask(std::shared_ptr<HeterContext> gpu_task) {
Y
yaoxuefeng 已提交
56
  VLOG(3) << "PSGPUWrapper::BuildGPUPSTask begin";
T
Thunderbrook 已提交
57 58
  platform::Timer timeline;
  timeline.Start();
59
  int device_num = heter_devices_.size();
60 61 62 63 64
  if (!multi_mf_dim_) {
    gpu_task->init(thread_keys_shard_num_, device_num);
  } else {
    gpu_task->init(thread_keys_shard_num_, device_num, multi_mf_dim_);
  }
Y
yaoxuefeng 已提交
65 66
  auto& local_keys = gpu_task->feature_keys_;
  auto& local_ptr = gpu_task->value_ptr_;
67

Y
yaoxuefeng 已提交
68 69 70
  std::vector<std::thread> threads;

  // data should be in input channel
71 72 73 74 75 76 77 78 79 80 81 82 83
  if (!multi_mf_dim_) {
    thread_keys_.resize(thread_keys_thread_num_);
    for (int i = 0; i < thread_keys_thread_num_; i++) {
      thread_keys_[i].resize(thread_keys_shard_num_);
    }
  } else {
    thread_dim_keys_.resize(thread_keys_thread_num_);
    for (int i = 0; i < thread_keys_thread_num_; i++) {
      thread_dim_keys_[i].resize(thread_keys_shard_num_);
      for (int j = 0; j < thread_keys_shard_num_; j++) {
        thread_dim_keys_[i][j].resize(multi_mf_dim_);
      }
    }
Y
yaoxuefeng 已提交
84
  }
Y
yaoxuefeng 已提交
85 86 87 88

  size_t total_len = 0;
  size_t len_per_thread = 0;
  int remain = 0;
Y
yaoxuefeng 已提交
89
  size_t begin = 0;
Y
yaoxuefeng 已提交
90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113

  std::string data_set_name = std::string(typeid(*dataset_).name());

  if (data_set_name.find("SlotRecordDataset") != std::string::npos) {
    VLOG(0) << "ps_gpu_wrapper use SlotRecordDataset";
    SlotRecordDataset* dataset = dynamic_cast<SlotRecordDataset*>(dataset_);
    auto input_channel = dataset->GetInputChannel();
    VLOG(0) << "yxf::buildtask::inputslotchannle size: "
            << input_channel->Size();
    const std::deque<SlotRecord>& vec_data = input_channel->GetData();
    total_len = vec_data.size();
    len_per_thread = total_len / thread_keys_thread_num_;
    remain = total_len % thread_keys_thread_num_;
    VLOG(0) << "total len: " << total_len;
    auto gen_func = [this](const std::deque<SlotRecord>& total_data,
                           int begin_index, int end_index, int i) {
      for (auto iter = total_data.begin() + begin_index;
           iter != total_data.begin() + end_index; iter++) {
        const auto& ins = *iter;
        const auto& feasign_v = ins->slot_uint64_feasigns_.slot_values;
        for (const auto feasign : feasign_v) {
          int shard_id = feasign % thread_keys_shard_num_;
          this->thread_keys_[i][shard_id].insert(feasign);
        }
Y
yaoxuefeng 已提交
114
      }
Y
yaoxuefeng 已提交
115
    };
116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144
    auto gen_dynamic_mf_func = [this](const std::deque<SlotRecord>& total_data,
                                      int begin_index, int end_index, int i) {
      for (auto iter = total_data.begin() + begin_index;
           iter != total_data.begin() + end_index; iter++) {
        const auto& ins = *iter;
        const auto& feasign_v = ins->slot_uint64_feasigns_.slot_values;
        const auto& slot_offset = ins->slot_uint64_feasigns_.slot_offsets;
        for (size_t slot_idx = 0; slot_idx < slot_offset_vector_.size();
             slot_idx++) {
          for (size_t j = slot_offset[slot_offset_vector_[slot_idx]];
               j < slot_offset[slot_offset_vector_[slot_idx] + 1]; j++) {
            int shard_id = feasign_v[j] % thread_keys_shard_num_;
            int dim_id = slot_index_vec_[slot_idx];
            this->thread_dim_keys_[i][shard_id][dim_id].insert(feasign_v[j]);
          }
        }
      }
      /*
      for (auto iter = total_data.begin() + begin_index;
           iter != total_data.begin() + end_index; iter++) {
        const auto& ins = *iter;
        const auto& feasign_v = ins->slot_uint64_feasigns_.slot_values;
        for (const auto feasign : feasign_v) {
          int shard_id = feasign % thread_keys_shard_num_;
          this->thread_dim_keys_[i][shard_id][0].insert(feasign);
        }
      }
      */
    };
Y
yaoxuefeng 已提交
145
    for (int i = 0; i < thread_keys_thread_num_; i++) {
146 147 148 149 150 151 152 153 154 155 156
      if (!multi_mf_dim_) {
        VLOG(0) << "yxf::psgpu wrapper genfunc";
        threads.push_back(
            std::thread(gen_func, std::ref(vec_data), begin,
                        begin + len_per_thread + (i < remain ? 1 : 0), i));
      } else {
        VLOG(0) << "yxf::psgpu wrapper genfunc with dynamic mf";
        threads.push_back(
            std::thread(gen_dynamic_mf_func, std::ref(vec_data), begin,
                        begin + len_per_thread + (i < remain ? 1 : 0), i));
      }
Y
yaoxuefeng 已提交
157
      begin += len_per_thread + (i < remain ? 1 : 0);
Y
yaoxuefeng 已提交
158
    }
Y
yaoxuefeng 已提交
159 160 161 162
    for (std::thread& t : threads) {
      t.join();
    }
    timeline.Pause();
T
Thunderbrook 已提交
163
    VLOG(0) << "GpuPs build task cost " << timeline.ElapsedSec() << " seconds.";
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 196
  } else {
    CHECK(data_set_name.find("MultiSlotDataset") != std::string::npos);
    VLOG(0) << "ps_gpu_wrapper use MultiSlotDataset";
    MultiSlotDataset* dataset = dynamic_cast<MultiSlotDataset*>(dataset_);
    auto input_channel = dataset->GetInputChannel();

    const std::deque<Record>& vec_data = input_channel->GetData();
    total_len = vec_data.size();
    len_per_thread = total_len / thread_keys_thread_num_;
    remain = total_len % thread_keys_thread_num_;
    auto gen_func = [this](const std::deque<Record>& total_data,
                           int begin_index, int end_index, int i) {
      for (auto iter = total_data.begin() + begin_index;
           iter != total_data.begin() + end_index; iter++) {
        const auto& ins = *iter;
        const auto& feasign_v = ins.uint64_feasigns_;
        for (const auto feasign : feasign_v) {
          uint64_t cur_key = feasign.sign().uint64_feasign_;
          int shard_id = cur_key % thread_keys_shard_num_;
          this->thread_keys_[i][shard_id].insert(cur_key);
        }
      }
    };
    for (int i = 0; i < thread_keys_thread_num_; i++) {
      threads.push_back(
          std::thread(gen_func, std::ref(vec_data), begin,
                      begin + len_per_thread + (i < remain ? 1 : 0), i));
      begin += len_per_thread + (i < remain ? 1 : 0);
    }
    for (std::thread& t : threads) {
      t.join();
    }
    timeline.Pause();
T
Thunderbrook 已提交
197
    VLOG(0) << "GpuPs build task cost " << timeline.ElapsedSec() << " seconds.";
Y
yaoxuefeng 已提交
198 199 200 201
  }

  timeline.Start();

202
  threads.clear();
Y
yaoxuefeng 已提交
203
  // merge thread_keys to shard_keys
204 205 206 207
  auto merge_ins_func = [this, gpu_task](int shard_num) {
    for (int i = 0; i < thread_keys_thread_num_; ++i) {
      gpu_task->batch_add_keys(shard_num, thread_keys_[i][shard_num]);
      thread_keys_[i][shard_num].clear();
Y
yaoxuefeng 已提交
208
    }
209
  };
210 211 212 213 214 215 216
  auto merge_ins_dynamic_mf_func = [this, gpu_task](int shard_num, int dim_id) {
    for (int i = 0; i < thread_keys_thread_num_; ++i) {
      gpu_task->batch_add_keys(shard_num, dim_id,
                               thread_dim_keys_[i][shard_num][dim_id]);
      thread_dim_keys_[i][shard_num][dim_id].clear();
    }
  };
217 218 219 220 221 222 223
  // for (size_t i = 0; i < thread_keys_.size(); i++) {
  //  gpu_task->batch_add_keys(thread_keys_[i]);
  //  for (int j = 0; j < thread_keys_thread_num_; j++) {
  //    thread_keys_[i][j].clear();
  //  }
  //}
  for (int i = 0; i < thread_keys_shard_num_; ++i) {
224 225 226 227 228 229 230
    if (!multi_mf_dim_) {
      threads.push_back(std::thread(merge_ins_func, i));
    } else {
      for (int j = 0; j < multi_mf_dim_; j++) {
        threads.push_back(std::thread(merge_ins_dynamic_mf_func, i, j));
      }
    }
231 232 233
  }
  for (auto& t : threads) {
    t.join();
Y
yaoxuefeng 已提交
234 235 236
  }
  timeline.Pause();

Y
yaoxuefeng 已提交
237
  VLOG(1) << "GpuPs task add keys cost " << timeline.ElapsedSec()
Y
yaoxuefeng 已提交
238 239 240 241 242
          << " seconds.";
  timeline.Start();
  gpu_task->UniqueKeys();
  timeline.Pause();

243
  VLOG(1) << "GpuPs task unique cost " << timeline.ElapsedSec() << " seconds.";
Y
yaoxuefeng 已提交
244

245 246 247 248 249 250 251 252 253 254 255 256 257 258
  if (!multi_mf_dim_) {
    for (int i = 0; i < thread_keys_shard_num_; i++) {
      VLOG(0) << "GpuPs shard: " << i << " key len: " << local_keys[i].size();
      local_ptr[i].resize(local_keys[i].size());
    }
  } else {
    for (int i = 0; i < thread_keys_shard_num_; i++) {
      for (int j = 0; j < multi_mf_dim_; j++) {
        VLOG(0) << "GpuPs shard: " << i << "mf dim: " << index_dim_vec_[j]
                << " key len: " << gpu_task->feature_dim_keys_[i][j].size();
        gpu_task->value_dim_ptr_[i][j].resize(
            gpu_task->feature_dim_keys_[i][j].size());
      }
    }
Y
yaoxuefeng 已提交
259
  }
260 261 262 263 264 265 266 267
}

void PSGPUWrapper::BuildPull(std::shared_ptr<HeterContext> gpu_task) {
  platform::Timer timeline;
  int device_num = heter_devices_.size();
  auto& local_keys = gpu_task->feature_keys_;
  auto& local_ptr = gpu_task->value_ptr_;

268 269 270
  auto& local_dim_keys = gpu_task->feature_dim_keys_;
  auto& local_dim_ptr = gpu_task->value_dim_ptr_;

271 272
  auto& device_keys = gpu_task->device_keys_;
  auto& device_vals = gpu_task->device_values_;
273 274 275 276 277 278 279 280 281
  auto& device_dim_keys = gpu_task->device_dim_keys_;
  auto& device_dim_ptr = gpu_task->device_dim_ptr_;
  auto& device_dim_mutex = gpu_task->dim_mutex_;
  if (multi_mf_dim_) {
    for (size_t dev = 0; dev < device_dim_keys.size(); dev++) {
      device_dim_keys[dev].resize(multi_mf_dim_);
      device_dim_ptr[dev].resize(multi_mf_dim_);
    }
  }
282 283 284 285 286 287 288
  auto& device_mutex = gpu_task->mutex_;

  std::vector<std::thread> threads(thread_keys_shard_num_);
#ifdef PADDLE_WITH_PSLIB
  auto fleet_ptr = FleetWrapper::GetInstance();
#endif
#ifdef PADDLE_WITH_PSCORE
289
  auto fleet_ptr = paddle::distributed::FleetWrapper::GetInstance();
290
#endif
291

292
#if (defined PADDLE_WITH_PSLIB) && (defined PADDLE_WITH_HETERPS)
293 294 295 296 297 298 299 300 301 302 303
  // get day_id: day nums from 1970
  struct std::tm b;
  b.tm_year = year_ - 1900;
  b.tm_mon = month_ - 1;
  b.tm_mday = day_;
  b.tm_min = b.tm_hour = b.tm_sec = 0;
  std::time_t seconds_from_1970 = std::mktime(&b);
  int day_id = seconds_from_1970 / 86400;
  fleet_ptr->pslib_ptr_->_worker_ptr->set_day_id(table_id_, day_id);
#endif

304
  timeline.Start();
305
  auto ptl_func = [this, &local_keys, &local_ptr, &fleet_ptr](int i) {
Y
yaoxuefeng 已提交
306
    size_t key_size = local_keys[i].size();
Y
yaoxuefeng 已提交
307
    int32_t status = -1;
T
Thunderbrook 已提交
308
#ifdef PADDLE_WITH_PSLIB
Y
yaoxuefeng 已提交
309 310 311 312 313 314
    // auto tt = fleet_ptr->pslib_ptr_->_worker_ptr->pull_sparse_ptr(
    //    reinterpret_cast<char**>(local_ptr[i].data()), this->table_id_,
    //    local_keys[i].data(), key_size);
    int32_t cnt = 0;
    while (true) {
      auto tt = fleet_ptr->pslib_ptr_->_worker_ptr->pull_sparse_ptr(
T
Thunderbrook 已提交
315
          i, reinterpret_cast<char**>(local_ptr[i].data()), this->table_id_,
Y
yaoxuefeng 已提交
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
          local_keys[i].data(), key_size);
      bool flag = true;

      tt.wait();

      try {
        status = tt.get();
      } catch (const std::future_error& e) {
        VLOG(0) << "Caught a future_error with code" << e.code()
                << ", Message:" << e.what();
      }
      if (status != 0) {
        VLOG(0) << "fleet pull sparse failed, status[" << status << "]";
        sleep(sleep_seconds_before_fail_exit_);
        flag = false;
        cnt++;
      }
      if (cnt > 3) {
        VLOG(0) << "fleet pull sparse failed, retry 3 times";
        exit(-1);
      }

      if (flag) {
        break;
      }
    }
T
Thunderbrook 已提交
342 343
#endif
#ifdef PADDLE_WITH_PSCORE
Y
yaoxuefeng 已提交
344 345
    int32_t cnt = 0;
    while (true) {
346
      auto tt = fleet_ptr->worker_ptr_->pull_sparse_ptr(
Y
yaoxuefeng 已提交
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
          reinterpret_cast<char**>(local_ptr[i].data()), this->table_id_,
          local_keys[i].data(), key_size);
      bool flag = true;

      tt.wait();

      try {
        status = tt.get();
      } catch (const std::future_error& e) {
        VLOG(0) << "Caught a future_error with code" << e.code()
                << ", Message:" << e.what();
      }
      if (status != 0) {
        VLOG(0) << "fleet pull sparse failed, status[" << status << "]";
        sleep(sleep_seconds_before_fail_exit_);
        flag = false;
        cnt++;
      }
      if (cnt > 3) {
        VLOG(0) << "fleet pull sparse failed, retry 3 times";
        exit(-1);
      }

      if (flag) {
        break;
      }
    }
T
Thunderbrook 已提交
374
#endif
Y
yaoxuefeng 已提交
375 376 377 378 379 380 381 382
    if (status != 0) {
      LOG(ERROR) << "fleet pull sparse failed, status[" << status << "]";
      sleep(300);
      exit(-1);
    } else {
      VLOG(3) << "FleetWrapper Pull sparse to local done with table size: "
              << local_keys[i].size();
    }
383
  };
384 385 386 387 388 389 390 391 392

  auto ptl_dynamic_mf_func = [this, &local_dim_keys, &local_dim_ptr,
                              &fleet_ptr](int i, int j) {
#ifdef PADDLE_WITH_PSLIB
    size_t key_size = local_dim_keys[i][j].size();
    int32_t status = -1;
    int32_t cnt = 0;
    while (true) {
      auto tt = fleet_ptr->pslib_ptr_->_worker_ptr->pull_sparse_ptr(
T
Thunderbrook 已提交
393 394
          i, reinterpret_cast<char**>(local_dim_ptr[i][j].data()),
          this->table_id_, local_dim_keys[i][j].data(), key_size);
395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440
      bool flag = true;

      tt.wait();

      try {
        status = tt.get();
      } catch (const std::future_error& e) {
        VLOG(0) << "Caught a future_error with code" << e.code()
                << ", Message:" << e.what();
      }
      if (status != 0) {
        VLOG(0) << "fleet pull sparse failed, status[" << status << "]";
        sleep(sleep_seconds_before_fail_exit_);
        flag = false;
        cnt++;
      }
      if (cnt > 3) {
        VLOG(0) << "fleet pull sparse failed, retry 3 times";
        exit(-1);
      }

      if (flag) {
        break;
      }
    }
    if (status != 0) {
      LOG(ERROR) << "fleet pull sparse failed, status[" << status << "]";
      sleep(300);
      exit(-1);
    } else {
      VLOG(0) << "FleetWrapper Pull sparse to local done with table size: "
              << local_dim_keys[i][j].size();
    }
#endif
  };
  if (!multi_mf_dim_) {
    for (size_t i = 0; i < threads.size(); i++) {
      threads[i] = std::thread(ptl_func, i);
    }
  } else {
    threads.resize(thread_keys_shard_num_ * multi_mf_dim_);
    for (int i = 0; i < thread_keys_shard_num_; i++) {
      for (int j = 0; j < multi_mf_dim_; j++) {
        threads[i * multi_mf_dim_ + j] = std::thread(ptl_dynamic_mf_func, i, j);
      }
    }
441 442 443 444 445
  }
  for (std::thread& t : threads) {
    t.join();
  }
  timeline.Pause();
T
Thunderbrook 已提交
446
  VLOG(0) << "pull sparse from CpuPS into GpuPS cost " << timeline.ElapsedSec()
447
          << " seconds.";
Y
yaoxuefeng 已提交
448 449 450 451 452 453 454 455
  if (multi_node_) {
    auto gloo_wrapper = paddle::framework::GlooWrapper::GetInstance();
    if (!gloo_wrapper->IsInitialized()) {
      VLOG(0) << "GLOO is not inited";
      gloo_wrapper->Init();
    }
    gloo_wrapper->Barrier();
  }
456 457

  timeline.Start();
Y
yaoxuefeng 已提交
458 459 460
  std::vector<std::vector<std::pair<uint64_t, char*>>> pass_values;

  bool record_status = false;
461 462
#ifdef PADDLE_WITH_PSLIB
  uint16_t pass_id = 0;
Y
yaoxuefeng 已提交
463 464 465 466
  if (multi_node_) {
    record_status = fleet_ptr->pslib_ptr_->_worker_ptr->take_sparse_record(
        table_id_, pass_id, pass_values);
  }
467
#endif
468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498
  auto build_dynamic_mf_func = [this, device_num, &local_dim_keys,
                                &local_dim_ptr, &device_dim_keys,
                                &device_dim_ptr,
                                &device_dim_mutex](int i, int j) {
#ifdef PADDLE_WITH_PSLIB
    std::vector<std::vector<FeatureKey>> task_keys(device_num);
    std::vector<std::vector<paddle::ps::DownpourFixedFeatureValue*>> task_ptrs(
        device_num);
    for (size_t k = 0; k < local_dim_keys[i][j].size(); k++) {
      int shard = local_dim_keys[i][j][k] % device_num;
      task_keys[shard].push_back(local_dim_keys[i][j][k]);
      task_ptrs[shard].push_back(local_dim_ptr[i][j][k]);
    }
    for (int dev = 0; dev < device_num; dev++) {
      for (int dim = 0; dim < multi_mf_dim_; dim++) {
        device_dim_mutex[dev][dim]->lock();

        int len = task_keys[dev].size();
        int cur = device_dim_keys[dev][dim].size();
        device_dim_keys[dev][dim].resize(device_dim_keys[dev][dim].size() +
                                         len);
        device_dim_ptr[dev][dim].resize(device_dim_ptr[dev][dim].size() + len);
        for (int k = 0; k < len; ++k) {
          device_dim_keys[dev][dim][cur + k] = task_keys[dev][k];
          device_dim_ptr[dev][dim][cur + k] = task_ptrs[dev][k];
        }
        device_dim_mutex[dev][dim]->unlock();
      }
    }
#endif
  };
Y
yaoxuefeng 已提交
499 500 501
  auto build_func = [device_num, record_status, &pass_values, &local_keys,
                     &local_ptr, &device_keys, &device_vals,
                     &device_mutex](int i) {
502
    std::vector<std::vector<FeatureKey>> task_keys(device_num);
T
Thunderbrook 已提交
503
#ifdef PADDLE_WITH_PSLIB
504 505
    std::vector<std::vector<paddle::ps::DownpourFixedFeatureValue*>> task_ptrs(
        device_num);
T
Thunderbrook 已提交
506 507 508
#endif

#ifdef PADDLE_WITH_PSCORE
509 510
    std::vector<std::vector<paddle::distributed::FixedFeatureValue*>> task_ptrs(
        device_num);
T
Thunderbrook 已提交
511
#endif
512 513 514 515 516 517

    for (size_t j = 0; j < local_keys[i].size(); j++) {
      int shard = local_keys[i][j] % device_num;
      task_keys[shard].push_back(local_keys[i][j]);
      task_ptrs[shard].push_back(local_ptr[i][j]);
    }
518
#ifdef PADDLE_WITH_PSLIB
Y
yaoxuefeng 已提交
519 520 521 522 523 524 525 526 527 528 529 530 531 532 533
    if (record_status) {
      size_t local_keys_size = local_keys.size();
      size_t pass_values_size = pass_values.size();
      for (size_t j = 0; j < pass_values_size; j += local_keys_size) {
        auto& shard_values = pass_values[j];
        for (size_t pair_idx = 0; pair_idx < pass_values[j].size();
             pair_idx++) {
          auto& cur_pair = shard_values[pair_idx];
          int shard = cur_pair.first % device_num;
          task_keys[shard].push_back(cur_pair.first);
          task_ptrs[shard].push_back(
              (paddle::ps::DownpourFixedFeatureValue*)cur_pair.second);
        }
      }
    }
534
#endif
535 536 537 538 539 540 541
    for (int dev = 0; dev < device_num; dev++) {
      device_mutex[dev]->lock();

      int len = task_keys[dev].size();
      int cur = device_keys[dev].size();
      device_keys[dev].resize(device_keys[dev].size() + len);
      device_vals[dev].resize(device_vals[dev].size() + len);
T
Thunderbrook 已提交
542
#ifdef PADDLE_WITH_PSLIB
543 544 545 546 547 548 549 550 551 552 553 554
      for (int j = 0; j < len; ++j) {
        device_keys[dev][cur + j] = task_keys[dev][j];
        float* ptr_val = task_ptrs[dev][j]->data();
        FeatureValue& val = device_vals[dev][cur + j];
        size_t dim = task_ptrs[dev][j]->size();

        val.delta_score = ptr_val[1];
        val.show = ptr_val[2];
        val.clk = ptr_val[3];
        val.slot = ptr_val[6];
        val.lr = ptr_val[4];
        val.lr_g2sum = ptr_val[5];
T
Thunderbrook 已提交
555
        val.cpu_ptr = (uint64_t)(task_ptrs[dev][j]);
556 557 558 559 560 561 562 563 564 565 566

        if (dim > 7) {
          val.mf_size = MF_DIM + 1;
          for (int x = 0; x < val.mf_size; x++) {
            val.mf[x] = ptr_val[x + 7];
          }
        } else {
          val.mf_size = 0;
          for (int x = 0; x < MF_DIM + 1; x++) {
            val.mf[x] = 0;
          }
Y
yaoxuefeng 已提交
567 568
        }
      }
T
Thunderbrook 已提交
569 570 571 572
#endif
#ifdef PADDLE_WITH_PSCORE
      for (int j = 0; j < len; ++j) {
        device_keys[dev][cur + j] = task_keys[dev][j];
573
        float* ptr_val = task_ptrs[dev][j]->data();
T
Thunderbrook 已提交
574
        FeatureValue& val = device_vals[dev][cur + j];
575 576 577 578 579 580 581
        size_t dim = task_ptrs[dev][j]->size();
        val.delta_score = ptr_val[2];
        val.show = ptr_val[3];
        val.clk = ptr_val[4];
        val.slot = ptr_val[0];
        val.lr = ptr_val[5];
        val.lr_g2sum = ptr_val[6];
T
Thunderbrook 已提交
582 583
        val.cpu_ptr = (uint64_t)(task_ptrs[dev][j]);

584
        if (dim > 7) {
T
Thunderbrook 已提交
585 586
          val.mf_size = MF_DIM + 1;
          for (int x = 0; x < val.mf_size; x++) {
587
            val.mf[x] = ptr_val[x + 7];
T
Thunderbrook 已提交
588 589 590 591 592 593 594 595 596
          }
        } else {
          val.mf_size = 0;
          for (int x = 0; x < MF_DIM + 1; x++) {
            val.mf[x] = 0;
          }
        }
      }
#endif
597
      VLOG(3) << "GpuPs build hbmps done";
598 599

      device_mutex[dev]->unlock();
Y
yaoxuefeng 已提交
600 601
    }
  };
602

603 604 605 606 607 608 609 610 611 612 613
  if (!multi_mf_dim_) {
    for (size_t i = 0; i < threads.size(); i++) {
      threads[i] = std::thread(build_func, i);
    }
  } else {
    for (int i = 0; i < thread_keys_shard_num_; i++) {
      for (int j = 0; j < multi_mf_dim_; j++) {
        threads[i * multi_mf_dim_ + j] =
            std::thread(build_dynamic_mf_func, i, j);
      }
    }
Y
yaoxuefeng 已提交
614 615 616 617 618
  }
  for (std::thread& t : threads) {
    t.join();
  }
  timeline.Pause();
T
Thunderbrook 已提交
619
  VLOG(0) << "GpuPs prepare for build hbm cost " << timeline.ElapsedSec()
620
          << " seconds.";
Y
yaoxuefeng 已提交
621 622
}

623
void PSGPUWrapper::BuildGPUTask(std::shared_ptr<HeterContext> gpu_task) {
624
  int device_num = heter_devices_.size();
Y
yaoxuefeng 已提交
625 626
  platform::Timer timeline;
  timeline.Start();
T
Thunderbrook 已提交
627

628
  std::vector<size_t> feature_keys_count(device_num);
T
Thunderbrook 已提交
629
  size_t size_max = 0;
630 631 632 633 634 635 636 637 638 639 640 641 642 643 644
  if (!multi_mf_dim_) {
    for (int i = 0; i < device_num; i++) {
      feature_keys_count[i] = gpu_task->device_keys_[i].size();
      VLOG(1) << i << " card contains feasign nums: " << feature_keys_count[i];
      size_max = std::max(size_max, feature_keys_count[i]);
    }
  } else {
    for (int i = 0; i < device_num; i++) {
      for (int j = 0; j < multi_mf_dim_; j++) {
        feature_keys_count[i] += gpu_task->device_dim_ptr_[i][j].size();
      }
      VLOG(1) << i << " card with dynamic mf contains feasign nums: "
              << feature_keys_count[i];
      size_max = std::max(size_max, feature_keys_count[i]);
    }
T
Thunderbrook 已提交
645 646
  }
  if (HeterPs_) {
647 648
    delete HeterPs_;
    HeterPs_ = nullptr;
T
Thunderbrook 已提交
649
  }
650 651 652 653
  if (size_max <= 0) {
    VLOG(1) << "Skip build gpu ps cause feasign nums = " << size_max;
    return;
  }
654
  std::vector<std::thread> threads(device_num);
T
Thunderbrook 已提交
655
  HeterPs_ = HeterPsBase::get_instance(size_max, resource_);
656
  HeterPs_->set_nccl_comm_and_size(inner_comms_, inter_comms_, node_size_);
Y
yaoxuefeng 已提交
657
  auto build_func = [this, &gpu_task, &feature_keys_count](int i) {
658
    VLOG(3) << "building table: " << i;
659 660 661
    this->HeterPs_->build_ps(i, gpu_task->device_keys_[i].data(),
                             gpu_task->device_values_[i].data(),
                             feature_keys_count[i], 500000, 2);
662 663 664
    // if (feature_keys_count[i] > 0) {
    //   HeterPs_->show_one_table(i);
    // }
Y
yaoxuefeng 已提交
665 666 667 668 669 670
  };
  for (size_t i = 0; i < threads.size(); i++) {
    threads[i] = std::thread(build_func, i);
  }
  for (std::thread& t : threads) {
    t.join();
T
Thunderbrook 已提交
671 672
  }
  timeline.Pause();
673
  VLOG(1) << "GpuPs build table total costs: " << timeline.ElapsedSec()
T
Thunderbrook 已提交
674
          << " s.";
675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697
}

void PSGPUWrapper::LoadIntoMemory(bool is_shuffle) {
  platform::Timer timer;
  VLOG(3) << "Begin LoadIntoMemory(), dataset[" << dataset_ << "]";
  timer.Start();
  dataset_->LoadIntoMemory();
  timer.Pause();
  VLOG(0) << "LoadIntoMemory cost: " << timer.ElapsedSec() << "s";

  // local shuffle
  if (is_shuffle) {
    dataset_->LocalShuffle();
  }

  std::shared_ptr<HeterContext> gpu_task = gpu_task_pool_.Get();
  gpu_task->Reset();
  data_ready_channel_->Put(gpu_task);
  VLOG(3) << "End LoadIntoMemory(), dataset[" << dataset_ << "]";
}

void PSGPUWrapper::start_build_thread() {
  running_ = true;
698
  VLOG(3) << "start build CPU ps thread.";
699
  pre_build_threads_ = std::thread([this] { pre_build_thread(); });
700 701
}

702 703
void PSGPUWrapper::pre_build_thread() {
  // prebuild: process load_data
704 705 706 707 708
  while (running_) {
    std::shared_ptr<HeterContext> gpu_task = nullptr;
    if (!data_ready_channel_->Get(gpu_task)) {
      continue;
    }
709
    VLOG(3) << "thread PreBuildTask start.";
710 711 712
    platform::Timer timer;
    timer.Start();
    // build cpu ps data process
713
    PreBuildTask(gpu_task);
714
    timer.Pause();
715 716
    VLOG(1) << "thread PreBuildTask end, cost time: " << timer.ElapsedSec()
            << "s";
717 718 719 720 721
    buildcpu_ready_channel_->Put(gpu_task);
  }
  VLOG(3) << "build cpu thread end";
}

722 723 724 725 726 727 728 729 730 731
void PSGPUWrapper::build_task() {
  // build_task: build_pull + build_gputask
  std::shared_ptr<HeterContext> gpu_task = nullptr;
  // train end, gpu free
  if (!gpu_free_channel_->Get(gpu_task)) {
    return;
  }
  // ins and pre_build end
  if (!buildcpu_ready_channel_->Get(gpu_task)) {
    return;
732
  }
733 734 735 736 737 738 739 740 741 742 743

  VLOG(1) << "BuildPull start.";
  platform::Timer timer;
  timer.Start();
  BuildPull(gpu_task);
  BuildGPUTask(gpu_task);
  timer.Pause();
  VLOG(1) << "BuildPull + BuildGPUTask end, cost time: " << timer.ElapsedSec()
          << "s";

  current_task_ = gpu_task;
744 745 746 747 748 749 750 751 752
}

void PSGPUWrapper::BeginPass() {
  platform::Timer timer;
  timer.Start();
  if (current_task_) {
    PADDLE_THROW(
        platform::errors::Fatal("[BeginPass] current task is not ended."));
  }
753 754

  build_task();
755
  timer.Pause();
756 757 758 759 760 761

  if (current_task_ == nullptr) {
    PADDLE_THROW(platform::errors::Fatal(
        "[BeginPass] after build_task, current task is not null."));
  }

T
Thunderbrook 已提交
762
  VLOG(0) << "BeginPass end, cost time: " << timer.ElapsedSec() << "s";
763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779
}

void PSGPUWrapper::EndPass() {
  if (!current_task_) {
    PADDLE_THROW(
        platform::errors::Fatal("[EndPass] current task has been ended."));
  }
  platform::Timer timer;
  timer.Start();
  size_t keysize_max = 0;
  // in case of feasign_num = 0, skip dump_to_cpu
  for (size_t i = 0; i < heter_devices_.size(); i++) {
    keysize_max = std::max(keysize_max, current_task_->device_keys_[i].size());
  }
  if (keysize_max != 0) {
    HeterPs_->end_pass();
  }
780 781

  gpu_task_pool_.Push(current_task_);
782 783 784
  current_task_ = nullptr;
  gpu_free_channel_->Put(current_task_);
  timer.Pause();
T
Thunderbrook 已提交
785
  VLOG(0) << "EndPass end, cost time: " << timer.ElapsedSec() << "s";
T
Thunderbrook 已提交
786 787 788 789 790 791 792 793 794 795 796 797 798 799
}

void PSGPUWrapper::PullSparse(const paddle::platform::Place& place,
                              const int table_id,
                              const std::vector<const uint64_t*>& keys,
                              const std::vector<float*>& values,
                              const std::vector<int64_t>& slot_lengths,
                              const int hidden_size) {
  VLOG(3) << "Begine Gpu Ps PullSparse";
  platform::Timer all_timer;
  platform::Timer pull_gpups_timer;
  all_timer.Start();
  int64_t total_length =
      std::accumulate(slot_lengths.begin(), slot_lengths.end(), 0UL);
800
  auto buf = memory::Alloc(place, total_length * sizeof(FeatureValue));
T
Thunderbrook 已提交
801 802 803 804 805 806
  FeatureValue* total_values_gpu = reinterpret_cast<FeatureValue*>(buf->ptr());
  if (platform::is_cpu_place(place)) {
    PADDLE_THROW(platform::errors::Unimplemented(
        "Warning:: CPUPlace is not supported in GpuPs now."));
  } else if (platform::is_gpu_place(place)) {
    VLOG(3) << "Begin copy keys, key_num[" << total_length << "]";
807
    int device_id = place.GetDeviceId();
T
Thunderbrook 已提交
808 809 810 811 812 813 814 815 816 817
    int devid_2_index = HeterPs_->get_index_by_devid(device_id);
    LoDTensor& total_keys_tensor = keys_tensor[devid_2_index];
    uint64_t* total_keys = reinterpret_cast<uint64_t*>(
        total_keys_tensor.mutable_data<int64_t>({total_length, 1}, place));

    // construct slot_level lod info
    auto slot_lengths_lod = slot_lengths;
    for (size_t i = 1; i < slot_lengths_lod.size(); i++) {
      slot_lengths_lod[i] += slot_lengths_lod[i - 1];
    }
818
    auto buf_key = memory::Alloc(place, keys.size() * sizeof(uint64_t*));
T
Thunderbrook 已提交
819
    auto buf_length =
820
        memory::Alloc(place, slot_lengths.size() * sizeof(int64_t));
T
Thunderbrook 已提交
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
    uint64_t** gpu_keys = reinterpret_cast<uint64_t**>(buf_key->ptr());
    int64_t* gpu_len = reinterpret_cast<int64_t*>(buf_length->ptr());
    cudaMemcpy(gpu_keys, keys.data(), keys.size() * sizeof(uint64_t*),
               cudaMemcpyHostToDevice);
    cudaMemcpy(gpu_len, slot_lengths_lod.data(),
               slot_lengths.size() * sizeof(int64_t), cudaMemcpyHostToDevice);

    this->CopyKeys(place, gpu_keys, total_keys, gpu_len,
                   static_cast<int>(slot_lengths.size()),
                   static_cast<int>(total_length));
    VLOG(3) << "Begin call PullSparseGPU in GPUPS, dev: " << devid_2_index
            << " len: " << total_length;
    pull_gpups_timer.Start();
    HeterPs_->pull_sparse(devid_2_index, total_keys, total_values_gpu,
                          static_cast<int>(total_length));
    // PADDLE_ENFORCE_EQ(ret, 0, platform::errors::PreconditionNotMet(
    //                              "PullSparseGPU failed in GPUPS."));
    pull_gpups_timer.Pause();

    VLOG(3) << "Begin Copy result to tensor, total_length[" << total_length
            << "]";
    this->CopyForPull(place, gpu_keys, values, total_values_gpu, gpu_len,
                      static_cast<int>(slot_lengths.size()), hidden_size,
                      total_length);
  } else {
    PADDLE_THROW(platform::errors::PreconditionNotMet(
        "GpuPs: PullSparse Only Support CUDAPlace Now."));
  }
  all_timer.Pause();
850
  VLOG(3) << "GpuPs PullSparse total costs: " << all_timer.ElapsedSec()
T
Thunderbrook 已提交
851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867
          << " s, of which GPUPS costs: " << pull_gpups_timer.ElapsedSec()
          << " s";
  VLOG(3) << "End PullSparse";
}

void PSGPUWrapper::PushSparseGrad(const paddle::platform::Place& place,
                                  const int table_id,
                                  const std::vector<const uint64_t*>& keys,
                                  const std::vector<const float*>& grad_values,
                                  const std::vector<int64_t>& slot_lengths,
                                  const int hidden_size, const int batch_size) {
  VLOG(3) << "Begin GPUPS PushSparseGrad";
  platform::Timer all_timer;
  platform::Timer push_gpups_timer;
  all_timer.Start();
  int64_t total_length =
      std::accumulate(slot_lengths.begin(), slot_lengths.end(), 0UL);
868
  auto buf = memory::Alloc(place, total_length * sizeof(FeaturePushValue));
T
Thunderbrook 已提交
869 870 871 872 873 874
  FeaturePushValue* total_grad_values_gpu =
      reinterpret_cast<FeaturePushValue*>(buf->ptr());
  if (platform::is_cpu_place(place)) {
    PADDLE_THROW(platform::errors::Unimplemented(
        "Warning:: CPUPlace is not supported in GPUPS now."));
  } else if (platform::is_gpu_place(place)) {
875
    int device_id = place.GetDeviceId();
T
Thunderbrook 已提交
876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894
    int devid_2_index = HeterPs_->get_index_by_devid(device_id);
    LoDTensor& cached_total_keys_tensor = keys_tensor[devid_2_index];
    uint64_t* total_keys =
        reinterpret_cast<uint64_t*>(cached_total_keys_tensor.data<int64_t>());
    VLOG(3) << "Begin copy grad tensor to gpups struct";
    this->CopyForPush(place, grad_values, total_grad_values_gpu, slot_lengths,
                      hidden_size, total_length, batch_size);

    VLOG(3) << "Begin call PushSparseGPU in GPUPS, dev: " << devid_2_index
            << " len: " << total_length;
    push_gpups_timer.Start();
    HeterPs_->push_sparse(devid_2_index, total_keys, total_grad_values_gpu,
                          static_cast<int>(total_length));
    push_gpups_timer.Pause();
  } else {
    PADDLE_THROW(platform::errors::PreconditionNotMet(
        "GPUPS: PushSparseGrad Only Support CUDAPlace Now."));
  }
  all_timer.Pause();
895
  VLOG(3) << "PushSparseGrad total cost: " << all_timer.ElapsedSec()
T
Thunderbrook 已提交
896 897 898 899 900 901 902 903
          << " s, of which GPUPS cost: " << push_gpups_timer.ElapsedSec()
          << " s";
  VLOG(3) << "End PushSparseGrad";
}

}  // end namespace framework
}  // end namespace paddle
#endif