Implement HIP build version using TORCH_HIP_VERSION

src/libtorchaudio/utils.cpp

@@ -15,7 +15,10 @@ bool is_align_available() {
 }
 
 std::optional<int64_t> cuda_version() {
-#ifdef USE_CUDA
+#if defined(TORCH_HIP_VERSION)
+  // TORCH_HIP_VERSION = {ROCM_VERSION[0] * 100 + ROCM_VERSION[1]}
+  return static_cast<int64_t>(TORCH_HIP_VERSION);
+#elif defined(USE_CUDA)
   return CUDA_VERSION;
 #else
   return {};

src/torchaudio/_extension/utils.py

@@ -126,4 +126,16 @@ def _check_cuda_version():
                 f"PyTorch has CUDA version {t_version} whereas TorchAudio has CUDA version {ta_version}. "
                 "Please install the TorchAudio version that matches your PyTorch version."
             )
+    elif version is not None and torch.version.hip is not None:
+        ta_major = int(version) // 100
+        ta_minor = int(version) % 100
+        ta_version = f"{ta_major}.{ta_minor}"
+        hip_parts = torch.version.hip.split(".")
+        t_version = f"{hip_parts[0]}.{hip_parts[1]}"
+        if ta_version != t_version:
+            raise RuntimeError(
+                "Detected that PyTorch and TorchAudio were compiled with different ROCm HIP versions. "
+                f"PyTorch has HIP version {t_version} whereas TorchAudio has HIP version {ta_version}. "
+                "Please install the TorchAudio version that matches your PyTorch version."
+            )
     return version

tools/setup_helpers/extension.py

@@ -3,7 +3,12 @@
 from pathlib import Path
 
 import torch
-from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension
+from torch.utils.cpp_extension import (
+    TORCH_HIP_VERSION,
+    BuildExtension,
+    CppExtension,
+    CUDAExtension,
+)
 
 __all__ = [
     "get_ext_modules",
@@ -74,14 +79,10 @@ def get_ext_modules():
     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.
-        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
-            extra_compile_args["cxx"].append("-DTORCH_HIP_VERSION=" + str(torch_hip_version))
-            extra_compile_args["nvcc"].append("-DTORCH_HIP_VERSION=" + str(torch_hip_version))
+        if TORCH_HIP_VERSION is not None:
+            flag = f"-DTORCH_HIP_VERSION={int(TORCH_HIP_VERSION)}"
+            extra_compile_args["cxx"].append(flag)
+            extra_compile_args["nvcc"].append(flag)
 
     sources = [
         "utils.cpp",