Skip to content

Commit 634a206

Browse files
manuelcandalesfacebook-github-bot
manuelcandales
authored and
facebook-github-bot
committed
Calibrate portable ops (#2537)
Summary: Pull Request resolved: #2537 Reviewed By: SS-JIA Differential Revision: D55142429 fbshipit-source-id: 77f1d145d4fe5e45ee90a6803f2f2956179b22cf
1 parent 3270d22 commit 634a206

12 files changed

Lines changed: 135 additions & 91 deletions

‎kernels/portable/cpu/op_any.cpp‎

Lines changed: 31 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ Tensor& any_all_out(RuntimeContext& ctx, const Tensor& in, Tensor& out) {
3232
auto data_out = out.mutable_data_ptr<CTYPE_OUT>();
3333
data_out[0] = static_cast<CTYPE_OUT>(false);
3434
for (auto i = 0; i < in.numel(); ++i) {
35-
if (static_cast<CTYPE_OUT>(data_in[i])) {
35+
if (static_cast<bool>(data_in[i])) {
3636
data_out[0] = static_cast<CTYPE_OUT>(true);
3737
break;
3838
}
@@ -57,11 +57,16 @@ Tensor& any_dims_out(
5757
InvalidArgument,
5858
out);
5959

60-
ET_KERNEL_CHECK(
61-
ctx,
62-
resize_reduction_out(in, dim_list, keepdim, out) == Error::Ok,
63-
InvalidArgument,
64-
out);
60+
if (dim_list.has_value() && dim_list.value().empty()) {
61+
ET_KERNEL_CHECK(
62+
ctx, resize_tensor(out, in.sizes()) == Error::Ok, InvalidArgument, out);
63+
} else {
64+
ET_KERNEL_CHECK(
65+
ctx,
66+
resize_reduction_out(in, dim_list, keepdim, out) == Error::Ok,
67+
InvalidArgument,
68+
out);
69+
}
6570

6671
ScalarType in_type = in.scalar_type();
6772
ScalarType out_type = out.scalar_type();
@@ -70,17 +75,25 @@ Tensor& any_dims_out(
7075
ET_SWITCH_REALHB_TYPES(in_type, ctx, name, CTYPE_IN, [&] {
7176
ET_SWITCH_TWO_TYPES(Bool, Byte, out_type, ctx, name, CTYPE_OUT, [&] {
7277
CTYPE_OUT* out_data = out.mutable_data_ptr<CTYPE_OUT>();
73-
for (size_t out_ix = 0; out_ix < out.numel(); ++out_ix) {
74-
CTYPE_OUT any = false;
75-
if (in.numel() > 0) {
76-
any = map_reduce_over_dim_list<CTYPE_IN, CTYPE_OUT>(
77-
[](CTYPE_IN v) { return static_cast<bool>(v); },
78-
[](CTYPE_OUT outv, CTYPE_OUT acc) { return acc || outv; },
79-
in,
80-
dim_list,
81-
out_ix);
78+
if (dim_list.has_value() && dim_list.value().empty()) {
79+
const CTYPE_IN* in_data = in.const_data_ptr<CTYPE_IN>();
80+
for (size_t out_ix = 0; out_ix < out.numel(); ++out_ix) {
81+
out_data[out_ix] =
82+
static_cast<CTYPE_OUT>(static_cast<bool>(in_data[out_ix]));
83+
}
84+
} else {
85+
for (size_t out_ix = 0; out_ix < out.numel(); ++out_ix) {
86+
bool any = false;
87+
if (in.numel() > 0) {
88+
any = map_reduce_over_dim_list<CTYPE_IN, bool>(
89+
[](CTYPE_IN v) { return static_cast<bool>(v); },
90+
[](bool outv, bool acc) { return acc || outv; },
91+
in,
92+
dim_list,
93+
out_ix);
94+
}
95+
out_data[out_ix] = static_cast<CTYPE_OUT>(any);
8296
}
83-
out_data[out_ix] = any;
8497
}
8598
});
8699
});
@@ -98,7 +111,8 @@ Tensor& any_out(
98111

99112
ET_KERNEL_CHECK(
100113
ctx,
101-
check_reduction_args_single_dim(in, dim, keepdim, {}, out),
114+
check_reduction_args_single_dim(
115+
in, dim, keepdim, {}, out, /*allow_empty_dim*/ true),
102116
InvalidArgument,
103117
out);
104118

‎kernels/portable/cpu/op_cdist_forward.cpp‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,9 +25,23 @@ inline ArrayRef<Tensor::SizesType> get_batch_sizes(const Tensor& tensor) {
2525

2626
template <typename CTYPE, typename Norm>
2727
void cdist(const Tensor& x1, const Tensor& x2, Tensor& out, double p) {
28+
if (out.numel() == 0) {
29+
return;
30+
}
31+
32+
CTYPE* out_data = out.mutable_data_ptr<CTYPE>();
33+
34+
// If the last dimension of x1 (which is equal to the last dimension of x2)
35+
// has size 0, then the output is filled with 0s.
36+
if (x1.numel() == 0) {
37+
for (size_t out_ix = 0; out_ix < out.numel(); ++out_ix) {
38+
out_data[out_ix] = 0;
39+
}
40+
return;
41+
}
42+
2843
const CTYPE* x1_data = x1.const_data_ptr<CTYPE>();
2944
const CTYPE* x2_data = x2.const_data_ptr<CTYPE>();
30-
CTYPE* out_data = out.mutable_data_ptr<CTYPE>();
3145

3246
const ArrayRef<Tensor::SizesType> x1_batch_sizes = get_batch_sizes(x1);
3347
const ArrayRef<Tensor::SizesType> x2_batch_sizes = get_batch_sizes(x2);

‎kernels/portable/cpu/op_diagonal_copy.cpp‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,10 @@ void diagonal_copy_impl(
2323
int64_t dim1,
2424
int64_t dim2,
2525
Tensor& out) {
26+
if (out.numel() == 0) {
27+
return;
28+
}
29+
2630
int64_t storage_offset = 0;
2731
size_t diag_size = out.size(out.dim() - 1);
2832

@@ -89,7 +93,7 @@ Tensor& diagonal_copy_out(
8993

9094
constexpr auto name = "diagonal_copy.out";
9195

92-
ET_SWITCH_REAL_TYPES(in.scalar_type(), ctx, name, CTYPE, [&] {
96+
ET_SWITCH_REALHB_TYPES(in.scalar_type(), ctx, name, CTYPE, [&] {
9397
diagonal_copy_impl<CTYPE>(in, offset, dim1, dim2, out);
9498
});
9599

‎kernels/portable/cpu/op_div.cpp‎

Lines changed: 35 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -163,25 +163,21 @@ Tensor& div_scalar_out(
163163

164164
ET_SWITCH_REAL_TYPES_AND(Bool, a_type, ctx, "div.Scalar_out", CTYPE_A, [&]() {
165165
ET_SWITCH_SCALAR_OBJ_TYPES(b_type, ctx, "div.Scalar_out", CTYPE_B, [&]() {
166-
ET_SWITCH_FLOAT_TYPES(
167-
common_type, ctx, "div.Scalar_out", CTYPE_IN, [&]() {
168-
ET_SWITCH_FLOAT_TYPES(
169-
out_type, ctx, "div.Scalar_out", CTYPE_OUT, [&]() {
170-
CTYPE_B b_val;
171-
utils::extract_scalar(b, &b_val);
172-
CTYPE_IN b_casted = static_cast<CTYPE_IN>(b_val);
173-
174-
apply_unary_map_fn(
175-
[b_casted](const CTYPE_A val_a) {
176-
CTYPE_IN a_casted = static_cast<CTYPE_IN>(val_a);
177-
CTYPE_IN value = a_casted / b_casted;
178-
return static_cast<CTYPE_OUT>(value);
179-
},
180-
a.const_data_ptr<CTYPE_A>(),
181-
out.mutable_data_ptr<CTYPE_OUT>(),
182-
out.numel());
183-
});
184-
});
166+
ET_SWITCH_FLOAT_TYPES(out_type, ctx, "div.Scalar_out", CTYPE, [&]() {
167+
CTYPE_B b_val;
168+
utils::extract_scalar(b, &b_val);
169+
CTYPE b_casted = static_cast<CTYPE>(b_val);
170+
171+
apply_unary_map_fn(
172+
[b_casted](const CTYPE_A val_a) {
173+
CTYPE a_casted = static_cast<CTYPE>(val_a);
174+
CTYPE value = a_casted / b_casted;
175+
return static_cast<CTYPE>(value);
176+
},
177+
a.const_data_ptr<CTYPE_A>(),
178+
out.mutable_data_ptr<CTYPE>(),
179+
out.numel());
180+
});
185181
});
186182
});
187183

@@ -206,7 +202,7 @@ Tensor& div_scalar_mode_out(
206202

207203
ScalarType a_type = a.scalar_type();
208204
ScalarType b_type = utils::get_scalar_dtype(b);
209-
ScalarType common_type = isFloatingType(a_type) ? a_type : ScalarType::Float;
205+
ScalarType common_type = utils::promote_type_with_scalar(a_type, b);
210206
ScalarType out_type = out.scalar_type();
211207

212208
ET_KERNEL_CHECK(ctx, common_type == out_type, InvalidArgument, out);
@@ -215,27 +211,25 @@ Tensor& div_scalar_mode_out(
215211

216212
ET_SWITCH_REALB_TYPES(a_type, ctx, name, CTYPE_A, [&]() {
217213
ET_SWITCH_SCALAR_OBJ_TYPES(b_type, ctx, name, CTYPE_B, [&]() {
218-
ET_SWITCH_FLOAT_TYPES(common_type, ctx, name, CTYPE_IN, [&]() {
219-
ET_SWITCH_FLOAT_TYPES(out_type, ctx, name, CTYPE_OUT, [&]() {
220-
CTYPE_B b_val;
221-
utils::extract_scalar(b, &b_val);
222-
CTYPE_IN b_casted = static_cast<CTYPE_IN>(b_val);
223-
224-
apply_unary_map_fn(
225-
[b_casted, mode](const CTYPE_A val_a) {
226-
CTYPE_IN a_casted = static_cast<CTYPE_IN>(val_a);
227-
CTYPE_IN value = a_casted / b_casted;
228-
if (mode.has_value() && mode.value() == "trunc") {
229-
value = std::trunc(value);
230-
} else if (mode.has_value() && mode.value() == "floor") {
231-
value = utils::floor_divide(a_casted, b_casted);
232-
}
233-
return static_cast<CTYPE_OUT>(value);
234-
},
235-
a.const_data_ptr<CTYPE_A>(),
236-
out.mutable_data_ptr<CTYPE_OUT>(),
237-
out.numel());
238-
});
214+
ET_SWITCH_REAL_TYPES(out_type, ctx, name, CTYPE, [&]() {
215+
CTYPE_B b_val;
216+
utils::extract_scalar(b, &b_val);
217+
CTYPE b_casted = static_cast<CTYPE>(b_val);
218+
219+
apply_unary_map_fn(
220+
[b_casted, mode](const CTYPE_A val_a) {
221+
CTYPE a_casted = static_cast<CTYPE>(val_a);
222+
CTYPE value = a_casted / b_casted;
223+
if (mode.has_value() && mode.value() == "trunc") {
224+
value = std::trunc(value);
225+
} else if (mode.has_value() && mode.value() == "floor") {
226+
value = utils::floor_divide(a_casted, b_casted);
227+
}
228+
return value;
229+
},
230+
a.const_data_ptr<CTYPE_A>(),
231+
out.mutable_data_ptr<CTYPE>(),
232+
out.numel());
239233
});
240234
});
241235
});

‎kernels/portable/cpu/op_flip.cpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -52,15 +52,15 @@ flip_out(RuntimeContext& ctx, const Tensor& in, IntArrayRef dims, Tensor& out) {
5252
flip_dim_data[i] = false;
5353
}
5454
for (size_t i = 0; i < dims.size(); i++) {
55-
const auto d = dims[i] < 0 ? dims[i] + in.dim() : dims[i];
55+
const auto d = dims[i] < 0 ? dims[i] + nonzero_dim(in) : dims[i];
5656
flip_dim_data[d] = true;
5757
}
5858
size_t flip_dim_length = static_cast<size_t>(in.dim()); // NOLINT
5959
ArrayRef<bool> flip_dim(flip_dim_data, flip_dim_length);
6060

6161
constexpr auto name = "flip.out";
6262

63-
ET_SWITCH_REAL_TYPES(in.scalar_type(), ctx, name, CTYPE, [&] {
63+
ET_SWITCH_REALHB_TYPES(in.scalar_type(), ctx, name, CTYPE, [&] {
6464
const CTYPE* in_data = in.const_data_ptr<CTYPE>();
6565
CTYPE* out_data = out.mutable_data_ptr<CTYPE>();
6666

‎kernels/portable/cpu/op_prod.cpp‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,8 @@ Tensor& prod_int_out(
5858

5959
ET_KERNEL_CHECK(
6060
ctx,
61-
check_reduction_args_single_dim(in, dim, keepdim, dtype, out),
61+
check_reduction_args_single_dim(
62+
in, dim, keepdim, dtype, out, /*allow_empty_dim=*/true),
6263
InvalidArgument,
6364
out);
6465

‎kernels/portable/cpu/op_roll.cpp‎

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,9 @@ bool check_roll_args(
1919
IntArrayRef shifts,
2020
IntArrayRef dims,
2121
const Tensor& out) {
22-
for (const auto& d : dims) {
23-
if (in.dim() == 0) {
24-
ET_LOG_AND_RETURN_IF_FALSE(d == 0 || d == -1);
25-
} else {
22+
ET_LOG_AND_RETURN_IF_FALSE(tensor_has_rank_greater_or_equal_to(in, 1));
23+
if (in.numel() > 0) {
24+
for (const auto& d : dims) {
2625
ET_LOG_AND_RETURN_IF_FALSE(dim_is_valid(d, in.dim()));
2726
}
2827
}
@@ -38,7 +37,8 @@ size_t unshift_flat_ix(size_t ix, const Tensor& in, IntArrayRef dim_shifts) {
3837

3938
size_t shifted_coord[kTensorDimensionLimit];
4039
for (size_t d = 0; d < in.dim(); d++) {
41-
shifted_coord[d] = (ix_coord[d] - dim_shifts[d]) % in.size(d);
40+
shifted_coord[d] =
41+
(ix_coord[d] + in.size(d) - dim_shifts[d] % in.size(d)) % in.size(d);
4242
}
4343

4444
return coordinateToIndex(in, shifted_coord);
@@ -60,7 +60,9 @@ Tensor& roll_out(
6060
ET_KERNEL_CHECK(
6161
ctx, check_roll_args(in, shifts, dims, out), InvalidArgument, out);
6262

63-
constexpr auto name = "roll.out";
63+
if (in.numel() == 0) {
64+
return out;
65+
}
6466

6567
int64_t dim_shift_array[kTensorDimensionLimit];
6668
for (size_t i = 0; i < in.dim(); i++) {
@@ -74,7 +76,9 @@ Tensor& roll_out(
7476
size_t dim_shift_array_length = static_cast<size_t>(in.dim()); // NOLINT
7577
IntArrayRef dim_shifts(dim_shift_array, dim_shift_array_length);
7678

77-
ET_SWITCH_REAL_TYPES(in.scalar_type(), ctx, name, CTYPE, [&] {
79+
constexpr auto name = "roll.out";
80+
81+
ET_SWITCH_REALHB_TYPES(in.scalar_type(), ctx, name, CTYPE, [&] {
7882
const CTYPE* in_data = in.const_data_ptr<CTYPE>();
7983
CTYPE* out_data = out.mutable_data_ptr<CTYPE>();
8084

‎kernels/portable/cpu/op_var.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ void compute_variance(
2626
const size_t num,
2727
const double denominator) {
2828
CTYPE_OUT* out_data = out.mutable_data_ptr<CTYPE_OUT>();
29-
if (num == 0 || denominator == 0) {
29+
if (num == 0 || denominator <= 0) {
3030
for (size_t out_ix = 0; out_ix < out.numel(); ++out_ix) {
3131
out_data[out_ix] = NAN;
3232
}

‎kernels/portable/cpu/util/copy_ops_util.cpp‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -886,11 +886,18 @@ void get_diagonal_copy_out_target_size(
886886

887887
size_t diagonal_size = 0;
888888
if (offset >= 0) {
889-
diagonal_size = std::min<size_t>(in.size(dim1), in.size(dim2) - offset);
889+
if (in.size(dim2) <= offset) {
890+
diagonal_size = 0;
891+
} else {
892+
diagonal_size = std::min<size_t>(in.size(dim1), in.size(dim2) - offset);
893+
}
890894
} else {
891-
diagonal_size = std::min<size_t>(in.size(dim1) + offset, in.size(dim2));
895+
if (in.size(dim1) <= -offset) {
896+
diagonal_size = 0;
897+
} else {
898+
diagonal_size = std::min<size_t>(in.size(dim1) + offset, in.size(dim2));
899+
}
892900
}
893-
diagonal_size = std::max<size_t>(diagonal_size, 0);
894901

895902
size_t shift = 0;
896903
for (size_t d = 0; d < in.dim(); ++d) {

‎kernels/portable/cpu/util/reduce_util.cpp‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -336,8 +336,6 @@ bool check_reduction_args(
336336
ET_LOG_AND_RETURN_IF_FALSE(dtype.value() == out.scalar_type());
337337
}
338338
ET_LOG_AND_RETURN_IF_FALSE(check_dim_list_is_valid(in, dim_list));
339-
ET_LOG_AND_RETURN_IF_FALSE(
340-
out.dim() == compute_reduced_out_dim(in, dim_list, keepdim));
341339
ET_LOG_AND_RETURN_IF_FALSE(tensor_is_default_or_channels_last_dim_order(in));
342340
ET_LOG_AND_RETURN_IF_FALSE(tensor_is_default_or_channels_last_dim_order(out));
343341

@@ -353,7 +351,8 @@ bool check_reduction_args_single_dim(
353351
optional<int64_t> dim,
354352
bool keepdim,
355353
optional<ScalarType> dtype,
356-
Tensor& out) {
354+
Tensor& out,
355+
bool allow_empty_dim) {
357356
if (dtype.has_value()) {
358357
ET_LOG_AND_RETURN_IF_FALSE(dtype.value() == out.scalar_type());
359358
}
@@ -366,11 +365,11 @@ bool check_reduction_args_single_dim(
366365

367366
if (dim.has_value()) {
368367
ET_LOG_AND_RETURN_IF_FALSE(dim_is_valid(dim.value(), in.dim()));
369-
ET_LOG_AND_RETURN_IF_FALSE(tensor_has_non_empty_dim(in, dim.value()));
368+
if (!allow_empty_dim) {
369+
ET_LOG_AND_RETURN_IF_FALSE(tensor_has_non_empty_dim(in, dim.value()));
370+
}
370371
}
371372

372-
ET_LOG_AND_RETURN_IF_FALSE(
373-
out.dim() == compute_reduced_out_dim(in, dim, keepdim));
374373
ET_LOG_AND_RETURN_IF_FALSE(tensor_is_default_or_channels_last_dim_order(in));
375374
ET_LOG_AND_RETURN_IF_FALSE(tensor_is_default_or_channels_last_dim_order(out));
376375

0 commit comments

Comments
 (0)