From f9437d2cf1ef6e91df0557795f7e35fb11cdecdc Mon Sep 17 00:00:00 2001 From: ethanwee1 Date: Wed, 1 Jul 2026 18:13:28 +0000 Subject: [PATCH 1/4] 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/4] 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 = [ From 133b66870860350b69fbb45869afe1c89b85709a Mon Sep 17 00:00:00 2001 From: ethanwee1 Date: Wed, 8 Jul 2026 15:00:31 +0000 Subject: [PATCH 3/4] HIPIFY v3: fix header_include_dirs so includes are rewritten hipify resolves angle-bracket includes against header_include_dirs. The sources include headers as via the extension -I root (src/), but v2 passed src/libtorchaudio, so hipify looked for src/libtorchaudio/libtorchaudio/... , failed to resolve, and left #include un-rewritten -> compiled the original cuda_utils.h with cuda_runtime_api.h. Point header_include_dirs at src/ (csrc.parent) so the includes resolve and get rewritten to the hipified header paths. --- tools/setup_helpers/extension.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tools/setup_helpers/extension.py b/tools/setup_helpers/extension.py index 48765d397d..c3463fc1da 100644 --- a/tools/setup_helpers/extension.py +++ b/tools/setup_helpers/extension.py @@ -70,7 +70,11 @@ def _hipify_libtorchaudio(root_dir: Path, csrc_dir: Path) -> dict: csrc = os.path.realpath(str(csrc_dir)) includes = _hipify_glob_patterns(csrc, depth=4) - header_include_dirs = [os.path.relpath(csrc, root)] + # Header includes are written as , resolved against the + # extension's -I root (src/), so hipify must search from src/ (csrc.parent) + # -- not src/libtorchaudio -- to find them and rewrite the includes to the + # hipified header paths. + header_include_dirs = [os.path.relpath(os.path.dirname(csrc), root)] cuctc = os.path.join(csrc, "cuctc") if os.path.isdir(cuctc): From 9808bfc9565dd298cec0b68d5787fcfd6ccad4eb Mon Sep 17 00:00:00 2001 From: ethanwee1 Date: Fri, 17 Jul 2026 15:10:26 +0000 Subject: [PATCH 4/4] Rework: drop manual setup hipify; port via torchaudio source (compat header + guards) Addresses review feedback (jithunnair-amd) to not invoke hipify manually in the torchaudio setup file. Instead this ports libtorchaudio to ROCm via fork source updates only, relying on the built-in CUDAExtension hipify for the .cu source bodies: - cuda_compat.h maps the CUDA runtime symbols used in headers to their HIP equivalents under USE_ROCM. - Widen the rnnt/forced_align GPU #ifdef USE_CUDA guards, macros.h, and Options::stream_ to also cover USE_ROCM. - Route cub through hipcub under USE_ROCM. The Windows B:/C: subst cross-drive HIPIFY fix lives in ROCm/pytorch#3382 (torch/utils/cpp_extension.py). --- .../cuctc/include/ctc_prefix_decoder_host.h | 2 + .../cuctc/src/ctc_prefix_decoder.cpp | 2 +- .../cuctc/src/ctc_prefix_decoder_kernel_v2.cu | 5 ++ .../cuctc/src/device_data_wrap.h | 2 +- src/libtorchaudio/cuda_compat.h | 41 ++++++++++ src/libtorchaudio/cuda_utils.h | 2 +- src/libtorchaudio/forced_align/gpu/compute.cu | 5 ++ .../rnnt/gpu/gpu_kernel_utils.cuh | 2 +- src/libtorchaudio/rnnt/gpu/gpu_kernels.cuh | 2 +- src/libtorchaudio/rnnt/gpu/gpu_transducer.h | 3 +- src/libtorchaudio/rnnt/gpu/math.cuh | 4 +- src/libtorchaudio/rnnt/macros.h | 9 ++- src/libtorchaudio/rnnt/options.h | 8 +- tools/setup_helpers/extension.py | 75 +++---------------- 14 files changed, 83 insertions(+), 79 deletions(-) create mode 100644 src/libtorchaudio/cuda_compat.h diff --git a/src/libtorchaudio/cuctc/include/ctc_prefix_decoder_host.h b/src/libtorchaudio/cuctc/include/ctc_prefix_decoder_host.h index 2d6574e36b..33d3675b5b 100644 --- a/src/libtorchaudio/cuctc/include/ctc_prefix_decoder_host.h +++ b/src/libtorchaudio/cuctc/include/ctc_prefix_decoder_host.h @@ -26,6 +26,8 @@ #ifndef __ctc_prefix_decoder_host_h_ #define __ctc_prefix_decoder_host_h_ +#include + #define CHECK(X, ERROR_INFO) \ do { \ auto result = (X); \ diff --git a/src/libtorchaudio/cuctc/src/ctc_prefix_decoder.cpp b/src/libtorchaudio/cuctc/src/ctc_prefix_decoder.cpp index 70fc801d0f..fd2ada7e1f 100644 --- a/src/libtorchaudio/cuctc/src/ctc_prefix_decoder.cpp +++ b/src/libtorchaudio/cuctc/src/ctc_prefix_decoder.cpp @@ -23,7 +23,7 @@ // OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -#include +#include #include "../include/ctc_prefix_decoder.h" #include "../include/ctc_prefix_decoder_host.h" diff --git a/src/libtorchaudio/cuctc/src/ctc_prefix_decoder_kernel_v2.cu b/src/libtorchaudio/cuctc/src/ctc_prefix_decoder_kernel_v2.cu index 92edcc83c1..14810e6db3 100644 --- a/src/libtorchaudio/cuctc/src/ctc_prefix_decoder_kernel_v2.cu +++ b/src/libtorchaudio/cuctc/src/ctc_prefix_decoder_kernel_v2.cu @@ -28,7 +28,12 @@ #include #include "../include/ctc_prefix_decoder_host.h" #include "ctc_fast_divmod.cuh" +#if defined(USE_ROCM) +#include +namespace cub = hipcub; +#else #include "cub/cub.cuh" +#endif #include "device_data_wrap.h" #include "device_log_prob.cuh" diff --git a/src/libtorchaudio/cuctc/src/device_data_wrap.h b/src/libtorchaudio/cuctc/src/device_data_wrap.h index 0b2ca6f1e0..fa3307e3f0 100644 --- a/src/libtorchaudio/cuctc/src/device_data_wrap.h +++ b/src/libtorchaudio/cuctc/src/device_data_wrap.h @@ -24,7 +24,7 @@ // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. #pragma once -#include +#include #include #include #include "../include/ctc_prefix_decoder_host.h" diff --git a/src/libtorchaudio/cuda_compat.h b/src/libtorchaudio/cuda_compat.h new file mode 100644 index 0000000000..93b6a8e189 --- /dev/null +++ b/src/libtorchaudio/cuda_compat.h @@ -0,0 +1,41 @@ +#pragma once + +// CUDA/HIP compatibility shim for building libtorchaudio GPU sources under ROCm. +// +// The torch CUDAExtension HIPIFY path does not reliably rewrite every CUDA +// symbol in torchaudio's sources/headers (torchaudio has no whole-tree +// build_amd.py hipify step). Map the CUDA runtime names used by libtorchaudio +// to their HIP equivalents under USE_ROCM so the sources compile with the HIP +// toolchain regardless of HIPIFY coverage. HIP natively supports the <<<>>> +// launch syntax and the warp shuffle intrinsics already guarded in the kernels. +#if defined(USE_ROCM) +#include +#include + +// Types +using cudaStream_t = hipStream_t; +using cudaError_t = hipError_t; + +// Enums / values +#define cudaSuccess hipSuccess +#define cudaMemcpyHostToDevice hipMemcpyHostToDevice +#define cudaMemcpyDeviceToHost hipMemcpyDeviceToHost +#define cudaMemcpyDeviceToDevice hipMemcpyDeviceToDevice + +// Error helpers +#define cudaGetLastError hipGetLastError +#define cudaGetErrorString hipGetErrorString +#define cudaGetErrorName hipGetErrorName + +// Device / stream / memory runtime API +#define cudaSetDevice hipSetDevice +#define cudaStreamSynchronize hipStreamSynchronize +#define cudaMemcpy hipMemcpy +#define cudaMemcpyAsync hipMemcpyAsync +#define cudaMemcpy2DAsync hipMemcpy2DAsync +#define cudaMemset hipMemset +#define cudaMemsetAsync hipMemsetAsync +#elif defined(USE_CUDA) +#include +#include +#endif diff --git a/src/libtorchaudio/cuda_utils.h b/src/libtorchaudio/cuda_utils.h index 4d122e5bdb..85e4cd0131 100644 --- a/src/libtorchaudio/cuda_utils.h +++ b/src/libtorchaudio/cuda_utils.h @@ -1,6 +1,6 @@ #pragma once -#include +#include #include #include diff --git a/src/libtorchaudio/forced_align/gpu/compute.cu b/src/libtorchaudio/forced_align/gpu/compute.cu index 444a4f8f6d..1557c6d22c 100644 --- a/src/libtorchaudio/forced_align/gpu/compute.cu +++ b/src/libtorchaudio/forced_align/gpu/compute.cu @@ -5,7 +5,12 @@ #include #include +#if defined(USE_ROCM) +#include +namespace cub = hipcub; +#else #include +#endif #include namespace { 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..23b4ce8fed 100644 --- a/src/libtorchaudio/rnnt/gpu/gpu_transducer.h +++ b/src/libtorchaudio/rnnt/gpu/gpu_transducer.h @@ -1,7 +1,8 @@ #pragma once -#ifdef USE_CUDA +#if defined(USE_CUDA) || defined(USE_ROCM) +#include #include #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..5e4530c806 100644 --- a/src/libtorchaudio/rnnt/macros.h +++ b/src/libtorchaudio/rnnt/macros.h @@ -1,14 +1,19 @@ #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 #define HOST_AND_DEVICE __host__ __device__ #define FORCE_INLINE __forceinline__ +#if defined(USE_ROCM) +#include +#include +#else #include #include +#endif #else #define HOST_AND_DEVICE #define FORCE_INLINE inline -#endif // USE_CUDA +#endif // USE_CUDA || USE_ROCM diff --git a/src/libtorchaudio/rnnt/options.h b/src/libtorchaudio/rnnt/options.h index 8a8fed1116..f854d4b4cb 100644 --- a/src/libtorchaudio/rnnt/options.h +++ b/src/libtorchaudio/rnnt/options.h @@ -1,8 +1,8 @@ #pragma once -#ifdef USE_CUDA -#include -#endif // USE_CUDA +#if defined(USE_CUDA) || defined(USE_ROCM) +#include +#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 c3463fc1da..a8d0689cb5 100644 --- a/tools/setup_helpers/extension.py +++ b/tools/setup_helpers/extension.py @@ -55,42 +55,6 @@ 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 includes are written as , resolved against the - # extension's -I root (src/), so hipify must search from src/ (csrc.parent) - # -- not src/libtorchaudio -- to find them and rewrite the includes to the - # hipified header paths. - header_include_dirs = [os.path.relpath(os.path.dirname(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": [ @@ -109,14 +73,13 @@ def get_ext_modules(): extra_compile_args["nvcc"] = ["-O2", "-DUSE_CUDA"] if _USE_ROCM: extension = CUDAExtension - # 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"] + 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. 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 + torch_hip_version = major * 100 + minor # e.g. 7.1.x -> 701 extra_compile_args["cxx"].append("-DTORCH_HIP_VERSION=" + str(torch_hip_version)) extra_compile_args["nvcc"].append("-DTORCH_HIP_VERSION=" + str(torch_hip_version)) @@ -150,36 +113,20 @@ 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=_srcs( + sources=[ _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=_srcs(*[_CSRC_DIR / s for s in sources]), + sources=[_CSRC_DIR / s for s in sources], py_limited_api=True, extra_compile_args=extra_compile_args, include_dirs=[_CSRC_DIR.parent], @@ -190,12 +137,10 @@ def _srcs(*paths): [ extension( name="torchaudio.lib.torchaudio_prefixctc", - sources=_srcs( - *[ - _CSRC_DIR / "cuctc" / "src" / s - for s in ["ctc_prefix_decoder.cpp", "ctc_prefix_decoder_kernel_v2.cu", "python_binding.cpp"] - ] - ), + sources=[ + _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],