Build libtorchaudio GPU extensions on ROCm via HIPIFY (torchvision-style)

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],