AIO on ROCM (#7023)

Adding compile support for AIO library on AMD GPUs. --------- Co-authored-by: Olatunji Ruwase <[email protected]> Co-authored-by: Logan Adams <[email protected]>
deepspeedai · Feb 14, 2025 · 14b3cce · 14b3cce
1 parent e637677
commit 14b3cce
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/op_builder/async_io.py b/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):