[Python STABLE ABI] Port cuctc extension to libtorch stable ABI

pearu · pearu · commit 03151f6ecfae · 2025-12-14T00:49:08.000+02:00
diff --git a/src/libtorchaudio/cuctc/CMakeLists.txt b/src/libtorchaudio/cuctc/CMakeLists.txt
@@ -34,10 +34,10 @@ torchaudio_library(
 
 if (BUILD_TORCHAUDIO_PYTHON_EXTENSION)
   torchaudio_extension(
-    pybind11_prefixctc
+    torchaudio_prefixctc
     src/python_binding.cpp
     "${CMAKE_CURRENT_SOURCE_DIR}"
-    "libctc_prefix_decoder;${additional_libs}"
+    "libctc_prefix_decoder;torch;${additional_libs}"
     ""
     )
 endif()
diff --git a/src/libtorchaudio/cuctc/src/python_binding.cpp b/src/libtorchaudio/cuctc/src/python_binding.cpp
@@ -24,6 +24,11 @@
 // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+#include <torch/csrc/stable/library.h>
+#include <torch/csrc/stable/ops.h>
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/ScalarType.h>
+
 // Stuff in third_party/ is considered non-ABI-stable
 // (https://github.com/pytorch/pytorch/issues/169893), but
 // it is safe to temporarily disable TORCH_TARGET_VERSION for pybind11
@@ -44,23 +49,43 @@
 #include <vector>
 #include "include/ctc_prefix_decoder.h"
 
-namespace py = pybind11;
+using torch::headeronly::ScalarType;
+using torch::stable::Tensor;
 
-std::tuple<size_t, std::vector<std::vector<std::pair<float, std::vector<int>>>>>
-ctc_prefix_decoder_batch_wrapper(
+std::tuple<size_t, Tensor, Tensor, Tensor> ctc_prefix_decoder_batch_wrapper(
     std::uintptr_t n_inter_data,
-    std::uintptr_t buff_ptr,
-    size_t buff_size,
-    std::uintptr_t pp,
-    std::uintptr_t seq_len_ptr,
-    const std::vector<int>& pp_sizes,
-    const std::vector<int>& pp_strides,
+    Tensor buff,
+    Tensor log_prob,
+    Tensor encoder_out_lens,
     int beam,
     int blid,
     int spid,
     float thresold) {
-  using SCORE_TYPE =
-      std::vector<std::vector<std::pair<float, std::vector<int>>>>;
+  STD_TORCH_CHECK(
+      encoder_out_lens.scalar_type() == ScalarType::Int,
+      "encoder_out_lens must be torch.int32");
+  STD_TORCH_CHECK(
+      log_prob.scalar_type() == ScalarType::Float,
+      "log_prob must be torch.float32");
+  STD_TORCH_CHECK(log_prob.is_cuda(), "log_prob must be cuda tensor");
+  STD_TORCH_CHECK(
+      encoder_out_lens.is_cuda(), "encoder_out_lens must be cuda tensor");
+  STD_TORCH_CHECK(
+      log_prob.get_device_index() == encoder_out_lens.get_device_index(),
+      "log_prob and encoder_out_lens must be on the same device");
+  STD_TORCH_CHECK(log_prob.is_contiguous(), "log_prob must be contiguous");
+  STD_TORCH_CHECK(
+      encoder_out_lens.is_contiguous(), "encoder_out_lens must be contiguous");
+  if (buff.numel() > 0) {
+    STD_TORCH_CHECK(buff.is_cuda(), "buff must be cuda tensor");
+    STD_TORCH_CHECK(
+        log_prob.get_device_index() == buff.get_device_index(),
+        "log_prob and buff must be on the same device");
+  }
+  auto pp_sizes_ = log_prob.sizes();
+  auto pp_strides_ = log_prob.strides();
+  std::vector<int> pp_sizes(std::begin(pp_sizes_), std::end(pp_sizes_));
+  std::vector<int> pp_strides(std::begin(pp_strides_), std::end(pp_strides_));
   cu_ctc::InternalData* inter_data = (cu_ctc::InternalData*)(n_inter_data);
   auto [require_size, max_select_seq_len] =
       cu_ctc::calculate_require_buff_and_init_internal_data(
@@ -69,46 +94,52 @@ ctc_prefix_decoder_batch_wrapper(
           pp_sizes[1],
           pp_sizes[2],
           beam,
-          buff_ptr,
-          buff_size,
-          (float*)pp,
-          (int*)seq_len_ptr,
+          reinterpret_cast<std::uintptr_t>(buff.data_ptr()),
+          buff.size(0),
+          log_prob.mutable_data_ptr<float>(),
+          encoder_out_lens.mutable_data_ptr<int>(),
           pp_sizes,
           pp_strides,
           blid,
           thresold);
   if (require_size > 0) {
-    return std::make_tuple(require_size, SCORE_TYPE{});
+    return std::make_tuple(require_size, Tensor{}, Tensor{}, Tensor{});
   }
   int batch_size = pp_sizes[0];
-  std::vector<int> list_data(batch_size * beam * max_select_seq_len);
-  std::vector<int> len_data(batch_size * beam);
-  std::vector<float> score(batch_size * beam);
+  Tensor list_data = torch::stable::empty(
+      {batch_size, beam, max_select_seq_len}, ScalarType::Int);
+  Tensor len_data = torch::stable::empty({batch_size, beam}, ScalarType::Int);
+  Tensor score = torch::stable::empty({batch_size, beam}, ScalarType::Float);
   cu_ctc::ctc_beam_search_decoder_batch_gpu(
-      inter_data, blid, spid, list_data.data(), len_data.data(), score.data());
-  SCORE_TYPE score_hyps{};
-  score_hyps.reserve(batch_size);
-  for (int b = 0; b < batch_size; b++) {
-    score_hyps.push_back(std::vector<std::pair<float, std::vector<int>>>{});
-    score_hyps.back().reserve(beam);
-    for (int beam_id = 0; beam_id < beam; beam_id++) {
-      int len = len_data[b * beam + beam_id];
-      int offset = b * beam * max_select_seq_len + beam_id * max_select_seq_len;
-      std::vector<int> clist(
-          list_data.data() + offset, list_data.data() + offset + len);
-      score_hyps.back().push_back(
-          std::pair{score[b * beam + beam_id], std::move(clist)});
-    }
-  }
-  return std::make_tuple(require_size, std::move(score_hyps));
+      inter_data,
+      blid,
+      spid,
+      list_data.mutable_data_ptr<int>(),
+      len_data.mutable_data_ptr<int>(),
+      score.mutable_data_ptr<float>());
+  return std::make_tuple(
+      require_size,
+      std::move(list_data),
+      std::move(len_data),
+      std::move(score));
 }
 
-PYBIND11_MODULE(pybind11_prefixctc, m) {
-  m.doc() = "none";
+PYBIND11_MODULE(torchaudio_prefixctc, m) {}
+
+STABLE_TORCH_LIBRARY_FRAGMENT(torchaudio_prefixctc, m) {
   m.def(
+      "ctc_beam_search_decoder_batch_gpu_v2(int interal_data_ptr, Tensor memory, Tensor log_prob, Tensor encoder_out_lens, int beam, int blid, int spid, float thresold) -> (int, Tensor, Tensor, Tensor)");
+  m.def("prefixCTC_alloc(int stream) -> int");
+  m.def("prefixCTC_free(int interal_data_ptr) -> ()");
+}
+
+STABLE_TORCH_LIBRARY_IMPL(torchaudio_prefixctc, CompositeExplicitAutograd, m) {
+  m.impl("prefixCTC_alloc", TORCH_BOX(&cu_ctc::prefixCTC_alloc));
+  m.impl("prefixCTC_free", TORCH_BOX(&cu_ctc::prefixCTC_free));
+}
+
+STABLE_TORCH_LIBRARY_IMPL(torchaudio_prefixctc, CUDA, m) {
+  m.impl(
       "ctc_beam_search_decoder_batch_gpu_v2",
-      &ctc_prefix_decoder_batch_wrapper,
-      "ctc prefix decoder  v2 computing on GPU");
-  m.def("prefixCTC_alloc", &cu_ctc::prefixCTC_alloc, "allocate internal data");
-  m.def("prefixCTC_free", &cu_ctc::prefixCTC_free, "free internal data");
+      TORCH_BOX(&ctc_prefix_decoder_batch_wrapper));
 }
diff --git a/src/torchaudio/models/decoder/_cuda_ctc_decoder.py b/src/torchaudio/models/decoder/_cuda_ctc_decoder.py
@@ -8,7 +8,9 @@
 import torchaudio
 
 torchaudio._extension._load_lib("libctc_prefix_decoder")
-import torchaudio.lib.pybind11_prefixctc as cuctc
+import torchaudio.lib.torchaudio_prefixctc
+
+cuctc = torch.ops.torchaudio_prefixctc
 
 
 __all__ = ["CUCTCHypothesis", "CUCTCDecoder", "cuda_ctc_decoder"]
@@ -102,51 +104,36 @@ def __call__(self, log_prob: torch.Tensor, encoder_out_lens: torch.Tensor):
             List[List[CUCTCHypothesis]]:
                 List of sorted best hypotheses for each audio sequence in the batch.
         """
-        if not encoder_out_lens.dtype == torch.int32:
-            raise AssertionError("encoder_out_lens must be torch.int32")
-        if not log_prob.dtype == torch.float32:
-            raise AssertionError("log_prob must be torch.float32")
-        if not (log_prob.is_cuda and encoder_out_lens.is_cuda):
-            raise AssertionError("inputs must be cuda tensors")
-        if not (log_prob.is_contiguous() and encoder_out_lens.is_contiguous()):
-            raise AssertionError("input tensors must be contiguous")
-        required_size, score_hyps = cuctc.ctc_beam_search_decoder_batch_gpu_v2(
+        required_size, list_data, len_data, score = cuctc.ctc_beam_search_decoder_batch_gpu_v2(
             self.internal_data,
-            self.memory.data_ptr(),
-            self.memory.size(0),
-            log_prob.data_ptr(),
-            encoder_out_lens.data_ptr(),
-            log_prob.size(),
-            log_prob.stride(),
+            self.memory,
+            log_prob,
+            encoder_out_lens,
             self.beam_size,
             self.blank_id,
             self.space_id,
             self.blank_skip_threshold,
         )
         if required_size > 0:
-            self.memory = torch.empty(required_size, dtype=torch.int8, device=log_prob.device).contiguous()
-            _, score_hyps = cuctc.ctc_beam_search_decoder_batch_gpu_v2(
+            self.memory = torch.empty(required_size, dtype=torch.int8, device=log_prob.device)
+            required_size, list_data, len_data, score = cuctc.ctc_beam_search_decoder_batch_gpu_v2(
                 self.internal_data,
-                self.memory.data_ptr(),
-                self.memory.size(0),
-                log_prob.data_ptr(),
-                encoder_out_lens.data_ptr(),
-                log_prob.size(),
-                log_prob.stride(),
+                self.memory,
+                log_prob,
+                encoder_out_lens,
                 self.beam_size,
                 self.blank_id,
                 self.space_id,
                 self.blank_skip_threshold,
             )
-        batch_size = len(score_hyps)
         hypos = []
-        for i in range(batch_size):
+        for i in range(log_prob.shape[0]):
             hypos.append(
                 [
                     CUCTCHypothesis(
-                        tokens=score_hyps[i][j][1],
-                        words=[self.vocab_list[word_id] for word_id in score_hyps[i][j][1]],
-                        score=score_hyps[i][j][0],
+                        tokens=list_data[i][j][: len_data[i][j]],
+                        words=[self.vocab_list[word_id] for word_id in list_data[i][j][: len_data[i][j]]],
+                        score=score[i][j],
                     )
                     for j in range(self.nbest)
                 ]
diff --git a/tools/setup_helpers/extension.py b/tools/setup_helpers/extension.py
@@ -52,7 +52,7 @@ def get_ext_modules():
         modules.extend(
             [
                 Extension(name="torchaudio.lib.libctc_prefix_decoder", sources=[]),
-                Extension(name="torchaudio.lib.pybind11_prefixctc", sources=[]),
+                Extension(name="torchaudio.lib.torchaudio_prefixctc", sources=[]),
             ]
         )
     return modules

Original file line number	Diff line number	Diff line change
`@@ -52,7 +52,7 @@ def get_ext_modules():`
`52`	`52`	`modules.extend(`
`53`	`53`	`[`
`54`	`54`	`Extension(name="torchaudio.lib.libctc_prefix_decoder", sources=[]),`
`55`		`- Extension(name="torchaudio.lib.pybind11_prefixctc", sources=[]),`
	`55`	`+ Extension(name="torchaudio.lib.torchaudio_prefixctc", sources=[]),`
`56`	`56`	`]`
`57`	`57`	`)`
`58`	`58`	`return modules`