Optimizations for index_select_scalar_cumsum_kernel

[amd-wsung102]

Optimization Changes

• Vectorized gather/store on ROCm: process 4/2/1 indices per thread, improving memory coalescing and bandwidth
• Fast path for single-block workloads, which uses a simple block scan and early return
• Multi-block cumsum uses a shared block prefix, reducing global synchronization overhead
• Auto-tuned vector width and block entry count at launch time to match indices length
• Guarded new ROCm changes using #ifdef USE_ROCM

Test Result

Reduced the duration of index_select_scalar_cumsum_kernel by 1.11 us, yielding a 1.3x speedup.

Test Plan

Unit test passed

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[avbokovoy]

Minor tweaks are needed, but overall LGTM

[avbokovoy]

We can pass that to the kernel as a template parameter with default value for easier tweaking if needed. Also VEC variable name is not self-explanatory

[amd-wsung102]

I changed the name of VEC to ENTRIES_PER_THREAD so it sounds more intuitive. It was originally passed to the kernel as a new template parameter, but Li said to avoid changing the template and kernel API. Should I change it back to a template parameter?

[amd-wsung102]

ENTRIES_PER_THREAD (previously VEC) is now passed to the kernel as a template parameter.

[avbokovoy]

@liligwu Could you elaborate why we should avoid changing template and kernel API?

[aryaman-gupta]

I took a glance at the code and left a couple of comments. Looks good to me otherwise

[aryaman-gupta]

As in the other file, you may consider passing multi_block as a compile-time parameter or splitting the function and dispatching the appropriate one at runtime.

[amd-wsung102]

Got it, I will be sure to try it to test its results.