onnxruntime
e3655b35 - Parallelize CPU ScatterElements kernel via ThreadPool (#28588)

Commit
67 days ago
Parallelize CPU ScatterElements kernel via ThreadPool (#28588) ### Description Parallelizes both `GetIndices` and `ScatterData` in the CPU `ScatterElements` implementation using `ThreadPool::TryParallelFor`. **Key insight**: For ScatterElements with `axis=a`, work units identified by coordinates orthogonal to the axis (`outer_size × inner_size`) are guaranteed to write to disjoint output elements—even with reductions. This enables lock-free parallelization without correctness concerns. Changes: - **`GetIndices`**: Index validation/normalization parallelized over the flat index array - **`ScatterData`**: Rewritten to decompose into `outer_size * inner_size` independent work units, each processing `axis_size` sequential scatter operations along the axis dimension - Thread pool plumbed through `ScatterDataDispatchTarget` from `OpKernelContext::GetOperatorThreadPool()` - Training `GatherElementsGradImpl` passes `nullptr` (sequential fallback preserved) For the reported workload (`axis=0`, indices shape `[481385, 80]`): 80 independent parallel streams, each processing 481385 elements—well-suited for multi-core execution. ### Motivation and Context The CPU `ScatterElements` kernel was entirely sequential—single-threaded index conversion followed by single-threaded scatter—yielding ~761ms on a 24-core ARM system for a workload that an optimized parallel implementation handles in ~6ms (129× gap). The kernel showed zero intra-op thread utilization in ORT profiling. --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: tianleiwu <30328909+tianleiwu@users.noreply.github.com> Co-authored-by: Tianlei Wu <tlwu@microsoft.com>
Author
Parents
Loading