修改默认编译器为hipcc

tools/setup_helpers/extension.py

@@ -7,7 +7,7 @@ from pathlib import Path
 import torch
 from setuptools import Extension
 from setuptools.command.build_ext import build_ext
-from torch.utils.cpp_extension import CUDA_HOME
+from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME
 
 __all__ = [
     "get_ext_modules",
@@ -124,6 +124,10 @@ class CMakeBuild(build_ext):
             cmake_args += [f"-DCMAKE_CUDA_COMPILER='{CUDA_HOME}/bin/nvcc'"]
             cmake_args += [f"-DCUDA_TOOLKIT_ROOT_DIR='{CUDA_HOME}'"]
         
+        if platform.system() != "Windows" and ROCM_HOME is not None:
+            cmake_args += [f"-DCMAKE_HIP_COMPILER='{ROCM_HOME}/bin/hipcc'"]
+            cmake_args += [f"-DCUDA_TOOLKIT_ROOT_DIR='{ROCM_HOME}'"]
+
         # Default to Ninja
         if "CMAKE_GENERATOR" not in os.environ or platform.system() == "Windows":
             cmake_args += ["-GNinja"]