diff --git a/tools/setup_helpers/extension.py b/tools/setup_helpers/extension.py index a8d0689cb5..48765d397d 100644 --- a/tools/setup_helpers/extension.py +++ b/tools/setup_helpers/extension.py @@ -55,6 +55,38 @@ def get_build_ext(): ) +def _hipify_glob_patterns(base_dir: str, depth: int = 4) -> list[str]: + """Build absolute-path fnmatch patterns for hipify (one * per directory level).""" + return [os.path.join(base_dir, *["*"] * level) for level in range(1, depth + 1)] + + +def _hipify_libtorchaudio(root_dir: Path, csrc_dir: Path) -> dict: + """Hipify libtorchaudio CUDA sources and headers before building on ROCm.""" + from torch.utils.hipify import hipify_python + + # Resolve through realpath so hipify updates the same files the compiler + # reads on Windows runners where B: is a subst alias of a C: checkout path. + root = os.path.realpath(str(root_dir)) + csrc = os.path.realpath(str(csrc_dir)) + + includes = _hipify_glob_patterns(csrc, depth=4) + header_include_dirs = [os.path.relpath(csrc, root)] + + cuctc = os.path.join(csrc, "cuctc") + if os.path.isdir(cuctc): + includes.extend(_hipify_glob_patterns(cuctc, depth=4)) + header_include_dirs.append(os.path.relpath(cuctc, root)) + + return hipify_python.hipify( + project_directory=root, + output_directory=root, + includes=includes, + header_include_dirs=header_include_dirs, + show_detailed=True, + is_pytorch_extension=True, + ) + + def get_ext_modules(): extra_compile_args = { "cxx": [ @@ -73,13 +105,14 @@ def get_ext_modules(): extra_compile_args["nvcc"] = ["-O2", "-DUSE_CUDA"] if _USE_ROCM: extension = CUDAExtension - extra_compile_args["nvcc"] = ["-O3"] - # TORCH_HIP_VERSION is used by hipified C++ (e.g. utils_hip.cpp); PyTorch only defines it when building PyTorch. + # GPU .hip sources still use USE_CUDA-guarded headers; hipify converts + # CUDA includes/APIs to HIP. Do not define USE_CUDA on cxx-only CPU files. + extra_compile_args["nvcc"] = ["-O3", "-DUSE_CUDA"] if torch.version.hip: parts = torch.version.hip.split(".") major = int(parts[0]) if len(parts) > 0 else 0 minor = int(parts[1]) if len(parts) > 1 else 0 - torch_hip_version = major * 100 + minor # e.g. 7.1.x -> 701 + torch_hip_version = major * 100 + minor extra_compile_args["cxx"].append("-DTORCH_HIP_VERSION=" + str(torch_hip_version)) extra_compile_args["nvcc"].append("-DTORCH_HIP_VERSION=" + str(torch_hip_version)) @@ -113,20 +146,36 @@ def get_ext_modules(): if _USE_CUDA or _USE_ROCM: sources.append("forced_align/gpu/compute.cu") + hipify_result = {} + if _USE_ROCM: + hipify_result = _hipify_libtorchaudio(_ROOT_DIR, _CSRC_DIR) + + def _srcs(*paths): + resolved = [] + for p in paths: + abs_path = os.path.realpath(str(p)) + result = hipify_result.get(abs_path) + hipified = getattr(result, "hipified_path", None) if result else None + if hipified and os.path.exists(hipified): + resolved.append(hipified) + else: + resolved.append(abs_path) + return resolved + modules = [ extension( name="torchaudio.lib._torchaudio", - sources=[ + sources=_srcs( _CSRC_DIR / "_torchaudio.cpp", _CSRC_DIR / "utils.cpp", - ], + ), py_limited_api=True, extra_compile_args=extra_compile_args, include_dirs=[_CSRC_DIR.parent], ), extension( name="torchaudio.lib.libtorchaudio", - sources=[_CSRC_DIR / s for s in sources], + sources=_srcs(*[_CSRC_DIR / s for s in sources]), py_limited_api=True, extra_compile_args=extra_compile_args, include_dirs=[_CSRC_DIR.parent], @@ -137,10 +186,12 @@ def get_ext_modules(): [ extension( name="torchaudio.lib.torchaudio_prefixctc", - sources=[ - _CSRC_DIR / "cuctc" / "src" / s - for s in ["ctc_prefix_decoder.cpp", "ctc_prefix_decoder_kernel_v2.cu", "python_binding.cpp"] - ], + sources=_srcs( + *[ + _CSRC_DIR / "cuctc" / "src" / s + for s in ["ctc_prefix_decoder.cpp", "ctc_prefix_decoder_kernel_v2.cu", "python_binding.cpp"] + ] + ), py_limited_api=True, extra_compile_args=extra_compile_args, include_dirs=[_CSRC_DIR / "cuctc", _CSRC_DIR.parent],