Skip to content

Commit 9cb7cf6

Browse files
swolchokswolchok
committed
Reapply "Add vectorized_math.h (#11204)", "Add optimized_portable_kernels test (#11205)", and "Add vectorization in elementwise_util (#9432)"
Stack was reverted due to internal CI failures. Reapplying as an exported internal diff so that we make sure to catch any more of those. New fixes: - straightforward op_sub build fixes - s/EXPECT_EQ/EXPECT_FLOAT_EQ/ in vectorized_math_test - define ET_USE_PYTORCH_HEADERS to detect whether exceptions are enabled, and use `#if` instead of `#ifdef` to check the macro so that we don't use PyTorch headers if exceptions are disabled. (otherwise, we might have problems with e.g. TORCH_CHECK) Original summary for #11204: Set of math functions that work on both scalars and at::vec::Vectorized, to be used in #9432. Original summary for #11205: Make sure we test the optimized versions of portable kernels even if they are shadowed by optimized implementations. Intended to support #9432. Original summary for #9432: This is a first cut at #9241 . In this PR I've vectorized a small initial set of ops: atan2, clamp, fmod_Scalar, maximum, minimum, mul, pow, and sigmoid. In addition, the following ops should have gotten vectorized automatically because they already used generic lambdas: add, div, rsub, sub. I've left covering ops that use the `unary_ufunc_*` utilities in [pattern.h](https://github.com/pytorch/executorch/blob/main/kernels/portable/cpu/pattern/pattern.h) for a follow-up push, because pattern.h and elementwise_util need some work before we can migrate pattern.h's utilities to be backed by elementwise_util. This PR adds an interesting testing problem: in theory, *all* operators might need test cases long enough to tickle vectorization, because we might accidentally vectorize ops unexpectedly and break their lambdas due to anticipated differences in semantics. I address this issue by using Vectorized for the scalar prologue/epilogue in debug mode (we run tests in both debug and release) so that we can detect broken lambdas. I additionally intentionally introduced a bug in the vectorized path in elementwise_util and manually verified that we saw test failures for each vectorized op called out above. Differential Revision: [D76467389](https://our.internmc.facebook.com/intern/diff/D76467389/) ghstack-source-id: 289985405 Pull Request resolved: #11604
1 parent 0e671b0 commit 9cb7cf6

35 files changed

+694
-74
lines changed

‎.lintrunner.toml

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -271,6 +271,10 @@ exclude_patterns = [
271271
'examples/**',
272272
'exir/verification/bindings.cpp',
273273
'extension/**',
274+
# Uses properly-gated (ET_USE_PYTORCH_HEADERS) ATen include.
275+
'kernels/portable/cpu/util/elementwise_util.h',
276+
'kernels/portable/cpu/util/math_util.h',
277+
'kernels/portable/cpu/util/vectorized_math.h',
274278
'kernels/optimized/**',
275279
'runtime/core/exec_aten/**',
276280
# Want to be able to keep c10 in sync with PyTorch core.

‎kernels/optimized/CMakeLists.txt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ message("Generated files ${gen_command_sources}")
6060
list(TRANSFORM _optimized_kernels__srcs PREPEND "${EXECUTORCH_ROOT}/")
6161
add_library(optimized_kernels ${_optimized_kernels__srcs})
6262
target_include_directories(optimized_kernels PRIVATE ${TORCH_INCLUDE_DIRS} "${EXECUTORCH_ROOT}/third-party/pocketfft")
63-
target_compile_definitions(optimized_kernels PRIVATE ET_USE_PYTORCH_HEADERS)
63+
target_compile_definitions(optimized_kernels PRIVATE "ET_USE_PYTORCH_HEADERS=ET_HAS_EXCEPTIONS")
6464
target_link_libraries(
6565
optimized_kernels PUBLIC executorch_core cpublas extension_threadpool kernels_util_all_deps
6666
)

‎kernels/portable/CMakeLists.txt

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -68,9 +68,16 @@ if(EXECUTORCH_BUILD_PTHREADPOOL AND EXECUTORCH_BUILD_KERNELS_OPTIMIZED)
6868
target_link_libraries(optimized_portable_kernels PUBLIC extension_threadpool)
6969
target_compile_options(optimized_portable_kernels PUBLIC ${_common_compile_options})
7070
target_include_directories(optimized_portable_kernels PRIVATE ${TORCH_INCLUDE_DIRS})
71-
target_compile_definitions(optimized_portable_kernels PRIVATE ET_USE_PYTORCH_HEADERS)
71+
target_compile_definitions(optimized_portable_kernels PRIVATE "ET_USE_PYTORCH_HEADERS=ET_HAS_EXCEPTIONS")
72+
gen_selected_ops(LIB_NAME "optimized_portable_ops_lib" OPS_SCHEMA_YAML "${_yaml}")
73+
generate_bindings_for_kernels(
74+
LIB_NAME "optimized_portable_ops_lib" FUNCTIONS_YAML "${_yaml}"
75+
)
76+
gen_operators_lib(
77+
LIB_NAME "optimized_portable_ops_lib" KERNEL_LIBS optimized_portable_kernels DEPS executorch_core
78+
)
7279
install(
73-
TARGETS optimized_portable_kernels
80+
TARGETS optimized_portable_kernels optimized_portable_ops_lib
7481
DESTINATION lib
7582
)
7683
endif()

‎kernels/portable/cpu/op_add.cpp

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -102,14 +102,18 @@ Tensor& add_scalar_out(
102102
static constexpr const char op_name[] = "add.Scalar_out";
103103

104104
ET_SWITCH_REALB_TYPES(compute_type, ctx, op_name, CTYPE_COMPUTE, [&]() {
105+
CTYPE_COMPUTE val_b = utils::scalar_to<CTYPE_COMPUTE>(b);
106+
CTYPE_COMPUTE val_alpha = utils::scalar_to<CTYPE_COMPUTE>(alpha);
107+
auto val_alpha_times_b = val_alpha * val_b;
105108
utils::apply_unitensor_elementwise_fn<
106109
CTYPE_COMPUTE,
107110
op_name,
108111
utils::SupportedTensorDtypes::SAME_AS_COMMON>(
109-
[b, alpha](const auto val_a) {
110-
CTYPE_COMPUTE val_b = utils::scalar_to<CTYPE_COMPUTE>(b);
111-
CTYPE_COMPUTE val_alpha = utils::scalar_to<CTYPE_COMPUTE>(alpha);
112-
return val_a + val_alpha * val_b;
112+
[val_alpha_times_b](const auto val_a) {
113+
// Cast here supports vectorization; either it does nothing
114+
// or it casts from CTYPE_COMPUTE to
115+
// Vectorized<CTYPE_COMPUTE>.
116+
return val_a + decltype(val_a)(val_alpha_times_b);
113117
},
114118
ctx,
115119
a,

‎kernels/portable/cpu/op_atan2.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ Tensor& atan2_out(
6060
op_name,
6161
utils::SupportedTensorDtypes::FLOATHBF16>(
6262
[](const auto val_a, const auto val_b) {
63-
return std::atan2(val_a, val_b);
63+
return executorch::math::atan2(val_a, val_b);
6464
},
6565
ctx,
6666
a,

‎kernels/portable/cpu/op_clamp.cpp

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -138,9 +138,8 @@ Tensor& clamp_out(
138138
CTYPE_COMPUTE,
139139
op_name,
140140
utils::SupportedTensorDtypes::SAME_AS_COMMON>(
141-
[has_min, min_opt, has_max, max_opt](const CTYPE_COMPUTE val_in) {
142-
// TODO: rewrite this to be vectorization-capable.
143-
CTYPE_COMPUTE val_out = val_in;
141+
[has_min, min_opt, has_max, max_opt](const auto val_in) {
142+
auto val_out = val_in;
144143
if (has_min) {
145144
val_out = utils::max_override(
146145
val_out, utils::scalar_to<CTYPE_COMPUTE>(min_opt.value()));

‎kernels/portable/cpu/op_elu.cpp

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -48,8 +48,7 @@ Tensor& elu_out(
4848
CTYPE,
4949
op_name,
5050
utils::SupportedTensorDtypes::SAME_AS_COMMON>(
51-
[negcoef, math_scale, math_input_scale](const auto x) {
52-
// TODO: rewrite this to be vectorization-capable.
51+
[negcoef, math_scale, math_input_scale](const CTYPE x) {
5352
return MathT(x) <= MathT(0)
5453
? std::expm1(MathT(x) * math_input_scale) * negcoef
5554
: MathT(x) * math_scale;

‎kernels/portable/cpu/op_fmod.cpp

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@ Tensor& fmod_Tensor_out(
6161
utils::SupportedTensorDtypes::REALHBF16>(
6262
[&div_by_zero_error](
6363
const CTYPE_COMPUTE val_a, const CTYPE_COMPUTE val_b) {
64-
// TODO: rewrite this to be vectorization-capable.
64+
// TODO: rewrite this to be vectorization-capable?
6565
CTYPE_COMPUTE value = 0;
6666
if (is_integral_type<CTYPE_COMPUTE, /*includeBool=*/true>::value) {
6767
if (val_b == 0) {
@@ -138,10 +138,8 @@ Tensor& fmod_Scalar_out(
138138
CTYPE_COMPUTE,
139139
op_name,
140140
utils::SupportedTensorDtypes::REALHBF16>(
141-
[val_b](const CTYPE_COMPUTE val_a) {
142-
// TODO: rewrite this to be vectorization-capable.
143-
CTYPE_COMPUTE value = std::fmod(val_a, val_b);
144-
return value;
141+
[val_b](const auto val_a) {
142+
return executorch::math::fmod(val_a, (decltype(val_a))val_b);
145143
},
146144
ctx,
147145
a,

‎kernels/portable/cpu/op_maximum.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ Tensor& maximum_out(
4949
CTYPE_COMPUTE,
5050
op_name,
5151
utils::SupportedTensorDtypes::REALHBBF16>(
52-
[](const CTYPE_COMPUTE val_a, const CTYPE_COMPUTE val_b) {
52+
[](const auto val_a, const auto val_b) {
5353
return utils::max_override(val_a, val_b);
5454
},
5555
ctx,

‎kernels/portable/cpu/op_minimum.cpp

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,8 +49,7 @@ Tensor& minimum_out(
4949
CTYPE_COMPUTE,
5050
op_name,
5151
utils::SupportedTensorDtypes::REALHBBF16>(
52-
[](const CTYPE_COMPUTE val_a, const CTYPE_COMPUTE val_b) {
53-
// TODO: rewrite this to be vectorization-capable.
52+
[](const auto val_a, const auto val_b) {
5453
return utils::min_override(val_a, val_b);
5554
},
5655
ctx,

0 commit comments

Comments
 (0)