Add PackPaddedSequence (#85)

@@ -16,6 +16,7 @@ set(sources
   online-transducer-modified-beam-search-decoder.cc
   online-zipformer-transducer-model.cc
   onnx-utils.cc
+  packed-sequence.cc
   pad-sequence.cc
   parse-options.cc
   resample.cc
@@ -123,6 +124,7 @@ endif()
 if(SHERPA_ONNX_ENABLE_TESTS)
   set(sherpa_onnx_test_srcs
     cat-test.cc
+    packed-sequence-test.cc
     pad-sequence-test.cc
     slice-test.cc
     transpose-test.cc

+// sherpa-onnx/csrc/packed-sequence-test.cc
+//
+// Copyright (c)  2023  Xiaomi Corporation
+
+#include "sherpa-onnx/csrc/packed-sequence.h"
+
+#include <numeric>
+
+#include "gtest/gtest.h"
+#include "sherpa-onnx/csrc/onnx-utils.h"
+
+namespace sherpa_onnx {
+
+TEST(PackedSequence, Case1) {
+  Ort::AllocatorWithDefaultOptions allocator;
+  std::array<int64_t, 3> shape{5, 5, 4};
+  Ort::Value v =
+      Ort::Value::CreateTensor<float>(allocator, shape.data(), shape.size());
+  float *p = v.GetTensorMutableData<float>();
+
+  std::iota(p, p + shape[0] * shape[1] * shape[2], 0);
+
+  Ort::Value length =
+      Ort::Value::CreateTensor<int64_t>(allocator, shape.data(), 1);
+  int64_t *p_length = length.GetTensorMutableData<int64_t>();
+  p_length[0] = 1;
+  p_length[1] = 2;
+  p_length[2] = 3;
+  p_length[3] = 5;
+  p_length[4] = 2;
+
+  auto packed_seq = PackPaddedSequence(allocator, &v, &length);
+  fprintf(stderr, "sorted indexes: ");
+  for (auto i : packed_seq.sorted_indexes) {
+    fprintf(stderr, "%d ", static_cast<int32_t>(i));
+  }
+  fprintf(stderr, "\n");
+  // output index:   0 1 2 3 4
+  // sorted indexes: 3 2 1 4 0
+  // length:         5 3 2 2 1
+  Print3D(&v);
+  Print2D(&packed_seq.data);
+  fprintf(stderr, "batch sizes per time step: ");
+  for (auto i : packed_seq.batch_sizes) {
+    fprintf(stderr, "%d ", static_cast<int32_t>(i));
+  }
+  fprintf(stderr, "\n");
+
+  // TODO(fangjun): Check that the return value is correct
+}
+
+}  // namespace sherpa_onnx

+// sherpa-onnx/csrc/packed-sequence.cc
+//
+// Copyright (c)  2023  Xiaomi Corporation
+
+#include "sherpa-onnx/csrc/packed-sequence.h"
+
+#include <assert.h>
+
+#include <algorithm>
+#include <numeric>
+#include <utility>
+
+#include "sherpa-onnx/csrc/slice.h"
+#include "sherpa-onnx/csrc/transpose.h"
+
+namespace sherpa_onnx {
+
+static Ort::Value IndexSelect(OrtAllocator *allocator, const Ort::Value *value,
+                              const std::vector<int32_t> &sorted_indexes) {
+  auto shape = value->GetTensorTypeAndShapeInfo().GetShape();
+  assert(shape.size() == 3);
+  std::array<int64_t, 3> ans_shape{static_cast<int64_t>(sorted_indexes.size()),
+                                   shape[1], shape[2]};
+
+  Ort::Value ans = Ort::Value::CreateTensor<float>(allocator, ans_shape.data(),
+                                                   ans_shape.size());
+  float *dst = ans.GetTensorMutableData<float>();
+  const float *src = value->GetTensorData<float>();
+
+  for (auto i : sorted_indexes) {
+    const float *start = src + i * shape[1] * shape[2];
+    std::copy(start, start + shape[1] * shape[2], dst);
+    dst += shape[1] * shape[2];
+  }
+  return ans;
+}
+
+PackedSequence PackPaddedSequence(OrtAllocator *allocator,
+                                  const Ort::Value *value, Ort::Value *length) {
+  std::vector<int64_t> v_shape = value->GetTensorTypeAndShapeInfo().GetShape();
+  std::vector<int64_t> l_shape = length->GetTensorTypeAndShapeInfo().GetShape();
+
+  assert(v_shape.size() == 3);
+  assert(l_shape.size() == 3);
+  assert(v_shape[0] == l_shape[0]);
+
+  std::vector<int32_t> indexes(v_shape[0]);
+  std::iota(indexes.begin(), indexes.end(), 0);
+
+  const int64_t *p_length = length->GetTensorData<int64_t>();
+  // sort in descending order
+  std::sort(indexes.begin(), indexes.end(), [p_length](int32_t i, int32_t j) {
+    return p_length[i] > p_length[j];
+  });
+
+  int32_t n = static_cast<int32_t>(v_shape[0]);
+
+  int64_t max_T = p_length[indexes[0]];
+
+  int32_t sum_T = std::accumulate(p_length, p_length + n, 0);
+
+  std::array<int64_t, 2> data_shape{sum_T, v_shape[2]};
+
+  Ort::Value data = Ort::Value::CreateTensor<float>(
+      allocator, data_shape.data(), data_shape.size());
+  float *dst = data.GetTensorMutableData<float>();
+
+  Ort::Value tensor = IndexSelect(allocator, value, indexes);
+  tensor = Transpose01(allocator, &tensor);
+
+  // batch size at each time step
+  std::vector<int32_t> batch_sizes;
+  batch_sizes.reserve(max_T);
+
+  int64_t prev_l = 0;
+  for (int32_t i = 0; i != n; ++i) {
+    auto cur_l = p_length[indexes[n - 1 - i]];
+    assert(cur_l >= prev_l);
+    if (cur_l == prev_l) {
+      continue;
+    }
+
+    auto cur_batch_size = n - i;
+
+    Ort::Value cur_batch =
+        Slice(allocator, &tensor, prev_l, cur_l, 0, cur_batch_size);
+    auto count = cur_batch.GetTensorTypeAndShapeInfo().GetElementCount();
+    const float *src = cur_batch.GetTensorData<float>();
+    std::copy(src, src + count, dst);
+    dst += count;
+
+    for (int32_t j = prev_l; j < cur_l; ++j) {
+      batch_sizes.push_back(cur_batch_size);
+    }
+
+    prev_l = cur_l;
+  }
+
+  PackedSequence packed_seq;
+  packed_seq.sorted_indexes = std::move(indexes);
+  packed_seq.data = std::move(data);
+  packed_seq.batch_sizes = std::move(batch_sizes);
+
+  return packed_seq;
+}
+
+}  // namespace sherpa_onnx

+// sherpa-onnx/csrc/packed-sequence.h
+//
+// Copyright (c)  2023  Xiaomi Corporation
+#ifndef SHERPA_ONNX_CSRC_PACKED_SEQUENCE_H_
+#define SHERPA_ONNX_CSRC_PACKED_SEQUENCE_H_
+
+#include <vector>
+
+#include "onnxruntime_cxx_api.h"  // NOLINT
+
+namespace sherpa_onnx {
+
+struct PackedSequence {
+  std::vector<int32_t> sorted_indexes;
+  std::vector<int32_t> batch_sizes;
+  Ort::Value data{nullptr};
+};
+
+/** Similar to torch.nn.utils.rnn.pad_sequence but it supports only
+ * batch_first=true.
+ *
+ * @param allocator
+ * @param value  A 3-D tensor of shape (B, T, C). Its dtype is float.
+ * @param length A 1-D tensor of shape (B,). Its dtype is int64_t. Each
+ *               element in it specifies the valid length of the corresponding
+ *               entry in value before padding.
+ */
+PackedSequence PackPaddedSequence(OrtAllocator *allocator,
+                                  const Ort::Value *value, Ort::Value *length);
+
+}  // namespace sherpa_onnx
+
+#endif  // SHERPA_ONNX_CSRC_PACKED_SEQUENCE_H_

@@ -13,19 +13,19 @@ namespace sherpa_onnx {
 
 TEST(Slice, Slice3D) {
   Ort::AllocatorWithDefaultOptions allocator;
-  std::array<int64_t, 3> shape{3, 5, 4};
+  std::array<int64_t, 3> shape{5, 5, 4};
   Ort::Value v =
       Ort::Value::CreateTensor<float>(allocator, shape.data(), shape.size());
   float *p = v.GetTensorMutableData<float>();
 
   std::iota(p, p + shape[0] * shape[1] * shape[2], 0);
 
-  auto v1 = Slice(&v, 0, 2, 5);
-  auto v2 = Slice(&v, 1, 2, 4);
+  auto v1 = Slice(allocator, &v, 2, 4, 0, 2);
+  auto v2 = Slice(allocator, &v, 1, 3, 1, 3);
 
   Print3D(&v);
-  Print2D(&v1);
-  Print2D(&v2);
+  Print3D(&v1);
+  Print3D(&v2);
 
   // TODO(fangjun): Check that the results are correct
 }

@@ -6,29 +6,48 @@
 
 #include <assert.h>
 
+#include <algorithm>
 #include <vector>
 
 namespace sherpa_onnx {
 
 template <typename T /*=float*/>
-Ort::Value Slice(const Ort::Value *v, int32_t dim0, int32_t dim1_start,
+Ort::Value Slice(OrtAllocator *allocator, const Ort::Value *v,
+                 int32_t dim0_start, int32_t dim0_end, int32_t dim1_start,
                  int32_t dim1_end) {
   std::vector<int64_t> shape = v->GetTensorTypeAndShapeInfo().GetShape();
   assert(shape.size() == 3);
 
-  auto memory_info =
-      Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault);
+  assert(0 <= dim0_start);
+  assert(dim0_start < dim0_end);
+  assert(dim0_end <= shape[0]);
+
+  assert(0 <= dim1_start);
+  assert(dim1_start < dim1_end);
+  assert(dim1_end < shape[1]);
 
-  std::array<int64_t, 2> ans_shape{dim1_end - dim1_start, shape[2]};
   const T *src = v->GetTensorData<T>();
-  src += dim0 * shape[1] * shape[2] + dim1_start * shape[2];
 
-  return Ort::Value::CreateTensor(memory_info, const_cast<T *>(src),
-                                  ans_shape[0] * ans_shape[1], ans_shape.data(),
+  std::array<int64_t, 3> ans_shape{dim0_end - dim0_start, dim1_end - dim1_start,
+                                   shape[2]};
+
+  Ort::Value ans = Ort::Value::CreateTensor<T>(allocator, ans_shape.data(),
                                                ans_shape.size());
+  T *dst = ans.GetTensorMutableData<T>();
+  for (int32_t i = dim0_start; i != dim0_end; ++i) {
+    const T *src = v->GetTensorData<T>() + i * shape[1] * shape[2];
+    const T *start = src + dim1_start * shape[2];
+    const T *end = src + dim1_end * shape[2];
+
+    std::copy(start, end, dst);
+    dst += ans_shape[1] * ans_shape[2];
+  }
+
+  return ans;
 }
 
-template Ort::Value Slice<float>(const Ort::Value *v, int32_t dim0,
+template Ort::Value Slice<float>(OrtAllocator *allocator, const Ort::Value *v,
+                                 int32_t dim0_start, int32_t dim0_end,
                                  int32_t dim1_start, int32_t dim1_end);
 
 }  // namespace sherpa_onnx

@@ -8,21 +8,23 @@
 
 namespace sherpa_onnx {
 
-/** Get a shallow copy by slicing v.
+/** Get a deep copy by slicing v.
  *
- * It returns v[dim0, dim1_start:dim1_end]
+ * It returns v[dim0_start:dim0_end, dim1_start:dim1_end]
  *
+ * @param allocator
  * @param v A 3-D tensor. Its data type is T.
- * @param dim0  Start index of the first dimension..
+ * @param dim0_start  Start index of the first dimension..
+ * @param dim0_end    End index of the first dimension..
  * @param dim1_start Start index of the second dimension.
  * @param dim1_end  End index of the second dimension.
  *
- * @return Return a 2-D tensor of shape (dim1_end-dim1_start, v.shape[2])
- *
- * @caution: The returned tensor is a shallow copy of `v`!
+ * @return Return a 3-D tensor of shape
+ *         (dim0_end-dim0_start, dim1_end-dim1_start, v.shape[2])
  */
 template <typename T = float>
-Ort::Value Slice(const Ort::Value *v, int32_t dim0, int32_t dim1_start,
+Ort::Value Slice(OrtAllocator *allocator, const Ort::Value *v,
+                 int32_t dim0_start, int32_t dim0_end, int32_t dim1_start,
                  int32_t dim1_end);
 }  // namespace sherpa_onnx