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