AIO on ROCM (deepspeedai#7023)

Adding compile support for AIO library on AMD GPUs.

---------

Co-authored-by: Olatunji Ruwase <olruwase@microsoft.com>
Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com>

Author: 3 people

op_builder/async_io.py

@@ -65,8 +65,11 @@ def extra_ldflags(self):
 
         import torch.utils.cpp_extension
         CUDA_HOME = torch.utils.cpp_extension.CUDA_HOME
-        CUDA_LIB64 = os.path.join(CUDA_HOME, "lib64")
-        ldflags = [f'-L{CUDA_HOME}', f'-L{CUDA_LIB64}', '-laio', '-lcuda', '-lcudart']
+        if CUDA_HOME is None:
+            ldflags = ['-laio']  # the ROCM case
+        else:
+            CUDA_LIB64 = os.path.join(CUDA_HOME, "lib64")
+            ldflags = [f'-L{CUDA_HOME}', f'-L{CUDA_LIB64}', '-laio', '-lcuda', '-lcudart']
         return ldflags
 
     def check_for_libaio_pkg(self):