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
| #ifndef CAFFE2_UTILS_MATH_DETAIL_H_
| #define CAFFE2_UTILS_MATH_DETAIL_H_
| namespace caffe2 {
|
| class CPUContext;
|
| namespace math {
| namespace detail {
|
| // proxy to a class because of partial specialization limitations for functions
|
| template<typename T, class Context, int FixedSize>
| struct ScaleImpl {
| inline void operator()(
| const int N,
| const float alpha,
| const T* x,
| T* y,
| Context* context) {
| Scale(N, alpha, x, y, context);
| }
| };
|
| // Put light-weight implementations in .h file to enable inlining
| template<typename T>
| struct ScaleImpl<T, CPUContext, 1> {
| inline void operator()(
| const int N,
| const float alpha,
| const T* x,
| T* y,
| CPUContext* /*context*/) {
| DCHECK_EQ(N, 1);
| *y = *x * alpha;
| }
| };
|
| template<typename T, class Context, int FixedSize>
| struct AxpyImpl {
| inline void operator()(
| const int N,
| const float alpha,
| const T* x,
| T* y,
| Context* context) {
| Axpy(N, alpha, x, y, context);
| }
| };
|
| // Put light-weight implementations in .h file to enable inlining
| template<typename T>
| struct AxpyImpl<T, CPUContext, 1> {
| inline void operator()(
| const int N,
| const float alpha,
| const T* x,
| T* y,
| CPUContext* /*context*/) {
| DCHECK_EQ(N, 1);
| *y += *x * alpha;
| }
| };
|
|
| } // namespace detail
|
| template <typename T, class Context, int FixedSize>
| inline void ScaleFixedSize(
| const int N,
| const float alpha,
| const T* x,
| T* y,
| Context* context) {
| detail::ScaleImpl<T, Context, FixedSize>()(N, alpha, x, y, context);
| }
|
| template <typename T, class Context, int FixedSize>
| inline void AxpyFixedSize(
| const int N,
| const float alpha,
| const T* x,
| T* y,
| Context* context) {
| detail::AxpyImpl<T, Context, FixedSize>()(N, alpha, x, y, context);
| }
|
| } // namespace math
| } // namespace caffe2
|
| #endif // CAFFE2_UTILS_MATH_DETAIL_H_
|
|
