517 lines
16 KiB
C++
517 lines
16 KiB
C++
#include "ai_scheduler.h"
|
|
|
|
#include <cstdlib>
|
|
#include <cstring>
|
|
#include <fstream>
|
|
|
|
#include "utils/logger.h"
|
|
|
|
namespace rk3588 {
|
|
|
|
#if defined(RK3588_ENABLE_RKNN)
|
|
namespace {
|
|
|
|
int GetEnvInt(const char* name, int default_value) {
|
|
if (!name) return default_value;
|
|
const char* v = std::getenv(name);
|
|
if (!v || !*v) return default_value;
|
|
try {
|
|
return std::stoi(v);
|
|
} catch (...) {
|
|
return default_value;
|
|
}
|
|
}
|
|
|
|
int ClampInt(int v, int lo, int hi) {
|
|
if (v < lo) return lo;
|
|
if (v > hi) return hi;
|
|
return v;
|
|
}
|
|
|
|
int DefaultContextPoolSize() {
|
|
// Default to 3 contexts to better utilize RK3588 NPU (3 cores).
|
|
// Can be overridden by env: RK3588_RKNN_CTX_POOL_SIZE.
|
|
const int v = GetEnvInt("RK3588_RKNN_CTX_POOL_SIZE", 3);
|
|
return ClampInt(v, 1, 16);
|
|
}
|
|
|
|
uint32_t TensorTypeSizeBytes(rknn_tensor_type t) {
|
|
switch (t) {
|
|
case RKNN_TENSOR_INT8:
|
|
case RKNN_TENSOR_UINT8:
|
|
return 1;
|
|
case RKNN_TENSOR_INT16:
|
|
case RKNN_TENSOR_UINT16:
|
|
case RKNN_TENSOR_FLOAT16:
|
|
return 2;
|
|
case RKNN_TENSOR_INT32:
|
|
case RKNN_TENSOR_UINT32:
|
|
case RKNN_TENSOR_FLOAT32:
|
|
return 4;
|
|
case RKNN_TENSOR_INT64:
|
|
return 8;
|
|
default:
|
|
return 1;
|
|
}
|
|
}
|
|
|
|
} // namespace
|
|
#endif
|
|
|
|
AiScheduler& AiScheduler::Instance() {
|
|
static AiScheduler instance;
|
|
return instance;
|
|
}
|
|
|
|
AiScheduler::AiScheduler() {
|
|
LogInfo("[AiScheduler] initialized");
|
|
}
|
|
|
|
AiScheduler::~AiScheduler() {
|
|
Shutdown();
|
|
}
|
|
|
|
void AiScheduler::Shutdown() {
|
|
#if defined(RK3588_ENABLE_RKNN)
|
|
{
|
|
std::lock_guard<std::mutex> lock(models_mutex_);
|
|
models_by_handle_.clear();
|
|
models_by_path_.clear();
|
|
}
|
|
LogInfo("[AiScheduler] shutdown, total inferences: " + std::to_string(total_inferences_.load()) +
|
|
", errors: " + std::to_string(total_errors_.load()));
|
|
#endif
|
|
}
|
|
|
|
ModelHandle AiScheduler::LoadModel(const std::string& model_path, std::string& err) {
|
|
#if defined(RK3588_ENABLE_RKNN)
|
|
{
|
|
std::lock_guard<std::mutex> lock(models_mutex_);
|
|
auto it = models_by_path_.find(model_path);
|
|
if (it != models_by_path_.end()) {
|
|
if (auto existing = it->second.lock()) {
|
|
ModelHandle handle = next_handle_.fetch_add(1);
|
|
models_by_handle_[handle] = existing;
|
|
LogInfo("[AiScheduler] reused model: " + model_path + " (handle=" + std::to_string(handle) + ")");
|
|
return handle;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Read model file
|
|
std::ifstream file(model_path, std::ios::binary | std::ios::ate);
|
|
if (!file.is_open()) {
|
|
err = "Failed to open model file: " + model_path;
|
|
return kInvalidModelHandle;
|
|
}
|
|
|
|
size_t model_size = file.tellg();
|
|
file.seekg(0, std::ios::beg);
|
|
|
|
auto model_data = std::make_shared<std::vector<uint8_t>>();
|
|
model_data->resize(model_size);
|
|
if (!file.read(reinterpret_cast<char*>(model_data->data()), model_size)) {
|
|
err = "Failed to read model file: " + model_path;
|
|
return kInvalidModelHandle;
|
|
}
|
|
|
|
auto group = std::make_shared<ModelGroup>();
|
|
group->path = model_path;
|
|
|
|
const int pool_size = DefaultContextPoolSize();
|
|
group->contexts.reserve(static_cast<size_t>(pool_size));
|
|
|
|
for (int i = 0; i < pool_size; ++i) {
|
|
auto ctx = std::make_shared<ModelContext>();
|
|
ctx->model_data = model_data;
|
|
ctx->path = model_path;
|
|
|
|
int ret = rknn_init(&ctx->ctx, ctx->model_data->data(), model_size, 0, nullptr);
|
|
if (ret < 0) {
|
|
err = "rknn_init failed with code: " + std::to_string(ret);
|
|
return kInvalidModelHandle;
|
|
}
|
|
|
|
// If we create multiple contexts, bind them to different NPU cores when possible.
|
|
// This reduces contention and avoids the single-context serialization bottleneck.
|
|
{
|
|
int mask = RKNN_NPU_CORE_0_1_2;
|
|
#if defined(RKNN_NPU_CORE_0) && defined(RKNN_NPU_CORE_1) && defined(RKNN_NPU_CORE_2)
|
|
if (pool_size >= 3) {
|
|
const int idx = i % 3;
|
|
mask = (idx == 0) ? RKNN_NPU_CORE_0 : (idx == 1 ? RKNN_NPU_CORE_1 : RKNN_NPU_CORE_2);
|
|
}
|
|
#endif
|
|
ret = rknn_set_core_mask(ctx->ctx, mask);
|
|
if (ret < 0) {
|
|
LogWarn("[AiScheduler] rknn_set_core_mask failed: " + std::to_string(ret));
|
|
}
|
|
}
|
|
|
|
rknn_input_output_num io_num;
|
|
ret = rknn_query(ctx->ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
|
|
if (ret < 0) {
|
|
err = "rknn_query IO num failed";
|
|
return kInvalidModelHandle;
|
|
}
|
|
|
|
ctx->n_input = io_num.n_input;
|
|
ctx->n_output = io_num.n_output;
|
|
|
|
ctx->input_attrs.resize(ctx->n_input);
|
|
for (uint32_t j = 0; j < ctx->n_input; ++j) {
|
|
ctx->input_attrs[j].index = j;
|
|
rknn_query(ctx->ctx, RKNN_QUERY_INPUT_ATTR, &ctx->input_attrs[j], sizeof(rknn_tensor_attr));
|
|
}
|
|
|
|
ctx->output_attrs.resize(ctx->n_output);
|
|
for (uint32_t j = 0; j < ctx->n_output; ++j) {
|
|
ctx->output_attrs[j].index = j;
|
|
rknn_query(ctx->ctx, RKNN_QUERY_OUTPUT_ATTR, &ctx->output_attrs[j], sizeof(rknn_tensor_attr));
|
|
}
|
|
|
|
ctx->output_buffers.resize(ctx->n_output);
|
|
for (uint32_t j = 0; j < ctx->n_output; ++j) {
|
|
uint32_t out_sz = ctx->output_attrs[j].size;
|
|
if (out_sz == 0 && ctx->output_attrs[j].size_with_stride > 0) {
|
|
out_sz = ctx->output_attrs[j].size_with_stride;
|
|
}
|
|
if (out_sz == 0 && ctx->output_attrs[j].n_elems > 0) {
|
|
out_sz = ctx->output_attrs[j].n_elems * TensorTypeSizeBytes(ctx->output_attrs[j].type);
|
|
}
|
|
if (out_sz > 0) {
|
|
ctx->output_buffers[j].resize(out_sz);
|
|
} else {
|
|
ctx->output_buffers[j].clear();
|
|
}
|
|
}
|
|
|
|
if (!ctx->input_attrs.empty()) {
|
|
if (ctx->input_attrs[0].fmt == RKNN_TENSOR_NCHW) {
|
|
ctx->input_c = ctx->input_attrs[0].dims[1];
|
|
ctx->input_h = ctx->input_attrs[0].dims[2];
|
|
ctx->input_w = ctx->input_attrs[0].dims[3];
|
|
} else {
|
|
ctx->input_h = ctx->input_attrs[0].dims[1];
|
|
ctx->input_w = ctx->input_attrs[0].dims[2];
|
|
ctx->input_c = ctx->input_attrs[0].dims[3];
|
|
}
|
|
}
|
|
|
|
group->contexts.push_back(ctx);
|
|
}
|
|
|
|
if (group->contexts.empty()) {
|
|
err = "No RKNN contexts created";
|
|
return kInvalidModelHandle;
|
|
}
|
|
|
|
ModelHandle handle = next_handle_.fetch_add(1);
|
|
{
|
|
std::lock_guard<std::mutex> lock(models_mutex_);
|
|
models_by_handle_[handle] = group;
|
|
models_by_path_[model_path] = group;
|
|
}
|
|
|
|
const auto& first = group->contexts.front();
|
|
LogInfo("[AiScheduler] loaded model: " + model_path +
|
|
" (handle=" + std::to_string(handle) +
|
|
", ctx_pool=" + std::to_string(group->contexts.size()) +
|
|
", input=" + std::to_string(first->input_w) + "x" + std::to_string(first->input_h) +
|
|
"x" + std::to_string(first->input_c) +
|
|
", outputs=" + std::to_string(first->n_output) + ")");
|
|
|
|
return handle;
|
|
#else
|
|
(void)model_path;
|
|
err = "RKNN not enabled";
|
|
return kInvalidModelHandle;
|
|
#endif
|
|
}
|
|
|
|
void AiScheduler::UnloadModel(ModelHandle handle) {
|
|
#if defined(RK3588_ENABLE_RKNN)
|
|
bool erased = false;
|
|
{
|
|
std::lock_guard<std::mutex> lock(models_mutex_);
|
|
auto it = models_by_handle_.find(handle);
|
|
if (it != models_by_handle_.end()) {
|
|
models_by_handle_.erase(it);
|
|
erased = true;
|
|
}
|
|
}
|
|
if (erased) {
|
|
LogInfo("[AiScheduler] unloaded model handle=" + std::to_string(handle));
|
|
}
|
|
#else
|
|
(void)handle;
|
|
#endif
|
|
}
|
|
|
|
bool AiScheduler::GetModelInfo(ModelHandle handle, ModelInfo& info) const {
|
|
#if defined(RK3588_ENABLE_RKNN)
|
|
std::shared_ptr<ModelGroup> group;
|
|
{
|
|
std::lock_guard<std::mutex> lock(models_mutex_);
|
|
auto it = models_by_handle_.find(handle);
|
|
if (it == models_by_handle_.end() || !it->second) {
|
|
return false;
|
|
}
|
|
group = it->second;
|
|
}
|
|
|
|
if (!group || group->contexts.empty() || !group->contexts[0]) return false;
|
|
auto ctx = group->contexts[0];
|
|
|
|
info.input_width = ctx->input_w;
|
|
info.input_height = ctx->input_h;
|
|
info.input_channels = ctx->input_c;
|
|
info.n_input = ctx->n_input;
|
|
info.n_output = ctx->n_output;
|
|
info.name = ctx->path;
|
|
return true;
|
|
#else
|
|
(void)handle;
|
|
(void)info;
|
|
return false;
|
|
#endif
|
|
}
|
|
|
|
InferResult AiScheduler::Infer(ModelHandle handle, const InferInput& input) {
|
|
InferResult result;
|
|
|
|
#if defined(RK3588_ENABLE_RKNN)
|
|
std::shared_ptr<ModelGroup> group;
|
|
{
|
|
std::lock_guard<std::mutex> lock(models_mutex_);
|
|
auto it = models_by_handle_.find(handle);
|
|
if (it == models_by_handle_.end() || !it->second) {
|
|
result.error = "Invalid model handle";
|
|
total_errors_.fetch_add(1);
|
|
return result;
|
|
}
|
|
group = it->second;
|
|
}
|
|
|
|
if (!group || group->contexts.empty()) {
|
|
result.error = "Invalid model context group";
|
|
total_errors_.fetch_add(1);
|
|
return result;
|
|
}
|
|
|
|
auto ctx = group->contexts[group->rr.fetch_add(1) % group->contexts.size()];
|
|
if (!ctx) {
|
|
result.error = "Invalid model context";
|
|
total_errors_.fetch_add(1);
|
|
return result;
|
|
}
|
|
|
|
// Lock this specific model for inference.
|
|
std::lock_guard<std::mutex> infer_lock(ctx->infer_mutex);
|
|
|
|
if (!input.data || input.size == 0) {
|
|
result.error = "Invalid input data";
|
|
total_errors_.fetch_add(1);
|
|
return result;
|
|
}
|
|
|
|
// Setup input
|
|
rknn_input inputs[1];
|
|
memset(inputs, 0, sizeof(inputs));
|
|
inputs[0].index = 0;
|
|
inputs[0].type = input.type;
|
|
inputs[0].size = input.size;
|
|
inputs[0].fmt = input.is_nhwc ? RKNN_TENSOR_NHWC : RKNN_TENSOR_NCHW;
|
|
inputs[0].buf = const_cast<void*>(input.data);
|
|
inputs[0].pass_through = 0;
|
|
|
|
int ret = rknn_inputs_set(ctx->ctx, ctx->n_input, inputs);
|
|
if (ret < 0) {
|
|
result.error = "rknn_inputs_set failed: " + std::to_string(ret);
|
|
total_errors_.fetch_add(1);
|
|
return result;
|
|
}
|
|
|
|
// Run inference
|
|
ret = rknn_run(ctx->ctx, nullptr);
|
|
if (ret < 0) {
|
|
result.error = "rknn_run failed: " + std::to_string(ret);
|
|
total_errors_.fetch_add(1);
|
|
return result;
|
|
}
|
|
|
|
// Get outputs
|
|
std::vector<rknn_output> outputs(ctx->n_output);
|
|
memset(outputs.data(), 0, sizeof(rknn_output) * ctx->n_output);
|
|
for (uint32_t i = 0; i < ctx->n_output; ++i) {
|
|
outputs[i].want_float = 0; // Keep quantized output
|
|
}
|
|
|
|
ret = rknn_outputs_get(ctx->ctx, ctx->n_output, outputs.data(), nullptr);
|
|
if (ret < 0) {
|
|
result.error = "rknn_outputs_get failed: " + std::to_string(ret);
|
|
total_errors_.fetch_add(1);
|
|
return result;
|
|
}
|
|
|
|
// Copy outputs to result
|
|
result.outputs.resize(ctx->n_output);
|
|
for (uint32_t i = 0; i < ctx->n_output; ++i) {
|
|
auto& out = result.outputs[i];
|
|
out.index = i;
|
|
out.size = outputs[i].size;
|
|
out.type = ctx->output_attrs[i].type;
|
|
out.zp = ctx->output_attrs[i].zp;
|
|
out.scale = ctx->output_attrs[i].scale;
|
|
|
|
// Copy dimensions
|
|
out.dims.resize(ctx->output_attrs[i].n_dims);
|
|
for (uint32_t d = 0; d < ctx->output_attrs[i].n_dims; ++d) {
|
|
out.dims[d] = ctx->output_attrs[i].dims[d];
|
|
}
|
|
|
|
// Copy data
|
|
out.data.resize(outputs[i].size);
|
|
memcpy(out.data.data(), outputs[i].buf, outputs[i].size);
|
|
}
|
|
|
|
rknn_outputs_release(ctx->ctx, ctx->n_output, outputs.data());
|
|
|
|
result.success = true;
|
|
total_inferences_.fetch_add(1);
|
|
|
|
#else
|
|
result.error = "RKNN not enabled";
|
|
(void)handle;
|
|
(void)input;
|
|
#endif
|
|
|
|
return result;
|
|
}
|
|
|
|
AiScheduler::BorrowedInferResult AiScheduler::InferBorrowed(ModelHandle handle, const InferInput& input) {
|
|
BorrowedInferResult result;
|
|
|
|
#if defined(RK3588_ENABLE_RKNN)
|
|
std::shared_ptr<ModelGroup> group;
|
|
{
|
|
std::lock_guard<std::mutex> lock(models_mutex_);
|
|
auto it = models_by_handle_.find(handle);
|
|
if (it == models_by_handle_.end() || !it->second) {
|
|
result.error = "Invalid model handle";
|
|
total_errors_.fetch_add(1);
|
|
return result;
|
|
}
|
|
group = it->second;
|
|
}
|
|
|
|
if (!group || group->contexts.empty()) {
|
|
result.error = "Invalid model context group";
|
|
total_errors_.fetch_add(1);
|
|
return result;
|
|
}
|
|
|
|
auto ctx = group->contexts[group->rr.fetch_add(1) % group->contexts.size()];
|
|
if (!ctx) {
|
|
result.error = "Invalid model context";
|
|
total_errors_.fetch_add(1);
|
|
return result;
|
|
}
|
|
|
|
if (!input.data || input.size == 0) {
|
|
result.error = "Invalid input data";
|
|
total_errors_.fetch_add(1);
|
|
return result;
|
|
}
|
|
|
|
// Hold per-model inference lock for the lifetime of this result.
|
|
result.infer_lock = std::unique_lock<std::mutex>(ctx->infer_mutex);
|
|
result.keepalive = ctx;
|
|
|
|
// Setup input
|
|
rknn_input inputs[1];
|
|
memset(inputs, 0, sizeof(inputs));
|
|
inputs[0].index = 0;
|
|
inputs[0].type = input.type;
|
|
inputs[0].size = input.size;
|
|
inputs[0].fmt = input.is_nhwc ? RKNN_TENSOR_NHWC : RKNN_TENSOR_NCHW;
|
|
inputs[0].buf = const_cast<void*>(input.data);
|
|
inputs[0].pass_through = 0;
|
|
|
|
int ret = rknn_inputs_set(ctx->ctx, ctx->n_input, inputs);
|
|
if (ret < 0) {
|
|
result.error = "rknn_inputs_set failed: " + std::to_string(ret);
|
|
total_errors_.fetch_add(1);
|
|
return result;
|
|
}
|
|
|
|
ret = rknn_run(ctx->ctx, nullptr);
|
|
if (ret < 0) {
|
|
result.error = "rknn_run failed: " + std::to_string(ret);
|
|
total_errors_.fetch_add(1);
|
|
return result;
|
|
}
|
|
|
|
std::vector<rknn_output> outputs(ctx->n_output);
|
|
memset(outputs.data(), 0, sizeof(rknn_output) * ctx->n_output);
|
|
for (uint32_t i = 0; i < ctx->n_output; ++i) {
|
|
outputs[i].want_float = 0;
|
|
outputs[i].index = i;
|
|
if (i < ctx->output_buffers.size() && !ctx->output_buffers[i].empty()) {
|
|
outputs[i].is_prealloc = 1;
|
|
outputs[i].buf = ctx->output_buffers[i].data();
|
|
outputs[i].size = static_cast<uint32_t>(ctx->output_buffers[i].size());
|
|
} else {
|
|
outputs[i].is_prealloc = 0;
|
|
outputs[i].buf = nullptr;
|
|
outputs[i].size = 0;
|
|
}
|
|
}
|
|
|
|
ret = rknn_outputs_get(ctx->ctx, ctx->n_output, outputs.data(), nullptr);
|
|
if (ret < 0) {
|
|
result.error = "rknn_outputs_get failed: " + std::to_string(ret);
|
|
total_errors_.fetch_add(1);
|
|
return result;
|
|
}
|
|
|
|
result.outputs.resize(ctx->n_output);
|
|
for (uint32_t i = 0; i < ctx->n_output; ++i) {
|
|
auto& out = result.outputs[i];
|
|
out.index = static_cast<int>(i);
|
|
out.size = outputs[i].size;
|
|
out.data = reinterpret_cast<const uint8_t*>(outputs[i].buf);
|
|
out.type = ctx->output_attrs[i].type;
|
|
out.zp = ctx->output_attrs[i].zp;
|
|
out.scale = ctx->output_attrs[i].scale;
|
|
out.dims.resize(ctx->output_attrs[i].n_dims);
|
|
for (uint32_t d = 0; d < ctx->output_attrs[i].n_dims; ++d) {
|
|
out.dims[d] = ctx->output_attrs[i].dims[d];
|
|
}
|
|
}
|
|
|
|
rknn_outputs_release(ctx->ctx, ctx->n_output, outputs.data());
|
|
|
|
result.success = true;
|
|
total_inferences_.fetch_add(1);
|
|
return result;
|
|
|
|
#else
|
|
(void)handle;
|
|
(void)input;
|
|
result.error = "RKNN not enabled";
|
|
return result;
|
|
#endif
|
|
}
|
|
|
|
void AiScheduler::InferAsync(ModelHandle handle, const InferInput& input, InferCallback callback) {
|
|
// Simple implementation: just call sync Infer
|
|
// Future: use a thread pool for true async
|
|
InferResult result = Infer(handle, input);
|
|
if (callback) {
|
|
callback(result);
|
|
}
|
|
}
|
|
|
|
} // namespace rk3588
|