build only one extension.

Author: vincentqb

build_tools/setup_helpers/extension.py

@@ -18,6 +18,8 @@
 _CSRC_DIR = _ROOT_DIR / 'torchaudio' / 'csrc'
 _TP_BASE_DIR = _ROOT_DIR / 'third_party'
 _TP_INSTALL_DIR = _TP_BASE_DIR / 'install'
+_TRANSDUCER_BUILD_DIR = _TP_BASE_DIR / 'build' / 'warp_transducer'
+_TRANSDUCER_BASE_DIR = _TP_BASE_DIR / 'warp_transducer' / 'submodule'
 
 
 def _get_build_sox():
@@ -64,6 +66,7 @@ def _get_srcs():
 def _get_include_dirs():
     dirs = [
         str(_ROOT_DIR),
+        str(_TRANSDUCER_BASE_DIR / 'include'),
     ]
     if _BUILD_SOX:
         dirs.append(str(_TP_INSTALL_DIR / 'include'))
@@ -94,6 +97,7 @@ def _get_extra_objects():
         ]
         for lib in libs:
             objs.append(str(_TP_INSTALL_DIR / 'lib' / lib))
+    objs.append(str(_TRANSDUCER_BUILD_DIR / 'libwarprnnt.a'))
     return objs
 
 
@@ -132,60 +136,11 @@ def get_ext_modules(debug=False):
             extra_objects=_get_extra_objects(),
             extra_link_args=_get_ela(debug),
         ),
-        _get_transducer_module(),
     ]
 
 
 class BuildExtension(TorchBuildExtension):
     def build_extension(self, ext):
         if ext.name == _EXT_NAME and _BUILD_SOX:
             _build_third_party()
-        if ext.name == _TRANSDUCER_NAME:
-            _build_transducer()
         super().build_extension(ext)
-
-
-_TRANSDUCER_NAME = '_warp_transducer'
-_TP_TRANSDUCER_BASE_DIR = _ROOT_DIR / 'third_party' / 'warp_transducer'
-
-
-def _build_transducer():
-    build_dir = str(_TP_TRANSDUCER_BASE_DIR / 'submodule' / 'build')
-    os.makedirs(build_dir, exist_ok=True)
-    subprocess.run(
-        args=['cmake', str(_TP_TRANSDUCER_BASE_DIR), '-DWITH_OMP=OFF'],
-        cwd=build_dir,
-        check=True,
-    )
-    subprocess.run(
-        args=['cmake', '--build', '.'],
-        cwd=build_dir,
-        check=True,
-    )
-
-
-def _get_transducer_module():
-    extra_compile_args = [
-        '-fPIC',
-        '-std=c++14',
-    ]
-
-    libraries = ['warprnnt']
-
-    source_paths = [
-        _TP_TRANSDUCER_BASE_DIR / 'binding.cpp',
-        _TP_TRANSDUCER_BASE_DIR / 'submodule' / 'pytorch_binding' / 'src' / 'binding.cpp',
-    ]
-    build_path = _TP_TRANSDUCER_BASE_DIR / 'submodule' / 'build'
-    include_path = _TP_TRANSDUCER_BASE_DIR / 'submodule' / 'include'
-
-    return CppExtension(
-        name=_TRANSDUCER_NAME,
-        sources=[os.path.realpath(path) for path in source_paths],
-        libraries=libraries,
-        include_dirs=[os.path.realpath(include_path)],
-        library_dirs=[os.path.realpath(build_path)],
-        extra_compile_args=extra_compile_args,
-        extra_objects=[str(build_path / f'lib{lib}.a') for lib in libraries],
-        extra_link_args=['-Wl,-rpath,' + os.path.realpath(build_path)],
-    )

third_party/CMakeLists.txt

@@ -88,3 +88,5 @@ ExternalProject_Add(libsox
   # See https://github.com/pytorch/audio/pull/1026
   CONFIGURE_COMMAND ${CMAKE_CURRENT_SOURCE_DIR}/build_codec_helper.sh ${CMAKE_CURRENT_SOURCE_DIR}/src/libsox/configure ${COMMON_ARGS} --with-lame --with-flac --with-mad --with-oggvorbis --without-alsa --without-coreaudio --without-png --without-oss --without-sndfile --with-opus --with-amrwb --with-amrnb --disable-openmp
 )
+
+add_subdirectory(warp_transducer)

third_party/warp_transducer/binding.cpp

@@ -0,0 +1,85 @@
+#include <iostream>
+#include <numeric>
+
+#include <torch/extension.h>
+#include "rnnt.h"
+
+int cpu_rnnt(torch::Tensor acts,
+            torch::Tensor labels,
+            torch::Tensor input_lengths,
+            torch::Tensor label_lengths,
+            torch::Tensor costs,
+            torch::Tensor grads,
+            int blank_label,
+            int num_threads) {
+
+    int maxT = acts.size(1);
+    int maxU = acts.size(2);
+    int minibatch_size = acts.size(0);
+    int alphabet_size = acts.size(3);
+
+    rnntOptions options;
+    memset(&options, 0, sizeof(options));
+    options.maxT = maxT;
+    options.maxU = maxU;
+    options.blank_label = blank_label;
+    options.batch_first = true;
+    options.loc = RNNT_CPU;
+    options.num_threads = num_threads;
+
+    // have to use at least one
+    options.num_threads = std::max(options.num_threads, (unsigned int) 1);
+
+    size_t cpu_size_bytes = 0;
+    switch (acts.type().scalarType()) {
+      case torch::ScalarType::Float:
+        {
+        get_workspace_size(maxT, maxU, minibatch_size,
+                           false, &cpu_size_bytes);
+
+        float* cpu_workspace = (float*) new unsigned char[cpu_size_bytes];
+        compute_rnnt_loss(acts.data<float>(), grads.data<float>(),
+                         labels.data<int>(), label_lengths.data<int>(),
+                         input_lengths.data<int>(), alphabet_size,
+                         minibatch_size, costs.data<float>(),
+                         cpu_workspace, options);
+
+        delete cpu_workspace;
+        return 0;
+        }
+      case torch::ScalarType::Double:
+        {
+        get_workspace_size(maxT, maxU, minibatch_size,
+                           false, &cpu_size_bytes,
+                           sizeof(double));
+
+        double* cpu_workspace = (double*) new unsigned char[cpu_size_bytes];
+        compute_rnnt_loss_fp64(acts.data<double>(), grads.data<double>(),
+                         labels.data<int>(), label_lengths.data<int>(),
+                         input_lengths.data<int>(), alphabet_size,
+                         minibatch_size, costs.data<double>(),
+                         cpu_workspace, options);
+
+        delete cpu_workspace;
+        return 0;
+        }
+      default:
+        std::cerr << __FILE__ << ':' << __LINE__ << ": " << "unsupported data type" << std::endl;
+    }
+    return -1;
+}
+
+TORCH_LIBRARY(warprnnt_pytorch_warp_rnnt, m) {
+    m.def("rnnt(Tensor acts,"
+               "Tensor labels,"
+               "Tensor input_lengths,"
+               "Tensor label_lengths,"
+               "Tensor costs,"
+               "Tensor grads,"
+               "int blank_label,"
+               "int num_threads) -> int");
+}
+
+TORCH_LIBRARY_IMPL(warprnnt_pytorch_warp_rnnt, CPU, m) {
+    m.impl("rnnt", &cpu_rnnt);
+}