[Feature Request] Atomic Ops on share memory like red.shared.add

[tianr22]

Required prerequisites

• I have searched the Issue Tracker that this hasn't already been reported. (comment there if it has.)

Motivation

I have been writing an backward kernel which performs atomic reduction across hc_mult(4) dim. If I use atomic.global.add, which means I need to read the whole tensor into L2 which is unacceptable.

It is better to perform reduction in share memory level. I have searched the docs and atomic.h and haven't found it implemented yet. Could u please add this feature in early future?
Thanks!

Solution

No response

Alternatives

No response

Additional context

No response

[SiriusNEO]

Hi @tianr22 , TileLang's atomic add ops (src/tl_templates/cuda/atomic.h) are implemented currently as:

• memory_order == relaxed: mostly calls CUDA atomicAdd(...) on generic pointers
• memory_order != relaxed: several paths use inline PTX hardcoded to .global

So I think for your feature request, we can refactor the template such that the memory_order != relaxed case can use different instructions for shared?

[tianr22]

Thanks!
memory_order == relaxed: CUDA atomicAdd is enough for my usage