Split extension into custom impl and Python wrapper libraries (#1752)

* Split `libtorchaudio` and `_torchaudio`

This change extract the core implementation from `_torchaudio` to `libtorchaudio`,
so that `libtorchaudio` is reusable in TorchScript-based app.

`_torchaudio` is a wrapper around `libtorchaudio` and only provides PyBind11-based
features. (currently file-like object support in I/O)

* Removed `BUILD_LIBTORCHAUDIO` option

When invoking `cmake`, `libtorchaudio` is always built, so this option is removed.

The new assumptions around the library discoverability

- In regular OSS workflow (`pip`/`conda`-based binary installation), both `libtorchaudio` and `_torchaudio` are present.
    In this case,`libtorchaudio` has to be loaded manually with `torch.ops.load_library` and/or `torch.classes.load_library` otherwise importing `_torchaudio` would not be able to resolve the symbols defined in `libtorchaudio`.
- When `torchaudio` is deployed with PEX format (single zip file)
  - We expect that`libtorchaudio.so` exists as a file in some search path configured by client code.
  - `_torchaudio` is still importable and because we do not know where `libtorchaudio` will exist, we will let the dynamic loader resolve the dependency from `_torchaudio` to `libtorchaudio`, which should work as long as `libtorchaudio` is in a library search path (search path is not modifiable from already-running Python process).

Author: mthrok

CMakeLists.txt

@@ -50,7 +50,6 @@ endif()
 option(BUILD_SOX "Build libsox statically" ON)
 option(BUILD_KALDI "Build kaldi statically" ON)
 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)

build_tools/setup_helpers/extension.py

@@ -43,7 +43,10 @@ def _get_build(var, default=False):
 
 
 def get_ext_modules():
-    return [Extension(name='torchaudio._torchaudio', sources=[])]
+    return [
+        Extension(name='torchaudio.libtorchaudio', sources=[]),
+        Extension(name='torchaudio._torchaudio', sources=[]),
+    ]
 
 
 # Based off of
@@ -53,10 +56,19 @@ def run(self):
         try:
             subprocess.check_output(['cmake', '--version'])
         except OSError:
-            raise RuntimeError("CMake is not available.")
+            raise RuntimeError("CMake is not available.") from None
         super().run()
 
     def build_extension(self, ext):
+        # Since two library files (libtorchaudio and _torchaudio) need to be
+        # recognized by setuptools, we instantiate `Extension` twice. (see `get_ext_modules`)
+        # This leads to the situation where this `build_extension` method is called twice.
+        # However, the following `cmake` command will build all of them at the same time,
+        # so, we do not need to perform `cmake` twice.
+        # Therefore we call `cmake` only for `torchaudio._torchaudio`.
+        if ext.name != 'torchaudio._torchaudio':
+            return
+
         extdir = os.path.abspath(
             os.path.dirname(self.get_ext_fullpath(ext.name)))
 
@@ -76,7 +88,6 @@ def build_extension(self, ext):
             f"-DBUILD_KALDI:BOOL={'ON' if _BUILD_KALDI else 'OFF'}",
             f"-DBUILD_RNNT:BOOL={'ON' if _BUILD_RNNT else 'OFF'}",
             "-DBUILD_TORCHAUDIO_PYTHON_EXTENSION:BOOL=ON",
-            "-DBUILD_LIBTORCHAUDIO:BOOL=OFF",
             f"-DUSE_ROCM:BOOL={'ON' if _USE_ROCM else 'OFF'}",
             f"-DUSE_CUDA:BOOL={'ON' if _USE_CUDA else 'OFF'}",
         ]

examples/libtorchaudio/CMakeLists.txt

@@ -2,7 +2,6 @@ cmake_minimum_required(VERSION 3.5)
 
 project(libtorchaudio-cpp-example)
 
-SET(BUILD_LIBTORCHAUDIO ON CACHE BOOL "Build libtorchaudio")
 SET(BUILD_SOX ON CACHE BOOL "Build libsox into libtorchaudio")
 
 SET(BUILD_KALDI OFF CACHE BOOL "Build Kaldi into libtorchaudio")

torchaudio/__init__.py

@@ -1,39 +1,4 @@
-from torchaudio._internal import module_utils as _mod_utils  # noqa: F401
-
-if _mod_utils.is_module_available('torchaudio._torchaudio'):
-    # Note this import has two purposes
-    # 1. Make _torchaudio accessible by the other modules (regular import)
-    # 2. Register torchaudio's custom ops bound via TorchScript
-    #
-    # For 2, normally function calls `torch.ops.load_library` and `torch.classes.load_library`
-    # are used. However, in our cases, this is inconvenient and unnecessary.
-    #
-    # - Why inconvenient?
-    # When torchaudio is deployed with `pex` format, all the files are deployed as a single zip
-    # file, and the extension module is not present as a file with full path. Therefore it is not
-    # possible to pass the path to library to `torch.[ops|classes].load_library` functions.
-    #
-    # - Why unnecessary?
-    # When torchaudio extension module (C++ module) is available, it is assumed that
-    # the extension contains both TorchScript-based binding and PyBind11-based binding.*
-    # Under this assumption, simply performing `from torchaudio import _torchaudio` will load the
-    # library which contains TorchScript-based binding as well, and the functions/classes bound
-    # via TorchScript become accessible under `torch.ops` and `torch.classes`.
-    #
-    # *Note that this holds true even when these two bindings are split into two library files and
-    # the library that contains PyBind11-based binding (`_torchaudio.so` in the following diagram)
-    # depends on the other one (`libtorchaudio.so`), because when the process tries to load
-    # `_torchaudio.so` it detects undefined symbols from `libtorchaudio.so` and will automatically
-    # loads `libtorchaudio.so`. (given that the library is found in a search path)
-    #
-    # [libtorchaudio.so] <- [_torchaudio.so]
-    #
-    #
-    from torchaudio import _torchaudio  # noqa
-else:
-    import warnings
-    warnings.warn('torchaudio C++ extension is not available.')
-
+from torchaudio import _extension  # noqa: F401
 from torchaudio import (
     compliance,
     datasets,

torchaudio/_extension.py

@@ -0,0 +1,27 @@
+import os
+import warnings
+from pathlib import Path
+
+import torch
+from torchaudio._internal import module_utils as _mod_utils  # noqa: F401
+
+
+def _init_extension():
+    if not _mod_utils.is_module_available('torchaudio._torchaudio'):
+        warnings.warn('torchaudio C++ extension is not available.')
+        return
+
+    suffix = 'dll' if os.name == 'nt' else 'so'
+    path = Path(__file__).parent / f'libtorchaudio.{suffix}'
+    # In case `torchaudio` is deployed with `pex` format, this file does not exist.
+    # In this case, we expect that `libtorchaudio` is available somewhere
+    # in the search path of dynamic loading mechanism, and importing `_torchaudio`,
+    # which depends on `libtorchaudio` and dynamic loader will handle it for us.
+    if path.exists():
+        torch.ops.load_library(path)
+        torch.classes.load_library(path)
+    # This import is for initializing the methods registered via PyBind11
+    from torchaudio import _torchaudio  # noqa
+
+
+_init_extension()

torchaudio/csrc/CMakeLists.txt

@@ -1,7 +1,7 @@
 get_property(TORCHAUDIO_THIRD_PARTIES GLOBAL PROPERTY TORCHAUDIO_THIRD_PARTIES)
 
 ################################################################################
-# Stuff common to libtorchaudio and _torchaudio.so
+# libtorchaudio
 ################################################################################
 set(
   LIBTORCHAUDIO_SOURCES
@@ -11,75 +11,94 @@ set(
   )
 
 if(BUILD_RNNT)
-  set(
-    RNNT_SOURCES
+  list(
+    APPEND
+    LIBTORCHAUDIO_SOURCES
     rnnt/cpu/compute_alphas.cpp
     rnnt/cpu/compute_betas.cpp
     rnnt/cpu/compute.cpp
     rnnt/compute_alphas.cpp
     rnnt/compute_betas.cpp
     rnnt/compute.cpp
     rnnt/autograd.cpp
-  )
-
+    )
   if (USE_CUDA)
-    set(
-      CUDA_RNNT_SOURCES
+    list(
+      APPEND
+      LIBTORCHAUDIO_SOURCES
       rnnt/gpu/compute_alphas.cu
       rnnt/gpu/compute_betas.cu
       rnnt/gpu/compute.cu
-    )
-    list(APPEND RNNT_SOURCES ${CUDA_RNNT_SOURCES})
+      )
   endif()
-
-  list(APPEND LIBTORCHAUDIO_SOURCES ${RNNT_SOURCES})
 endif()
 
 if(BUILD_KALDI)
   list(APPEND LIBTORCHAUDIO_SOURCES kaldi.cpp)
 endif()
 
 if(BUILD_SOX)
-  set(
-    SOX_SOURCES
+  list(
+    APPEND
+    LIBTORCHAUDIO_SOURCES
     sox/io.cpp
     sox/utils.cpp
     sox/effects.cpp
     sox/effects_chain.cpp
     sox/types.cpp
+    )
+endif()
+
+add_library(
+  libtorchaudio
+  SHARED
+  ${LIBTORCHAUDIO_SOURCES}
+  )
+set_target_properties(libtorchaudio PROPERTIES PREFIX "")
+
+target_include_directories(
+  libtorchaudio
+  PRIVATE
+  ${PROJECT_SOURCE_DIR}
   )
-  list(APPEND LIBTORCHAUDIO_SOURCES ${SOX_SOURCES})
+
+target_link_libraries(
+  libtorchaudio
+  torch
+  ${TORCHAUDIO_THIRD_PARTIES}
+  )
+
+if (BUILD_SOX)
+  target_compile_definitions(libtorchaudio PUBLIC INCLUDE_SOX)
 endif()
 
-################################################################################
-# libtorchaudio.so
-################################################################################
-if(BUILD_LIBTORCHAUDIO)
-  add_library(
-    libtorchaudio
-    SHARED
-    ${LIBTORCHAUDIO_SOURCES}
-    )
-  set_target_properties(libtorchaudio PROPERTIES PREFIX "")
+if (BUILD_KALDI)
+  target_compile_definitions(libtorchaudio PUBLIC INCLUDE_KALDI)
+endif()
 
+if(USE_CUDA)
+  target_compile_definitions(libtorchaudio PRIVATE USE_CUDA)
   target_include_directories(
     libtorchaudio
-    PUBLIC
-    ${PROJECT_SOURCE_DIR}
+    PRIVATE
+    ${CUDA_TOOLKIT_INCLUDE}
     )
-
   target_link_libraries(
     libtorchaudio
-    ${TORCH_LIBRARIES}
-    ${TORCHAUDIO_THIRD_PARTIES}
+    ${C10_CUDA_LIBRARY}
+    ${CUDA_CUDART_LIBRARY}
     )
+endif()
 
-  install(
-    TARGETS
-    libtorchaudio
-    LIBRARY DESTINATION lib
-    )
+install(
+  TARGETS libtorchaudio
+  LIBRARY DESTINATION .
+  RUNTIME DESTINATION .
+  )
 
+if (APPLE)
+  set(TORCHAUDIO_LIBRARY libtorchaudio CACHE INTERNAL "")
+else()
   set(TORCHAUDIO_LIBRARY -Wl,--no-as-needed libtorchaudio -Wl,--as-needed CACHE INTERNAL "")
 endif()
 
@@ -104,7 +123,6 @@ if (BUILD_TORCHAUDIO_PYTHON_EXTENSION)
   add_library(
     _torchaudio
     SHARED
-    ${LIBTORCHAUDIO_SOURCES}
     ${EXTENSION_SOURCES}
     )
 
@@ -119,31 +137,12 @@ if (BUILD_TORCHAUDIO_PYTHON_EXTENSION)
     set_target_properties(_torchaudio PROPERTIES LINK_FLAGS "-undefined dynamic_lookup")
   endif()
 
-  if (BUILD_SOX)
-    target_compile_definitions(_torchaudio PRIVATE INCLUDE_SOX)
-  endif()
-
-  if (BUILD_KALDI)
-    target_compile_definitions(_torchaudio PRIVATE INCLUDE_KALDI)
-  endif()
-
-  if (USE_CUDA)
-    target_compile_definitions(_torchaudio PRIVATE USE_CUDA)
-  endif()
-
   target_include_directories(
     _torchaudio
     PRIVATE
     ${PROJECT_SOURCE_DIR}
     ${Python_INCLUDE_DIR}
     )
-  if(USE_CUDA)
-    target_include_directories(
-      _torchaudio
-      PRIVATE
-      ${CUDA_TOOLKIT_INCLUDE}
-      )
-  endif()
 
   # See https://github.com/pytorch/pytorch/issues/38122
   find_library(TORCH_PYTHON_LIBRARY torch_python PATHS "${TORCH_INSTALL_PREFIX}/lib")
@@ -155,20 +154,11 @@ if (BUILD_TORCHAUDIO_PYTHON_EXTENSION)
 
   target_link_libraries(
     _torchaudio
-    torch
+    libtorchaudio
     ${TORCH_PYTHON_LIBRARY}
-    ${TORCHAUDIO_THIRD_PARTIES}
     ${ADDITIONAL_ITEMS}
     )
 
-  if(USE_CUDA)
-    target_link_libraries(
-      _torchaudio
-      ${C10_CUDA_LIBRARY}
-      ${CUDA_CUDART_LIBRARY}
-      )
-  endif()
-
   install(
     TARGETS _torchaudio
     LIBRARY DESTINATION .