nccl_gpu_common.h 1.8 KB
Newer Older
D
Dong Zhihong 已提交
1 2
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.

L
Luo Tao 已提交
3 4 5
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
D
Dong Zhihong 已提交
6

L
Luo Tao 已提交
7
    http://www.apache.org/licenses/LICENSE-2.0
D
Dong Zhihong 已提交
8

L
Luo Tao 已提交
9 10 11 12 13
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. */
D
Dong Zhihong 已提交
14

D
dongzhihong 已提交
15 16
#pragma once

D
Dong Zhihong 已提交
17 18
#include <algorithm>
#include <condition_variable>
D
dzhwinter 已提交
19 20
#include <memory>
#include <mutex>
D
Dong Zhihong 已提交
21
#include <string>
D
dzhwinter 已提交
22
#include <unordered_map>
D
Dong Zhihong 已提交
23
#include <vector>
D
dzhwinter 已提交
24

D
dongzhihong 已提交
25
#include "paddle/platform/device_context.h"
D
Dong Zhihong 已提交
26
#include "paddle/platform/dynload/nccl.h"
D
Dong Zhihong 已提交
27
#include "paddle/platform/enforce.h"
D
Dong Zhihong 已提交
28
#include "paddle/platform/macros.h"
D
dongzhihong 已提交
29 30 31 32

namespace paddle {
namespace platform {

D
Dong Zhihong 已提交
33 34
constexpr int kInvalidGPUId = -1;

D
Dong Zhihong 已提交
35 36
struct Communicator {
  std::vector<ncclComm_t> comms_;
D
Dong Zhihong 已提交
37
  std::unordered_map<int, int> comm_id_map_;
38
  bool inited_;
D
Dong Zhihong 已提交
39

D
fix ci  
Dong Zhihong 已提交
40 41
  Communicator() {}

D
Dong Zhihong 已提交
42 43 44
  int GetCommId(int device_id) const { return comm_id_map_.at(device_id); }

  void InitAll(const std::vector<int>& gpus) {
D
Dong Zhihong 已提交
45
    comms_.resize(gpus.size());
46
    inited_ = false;
D
Dong Zhihong 已提交
47 48 49
    for (size_t i = 0; i < gpus.size(); ++i) {
      comm_id_map_[gpus[i]] = i;
    }
D
Dong Zhihong 已提交
50 51
    PADDLE_ENFORCE(
        dynload::ncclCommInitAll(comms_.data(), gpus.size(), gpus.data()));
52
    inited_ = true;
D
Dong Zhihong 已提交
53
  }
D
Dong Zhihong 已提交
54 55

  ~Communicator() {
56 57 58 59 60
    if (inited_) {
      for (size_t i = 0; i < comms_.size(); ++i) {
        // FIXME(dzh) : PADDLE_ENFORCE return void
        dynload::ncclCommDestroy(comms_[i]);
      }
D
Dong Zhihong 已提交
61 62
    }
  }
D
Dong Zhihong 已提交
63

D
Dong Zhihong 已提交
64
  DISABLE_COPY_AND_ASSIGN(Communicator);
D
Dong Zhihong 已提交
65
};
D
dzhwinter 已提交
66

D
Dong Zhihong 已提交
67
}  // namespace platform
D
dzhwinter 已提交
68
}  // namespace paddle