Support building GPU-capable sherpa-onnx on Linux aarch64. (#1500)

Thanks to @Peakyxh for providing pre-built onnxruntime libraries 
with CUDA support for Linux aarch64.

Tested on Jetson nano b01
This commit is contained in:
Fangjun Kuang
2024-11-01 11:16:28 +08:00
committed by GitHub
parent a3c89aa0d8
commit 9ab89c33bc
41 changed files with 537 additions and 291 deletions

View File

@@ -185,6 +185,9 @@ std::vector<Ort::Value> OnlineZipformer2TransducerModel::StackStates(
std::vector<const Ort::Value *> buf(batch_size);
auto allocator =
const_cast<OnlineZipformer2TransducerModel *>(this)->allocator_;
std::vector<Ort::Value> ans;
int32_t num_states = static_cast<int32_t>(states[0].size());
ans.reserve(num_states);
@@ -194,42 +197,42 @@ std::vector<Ort::Value> OnlineZipformer2TransducerModel::StackStates(
for (int32_t n = 0; n != batch_size; ++n) {
buf[n] = &states[n][6 * i];
}
auto v = Cat(allocator_, buf, 1);
auto v = Cat(allocator, buf, 1);
ans.push_back(std::move(v));
}
{
for (int32_t n = 0; n != batch_size; ++n) {
buf[n] = &states[n][6 * i + 1];
}
auto v = Cat(allocator_, buf, 1);
auto v = Cat(allocator, buf, 1);
ans.push_back(std::move(v));
}
{
for (int32_t n = 0; n != batch_size; ++n) {
buf[n] = &states[n][6 * i + 2];
}
auto v = Cat(allocator_, buf, 1);
auto v = Cat(allocator, buf, 1);
ans.push_back(std::move(v));
}
{
for (int32_t n = 0; n != batch_size; ++n) {
buf[n] = &states[n][6 * i + 3];
}
auto v = Cat(allocator_, buf, 1);
auto v = Cat(allocator, buf, 1);
ans.push_back(std::move(v));
}
{
for (int32_t n = 0; n != batch_size; ++n) {
buf[n] = &states[n][6 * i + 4];
}
auto v = Cat(allocator_, buf, 0);
auto v = Cat(allocator, buf, 0);
ans.push_back(std::move(v));
}
{
for (int32_t n = 0; n != batch_size; ++n) {
buf[n] = &states[n][6 * i + 5];
}
auto v = Cat(allocator_, buf, 0);
auto v = Cat(allocator, buf, 0);
ans.push_back(std::move(v));
}
}
@@ -238,7 +241,7 @@ std::vector<Ort::Value> OnlineZipformer2TransducerModel::StackStates(
for (int32_t n = 0; n != batch_size; ++n) {
buf[n] = &states[n][num_states - 2];
}
auto v = Cat(allocator_, buf, 0);
auto v = Cat(allocator, buf, 0);
ans.push_back(std::move(v));
}
@@ -246,7 +249,7 @@ std::vector<Ort::Value> OnlineZipformer2TransducerModel::StackStates(
for (int32_t n = 0; n != batch_size; ++n) {
buf[n] = &states[n][num_states - 1];
}
auto v = Cat<int64_t>(allocator_, buf, 0);
auto v = Cat<int64_t>(allocator, buf, 0);
ans.push_back(std::move(v));
}
return ans;
@@ -261,12 +264,15 @@ OnlineZipformer2TransducerModel::UnStackStates(
int32_t batch_size = states[0].GetTensorTypeAndShapeInfo().GetShape()[1];
auto allocator =
const_cast<OnlineZipformer2TransducerModel *>(this)->allocator_;
std::vector<std::vector<Ort::Value>> ans;
ans.resize(batch_size);
for (int32_t i = 0; i != m; ++i) {
{
auto v = Unbind(allocator_, &states[i * 6], 1);
auto v = Unbind(allocator, &states[i * 6], 1);
assert(static_cast<int32_t>(v.size()) == batch_size);
for (int32_t n = 0; n != batch_size; ++n) {
@@ -274,7 +280,7 @@ OnlineZipformer2TransducerModel::UnStackStates(
}
}
{
auto v = Unbind(allocator_, &states[i * 6 + 1], 1);
auto v = Unbind(allocator, &states[i * 6 + 1], 1);
assert(static_cast<int32_t>(v.size()) == batch_size);
for (int32_t n = 0; n != batch_size; ++n) {
@@ -282,7 +288,7 @@ OnlineZipformer2TransducerModel::UnStackStates(
}
}
{
auto v = Unbind(allocator_, &states[i * 6 + 2], 1);
auto v = Unbind(allocator, &states[i * 6 + 2], 1);
assert(static_cast<int32_t>(v.size()) == batch_size);
for (int32_t n = 0; n != batch_size; ++n) {
@@ -290,7 +296,7 @@ OnlineZipformer2TransducerModel::UnStackStates(
}
}
{
auto v = Unbind(allocator_, &states[i * 6 + 3], 1);
auto v = Unbind(allocator, &states[i * 6 + 3], 1);
assert(static_cast<int32_t>(v.size()) == batch_size);
for (int32_t n = 0; n != batch_size; ++n) {
@@ -298,7 +304,7 @@ OnlineZipformer2TransducerModel::UnStackStates(
}
}
{
auto v = Unbind(allocator_, &states[i * 6 + 4], 0);
auto v = Unbind(allocator, &states[i * 6 + 4], 0);
assert(static_cast<int32_t>(v.size()) == batch_size);
for (int32_t n = 0; n != batch_size; ++n) {
@@ -306,7 +312,7 @@ OnlineZipformer2TransducerModel::UnStackStates(
}
}
{
auto v = Unbind(allocator_, &states[i * 6 + 5], 0);
auto v = Unbind(allocator, &states[i * 6 + 5], 0);
assert(static_cast<int32_t>(v.size()) == batch_size);
for (int32_t n = 0; n != batch_size; ++n) {
@@ -316,7 +322,7 @@ OnlineZipformer2TransducerModel::UnStackStates(
}
{
auto v = Unbind(allocator_, &states[m * 6], 0);
auto v = Unbind(allocator, &states[m * 6], 0);
assert(static_cast<int32_t>(v.size()) == batch_size);
for (int32_t n = 0; n != batch_size; ++n) {
@@ -324,7 +330,7 @@ OnlineZipformer2TransducerModel::UnStackStates(
}
}
{
auto v = Unbind<int64_t>(allocator_, &states[m * 6 + 1], 0);
auto v = Unbind<int64_t>(allocator, &states[m * 6 + 1], 0);
assert(static_cast<int32_t>(v.size()) == batch_size);
for (int32_t n = 0; n != batch_size; ++n) {