14
14
15
15
#pragma once
16
16
17
- #include < thread>
17
+ #include < thread> // NOLINT
18
18
#include < typeindex>
19
+ #include < vector>
19
20
#include " paddle/fluid/platform/dynload/nccl.h"
20
21
#include " paddle/fluid/platform/enforce.h"
21
22
@@ -29,6 +30,8 @@ inline ncclDataType_t ToNCCLDataType(std::type_index type) {
29
30
return ncclDouble;
30
31
} else if (type == typeid (int )) { // NOLINT
31
32
return ncclInt;
33
+ } else if (type == typeid (int64_t )) { // NOLINT
34
+ return ncclInt64;
32
35
} else {
33
36
PADDLE_THROW (" Not supported" );
34
37
}
@@ -66,23 +69,23 @@ struct NCCLContext {
66
69
return boost::get<platform::CUDAPlace>(ctx_->GetPlace ()).device ;
67
70
}
68
71
69
- static void InitNCCLContext (std::unordered_map<int , NCCLContext> & contexts,
72
+ static void InitNCCLContext (std::unordered_map<int , NCCLContext> * contexts,
70
73
const std::vector<platform::Place> &places) {
71
74
std::vector<ncclComm_t> comms;
72
75
std::vector<int > devs;
73
- comms.resize (contexts. size ());
74
- devs.reserve (contexts. size ());
76
+ comms.resize (contexts-> size ());
77
+ devs.reserve (contexts-> size ());
75
78
76
79
for (auto &p : places) {
77
80
devs.push_back (boost::get<platform::CUDAPlace>(p).device );
78
81
}
79
82
80
83
PADDLE_ENFORCE (platform::dynload::ncclCommInitAll (
81
- &comms[0 ], static_cast <int >(contexts. size ()), &devs[0 ]));
84
+ &comms[0 ], static_cast <int >(contexts-> size ()), &devs[0 ]));
82
85
83
86
int i = 0 ;
84
87
for (auto &dev_id : devs) {
85
- contexts. at (dev_id).comm_ = comms[i++];
88
+ contexts-> at (dev_id).comm_ = comms[i++];
86
89
}
87
90
}
88
91
};
@@ -91,7 +94,7 @@ struct NCCLContextMap {
91
94
std::unordered_map<int , NCCLContext> contexts_;
92
95
std::vector<int > order_;
93
96
94
- NCCLContextMap (const std::vector<platform::Place> &places) {
97
+ explicit NCCLContextMap (const std::vector<platform::Place> &places) {
95
98
order_.reserve (places.size ());
96
99
for (auto &p : places) {
97
100
int dev_id = boost::get<CUDAPlace>(p).device ;
0 commit comments