274 lines
7.7 KiB
C++
274 lines
7.7 KiB
C++
#include "ai_scheduler.h"
|
|
|
|
#include <cstring>
|
|
#include <fstream>
|
|
#include <iostream>
|
|
|
|
namespace rk3588 {
|
|
|
|
AiScheduler& AiScheduler::Instance() {
|
|
static AiScheduler instance;
|
|
return instance;
|
|
}
|
|
|
|
AiScheduler::AiScheduler() {
|
|
std::cout << "[AiScheduler] initialized\n";
|
|
}
|
|
|
|
AiScheduler::~AiScheduler() {
|
|
Shutdown();
|
|
}
|
|
|
|
void AiScheduler::Shutdown() {
|
|
#if defined(RK3588_ENABLE_RKNN)
|
|
{
|
|
std::lock_guard<std::mutex> lock(models_mutex_);
|
|
models_.clear();
|
|
}
|
|
std::cout << "[AiScheduler] shutdown, total inferences: " << total_inferences_.load()
|
|
<< ", errors: " << total_errors_.load() << "\n";
|
|
#endif
|
|
}
|
|
|
|
ModelHandle AiScheduler::LoadModel(const std::string& model_path, std::string& err) {
|
|
#if defined(RK3588_ENABLE_RKNN)
|
|
// 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 ctx = std::make_shared<ModelContext>();
|
|
ctx->model_data.resize(model_size);
|
|
ctx->path = model_path;
|
|
|
|
if (!file.read(reinterpret_cast<char*>(ctx->model_data.data()), model_size)) {
|
|
err = "Failed to read model file: " + model_path;
|
|
return kInvalidModelHandle;
|
|
}
|
|
|
|
// Initialize RKNN context
|
|
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;
|
|
}
|
|
|
|
// Query input/output info
|
|
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";
|
|
rknn_destroy(ctx->ctx);
|
|
ctx->ctx = 0;
|
|
return kInvalidModelHandle;
|
|
}
|
|
|
|
ctx->n_input = io_num.n_input;
|
|
ctx->n_output = io_num.n_output;
|
|
|
|
// Query input attributes
|
|
ctx->input_attrs.resize(ctx->n_input);
|
|
for (uint32_t i = 0; i < ctx->n_input; ++i) {
|
|
ctx->input_attrs[i].index = i;
|
|
rknn_query(ctx->ctx, RKNN_QUERY_INPUT_ATTR, &ctx->input_attrs[i], sizeof(rknn_tensor_attr));
|
|
}
|
|
|
|
// Query output attributes
|
|
ctx->output_attrs.resize(ctx->n_output);
|
|
for (uint32_t i = 0; i < ctx->n_output; ++i) {
|
|
ctx->output_attrs[i].index = i;
|
|
rknn_query(ctx->ctx, RKNN_QUERY_OUTPUT_ATTR, &ctx->output_attrs[i], sizeof(rknn_tensor_attr));
|
|
}
|
|
|
|
// Extract input dimensions
|
|
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 {
|
|
// NHWC
|
|
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];
|
|
}
|
|
|
|
ModelHandle handle = next_handle_.fetch_add(1);
|
|
|
|
{
|
|
std::lock_guard<std::mutex> lock(models_mutex_);
|
|
models_[handle] = ctx;
|
|
}
|
|
|
|
std::cout << "[AiScheduler] loaded model: " << model_path
|
|
<< " (handle=" << handle << ", input=" << ctx->input_w
|
|
<< "x" << ctx->input_h << "x" << ctx->input_c
|
|
<< ", outputs=" << ctx->n_output << ")\n";
|
|
|
|
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_.find(handle);
|
|
if (it != models_.end()) {
|
|
models_.erase(it);
|
|
erased = true;
|
|
}
|
|
}
|
|
if (erased) {
|
|
std::cout << "[AiScheduler] unloaded model handle=" << handle << "\n";
|
|
}
|
|
#else
|
|
(void)handle;
|
|
#endif
|
|
}
|
|
|
|
bool AiScheduler::GetModelInfo(ModelHandle handle, ModelInfo& info) const {
|
|
#if defined(RK3588_ENABLE_RKNN)
|
|
std::shared_ptr<ModelContext> ctx;
|
|
{
|
|
std::lock_guard<std::mutex> lock(models_mutex_);
|
|
auto it = models_.find(handle);
|
|
if (it == models_.end() || !it->second) {
|
|
return false;
|
|
}
|
|
ctx = it->second;
|
|
}
|
|
|
|
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<ModelContext> ctx;
|
|
{
|
|
std::lock_guard<std::mutex> lock(models_mutex_);
|
|
auto it = models_.find(handle);
|
|
if (it == models_.end() || !it->second) {
|
|
result.error = "Invalid model handle";
|
|
total_errors_.fetch_add(1);
|
|
return result;
|
|
}
|
|
ctx = it->second;
|
|
}
|
|
|
|
// 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 = RKNN_TENSOR_UINT8;
|
|
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;
|
|
}
|
|
|
|
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
|