From ddfbc452bf1e5090776f1ad6747c68baeb67104a Mon Sep 17 00:00:00 2001 From: ethanwee1 Date: Wed, 1 Jul 2026 01:54:37 +0000 Subject: [PATCH 1/2] Port libtorchaudio GPU extensions to build under ROCm/HIP The torch CUDAExtension HIPIFY path only rewrites source files, not headers, so CUDA symbols used in headers survived un-hipified and broke the ROCm build (e.g. "cuda_runtime_api.h file not found", and the rnnt/forced_align GPU kernels were entirely gated behind #ifdef USE_CUDA). - Add cuda_compat.h mapping the CUDA runtime names used in headers (cudaStream_t, cudaError_t, cudaSuccess, cudaGetErrorString, cudaGetLastError) to their HIP equivalents under USE_ROCM. HIP natively supports the <<<>>> launch syntax and warp shuffle intrinsics already guarded in the kernels. - Extend the #ifdef USE_CUDA guards in the rnnt GPU headers, macros.h, options.h, and cuctc headers to also cover USE_ROCM, including HIP fp16/runtime headers. --- .../cuctc/include/ctc_prefix_decoder_host.h | 2 ++ .../cuctc/src/device_data_wrap.h | 2 +- src/libtorchaudio/cuda_compat.h | 23 +++++++++++++++++++ src/libtorchaudio/cuda_utils.h | 2 +- .../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 | 6 ++--- 10 files changed, 43 insertions(+), 12 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/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..991267d4d8 --- /dev/null +++ b/src/libtorchaudio/cuda_compat.h @@ -0,0 +1,23 @@ +#pragma once + +// CUDA/HIP compatibility shim for building libtorchaudio GPU sources under ROCm. +// +// The torch CUDAExtension HIPIFY path only rewrites source files (.cu/.cpp +// listed as extension sources), not headers. CUDA runtime symbols that appear +// in headers therefore survive un-hipified, so map them to their HIP +// equivalents explicitly here. HIP natively supports the triple-chevron kernel +// launch syntax and the warp shuffle intrinsics used elsewhere; only the +// runtime type/enum/function names need aliasing. +#if defined(USE_ROCM) +#include +#include + +using cudaStream_t = hipStream_t; +using cudaError_t = hipError_t; +#define cudaSuccess hipSuccess +#define cudaGetErrorString hipGetErrorString +#define cudaGetLastError hipGetLastError +#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/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..0dae561f65 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 From a890448a93de1a06b67f6469a75485360a956204 Mon Sep 17 00:00:00 2001 From: ethanwee1 Date: Wed, 1 Jul 2026 15:28:53 +0000 Subject: [PATCH 2/2] Expand ROCm/HIP port: cuda runtime shim, hipcub, Options::stream_ - cuda_compat.h: alias the CUDA runtime functions/enums used by the sources (cudaSetDevice, cudaMemcpy*, cudaMemset*, cudaStreamSynchronize, cudaGetErrorName, cudaMemcpy kinds) to HIP under USE_ROCM. - Route cub through hipcub (namespace cub = hipcub) under USE_ROCM in forced_align and cuctc kernels. - Enable Options::stream_ for USE_ROCM. - ctc_prefix_decoder.cpp includes cuda_compat.h instead of cuda_runtime.h. --- .../cuctc/src/ctc_prefix_decoder.cpp | 2 +- .../cuctc/src/ctc_prefix_decoder_kernel_v2.cu | 5 ++ src/libtorchaudio/cuda_compat.h | 64 ++++++++++++------- src/libtorchaudio/forced_align/gpu/compute.cu | 5 ++ src/libtorchaudio/rnnt/options.h | 2 +- 5 files changed, 53 insertions(+), 25 deletions(-) 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/cuda_compat.h b/src/libtorchaudio/cuda_compat.h index 991267d4d8..93b6a8e189 100644 --- a/src/libtorchaudio/cuda_compat.h +++ b/src/libtorchaudio/cuda_compat.h @@ -1,23 +1,41 @@ -#pragma once - -// CUDA/HIP compatibility shim for building libtorchaudio GPU sources under ROCm. -// -// The torch CUDAExtension HIPIFY path only rewrites source files (.cu/.cpp -// listed as extension sources), not headers. CUDA runtime symbols that appear -// in headers therefore survive un-hipified, so map them to their HIP -// equivalents explicitly here. HIP natively supports the triple-chevron kernel -// launch syntax and the warp shuffle intrinsics used elsewhere; only the -// runtime type/enum/function names need aliasing. -#if defined(USE_ROCM) -#include -#include - -using cudaStream_t = hipStream_t; -using cudaError_t = hipError_t; -#define cudaSuccess hipSuccess -#define cudaGetErrorString hipGetErrorString -#define cudaGetLastError hipGetLastError -#elif defined(USE_CUDA) -#include -#include -#endif +#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/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/options.h b/src/libtorchaudio/rnnt/options.h index 0dae561f65..f854d4b4cb 100644 --- a/src/libtorchaudio/rnnt/options.h +++ b/src/libtorchaudio/rnnt/options.h @@ -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