Skip to content

Commit

Permalink
Fix typo in sub & clean up
Browse files Browse the repository at this point in the history 
Differential Revision: D56255838
  • Loading branch information
@manuelcandales @facebook-github-bot
manuelcandales authored and facebook-github-bot committed Apr 17, 2024
1 parent ebde8e1 commit d9c3d38
Show file tree
Hide file tree
Showing 2 changed files with 62 additions and 60 deletions.
56 changes: 29 additions & 27 deletions kernels/portable/cpu/op_add.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,10 +41,12 @@ Tensor& add_out(
ET_KERNEL_CHECK(
ctx, check_alpha_type(alpha_type, common_type), InvalidArgument, out);

ET_SWITCH_REALHB_TYPES(a_type, ctx, "add.out", CTYPE_A, [&]() {
ET_SWITCH_REALHB_TYPES(b_type, ctx, "add.out", CTYPE_B, [&]() {
ET_SWITCH_REALB_TYPES(common_type, ctx, "add.out", CTYPE_IN, [&]() {
ET_SWITCH_REALHB_TYPES(out_type, ctx, "add.out", CTYPE_OUT, [&]() {
constexpr auto name = "add.out";

ET_SWITCH_REALHB_TYPES(a_type, ctx, name, CTYPE_A, [&]() {
ET_SWITCH_REALHB_TYPES(b_type, ctx, name, CTYPE_B, [&]() {
ET_SWITCH_REALB_TYPES(common_type, ctx, name, CTYPE_IN, [&]() {
ET_SWITCH_REALHB_TYPES(out_type, ctx, name, CTYPE_OUT, [&]() {
CTYPE_IN alpha_val;
utils::extract_scalar(alpha, &alpha_val);

Expand Down Expand Up @@ -99,29 +101,29 @@ Tensor& add_scalar_out(
common_type = ScalarType::Float;
}

ET_SWITCH_REALHB_TYPES(a_type, ctx, "add.Scalar_out", CTYPE_A, [&]() {
ET_SWITCH_SCALAR_OBJ_TYPES(b_type, ctx, "add.Scalar_out", CTYPE_B, [&]() {
ET_SWITCH_REALB_TYPES(
common_type, ctx, "add.Scalar_out", CTYPE_IN, [&]() {
ET_SWITCH_REALHB_TYPES(
out_type, ctx, "add.Scalar_out", CTYPE_OUT, [&]() {
CTYPE_B b_val;
utils::extract_scalar(b, &b_val);
CTYPE_IN b_casted = static_cast<CTYPE_IN>(b_val);
CTYPE_IN alpha_val;
utils::extract_scalar(alpha, &alpha_val);

apply_unary_map_fn(
[b_casted, alpha_val](const CTYPE_A val_a) {
CTYPE_IN a_casted = static_cast<CTYPE_IN>(val_a);
CTYPE_IN value = a_casted + alpha_val * b_casted;
return static_cast<CTYPE_OUT>(value);
},
a.const_data_ptr<CTYPE_A>(),
out.mutable_data_ptr<CTYPE_OUT>(),
out.numel());
});
});
constexpr auto name = "add.Scalar_out";

ET_SWITCH_REALHB_TYPES(a_type, ctx, name, CTYPE_A, [&]() {
ET_SWITCH_SCALAR_OBJ_TYPES(b_type, ctx, name, CTYPE_B, [&]() {
ET_SWITCH_REALB_TYPES(common_type, ctx, name, CTYPE_IN, [&]() {
ET_SWITCH_REALHB_TYPES(out_type, ctx, name, CTYPE_OUT, [&]() {
CTYPE_B b_val;
utils::extract_scalar(b, &b_val);
CTYPE_IN b_casted = static_cast<CTYPE_IN>(b_val);
CTYPE_IN alpha_val;
utils::extract_scalar(alpha, &alpha_val);

apply_unary_map_fn(
[b_casted, alpha_val](const CTYPE_A val_a) {
CTYPE_IN a_casted = static_cast<CTYPE_IN>(val_a);
CTYPE_IN value = a_casted + alpha_val * b_casted;
return static_cast<CTYPE_OUT>(value);
},
a.const_data_ptr<CTYPE_A>(),
out.mutable_data_ptr<CTYPE_OUT>(),
out.numel());
});
});
});
});

Expand Down
66 changes: 33 additions & 33 deletions kernels/portable/cpu/op_sub.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,23 +29,24 @@ Tensor& sub_out(
InvalidArgument,
out);

ET_KERNEL_CHECK(ctx, tensor_is_realhb_type(out), InvalidArgument, out);
ET_KERNEL_CHECK(ctx, tensor_is_realh_type(out), InvalidArgument, out);

ScalarType a_type = a.scalar_type();
ScalarType b_type = b.scalar_type();
ScalarType alpha_type = utils::get_scalar_dtype(alpha);
ScalarType common_type = promoteTypes(a_type, b_type, /*half_to_float*/ true);
ScalarType out_type = out.scalar_type();

ET_KERNEL_CHECK(ctx, canCast(common_type, out_type), InvalidArgument, out);
ET_KERNEL_CHECK(
ctx, check_alpha_type(alpha_type, common_type), InvalidArgument, out);
ET_KERNEL_CHECK(ctx, canCast(common_type, out_type), InvalidArgument, out);
ET_KERNEL_CHECK(ctx, tensor_is_realh_type(out), InvalidArgument, out);

ET_SWITCH_REALH_TYPES(a_type, ctx, "sub.out", CTYPE_A, [&]() {
ET_SWITCH_REALH_TYPES(b_type, ctx, "sub.out", CTYPE_B, [&]() {
ET_SWITCH_REAL_TYPES(common_type, ctx, "sub.out", CTYPE_IN, [&]() {
ET_SWITCH_REALH_TYPES(out_type, ctx, "sub.out", CTYPE_OUT, [&]() {
constexpr auto name = "sub.out";

ET_SWITCH_REALH_TYPES(a_type, ctx, name, CTYPE_A, [&]() {
ET_SWITCH_REALH_TYPES(b_type, ctx, name, CTYPE_B, [&]() {
ET_SWITCH_REAL_TYPES(common_type, ctx, name, CTYPE_IN, [&]() {
ET_SWITCH_REALH_TYPES(out_type, ctx, name, CTYPE_OUT, [&]() {
CTYPE_IN alpha_val;
utils::extract_scalar(alpha, &alpha_val);

Expand Down Expand Up @@ -84,11 +85,11 @@ Tensor& sub_scalar_out(
out,
"Failed to resize output tensor.");

ET_KERNEL_CHECK(ctx, tensor_is_realhb_type(out), InvalidArgument, out);
ET_KERNEL_CHECK(ctx, tensor_is_realh_type(out), InvalidArgument, out);

ScalarType a_type = a.scalar_type();
ScalarType b_type = utils::get_scalar_dtype(b);
ScalarType alpha_type = utils::get_scalar_dtype(b);
ScalarType alpha_type = utils::get_scalar_dtype(alpha);
ScalarType common_type =
utils::promote_type_with_scalar(a_type, b, /*half_to_float*/ false);
ScalarType out_type = out.scalar_type();
Expand All @@ -100,31 +101,30 @@ Tensor& sub_scalar_out(
common_type = ScalarType::Float;
}

ET_SWITCH_REALH_TYPES(a_type, ctx, "sub.Scalar_out", CTYPE_A, [&]() {
ET_SWITCH_SCALAR_OBJ_REAL_TYPES(
b_type, ctx, "sub.Scalar_out", CTYPE_B, [&]() {
ET_SWITCH_REAL_TYPES(
common_type, ctx, "sub.Scalar_out", CTYPE_IN, [&]() {
ET_SWITCH_REALH_TYPES(
out_type, ctx, "sub.Scalar_out", CTYPE_OUT, [&]() {
CTYPE_B b_val;
utils::extract_scalar(b, &b_val);
CTYPE_IN b_casted = static_cast<CTYPE_IN>(b_val);
CTYPE_IN alpha_val;
utils::extract_scalar(alpha, &alpha_val);

apply_unary_map_fn(
[b_casted, alpha_val](const CTYPE_A val_a) {
CTYPE_IN a_casted = static_cast<CTYPE_IN>(val_a);
CTYPE_IN value = a_casted - alpha_val * b_casted;
return static_cast<CTYPE_OUT>(value);
},
a.const_data_ptr<CTYPE_A>(),
out.mutable_data_ptr<CTYPE_OUT>(),
out.numel());
});
});
constexpr auto name = "sub.Scalar_out";

ET_SWITCH_REALH_TYPES(a_type, ctx, name, CTYPE_A, [&]() {
ET_SWITCH_SCALAR_OBJ_REAL_TYPES(b_type, ctx, name, CTYPE_B, [&]() {
ET_SWITCH_REAL_TYPES(common_type, ctx, name, CTYPE_IN, [&]() {
ET_SWITCH_REALH_TYPES(out_type, ctx, name, CTYPE_OUT, [&]() {
CTYPE_B b_val;
utils::extract_scalar(b, &b_val);
CTYPE_IN b_casted = static_cast<CTYPE_IN>(b_val);
CTYPE_IN alpha_val;
utils::extract_scalar(alpha, &alpha_val);

apply_unary_map_fn(
[b_casted, alpha_val](const CTYPE_A val_a) {
CTYPE_IN a_casted = static_cast<CTYPE_IN>(val_a);
CTYPE_IN value = a_casted - alpha_val * b_casted;
return static_cast<CTYPE_OUT>(value);
},
a.const_data_ptr<CTYPE_A>(),
out.mutable_data_ptr<CTYPE_OUT>(),
out.numel());
});
});
});
});

return out;
Expand Down

0 comments on commit d9c3d38

Please sign in to comment.