Atomic opeartion in TIR

I want to use tensorIR to write Pooling backward operation. It seems that this needs atomic add operation, ie. accumulating the same element not continuously in the output array. Does tensor IR support atomic operator or have better choice to achieve pooling backward operator? Thanks

@yzh119 has some exploration on atomic operations :slight_smile: essentially via intrinsics

1 Like

Hi @lhd1029 , yes I have added atomic intrinsic in my personal project. I can upstream this part if necessary. You can also check the implementation of nms: https://github.com/apache/tvm/blob/f4ec5fd4ae346dbdd8e915c048aeed94b44f6776/python/tvm/topi/cuda/nms.py

hi @yzh119 .Thank you for your helpful reply. I would appreciate it if you could upstream the part of atomic intrisic.