Use PyBind for binding utilities (#2956)

Summary:
Pull Request resolved: https://github.com/pytorch/audio/pull/2956

Merge utility binding

This commit updates the utility binding, so that we can use `is_module_available()`
for checking the existence of extension modules.

To ensure the existence of module, this commit migrates the binding of utility functions
to PyBind11.

Going forward, we should use TorchBind for ops that we want to support TorchScript,
otherwise default to PyBind11. (PyBind has advantage of not copying strings.)

Reviewed By: hwangjeff

Differential Revision: D42355992

fbshipit-source-id: 4c71d65b24a0882a38a80dc097d45ba72b4c4a6b

torchaudio/_extension.py

@@ -57,6 +57,7 @@ def _load_lib(lib: str) -> bool:
     return True
 
 
+_IS_TORCHAUDIO_EXT_AVAILABLE = is_module_available("torchaudio.lib._torchaudio")
 _FFMPEG_INITIALIZED = False
 
 
@@ -101,7 +102,7 @@ def _init_extension():
                 except Exception:
                     pass
 
-    if is_module_available("torchaudio.lib._torchaudio"):
+    if _IS_TORCHAUDIO_EXT_AVAILABLE:
         try:
             _load_lib("libtorchaudio")
             import torchaudio.lib._torchaudio  # noqa
@@ -125,9 +126,12 @@ def _init_extension():
 
 
 def _check_cuda_version():
-    if not _get_lib_path("libtorchaudio").exists():
+    if not _IS_TORCHAUDIO_EXT_AVAILABLE:
         return None
-    version = torch.ops.torchaudio.cuda_version()
+
+    import torchaudio.lib._torchaudio
+
+    version = torchaudio.lib._torchaudio.cuda_version()
     if version is not None and torch.version.cuda is not None:
         version_str = str(version)
         ta_version = f"{version_str[:-3]}.{version_str[-2]}"

torchaudio/csrc/pybind/pybind.cpp

@@ -4,15 +4,6 @@
 namespace torchaudio {
 namespace {
 
-// Note
-// These functions are not intended for a real usecase.
-// They are accessible via TorchBind.
-// It is beneficial to have _torchaudio that is linked to libtorchaudio,
-// when torchaudio is deployed with PEX format, where the library location
-// is not in torchaudio/lib. But somewhere in LD_LIBRARY_PATH.
-// In this case, attempt to import _torchaudio will automatically resolves
-// libtorchaudio, if _torchaudio is linked to libtorchaudio.
-
 PYBIND11_MODULE(_torchaudio, m) {
   m.def("is_kaldi_available", &is_kaldi_available, "");
   m.def("cuda_version", &cuda_version, "");

torchaudio/csrc/utils.cpp

@@ -22,11 +22,5 @@ c10::optional<int64_t> cuda_version() {
   return {};
 #endif
 }
-namespace {
-TORCH_LIBRARY_FRAGMENT(torchaudio, m) {
-  m.def("torchaudio::is_kaldi_available", &is_kaldi_available);
-  m.def("torchaudio::cuda_version", &cuda_version);
-}
-} // namespace
 
 } // namespace torchaudio