From b86b38045dc8407158e166d2309c462aa7e8c8d5 Mon Sep 17 00:00:00 2001 From: hjwei Date: Mon, 23 Dec 2024 23:19:47 -0800 Subject: [PATCH 1/2] [BUG FIX]:fix get torch.version.cuda error when cuda is None in rocm --- op_builder/builder.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/op_builder/builder.py b/op_builder/builder.py index 461281d4a569..f0ca90880d33 100644 --- a/op_builder/builder.py +++ b/op_builder/builder.py @@ -839,7 +839,9 @@ def cxx_args(self): CPU_ARCH = self.cpu_arch() SIMD_WIDTH = self.simd_width() - CUDA_ENABLE = self.is_cuda_enable() + CUDA_ENABLE = ( + "-D__DISABLE_CUDA__" if self.is_rocm_pytorch() else self.is_cuda_enable() + ) args += [ CPU_ARCH, '-fopenmp', From 309a666d898cffe98a4657ade6131017bf8d7931 Mon Sep 17 00:00:00 2001 From: hjwei Date: Tue, 24 Dec 2024 18:17:08 -0800 Subject: [PATCH 2/2] Move ROCm torch branch check into get_cuda_compile_flag --- op_builder/builder.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/op_builder/builder.py b/op_builder/builder.py index f0ca90880d33..ab26054bda7d 100644 --- a/op_builder/builder.py +++ b/op_builder/builder.py @@ -415,10 +415,11 @@ def cpu_arch(self): return '-mcpu=native' return '-march=native' - def is_cuda_enable(self): + def get_cuda_compile_flag(self): try: - assert_no_cuda_mismatch(self.name) - return '-D__ENABLE_CUDA__' + if not self.is_rocm_pytorch(): + assert_no_cuda_mismatch(self.name) + return "-D__ENABLE_CUDA__" except MissingCUDAException: print(f"{WARNING} {self.name} cuda is missing or is incompatible with installed torch, " "only cpu ops can be compiled!") @@ -839,9 +840,7 @@ def cxx_args(self): CPU_ARCH = self.cpu_arch() SIMD_WIDTH = self.simd_width() - CUDA_ENABLE = ( - "-D__DISABLE_CUDA__" if self.is_rocm_pytorch() else self.is_cuda_enable() - ) + CUDA_ENABLE = self.get_cuda_compile_flag() args += [ CPU_ARCH, '-fopenmp',