Parallelize portable ops if threadpool is available, with fallback to parallel_for-as-for-loop

[swolchok]

🚀 The feature, motivation and pitch

It seems suboptimal to me that we have to create separate optimized ops just to get basic stuff like parallelization (and vectorization, but let's start with parallelization). Here's what I'd like to do: (The timeline here is "ASAP", but I'm opening an issue because this got too long for chat and so that I can point to this issue on the PRs.)

1. Set up a proper CMake build for extension/parallel; right now it's free-riding on buck and getting automatically duplicated into 3 different targets per the generated executorch_srcs.cmake. (done; Add proper CMake build for extension_parallel #8938)
2. Make extension_threadpool itself export the -DET_USE_THREADPOOL macro we already use and define somewhat ad-hoc. (done; Properly export ET_USE_THREADPOOL from the threadpool extension #8947)
3. move extension/parallel/thread_parallel.h to core. (@larryliu0820 suggests runtime/kernel/thread_parallel.h) (Yes I will leave a stub header behind for backward compatibility.) Move thread_parallel.cpp to threadpool, since there will be no reason not to provide it when threads are available. Provide a default implementation of parallel_for if threadpool is not built (gated behind ET_USE_THREADPOOL) that is just an inlinable for loop. (Split & remove extension_parallel #8983)
4. use parallel_for in at least one portable op, either directly or via the workhorse "util" functions. (Add basic parallel_for support to reduce_util #8986)
5. Verify that, because the optimized library is built with threadpool, it gets parallelization. Adjust build configuration for optimized ops lib if necessary. (Build optimized_portable_kernels if threadpool is enabled #8987)
6. Roll out parallel_for across portable ops and workhorse "util" functions.

Thoughts? Blockers?

Alternatives

status quo -- slow portable ops

Additional context

No response

RFC (Optional)

No response

cc @larryliu0820 @manuelcandales

[kimishpatel]

One of the goals on portable op was readability as well. SO it would be good to at the least touch upon that.

move extension/parallel/thread_parallel.h to core

probably be better find a different location because it will start pulling in pthreadpool and other dependencies when you actually want threadpool. Keeping in extension would mean portable ops will depend on something outside of the core, but within the core will complicate the build system?

[swolchok]

Why would the header pull in threadpool? I didn't say I would move the .cpp file.

[digantdesai]

I am curious, since we have a dedicated CPU delegate for cases when we need CPU perf, why enhancing that to deal with perf issues is not an option, as opposed to pushing perf through portable or even optimized ops? We can also explore custom op interface if upstreaming to XNNPACK and XNNPACK rebase is too much effort.

[swolchok]

readability

parallel_for isn't much worse than a for loop

[swolchok]

since we have a dedicated CPU delegate for cases when we need CPU perf, why enhancing that to deal with perf issues is not an option, as opposed to pushing perf through portable or even optimized ops

It is much better if our ops are reasonably fast by default, rather than only performing well for specific cases we have looked at and done work on.

[swolchok]

XNNPACK

forgot to mention that it is also better if we don't have to write ops multiple times (once for portable and once for optimized and/or XNNPACK).

[mergennachin]

Do you have a prototype of what a portable op would look like with this? Maybe anchor with an example

[swolchok]

Do you have a prototype of what a portable op would look like with this?

Sure, I expect

executorch/kernels/portable/cpu/op_argmin.cpp

Line 50 in d897e73

for (const auto out_ix : c10::irange(out.numel())) {

to be replaced with parallel_for(0, out.numel(), internal::GRAIN_SIZE, [&]() {.
Relatedly, I expect

executorch/kernels/portable/cpu/util/reduce_util.h

Line 27 in d897e73

for (size_t i = start; i <= end; i++) {

and

executorch/kernels/portable/cpu/util/reduce_util.h

Line 39 in d897e73

for (size_t i = start; i <= end; i++) {

will read parallel_for(start, end, internal::GRAIN_SIZE, [&]() {.

[mergennachin]

That's looks ugly... A simple for loop now has three new concepts, namely (1) a new parallel_for util function (2) grain_size (3) lambda stuff

Can we do completely inline and using macros?

for (size_t i = start; i < end; i ++) {
   ...
}

could be replaced by directly

FOR(a, start, end) {
   ...
} END_FOR

and define #FOR and #END_FOR macros and switch based on a flag, and hide the grain size and lambdas etc underneath?

[swolchok]

Can we do completely inline and using macros?

That's not better. Now you have 1) macros, which are inherently a problem 2) a new FOR concept 3) a lambda that you can't see (overlaps with point (1)). The one thing we can easily get rid of is the grain size, by simply adding an overload of parallel_for that defaults it.

[swolchok]

#9197 is the top of a stack of rollouts for reductions. After that, need to roll out across other util functions (most notably elementwise_util) before closing.

[swolchok]

Looking through the other util functions to finish rolling out. Questions:

• should we parallelize apply_kernel_2d_reduce_then_map_fn? @manuelcandales
• ditto padding_util

Notes:

• parallelizing memcpy seems unlikely to be very effective on mobile/embedded devices, skipping repeat_util.cpp and select_copy_util.cpp and slice_util.cpp and transpose_util.h

[digantdesai]

it is also better if we don't have to write ops multiple times (once for portable and once for optimized and/or XNNPACK).

Portable ops with threads may be necessary (default out-of-box is not too bad) but not really sufficient if we care about performance. So, while we do this, it doesn't practically save us from implementing arm and/or x86 specific (which we mainly care about) SIMD + Multi-threaded ops in either optimized or XNNPACK.

That said, I agree that it may reduce pressure from us providing an optimized impl, and not having "guaranteed bad" perf when using portable. Do we have perf uplift data from this? Curious to see how far this + autovec + out-of-order CPUs can get.

[swolchok]

SIMD + Multi-threaded

we are also going to vectorize, at least for elementwise ops. #9241

perf uplift data

nothing particularly concrete, but I can vouch that it goes faster.