From 14b3cce4aaedac69120d386953e2b4cae8c2cf2c Mon Sep 17 00:00:00 2001 From: Joe Mayer <114769929+jomayeri@users.noreply.github.com> Date: Fri, 14 Feb 2025 09:17:08 -0800 Subject: [PATCH] AIO on ROCM (#7023) Adding compile support for AIO library on AMD GPUs. --------- Co-authored-by: Olatunji Ruwase Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com> --- op_builder/async_io.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/op_builder/async_io.py b/op_builder/async_io.py index 0f9e34106487..f29f88531a98 100644 --- 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):