Add OpenMP support

Author: mthrok

CMakeLists.txt

@@ -53,7 +53,7 @@ option(BUILD_RNNT "Enable RNN transducer" ON)
 option(BUILD_TORCHAUDIO_PYTHON_EXTENSION "Build Python extension" OFF)
 option(USE_CUDA "Enable CUDA support" OFF)
 option(USE_ROCM "Enable ROCM support" OFF)
-
+option(USE_OPENMP "Enable OpenMP support" OFF)
 
 # check that USE_CUDA and USE_ROCM are not set at the same time
 if(USE_CUDA AND USE_ROCM)
@@ -74,6 +74,13 @@ endif()
 
 find_package(Torch REQUIRED)
 
+if(USE_OPENMP)
+  find_package(OpenMP)
+  if (NOT OpenMP_CXX_FOUND)
+    message(WARNING "OpenMP support is disabled.")
+  endif()
+endif()
+
 # TORCH_CXX_FLAGS contains the same -D_GLIBCXX_USE_CXX11_ABI value as PyTorch
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall ${TORCH_CXX_FLAGS}")
 

build_tools/setup_helpers/extension.py

@@ -39,6 +39,7 @@ def _get_build(var, default=False):
 _BUILD_RNNT = _get_build("BUILD_RNNT", True)
 _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)
+_USE_OPENMP = _get_build("USE_OPENMP", True) and 'ATen parallel backend: OpenMP' in torch.__config__.parallel_info()
 _TORCH_CUDA_ARCH_LIST = os.environ.get('TORCH_CUDA_ARCH_LIST', None)
 
 
@@ -90,6 +91,7 @@ def build_extension(self, ext):
             "-DBUILD_TORCHAUDIO_PYTHON_EXTENSION:BOOL=ON",
             f"-DUSE_ROCM:BOOL={'ON' if _USE_ROCM else 'OFF'}",
             f"-DUSE_CUDA:BOOL={'ON' if _USE_CUDA else 'OFF'}",
+            f"-DUSE_OPENMP:BOOL={'ON' if _USE_OPENMP else 'OFF'}",
         ]
         build_args = [
             '--target', 'install'

torchaudio/csrc/CMakeLists.txt

@@ -90,6 +90,10 @@ if(USE_CUDA)
     )
 endif()
 
+if(OpenMP_CXX_FOUND)
+  target_link_libraries(libtorchaudio OpenMP::OpenMP_CXX)
+endif()
+
 install(
   TARGETS libtorchaudio
   LIBRARY DESTINATION .