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
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
|
#include "caffe2/operators/elementwise_ops.h"
#include "caffe2/utils/eigen_utils.h"
#include <algorithm>
namespace caffe2 {
REGISTER_CPU_OPERATOR(
Not,
UnaryElementwiseOp<BoolTypes, CPUContext, NotFunctor<CPUContext>>);
REGISTER_CPU_OPERATOR(
Sign,
UnaryElementwiseOp<NumericTypes, CPUContext, SignFunctor<CPUContext>>);
#define REGISTER_CPU_COMPARE_OPERATOR(Op) \
REGISTER_CPU_OPERATOR( \
Op, \
BinaryElementwiseOp< \
TensorTypes<bool, int32_t, int64_t, float, double>, \
CPUContext, \
Op##Functor<CPUContext>, \
FixedType<bool>>)
REGISTER_CPU_COMPARE_OPERATOR(EQ);
REGISTER_CPU_COMPARE_OPERATOR(NE);
REGISTER_CPU_COMPARE_OPERATOR(LT);
REGISTER_CPU_COMPARE_OPERATOR(LE);
REGISTER_CPU_COMPARE_OPERATOR(GT);
REGISTER_CPU_COMPARE_OPERATOR(GE);
#undef REGISTER_CPU_COMPARE_OPERATOR
#define REGISTER_CPU_LOGICAL_BINARY_OPERATOR(Op) \
REGISTER_CPU_OPERATOR( \
Op, BinaryElementwiseOp<BoolTypes, CPUContext, Op##Functor<CPUContext>>)
REGISTER_CPU_LOGICAL_BINARY_OPERATOR(And);
REGISTER_CPU_LOGICAL_BINARY_OPERATOR(Or);
REGISTER_CPU_LOGICAL_BINARY_OPERATOR(Xor);
#undef REGISTER_CPU_LOGICAL_BINARY_OPERATOR
#define REGISTER_CPU_BITWISE_BINARY_OPERATOR(Op) \
REGISTER_CPU_OPERATOR( \
Op, \
BinaryElementwiseOp<IntBoolTypes, CPUContext, Op##Functor<CPUContext>>)
REGISTER_CPU_BITWISE_BINARY_OPERATOR(BitwiseAnd);
REGISTER_CPU_BITWISE_BINARY_OPERATOR(BitwiseOr);
REGISTER_CPU_BITWISE_BINARY_OPERATOR(BitwiseXor);
#undef REGISTER_CPU_BITWISE_BINARY_OPERATOR
template <typename T>
void SRLHelper::sum2one(const T* x, T* y, size_t n) {
*y = ConstEigenArrayMap<T>(x, n, 1).sum();
}
template <typename T>
void SRLHelper::RunWithBroadcastFront(
const T* x,
T* y,
size_t pre,
size_t n,
CPUContext*) {
EigenArrayMap<T>(y, n, 1) = ConstEigenArrayMap<T>(x, n, pre).rowwise().sum();
}
template <typename T>
void SRLHelper::RunWithBroadcastBack(
const T* x,
T* y,
size_t post,
size_t n,
CPUContext*) {
EigenArrayMap<T>(y, 1, n) = ConstEigenArrayMap<T>(x, post, n).colwise().sum();
}
template <typename T>
void SRLHelper::RunWithBroadcast2(
const T* a,
T* y,
size_t pre,
size_t n,
size_t post,
CPUContext*) {
for (auto i = 0U; i < n; ++i) {
y[i] = 0;
for (auto j = 0U; j < pre; ++j) {
for (auto k = 0U; k < post; ++k) {
y[i] += a[(j * n + i) * post + k];
}
}
}
}
template <>
template <typename T>
bool SumReduceLikeOp<CPUContext>::DoRunWithType() {
const auto& A = Input(0);
const auto& B = Input(1);
CAFFE_ENFORCE(!IsInputOutputAlias(1, 0), "In-place is not allowed.");
auto* C = Output(0, B.sizes(), at::dtype<T>());
const T* Adata = A.template data<T>();
auto* Cdata = C->template mutable_data<T>();
if (B.numel() == 1) {
auto count = A.numel();
SRLHelper::sum2one<T>(Adata, Cdata, count);
} else {
// NOLINTNEXTLINE(cppcoreguidelines-init-variables)
size_t pre, n, post;
std::tie(pre, n, post) =
elementwise_ops_utils::ComputeLegacyBroadcastSizes(A, B, axis_);
if (post == 1) {
SRLHelper::RunWithBroadcastFront<T>(Adata, Cdata, pre, n, &context_);
} else if (pre == 1) {
SRLHelper::RunWithBroadcastBack<T>(Adata, Cdata, post, n, &context_);
} else {
SRLHelper::RunWithBroadcast2<T>(Adata, Cdata, pre, n, post, &context_);
}
}
return true;
}
REGISTER_CPU_OPERATOR(SumReduceLike, SumReduceLikeOp<CPUContext>);
} // namespace caffe2
|
