SIGN IN SIGN UP

Proof-of-concept: parallelize argmin (#9066)

I attempted to port `at::parallel_reduce` to ExecuTorch and use that
in reduce_util.h, but it turned out to be much trickier than expected.

(In brief: parallel reduction requires two steps: 1) split the input
range into chunks and reduce over them (easily done like
parallel_for), and then 2) combine the sub-results from chunks. The
reduction function accepted by reduce_over_dim is not well-suited to
step (2).)

Instead, I ported the parallelization strategy used by
binary_kernel_reduce_lastdim: just parallelize over the *non*-reduced
dimensions of the tensor. I don't understand why this strategy isn't
generally applicable and we aren't otherwise capable of parallelizing
reductions, so I haven't gated it to the case where we are reducing
over a contiguous last dimension.

I will send a follow-up that packages up this strategy nicely and uses
it in our reduction portable ops.
S
Scott Wolchok committed
243c8856348e18af5e3cf24016f5dd8a87dc0741
Parent: 907e97e
Committed by GitHub <noreply@github.com> on 3/12/2025, 11:17:21 PM