Fix build on AMD GPUs (related to DeepCompile) (#7224)

We should use `torch.utils.cpp_extension.ROCM_HOME` for ROCm pytorch.

```log
  Traceback (most recent call last):
    File "<string>", line 2, in <module>
    File "<pip-setuptools-caller>", line 34, in <module>
    File "DeepSpeed/setup.py", line 195, in <module>
      builder.hipify_extension()
    File "DeepSpeed/op_builder/builder.py", line 750, in hipify_extension
      header_include_dirs=self.include_paths(),
                          ^^^^^^^^^^^^^^^^^^^^
    File "DeepSpeed/op_builder/dc.py", line 32, in include_paths
      return ['csrc/includes', os.path.join(torch.utils.cpp_extension.CUDA_HOME, "include")]
                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "<frozen posixpath>", line 76, in join
  TypeError: expected str, bytes or os.PathLike object, not NoneType
```

Signed-off-by: Hollow Man <hollowman@opensuse.org>
This commit is contained in:
ℍ𝕠𝕝𝕝𝕠𝕨 𝕄𝕒𝕟
2025-04-17 19:27:10 +03:00
committed by GitHub
parent 1d022e05c3
commit 0f224d79c6

View File

@ -29,4 +29,12 @@ class DeepCompileBuilder(TorchCPUOpBuilder):
def include_paths(self):
import os
import torch
return ['csrc/includes', os.path.join(torch.utils.cpp_extension.CUDA_HOME, "include")]
if self.build_for_cpu:
CUDA_INCLUDE = []
elif not self.is_rocm_pytorch():
CUDA_INCLUDE = [os.path.join(torch.utils.cpp_extension.CUDA_HOME, "include")]
else:
CUDA_INCLUDE = [
os.path.join(torch.utils.cpp_extension.ROCM_HOME, "include"),
]
return ['csrc/includes'] + CUDA_INCLUDE