Make index_mul_2d extension backward compatible for Atomic header include

ROCm / apex

A PyTorch Extension: Tools for easy mixed precision and distributed training in Pytorch

BSD 3-Clause "New" or "Revised" License

17 stars 14 forks source link

Make index_mul_2d extension backward compatible for Atomic header include #96

Closed hubertlu-tw closed 1 year ago

hubertlu-tw commented 1 year ago

ATen/cuda/Atomic.cuh does not exist before Pytorch 1.11. Therefore, this PR is to make index_mul_2d extension compatible with various Pytorch releases.

New Pytorch releases (after 1.11): #include <ATen/cuda/Atomic.cuh>
Old Pytorch releases (before 1.11): #include <THC/THCAtomics.cuh>

jithunnair-amd commented 1 year ago

Tested with 1.10 and 1.11 locally, according to @hubertlu-tw. Merging