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
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
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
|
#include "allreduce_ops.h"
#include <math.h>
#include <gloo/allreduce_bcube.h>
#include <gloo/allreduce_halving_doubling.h>
#include <gloo/allreduce_ring.h>
#include <gloo/allreduce_ring_chunked.h>
#include <gloo/types.h>
namespace {
/**
* This is a helper function which attempts to get a base value depending on the
* # of nodes. Larger the base the better performance (up to 4) is what we have
* observed in gloo benchmarks. At the moment bcube works only if # nodes = base
* ^ x. Where x is some constant. So, if # node don't match our expectation
* simply return -1. This will indicate caller to switch to another algorithm
* like halving-doubling.
*/
static int getAllrduceBcubeBase(int nodes) {
auto getExponent = [](int n, int b) -> int {
float lg2n = log2(n);
float lg2b = log2(b);
return ceil(lg2n / lg2b);
};
auto baseCheck = [&](int n, int b) -> bool {
int e = getExponent(n, b);
return n == pow(b, e);
};
for (const auto base : {6, 5, 4, 3, 2}) {
if (baseCheck(nodes, base)) {
return base;
}
/*
* Base could work if # nodes is multiple of the base yet smaller than
* base^2
*/
if (nodes < base * base && 0 == nodes % base) {
return base;
}
}
return -1;
}
} // namespace
namespace caffe2 {
namespace gloo {
template <class Context>
void AllreduceOp<Context>::initializeBcube() {
int base = getAllrduceBcubeBase(init_.size);
if (-1 == base) {
return initializeHalvingDoubling();
}
init_.context->base = base;
if (init_.template IsType<float>()) {
algorithm_.reset(new ::gloo::AllreduceBcube<float>(
init_.context, init_.template getOutputs<float>(), init_.size));
} else if (init_.template IsType<::at::Half>()) {
algorithm_.reset(new ::gloo::AllreduceBcube<::gloo::float16>(
init_.context,
init_.template getOutputs<::gloo::float16>(),
init_.size));
} else {
CAFFE_ENFORCE(false, "Unhandled type: ", init_.meta.name());
}
}
template <class Context>
void AllreduceOp<Context>::initializeHalvingDoubling() {
if (init_.template IsType<float>()) {
algorithm_.reset(new ::gloo::AllreduceHalvingDoubling<float>(
init_.context, init_.template getOutputs<float>(), init_.size));
} else if (init_.template IsType<::at::Half>()) {
algorithm_.reset(new ::gloo::AllreduceHalvingDoubling<::gloo::float16>(
init_.context,
init_.template getOutputs<::gloo::float16>(),
init_.size));
} else {
CAFFE_ENFORCE(false, "Unhandled type: ", init_.meta.name());
}
}
// Used outside of the translation unit
template void AllreduceOp<CPUContext>::initializeHalvingDoubling();
template <class Context>
void AllreduceOp<Context>::initializeRingFull() {
if (init_.template IsType<float>()) {
algorithm_.reset(new ::gloo::AllreduceRing<float>(
init_.context, init_.template getOutputs<float>(), init_.size));
} else if (init_.template IsType<::at::Half>()) {
algorithm_.reset(new ::gloo::AllreduceRing<::gloo::float16>(
init_.context,
init_.template getOutputs<::gloo::float16>(),
init_.size));
} else {
CAFFE_ENFORCE(false, "Unhandled type: ", init_.meta.name());
}
}
template <class Context>
void AllreduceOp<Context>::initializeRingChunked() {
if (init_.template IsType<float>()) {
algorithm_.reset(new ::gloo::AllreduceRingChunked<float>(
init_.context, init_.template getOutputs<float>(), init_.size));
} else if (init_.template IsType<::at::Half>()) {
algorithm_.reset(new ::gloo::AllreduceRingChunked<::gloo::float16>(
init_.context,
init_.template getOutputs<::gloo::float16>(),
init_.size));
} else {
CAFFE_ENFORCE(false, "Unhandled type: ", init_.meta.name());
}
}
namespace {
REGISTER_CPU_OPERATOR_WITH_ENGINE(Allreduce, GLOO, AllreduceOp<CPUContext>);
} // namespace
} // namespace gloo
} // namespace caffe2
|
