Does expressions like "x.at[idx].add(y)" are executed in parallel on GPU?

[Dong-Jiahuan]

As described in title, I want to know if expressions like "x.at[idx].add(y)" are executed in parallel on GPU? Is it necessary to write CUDA to make custom operator to realize the same logic?

[hawkinsp]

Yes. That expression lowers to an HLO scatter operator, and if the scatter updates multiple indices in one operation, those updates may happen in parallel.

[Dong-Jiahuan]

Thanks for your reply! I'm not sure your meaning of "those updates may happen in parallel". Does it mean some times those updates may happen in parallel and some times may not happen in parallel? What's the standard of parallel computing happen or not?

[hawkinsp]

It's up to the compiler (XLA) what happens, but in practice, for a scatter-add on GPU: yes, the updates will occur in parallel.

[Dong-Jiahuan]

I see. Thanks a lot!