Port libtorchaudio GPU extensions to build under ROCm/HIP

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 <libtorchaudio/cuda_compat.h>
+
 #define CHECK(X, ERROR_INFO)                        \
   do {                                              \
     auto result = (X);                              \

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 <cuda_runtime.h>
+#include <libtorchaudio/cuda_compat.h>
 
 #include "../include/ctc_prefix_decoder.h"
 #include "../include/ctc_prefix_decoder_host.h"

src/libtorchaudio/cuctc/src/ctc_prefix_decoder_kernel_v2.cu

@@ -28,7 +28,12 @@
 #include <limits>
 #include "../include/ctc_prefix_decoder_host.h"
 #include "ctc_fast_divmod.cuh"
+#if defined(USE_ROCM)
+#include <hipcub/hipcub.hpp>
+namespace cub = hipcub;
+#else
 #include "cub/cub.cuh"
+#endif
 #include "device_data_wrap.h"
 #include "device_log_prob.cuh"
 

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 <cuda_runtime.h>
+#include <libtorchaudio/cuda_compat.h>
 #include <iostream>
 #include <vector>
 #include "../include/ctc_prefix_decoder_host.h"

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 <hip/hip_runtime.h>
+#include <hip/hip_runtime_api.h>
+
+// 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 <cuda_runtime.h>
+#include <cuda_runtime_api.h>
+#endif

src/libtorchaudio/cuda_utils.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include <cuda_runtime_api.h>
+#include <libtorchaudio/cuda_compat.h>
 #include <torch/csrc/stable/c/shim.h>
 #include <torch/csrc/stable/device.h>
 

src/libtorchaudio/forced_align/gpu/compute.cu

@@ -5,7 +5,12 @@
 #include <torch/headeronly/core/Dispatch_v2.h>
 #include <torch/headeronly/core/ScalarType.h>
 
+#if defined(USE_ROCM)
+#include <hipcub/hipcub.hpp>
+namespace cub = hipcub;
+#else
 #include <cub/cub.cuh>
+#endif
 #include <limits.h>
 
 namespace {

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 <libtorchaudio/rnnt/gpu/math.cuh>
 

src/libtorchaudio/rnnt/gpu/gpu_kernels.cuh

@@ -1,6 +1,6 @@
 #pragma once
 
-#ifdef USE_CUDA
+#if defined(USE_CUDA) || defined(USE_ROCM)
 
 #include <cassert>
 

src/libtorchaudio/rnnt/gpu/gpu_transducer.h

@@ -1,7 +1,8 @@
 #pragma once
 
-#ifdef USE_CUDA
+#if defined(USE_CUDA) || defined(USE_ROCM)
 
+#include <libtorchaudio/cuda_compat.h>
 #include <libtorchaudio/rnnt/workspace.h>
 #include <libtorchaudio/rnnt/gpu/gpu_kernel_utils.cuh>
 #include <libtorchaudio/rnnt/gpu/gpu_kernels.cuh>

src/libtorchaudio/rnnt/gpu/math.cuh

@@ -1,10 +1,10 @@
 #pragma once
 
-#ifdef USE_CUDA
+#if defined(USE_CUDA) || defined(USE_ROCM)
 
 #include <cmath>
 
-#endif // USE_CUDA
+#endif // USE_CUDA || USE_ROCM
 
 #include <libtorchaudio/rnnt/gpu/half.cuh>
 

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 <hip/hip_fp16.h>
+#include <hip/hip_runtime.h>
+#else
 #include <cuda_fp16.h>
 #include <cuda_runtime.h>
+#endif
 #else
 #define HOST_AND_DEVICE
 #define FORCE_INLINE inline
-#endif // USE_CUDA
+#endif // USE_CUDA || USE_ROCM

src/libtorchaudio/rnnt/options.h

@@ -1,8 +1,8 @@
 #pragma once
 
-#ifdef USE_CUDA
-#include <cuda_runtime.h>
-#endif // USE_CUDA
+#if defined(USE_CUDA) || defined(USE_ROCM)
+#include <libtorchaudio/cuda_compat.h>
+#endif // USE_CUDA || USE_ROCM
 
 #include <libtorchaudio/rnnt/types.h>
 #include <ostream>
@@ -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