forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
PointwiseOpsKernel.cpp
197 lines (188 loc) · 7.93 KB
/
PointwiseOpsKernel.cpp
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
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
// Ternary and higher-order pointwise operations
#define TORCH_ASSERT_NO_OPERATORS
#include <ATen/Dispatch.h>
#include <ATen/native/PointwiseOps.h>
#include <ATen/native/TensorIterator.h>
#include <ATen/native/cpu/Loops.h>
#include <c10/core/Scalar.h>
namespace at {
namespace native {
namespace {
static void addcmul_cpu_kernel(TensorIteratorBase& iter, const Scalar& value) {
ScalarType dtype = iter.dtype(0);
if (iter.dtype() == kBFloat16) {
float float_val = value.to<float>();
auto float_vec = Vectorized<float>(float_val);
cpu_kernel_vec(
iter,
[=](BFloat16 self_val, BFloat16 t1_val, BFloat16 t2_val) -> BFloat16 {
return float(self_val) + float_val * float(t1_val) * float(t2_val);
},
[=](Vectorized<BFloat16> self_vec,
Vectorized<BFloat16> t1_vec,
Vectorized<BFloat16> t2_vec) {
Vectorized<float> self_vec0, self_vec1;
std::tie(self_vec0, self_vec1) = convert_bfloat16_float(self_vec);
Vectorized<float> t1_vec0, t1_vec1, t2_vec0, t2_vec1;
std::tie(t1_vec0, t1_vec1) = convert_bfloat16_float(t1_vec);
std::tie(t2_vec0, t2_vec1) = convert_bfloat16_float(t2_vec);
self_vec0 = self_vec0 + float_vec * t1_vec0 * t2_vec0;
self_vec1 = self_vec1 + float_vec * t1_vec1 * t2_vec1;
return convert_float_bfloat16(self_vec0, self_vec1);
});
} else {
AT_DISPATCH_ALL_TYPES_AND_COMPLEX(dtype, "addcmul_cpu_out", [&] {
scalar_t scalar_val = value.to<scalar_t>();
auto scalar_vec = Vectorized<scalar_t>(scalar_val);
cpu_kernel_vec(
iter,
[=](scalar_t self_val, scalar_t t1_val, scalar_t t2_val) -> scalar_t {
return self_val + scalar_val * t1_val * t2_val;
},
[=](Vectorized<scalar_t> self_vec,
Vectorized<scalar_t> t1_vec,
Vectorized<scalar_t> t2_vec) {
return self_vec + scalar_vec * t1_vec * t2_vec;
});
});
}
}
static void addcdiv_cpu_kernel(TensorIteratorBase& iter, const Scalar& value) {
ScalarType dtype = iter.dtype(0);
if (dtype == kBFloat16) {
float float_val = value.to<float>();
auto float_vec = Vectorized<float>(float_val);
cpu_kernel_vec(
iter,
[=](BFloat16 self_val, BFloat16 t1_val, BFloat16 t2_val) -> BFloat16 {
return float(self_val) + float_val * float(t1_val) / float(t2_val);
},
[=](Vectorized<BFloat16> self_vec,
Vectorized<BFloat16> t1_vec,
Vectorized<BFloat16> t2_vec) {
Vectorized<float> self_vec0, self_vec1;
std::tie(self_vec0, self_vec1) = convert_bfloat16_float(self_vec);
Vectorized<float> t1_vec0, t1_vec1, t2_vec0, t2_vec1;
std::tie(t1_vec0, t1_vec1) = convert_bfloat16_float(t1_vec);
std::tie(t2_vec0, t2_vec1) = convert_bfloat16_float(t2_vec);
self_vec0 = self_vec0 + float_vec * t1_vec0 / t2_vec0;
self_vec1 = self_vec1 + float_vec * t1_vec1 / t2_vec1;
return convert_float_bfloat16(self_vec0, self_vec1);
});
} else {
AT_DISPATCH_ALL_TYPES_AND_COMPLEX(dtype, "addcdiv_cpu_out", [&] {
scalar_t scalar_val = value.to<scalar_t>();
auto scalar_vec = Vectorized<scalar_t>(scalar_val);
cpu_kernel_vec(
iter,
[=](scalar_t self_val, scalar_t t1_val, scalar_t t2_val) -> scalar_t {
return self_val + scalar_val * t1_val / t2_val;
},
[=](Vectorized<scalar_t> self_vec,
Vectorized<scalar_t> t1_vec,
Vectorized<scalar_t> t2_vec) {
return self_vec + scalar_vec * t1_vec / t2_vec;
});
});
}
}
static void smooth_l1_backward_cpu_kernel(TensorIterator& iter, const Scalar& norm, double beta) {
ScalarType dtype = iter.dtype(0);
AT_DISPATCH_ALL_TYPES(dtype, "smooth_l1_backward_cpu_out", [&] {
auto norm_val = norm.to<scalar_t>();
scalar_t beta_val(beta);
auto norm_val_vec = Vectorized<scalar_t>(norm_val);
auto beta_val_vec = Vectorized<scalar_t>(beta_val);
const auto neg_1_vec = Vectorized<scalar_t>(-1);
const auto zero_vec = Vectorized<scalar_t>(0);
const auto pos_1_vec = Vectorized<scalar_t>(1);
cpu_kernel_vec(iter,
[=](scalar_t input, scalar_t target, scalar_t grad_output) -> scalar_t {
const auto x = input - target;
if (x <= -beta)
return -norm_val * grad_output;
else if (x >= beta)
return norm_val * grad_output;
else
return norm_val * x * grad_output / beta;
},
[norm_val_vec, beta_val_vec, neg_1_vec, zero_vec, pos_1_vec](
Vectorized<scalar_t> input, Vectorized<scalar_t> target, Vectorized<scalar_t> grad_output) -> Vectorized<scalar_t> {
// using two blendv calls to simulate the 3 cases
// 1 if x >= beta
// -1 if x <= -beta
// x / beta if |x| < beta
const auto x = input - target;
const auto pos_or_neg_1_vec = Vectorized<scalar_t>::blendv(
neg_1_vec, pos_1_vec, x > zero_vec);
const auto x_abs = x.abs();
const auto output = Vectorized<scalar_t>::blendv(
x / beta_val_vec, pos_or_neg_1_vec, x_abs >= beta_val_vec);
return norm_val_vec * output * grad_output;
}
);
});
}
static void huber_backward_cpu_kernel(TensorIterator& iter, const Scalar& norm, double delta) {
ScalarType dtype = iter.dtype(0);
AT_DISPATCH_FLOATING_TYPES_AND2(kBFloat16, kHalf, dtype, "huber_backward_cpu_out", [&] {
auto norm_val = norm.to<scalar_t>();
scalar_t delta_val(delta);
auto norm_val_vec = Vectorized<scalar_t>(norm_val);
auto delta_val_vec = Vectorized<scalar_t>(delta_val);
const auto neg_1_vec = Vectorized<scalar_t>(-1);
const auto zero_vec = Vectorized<scalar_t>(0);
const auto pos_1_vec = Vectorized<scalar_t>(1);
cpu_kernel_vec(iter,
[=](scalar_t input, scalar_t target, scalar_t grad_output) -> scalar_t {
const auto x = input - target;
if (x <= -delta) {
return -norm_val * grad_output * delta;
} else if (x >= delta) {
return norm_val * grad_output * delta;
} else {
return norm_val * x * grad_output;
}
},
[norm_val_vec, delta_val_vec, neg_1_vec, zero_vec, pos_1_vec](
Vectorized<scalar_t> input, Vectorized<scalar_t> target, Vectorized<scalar_t> grad_output) -> Vectorized<scalar_t> {
// using two blendv calls to simulate the 3 cases
// delta if x >= delta
// -delta if x <= -delta
// x if |x| < delta
const auto x = input - target;
const auto pos_or_neg_1_vec = Vectorized<scalar_t>::blendv(
neg_1_vec, pos_1_vec, x > zero_vec);
const auto x_abs = x.abs();
const auto output = Vectorized<scalar_t>::blendv(
x, pos_or_neg_1_vec * delta_val_vec, x_abs >= delta_val_vec);
return norm_val_vec * output * grad_output;
}
);
});
}
static void mse_backward_cpu_kernel(TensorIterator& iter, const Scalar& value) {
ScalarType dtype = iter.dtype(0);
AT_DISPATCH_ALL_TYPES(dtype, "mse_backward_cpu_out", [&] {
scalar_t scalar_val = value.to<scalar_t>();
auto scalar_vec = Vectorized<scalar_t>(scalar_val);
cpu_kernel_vec(
iter,
[=](scalar_t self_val, scalar_t t1_val, scalar_t t2_val) -> scalar_t {
return scalar_val * (self_val - t1_val) * t2_val;
},
[=](Vectorized<scalar_t> self_vec,
Vectorized<scalar_t> t1_vec,
Vectorized<scalar_t> t2_vec) {
return scalar_vec * (self_vec - t1_vec) * t2_vec;
});
});
}
} // anonymous namespace
REGISTER_DISPATCH(addcmul_stub, &addcmul_cpu_kernel);
REGISTER_DISPATCH(addcdiv_stub, &addcdiv_cpu_kernel);
REGISTER_DISPATCH(smooth_l1_backward_stub, &smooth_l1_backward_cpu_kernel);
REGISTER_DISPATCH(huber_backward_stub, &huber_backward_cpu_kernel);
REGISTER_DISPATCH(mse_backward_stub, &mse_backward_cpu_kernel);
} // namespace native
} // namespace at