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
|
#include "caffe2/sgd/math_lp.h"
namespace caffe2 {
namespace internal {
template <>
void dot<float, float, float>(
const int N,
const float* x,
const float* y,
float* z,
CPUContext* ctx) {
math::Dot<float, CPUContext>(N, x, y, z, ctx);
}
template <>
void dot<float, at::Half, float>(
const int N,
const float* x,
const at::Half* y,
float* z,
CPUContext* ctx) {
#ifdef _MSC_VER
std::vector<float> tmp_y_vec(N);
float* tmp_y = tmp_y_vec.data();
#else
float tmp_y[N];
#endif
for (int i = 0; i < N; i++) {
#ifdef __F16C__
tmp_y[i] = _cvtss_sh(y[i], 0); // TODO: vectorize
#else
tmp_y[i] = y[i];
#endif
}
math::Dot<float, CPUContext>(N, x, tmp_y, z, ctx);
}
} // namespace internal
} // namespace caffe2
|
