From f9437d2cf1ef6e91df0557795f7e35fb11cdecdc Mon Sep 17 00:00:00 2001 From: ethanwee1 Date: Wed, 1 Jul 2026 18:13:28 +0000 Subject: [PATCH 1/2] Build libtorchaudio GPU extensions on ROCm via HIPIFY (torchvision-style) Mirror torchvision's setup.py approach: run hipify_python.hipify() over the libtorchaudio sources when building for ROCm and compile the generated .hip outputs, instead of relying solely on the torch CUDAExtension HIPIFY (which only processes listed source files and left torchaudio headers/symbols like cuda_runtime_api.h, cudaSetDevice and cub un-converted). HIPIFY cannot rewrite preprocessor guards, so the #ifdef USE_CUDA gates around the rnnt/forced_align GPU headers, macros.h, and Options::stream_ are widened to `#if defined(USE_CUDA) || defined(USE_ROCM)`. Reference: https://github.com/ROCm/vision/blob/release/0.27/setup.py --- .../rnnt/gpu/gpu_kernel_utils.cuh | 2 +- src/libtorchaudio/rnnt/gpu/gpu_kernels.cuh | 2 +- src/libtorchaudio/rnnt/gpu/gpu_transducer.h | 2 +- src/libtorchaudio/rnnt/gpu/math.cuh | 4 +- src/libtorchaudio/rnnt/macros.h | 2 +- src/libtorchaudio/rnnt/options.h | 6 +-- tools/setup_helpers/extension.py | 47 ++++++++++++++++--- 7 files changed, 49 insertions(+), 16 deletions(-) diff --git a/src/libtorchaudio/rnnt/gpu/gpu_kernel_utils.cuh b/src/libtorchaudio/rnnt/gpu/gpu_kernel_utils.cuh index f4ad3add2b..e58b9c185e 100644 --- a/src/libtorchaudio/rnnt/gpu/gpu_kernel_utils.cuh +++ b/src/libtorchaudio/rnnt/gpu/gpu_kernel_utils.cuh @@ -1,6 +1,6 @@ #pragma once -#ifdef USE_CUDA +#if defined(USE_CUDA) || defined(USE_ROCM) #include diff --git a/src/libtorchaudio/rnnt/gpu/gpu_kernels.cuh b/src/libtorchaudio/rnnt/gpu/gpu_kernels.cuh index 136e6844f2..1d29993d23 100644 --- a/src/libtorchaudio/rnnt/gpu/gpu_kernels.cuh +++ b/src/libtorchaudio/rnnt/gpu/gpu_kernels.cuh @@ -1,6 +1,6 @@ #pragma once -#ifdef USE_CUDA +#if defined(USE_CUDA) || defined(USE_ROCM) #include diff --git a/src/libtorchaudio/rnnt/gpu/gpu_transducer.h b/src/libtorchaudio/rnnt/gpu/gpu_transducer.h index 875c47974f..ad07b613fd 100644 --- a/src/libtorchaudio/rnnt/gpu/gpu_transducer.h +++ b/src/libtorchaudio/rnnt/gpu/gpu_transducer.h @@ -1,6 +1,6 @@ #pragma once -#ifdef USE_CUDA +#if defined(USE_CUDA) || defined(USE_ROCM) #include #include diff --git a/src/libtorchaudio/rnnt/gpu/math.cuh b/src/libtorchaudio/rnnt/gpu/math.cuh index 8a3c664505..731bd0de0f 100644 --- a/src/libtorchaudio/rnnt/gpu/math.cuh +++ b/src/libtorchaudio/rnnt/gpu/math.cuh @@ -1,10 +1,10 @@ #pragma once -#ifdef USE_CUDA +#if defined(USE_CUDA) || defined(USE_ROCM) #include -#endif // USE_CUDA +#endif // USE_CUDA || USE_ROCM #include diff --git a/src/libtorchaudio/rnnt/macros.h b/src/libtorchaudio/rnnt/macros.h index cdc83dd5d2..455e439974 100644 --- a/src/libtorchaudio/rnnt/macros.h +++ b/src/libtorchaudio/rnnt/macros.h @@ -1,6 +1,6 @@ #pragma once -#ifdef USE_CUDA +#if defined(USE_CUDA) || defined(USE_ROCM) #define WARP_SIZE 32 #define MAX_THREADS_PER_BLOCK 1024 #define REDUCE_THREADS 256 diff --git a/src/libtorchaudio/rnnt/options.h b/src/libtorchaudio/rnnt/options.h index 8a8fed1116..eae87deb2f 100644 --- a/src/libtorchaudio/rnnt/options.h +++ b/src/libtorchaudio/rnnt/options.h @@ -1,8 +1,8 @@ #pragma once -#ifdef USE_CUDA +#if defined(USE_CUDA) || defined(USE_ROCM) #include -#endif // USE_CUDA +#endif // USE_CUDA || USE_ROCM #include #include @@ -13,7 +13,7 @@ namespace rnnt { struct Options { // the device to compute transducer loss. device_t device_; -#ifdef USE_CUDA +#if defined(USE_CUDA) || defined(USE_ROCM) // the stream to launch kernels in when using GPU. cudaStream_t stream_; #endif diff --git a/tools/setup_helpers/extension.py b/tools/setup_helpers/extension.py index a8d0689cb5..6d4f353e3b 100644 --- a/tools/setup_helpers/extension.py +++ b/tools/setup_helpers/extension.py @@ -113,20 +113,51 @@ def get_ext_modules(): if _USE_CUDA or _USE_ROCM: sources.append("forced_align/gpu/compute.cu") + # For ROCm, hipify libtorchaudio's sources before building. This mirrors + # torchvision's setup.py, which runs hipify_python.hipify() and then + # compiles the generated .hip outputs. HIPIFY rewrites CUDA symbols in the + # .cu sources *and their headers* (which the torch CUDAExtension HIPIFY, + # limited to extra source files, does not fully cover) and renames .cu -> + # .hip, updating includes accordingly. + # See: https://github.com/ROCm/vision/blob/release/0.27/setup.py + hipify_result = {} + if _USE_ROCM: + from torch.utils.hipify import hipify_python + + hipify_result = hipify_python.hipify( + project_directory=str(_ROOT_DIR), + output_directory=str(_ROOT_DIR), + includes=[str(_CSRC_DIR / "*")], + show_detailed=True, + is_pytorch_extension=True, + ) + + def _srcs(*paths): + resolved = [] + for p in paths: + abs_path = os.path.abspath(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(str(p)) + 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 +168,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], From 0c7509ef5eae0ad822d2099ca5b42fce18db0148 Mon Sep 17 00:00:00 2001 From: ethanwee1 Date: Mon, 6 Jul 2026 02:53:04 +0000 Subject: [PATCH 2/2] Build libtorchaudio GPU extensions on ROCm via HIPIFY (torchvision-style) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mirror torchvision's setup.py approach: run hipify_python.hipify() over libtorchaudio CUDA sources and headers, then compile the generated .hip outputs. Fixes vs the first HIPIFY attempt (TheRock 28538325024): 1. Do not widen `#ifdef USE_CUDA` to `USE_ROCM` in headers — that pulled `cuda_fp16.h` into CPU `.cpp` compiles because PyTorch defines USE_ROCM globally on ROCm extension builds. Keep upstream USE_CUDA guards; pass `-DUSE_CUDA` on hipcc/nvcc only for GPU translation units. 2. Hipify with `os.path.realpath()` paths and depth-aware include globs so headers like `cuda_utils.h` are rewritten on the same inode Windows runners use when B: is a subst alias of the C: checkout (otherwise the compiler reads un-hipified headers from the underlying path). 3. Include `cuctc/` in the hipify scope. Reference: https://github.com/ROCm/vision/blob/release/0.27/setup.py Paired with ROCm/pytorch#3382 (torchaudio pin + Windows cross-drive HIPIFY fix in torch/utils/cpp_extension.py — that fix is in PyTorch, not here). --- .../rnnt/gpu/gpu_kernel_utils.cuh | 2 +- src/libtorchaudio/rnnt/gpu/gpu_kernels.cuh | 2 +- src/libtorchaudio/rnnt/gpu/gpu_transducer.h | 2 +- src/libtorchaudio/rnnt/gpu/math.cuh | 4 +- src/libtorchaudio/rnnt/macros.h | 2 +- src/libtorchaudio/rnnt/options.h | 6 +- tools/setup_helpers/extension.py | 60 ++++++++++++------- 7 files changed, 48 insertions(+), 30 deletions(-) diff --git a/src/libtorchaudio/rnnt/gpu/gpu_kernel_utils.cuh b/src/libtorchaudio/rnnt/gpu/gpu_kernel_utils.cuh index e58b9c185e..f4ad3add2b 100644 --- a/src/libtorchaudio/rnnt/gpu/gpu_kernel_utils.cuh +++ b/src/libtorchaudio/rnnt/gpu/gpu_kernel_utils.cuh @@ -1,6 +1,6 @@ #pragma once -#if defined(USE_CUDA) || defined(USE_ROCM) +#ifdef USE_CUDA #include diff --git a/src/libtorchaudio/rnnt/gpu/gpu_kernels.cuh b/src/libtorchaudio/rnnt/gpu/gpu_kernels.cuh index 1d29993d23..136e6844f2 100644 --- a/src/libtorchaudio/rnnt/gpu/gpu_kernels.cuh +++ b/src/libtorchaudio/rnnt/gpu/gpu_kernels.cuh @@ -1,6 +1,6 @@ #pragma once -#if defined(USE_CUDA) || defined(USE_ROCM) +#ifdef USE_CUDA #include diff --git a/src/libtorchaudio/rnnt/gpu/gpu_transducer.h b/src/libtorchaudio/rnnt/gpu/gpu_transducer.h index ad07b613fd..875c47974f 100644 --- a/src/libtorchaudio/rnnt/gpu/gpu_transducer.h +++ b/src/libtorchaudio/rnnt/gpu/gpu_transducer.h @@ -1,6 +1,6 @@ #pragma once -#if defined(USE_CUDA) || defined(USE_ROCM) +#ifdef USE_CUDA #include #include diff --git a/src/libtorchaudio/rnnt/gpu/math.cuh b/src/libtorchaudio/rnnt/gpu/math.cuh index 731bd0de0f..8a3c664505 100644 --- a/src/libtorchaudio/rnnt/gpu/math.cuh +++ b/src/libtorchaudio/rnnt/gpu/math.cuh @@ -1,10 +1,10 @@ #pragma once -#if defined(USE_CUDA) || defined(USE_ROCM) +#ifdef USE_CUDA #include -#endif // USE_CUDA || USE_ROCM +#endif // USE_CUDA #include diff --git a/src/libtorchaudio/rnnt/macros.h b/src/libtorchaudio/rnnt/macros.h index 455e439974..cdc83dd5d2 100644 --- a/src/libtorchaudio/rnnt/macros.h +++ b/src/libtorchaudio/rnnt/macros.h @@ -1,6 +1,6 @@ #pragma once -#if defined(USE_CUDA) || defined(USE_ROCM) +#ifdef USE_CUDA #define WARP_SIZE 32 #define MAX_THREADS_PER_BLOCK 1024 #define REDUCE_THREADS 256 diff --git a/src/libtorchaudio/rnnt/options.h b/src/libtorchaudio/rnnt/options.h index eae87deb2f..8a8fed1116 100644 --- a/src/libtorchaudio/rnnt/options.h +++ b/src/libtorchaudio/rnnt/options.h @@ -1,8 +1,8 @@ #pragma once -#if defined(USE_CUDA) || defined(USE_ROCM) +#ifdef USE_CUDA #include -#endif // USE_CUDA || USE_ROCM +#endif // USE_CUDA #include #include @@ -13,7 +13,7 @@ namespace rnnt { struct Options { // the device to compute transducer loss. device_t device_; -#if defined(USE_CUDA) || defined(USE_ROCM) +#ifdef USE_CUDA // the stream to launch kernels in when using GPU. cudaStream_t stream_; #endif diff --git a/tools/setup_helpers/extension.py b/tools/setup_helpers/extension.py index 6d4f353e3b..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,35 +146,20 @@ def get_ext_modules(): if _USE_CUDA or _USE_ROCM: sources.append("forced_align/gpu/compute.cu") - # For ROCm, hipify libtorchaudio's sources before building. This mirrors - # torchvision's setup.py, which runs hipify_python.hipify() and then - # compiles the generated .hip outputs. HIPIFY rewrites CUDA symbols in the - # .cu sources *and their headers* (which the torch CUDAExtension HIPIFY, - # limited to extra source files, does not fully cover) and renames .cu -> - # .hip, updating includes accordingly. - # See: https://github.com/ROCm/vision/blob/release/0.27/setup.py hipify_result = {} if _USE_ROCM: - from torch.utils.hipify import hipify_python - - hipify_result = hipify_python.hipify( - project_directory=str(_ROOT_DIR), - output_directory=str(_ROOT_DIR), - includes=[str(_CSRC_DIR / "*")], - show_detailed=True, - is_pytorch_extension=True, - ) + hipify_result = _hipify_libtorchaudio(_ROOT_DIR, _CSRC_DIR) def _srcs(*paths): resolved = [] for p in paths: - abs_path = os.path.abspath(str(p)) + 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(str(p)) + resolved.append(abs_path) return resolved modules = [