[ROCM] fix build error (#1729)

* fix build error on ROCM

* Update CMakeLists.txt
Co-authored-by: Nikita Shulga <nikita.shulga@gmail.com>

* address comments and fix cuda detction on rocm
Co-authored-by: Nikita Shulga <nikita.shulga@gmail.com>

CMakeLists.txt

@@ -18,13 +18,6 @@ endif()
 
 project(torchaudio)
 
-# Find the HIP package, set the HIP paths, load the HIP CMake.
-if(USE_ROCM)
-  include(cmake/LoadHIP.cmake)
-  if(NOT PYTORCH_FOUND_HIP)
-    set(USE_ROCM OFF)
-  endif()
-endif()
 
 # check and set CMAKE_CXX_STANDARD
 string(FIND "${CMAKE_CXX_FLAGS}" "-std=c++" env_cxx_standard)
@@ -60,6 +53,21 @@ option(BUILD_RNNT "Enable RNN transducer" ON)
 option(BUILD_LIBTORCHAUDIO "Build C++ Library" ON)
 option(BUILD_TORCHAUDIO_PYTHON_EXTENSION "Build Python extension" OFF)
 option(USE_CUDA "Enable CUDA support" OFF)
+option(USE_ROCM "Enable ROCM support" OFF)
+
+
+# check that USE_CUDA and USE_ROCM are not set at the same time
+if(USE_CUDA AND USE_ROCM)
+  message(FATAL "CUDA and ROCm are mutually exclusive")
+endif()
+
+if(USE_ROCM) 
+  # Find the HIP package, set the HIP paths, load the HIP CMake.
+  include(cmake/LoadHIP.cmake)
+  if(NOT PYTORCH_FOUND_HIP)
+    set(USE_ROCM OFF)
+  endif()
+endif()
 
 if(USE_CUDA)
   enable_language(CUDA)

build_tools/setup_helpers/extension.py

@@ -37,8 +37,8 @@ def _get_build(var, default=False):
 _BUILD_SOX = False if platform.system() == 'Windows' else _get_build("BUILD_SOX", True)
 _BUILD_KALDI = False if platform.system() == 'Windows' else _get_build("BUILD_KALDI", True)
 _BUILD_RNNT = _get_build("BUILD_RNNT", True)
-_USE_ROCM = _get_build("USE_ROCM")
-_USE_CUDA = _get_build("USE_CUDA", torch.cuda.is_available())
+_USE_ROCM = _get_build("USE_ROCM", torch.cuda.is_available() and torch.version.hip is not None)
+_USE_CUDA = _get_build("USE_CUDA", torch.cuda.is_available() and torch.version.hip is None)
 _TORCH_CUDA_ARCH_LIST = os.environ.get('TORCH_CUDA_ARCH_LIST', None)