Caffe
include
caffe
util
nccl.hpp
1
#ifndef CAFFE_UTIL_NCCL_H_
2
#define CAFFE_UTIL_NCCL_H_
3
#ifdef USE_NCCL
4
5
#include <nccl.h>
6
7
#include "caffe/common.hpp"
8
9
#define NCCL_CHECK(condition) \
10
{ \
11
ncclResult_t result = condition; \
12
CHECK_EQ(result, ncclSuccess) << " " \
13
<< ncclGetErrorString(result); \
14
}
15
16
namespace
caffe
{
17
18
namespace
nccl {
19
20
template
<
typename
Dtype>
class
dataType;
21
22
template
<>
class
dataType<float> {
23
public
:
24
static
const
ncclDataType_t type = ncclFloat;
25
};
26
template
<>
class
dataType<double> {
27
public
:
28
static
const
ncclDataType_t type = ncclDouble;
29
};
30
31
}
// namespace nccl
32
33
}
// namespace caffe
34
35
#endif // end USE_NCCL
36
37
#endif // CAFFE_UTIL_NCCL_H_
caffe
A layer factory that allows one to register layers. During runtime, registered layers can be called b...
Definition:
blob.hpp:14
Generated on Wed Dec 23 2020 16:08:21 for Caffe by
1.8.20