Caffe2 - C++ API
A deep learning, cross platform ML framework
typed_axpy.cc
1 
17 #include "caffe2/perfkernels/typed_axpy.h"
18 #include "caffe2/core/types.h"
19 #include "caffe2/perfkernels/common.h"
20 #include "caffe2/utils/cpuid.h"
21 #include "caffe2/utils/math.h"
22 
23 namespace caffe2 {
24 
25 template <>
26 void TypedAxpy<float, float>(int N, const float a, const float* x, float* y) {
27  // This uses a hack that axpy implementation actually does not use the
28  // CPUContext, so passing in a nullpointer works.
29  math::Axpy<float, CPUContext>(N, a, x, y, nullptr);
30 }
31 
32 void TypedAxpy_float16_float__base(
33  int N,
34  const float a,
35  const float16* x,
36  float* y) {
37  for (int i = 0; i < N; ++i) {
38  union {
39  uint32_t intval;
40  float floatval;
41  } t1;
42  uint32_t t2, t3;
43  t1.intval = x[i].x & 0x7fff; // Non-sign bits
44  t2 = x[i].x & 0x8000; // Sign bit
45  t3 = x[i].x & 0x7c00; // Exponent
46  t1.intval <<= 13; // Align mantissa on MSB
47  t2 <<= 16; // Shift sign bit into position
48  t1.intval += 0x38000000; // Adjust bias
49  t1.intval = (t3 == 0 ? 0 : t1.intval); // Denormals-as-zero
50  t1.intval |= t2; // Re-insert sign bit
51  y[i] += t1.floatval * a;
52  }
53 }
54 
55 template <>
56 void TypedAxpy<float16, float>(
57  int N,
58  const float a,
59  const float16* x,
60  float* y) {
61  AVX2_FMA_DO(TypedAxpy_float16_float, N, a, x, y);
62  AVX_F16C_DO(TypedAxpy_float16_float, N, a, x, y);
63  BASE_DO(TypedAxpy_float16_float, N, a, x, y);
64 }
65 
66 void TypedAxpy_uint8_float__base(
67  int N,
68  const float a,
69  const std::uint8_t* x,
70  float* y) {
71  for (int i = 0; i < N; ++i) {
72  y[i] += (float)(x[i]) * a;
73  }
74 }
75 
76 template <>
77 void TypedAxpy<std::uint8_t, float>(
78  int N,
79  const float a,
80  const std::uint8_t* x,
81  float* y) {
82  AVX2_FMA_DO(TypedAxpy_uint8_float, N, a, x, y);
83  BASE_DO(TypedAxpy_uint8_float, N, a, x, y);
84 }
85 
86 } // namespace caffe2
Copyright (c) 2016-present, Facebook, Inc.