Caffe2 - C++ API
A deep learning, cross platform ML framework
typed_axpy_avx.cc
1 
17 #include "caffe2/core/types.h"
18 #include "caffe2/perfkernels/cvtsh_ss_bugfix.h"
19 #include "caffe2/perfkernels/typed_axpy.h"
20 #include "caffe2/utils/math.h"
21 
22 #include <emmintrin.h>
23 #include <immintrin.h>
24 
25 namespace caffe2 {
26 
27 void TypedAxpy_float16_float__avx_f16c(
28  int N,
29  const float a,
30  const float16* x,
31  float* y) {
32  // if x does not start at the 16 byte boundary, we will process the first few.
33  // before we get to a real one.
34  while (N && (unsigned long)x % 16) {
35  *(y++) += _cvtsh_ss((*(x++)).x) * a;
36  --N;
37  }
38 
39  // From now on we can do vectorized additions using __m256, which is 8 floats,
40  // so we will vectorize every 8 element and then resort to cvtsh_ss.
41  __m256 mma = _mm256_set1_ps(a);
42  int current = 0;
43  const int bound = (N % 8) ? N - 8 : N;
44 
45  for (; current < bound; current += 8) {
46  __m128i mmx_16 =
47  _mm_loadu_si128(reinterpret_cast<const __m128i*>(x + current));
48  __m256 mmx_32 = _mm256_cvtph_ps(mmx_16);
49  __m256 mmy_in = _mm256_loadu_ps(y + current);
50  __m256 mmmul = _mm256_mul_ps(mmx_32, mma);
51  __m256 mmy_out = _mm256_add_ps(mmmul, mmy_in);
52  _mm256_storeu_ps(y + current, mmy_out);
53  }
54 
55  if (bound != N) {
56  while (current < N) {
57  y[current] += _cvtsh_ss(x[current].x) * a;
58  ++current;
59  }
60  }
61 }
62 
63 } // namespace caffe2
Copyright (c) 2016-present, Facebook, Inc.