Caffe2 - C++ API
A deep learning, cross platform ML framework
math-detail.h
1 
17 #ifndef CAFFE2_UTILS_MATH_DETAIL_H_
18 #define CAFFE2_UTILS_MATH_DETAIL_H_
19 namespace caffe2 {
20 
21 class CPUContext;
22 
23 namespace math {
24 namespace detail {
25 
26 // proxy to a class because of partial specialization limitations for functions
27 
28 template<typename T, class Context, int FixedSize>
29 struct ScaleImpl {
30  inline void operator()(
31  const int N,
32  const float alpha,
33  const T* x,
34  T* y,
35  Context* context) {
36  Scale(N, alpha, x, y, context);
37  }
38 };
39 
40 // Put light-weight implementations in .h file to enable inlining
41 template<typename T>
42 struct ScaleImpl<T, CPUContext, 1> {
43  inline void operator()(
44  const int N,
45  const float alpha,
46  const T* x,
47  T* y,
48  CPUContext* /*context*/) {
49  DCHECK_EQ(N, 1);
50  *y = *x * alpha;
51  }
52 };
53 
54 template<typename T, class Context, int FixedSize>
55 struct AxpyImpl {
56  inline void operator()(
57  const int N,
58  const float alpha,
59  const T* x,
60  T* y,
61  Context* context) {
62  Axpy(N, alpha, x, y, context);
63  }
64 };
65 
66 // Put light-weight implementations in .h file to enable inlining
67 template<typename T>
68 struct AxpyImpl<T, CPUContext, 1> {
69  inline void operator()(
70  const int N,
71  const float alpha,
72  const T* x,
73  T* y,
74  CPUContext* /*context*/) {
75  DCHECK_EQ(N, 1);
76  *y += *x * alpha;
77  }
78 };
79 
80 
81 } // namespace detail
82 
83 template <typename T, class Context, int FixedSize>
84 inline void ScaleFixedSize(
85  const int N,
86  const float alpha,
87  const T* x,
88  T* y,
89  Context* context) {
90  detail::ScaleImpl<T, Context, FixedSize>()(N, alpha, x, y, context);
91 }
92 
93 template <typename T, class Context, int FixedSize>
94 inline void AxpyFixedSize(
95  const int N,
96  const float alpha,
97  const T* x,
98  T* y,
99  Context* context) {
100  detail::AxpyImpl<T, Context, FixedSize>()(N, alpha, x, y, context);
101 }
102 
103 } // namespace math
104 } // namespace caffe2
105 
106 #endif // CAFFE2_UTILS_MATH_DETAIL_H_
The CPU Context, representing the bare minimum of what a Context class in Caffe2 should implement...
Definition: context.h:82
Copyright (c) 2016-present, Facebook, Inc.