diff --git a/PRD_01_Device_MediaServer_RemoteMgmt.md b/PRD_01_Device_MediaServer_RemoteMgmt.md index 0266bda..7d79418 100644 --- a/PRD_01_Device_MediaServer_RemoteMgmt.md +++ b/PRD_01_Device_MediaServer_RemoteMgmt.md @@ -158,12 +158,20 @@ Restart=always RestartSec=2 LimitNOFILE=65535 +# (可选)性能调参:多路推理并发 / RGA 并发 +Environment=RK3588_RKNN_CTX_POOL_SIZE=3 +Environment=RK3588_RGA_MAX_INFLIGHT=2 + [Install] WantedBy=multi-user.target ``` > V1 不要求 agent 控制 systemd,但要求两服务可开机自启。 +### 6.5 性能调参说明(新增) +- `RK3588_RKNN_CTX_POOL_SIZE`:同一模型的 RKNN context 池大小(默认 3)。 +- `RK3588_RGA_MAX_INFLIGHT`:RGA 全局并发上限(默认 2)。 + ### 6.4 模型目录结构 - `${models_dir}/manifest.json` - `${models_dir}/files/__.`(`ext` 默认 `rknn`,白名单控制) diff --git a/PRD_04_Manager_ConfigGUI_AgentAPI.md b/PRD_04_Manager_ConfigGUI_AgentAPI.md index ffc40ca..562eb2f 100644 --- a/PRD_04_Manager_ConfigGUI_AgentAPI.md +++ b/PRD_04_Manager_ConfigGUI_AgentAPI.md @@ -239,7 +239,28 @@ GUI 侧以 `instances[]` 为“通道列表”。 --- -## 8. 人脸库路径对齐建议(避免“上传了但识别不到”) +## 8. 性能调优参数(可选,建议 GUI 以“高级设置”方式暴露) + +### 8.1 `preprocess`(AI 分支减拷贝 / RGA 并发) +- `dst_packed`:bool,默认 `false`。 + - 说明:当 `dst_format` 为 `rgb/bgr` 时,若开启则输出紧凑 packed(`stride = width * 3`),可避免下游(如 `ai_yolo`)对齐 stride 导致的逐行 memcpy。 +- `rga_max_inflight`:int,默认 `0`(表示不在配置里覆盖全局值)。 + - 说明:RGA 全局并发上限(原先固定为 2)。 + - 备注:也可通过环境变量 `RK3588_RGA_MAX_INFLIGHT` 设置。 + +### 8.2 `publish/storage`(一次编码,多处复用) +- `publish.attach_encoded_meta`:bool,默认 `true`(当 publish 有下游 output_queues 时)。 + - 说明:将已编码视频包(含 codec extradata/pts/key)挂在 `frame.user_meta`,供下游复用(例如报警 clip、storage 复用码流)。 +- `storage.reuse_encoded_meta`:bool,默认 `true`。 + - 说明:若收到 `EncodedVideoFrameMeta`,storage 将直接 remux 写文件(避免再次 MPP 编码);若未收到则保持原逻辑(MPP 编码录制)。 + +### 8.3 推理并发(环境变量,部署侧配置) +- `RK3588_RKNN_CTX_POOL_SIZE`:默认 `3`。 + - 说明:同一模型创建的 RKNN context 数量(多路同模型推理可并发,避免被单 context 串行化)。 + +--- + +## 9. 人脸库路径对齐建议(避免“上传了但识别不到”) 1) 默认推荐(最省事): - media-server work_dir:`/opt/rk3588sys` @@ -251,7 +272,7 @@ GUI 侧以 `instances[]` 为“通道列表”。 --- -## 9. 验收标准 +## 10. 验收标准 1) GUI 新增通道 → apply 成功后:`GET /v1/graphs` 出现该通道;`GET /v1/graphs/{name}` 能看到节点链路。 2) 删除通道 → apply 成功后 graphs 中消失。 diff --git a/include/ai_scheduler.h b/include/ai_scheduler.h index b1a07ce..0450537 100644 --- a/include/ai_scheduler.h +++ b/include/ai_scheduler.h @@ -93,17 +93,16 @@ public: #endif }; - // Borrowed inference avoids per-call output allocations/copies by using per-model preallocated buffers. - // The returned result holds the per-model inference lock; it must be destroyed before the next InferBorrowed() - // on the same model handle. + // Borrowed inference avoids per-call output allocations/copies by using per-context preallocated buffers. + // The returned result holds the selected context's inference lock until it is destroyed. struct BorrowedInferResult { bool success = false; std::string error; std::vector outputs; #if defined(RK3588_ENABLE_RKNN) - std::shared_ptr keepalive; // keeps ModelContext alive - std::unique_lock infer_lock; // holds ModelContext::infer_mutex + std::shared_ptr keepalive; // keeps the selected ModelContext alive + std::unique_lock infer_lock; // holds ModelContext::infer_mutex for that context #endif BorrowedInferResult() = default; @@ -133,7 +132,8 @@ private: #if defined(RK3588_ENABLE_RKNN) struct ModelContext { rknn_context ctx = 0; - std::vector model_data; + // Shared model blob kept alive for RKNN runtime. Multiple contexts can share the same data. + std::shared_ptr> model_data; std::vector input_attrs; std::vector output_attrs; std::vector> output_buffers; // preallocated output buffers @@ -143,7 +143,7 @@ private: int input_h = 0; int input_c = 0; std::string path; - std::mutex infer_mutex; // Per-model lock for inference + std::mutex infer_mutex; // Per-context lock for inference ~ModelContext() { if (ctx) { @@ -153,8 +153,14 @@ private: } }; - std::unordered_map> models_by_handle_; - std::unordered_map> models_by_path_; + struct ModelGroup { + std::string path; + std::vector> contexts; + std::atomic rr{0}; // round-robin context selection + }; + + std::unordered_map> models_by_handle_; + std::unordered_map> models_by_path_; #endif mutable std::mutex models_mutex_; // Protects models_ map diff --git a/plugins/preprocess/preprocess_node.cpp b/plugins/preprocess/preprocess_node.cpp index c81a219..55c2094 100644 --- a/plugins/preprocess/preprocess_node.cpp +++ b/plugins/preprocess/preprocess_node.cpp @@ -2,6 +2,7 @@ #include #include #include +#include #include #include #include @@ -43,9 +44,30 @@ int ToRgaFormat(PixelFormat fmt) { } } -constexpr int GlobalRgaMaxInflight() { - // Fixed concurrency cap for RGA submissions. - return 2; +std::atomic& GlobalRgaMaxInflightRef() { + static std::atomic v{2}; + return v; +} + +int GlobalRgaMaxInflight() { + static std::once_flag once; + std::call_once(once, []() { + const char* s = std::getenv("RK3588_RGA_MAX_INFLIGHT"); + if (!s || !*s) return; + try { + int v = std::stoi(s); + if (v > 0 && v <= 32) GlobalRgaMaxInflightRef().store(v); + } catch (...) { + } + }); + int v = GlobalRgaMaxInflightRef().load(); + return v > 0 ? v : 1; +} + +void SetGlobalRgaMaxInflight(int v) { + if (v <= 0) return; + if (v > 32) v = 32; + GlobalRgaMaxInflightRef().store(v); } class RgaGate { @@ -254,10 +276,18 @@ public: dst_w_ = config.ValueOr("dst_w", 640); dst_h_ = config.ValueOr("dst_h", 640); keep_ratio_ = config.ValueOr("keep_ratio", false); + dst_packed_ = config.ValueOr("dst_packed", false); std::string fmt_str = config.ValueOr("dst_format", ""); if (!fmt_str.empty()) { dst_fmt_ = ParseFormat(fmt_str); } + +#if defined(RK3588_ENABLE_RGA) + const int rga_max_inflight = config.ValueOr("rga_max_inflight", 0); + if (rga_max_inflight > 0) { + SetGlobalRgaMaxInflight(rga_max_inflight); + } +#endif const bool requested_use_rga = config.ValueOr("use_rga", true); use_rga_ = requested_use_rga; @@ -415,6 +445,11 @@ private: } } int dst_wstride = Align16(out_w); + // For AI input (RGB/BGR), allow a tightly packed output to avoid an extra per-frame memcpy + // in downstream nodes (e.g. ai_yolo). + if (dst_packed_ && (out_fmt == PixelFormat::RGB || out_fmt == PixelFormat::BGR)) { + dst_wstride = out_w; + } int dst_hstride = Align16(out_h); if (src_fmt_rga == RK_FORMAT_UNKNOWN || dst_fmt_rga == RK_FORMAT_UNKNOWN) { @@ -802,6 +837,7 @@ private: int dst_h_ = 640; bool keep_ratio_ = false; PixelFormat dst_fmt_ = PixelFormat::UNKNOWN; + bool dst_packed_ = false; bool use_rga_ = true; bool stats_log_ = false; diff --git a/plugins/storage/storage_node.cpp b/plugins/storage/storage_node.cpp index 1ea1856..9301619 100644 --- a/plugins/storage/storage_node.cpp +++ b/plugins/storage/storage_node.cpp @@ -1,3 +1,4 @@ +#include #include #include #include @@ -12,6 +13,7 @@ #include #include "node.h" +#include "media/encoded_video_meta.h" #if defined(RK3588_ENABLE_FFMPEG) extern "C" { @@ -42,6 +44,118 @@ namespace rk3588 { namespace { inline int Align16(int v) { return (v + 15) & ~15; } +std::shared_ptr TryGetEncodedMeta(const FramePtr& frame) { + if (!frame || !frame->user_meta) return nullptr; + auto meta = std::static_pointer_cast(frame->user_meta); + if (!meta || meta->magic != EncodedVideoFrameMeta::kMagic) return nullptr; + if (!meta->codec || meta->pkt.data.empty() || meta->pkt.pts_ms <= 0) return nullptr; + return meta; +} + +#if HAS_FFMPEG +static bool HasAnnexBStartCode(const uint8_t* d, size_t n) { + if (!d || n < 4) return false; + for (size_t i = 0; i + 3 < n; ++i) { + if (d[i] == 0 && d[i + 1] == 0 && d[i + 2] == 1) return true; + if (i + 4 < n && d[i] == 0 && d[i + 1] == 0 && d[i + 2] == 0 && d[i + 3] == 1) return true; + } + return false; +} + +static size_t FindStartCode(const uint8_t* d, size_t n, size_t from, size_t* sc_len) { + for (size_t i = from; i + 3 < n; ++i) { + if (d[i] == 0 && d[i + 1] == 0 && d[i + 2] == 1) { + if (sc_len) *sc_len = 3; + return i; + } + if (i + 4 < n && d[i] == 0 && d[i + 1] == 0 && d[i + 2] == 0 && d[i + 3] == 1) { + if (sc_len) *sc_len = 4; + return i; + } + } + if (sc_len) *sc_len = 0; + return n; +} + +static std::vector> SplitAnnexBNals(const uint8_t* d, size_t n) { + std::vector> out; + if (!d || n < 4) return out; + size_t pos = 0; + while (true) { + size_t sc_len = 0; + size_t sc = FindStartCode(d, n, pos, &sc_len); + if (sc >= n) break; + size_t nal_start = sc + sc_len; + size_t next_sc = FindStartCode(d, n, nal_start, nullptr); + size_t nal_end = (next_sc >= n) ? n : next_sc; + + while (nal_end > nal_start && d[nal_end - 1] == 0) --nal_end; + if (nal_end > nal_start) { + out.emplace_back(d + nal_start, d + nal_end); + } + pos = nal_end; + } + return out; +} + +static bool BuildAvccFromAnnexB(const std::vector& annexb, std::vector& avcc_out) { + avcc_out.clear(); + if (annexb.empty()) return false; + auto nals = SplitAnnexBNals(annexb.data(), annexb.size()); + + const uint8_t* sps = nullptr; + size_t sps_len = 0; + const uint8_t* pps = nullptr; + size_t pps_len = 0; + + for (const auto& nal : nals) { + if (nal.empty()) continue; + const uint8_t t = nal[0] & 0x1F; + if (t == 7 && !sps) { + sps = nal.data(); + sps_len = nal.size(); + } else if (t == 8 && !pps) { + pps = nal.data(); + pps_len = nal.size(); + } + } + if (!sps || sps_len < 4 || !pps || pps_len < 1) return false; + + avcc_out.reserve(11 + sps_len + pps_len); + avcc_out.push_back(1); + avcc_out.push_back(sps[1]); + avcc_out.push_back(sps[2]); + avcc_out.push_back(sps[3]); + avcc_out.push_back(0xFF); + avcc_out.push_back(0xE1); + avcc_out.push_back(static_cast((sps_len >> 8) & 0xFF)); + avcc_out.push_back(static_cast((sps_len) & 0xFF)); + avcc_out.insert(avcc_out.end(), sps, sps + sps_len); + avcc_out.push_back(1); + avcc_out.push_back(static_cast((pps_len >> 8) & 0xFF)); + avcc_out.push_back(static_cast((pps_len) & 0xFF)); + avcc_out.insert(avcc_out.end(), pps, pps + pps_len); + return true; +} + +static std::vector AnnexBToLengthPrefixed(const uint8_t* d, size_t n) { + std::vector out; + auto nals = SplitAnnexBNals(d, n); + size_t total = 0; + for (const auto& nal : nals) total += 4 + nal.size(); + out.reserve(total); + for (const auto& nal : nals) { + uint32_t len = static_cast(nal.size()); + out.push_back(static_cast((len >> 24) & 0xFF)); + out.push_back(static_cast((len >> 16) & 0xFF)); + out.push_back(static_cast((len >> 8) & 0xFF)); + out.push_back(static_cast((len) & 0xFF)); + out.insert(out.end(), nal.begin(), nal.end()); + } + return out; +} +#endif + std::string FormatTime(const std::string& pattern) { auto now = std::chrono::system_clock::now(); auto time_t_now = std::chrono::system_clock::to_time_t(now); @@ -264,6 +378,7 @@ public: filename_pattern_ = config.ValueOr("filename_pattern", "%Y%m%d/%H%M%S"); fps_ = config.ValueOr("fps", 25); bitrate_kbps_ = config.ValueOr("bitrate_kbps", 2000); + reuse_encoded_meta_ = config.ValueOr("reuse_encoded_meta", true); input_queue_ = ctx.input_queue; if (!input_queue_) { @@ -373,6 +488,11 @@ private: void OpenNewFile(FramePtr frame) { #if HAS_FFMPEG + auto meta = reuse_encoded_meta_ ? TryGetEncodedMeta(frame) : nullptr; + using_encoded_meta_ = (meta != nullptr); + segment_base_pts_ms_ = using_encoded_meta_ ? meta->pkt.pts_ms : 0; + last_pts90k_ = -1; + std::string filename = FormatTime(filename_pattern_) + "." + format_; current_path_ = base_path_ + "/" + filename; @@ -390,41 +510,78 @@ private: return; } - AVCodecID codec_id = (codec_ == "h265" || codec_ == "hevc") - ? AV_CODEC_ID_HEVC : AV_CODEC_ID_H264; - const AVCodec* encoder = avcodec_find_encoder(codec_id); - stream_ = avformat_new_stream(fmt_ctx_, encoder); + AVCodecID codec_id = AV_CODEC_ID_H264; + if (using_encoded_meta_ && meta && meta->codec) { + codec_id = (meta->codec->codec == VideoCodec::H265) ? AV_CODEC_ID_HEVC : AV_CODEC_ID_H264; + } else { + codec_id = (codec_ == "h265" || codec_ == "hevc") ? AV_CODEC_ID_HEVC : AV_CODEC_ID_H264; + } + + stream_ = avformat_new_stream(fmt_ctx_, nullptr); if (!stream_) { avformat_free_context(fmt_ctx_); fmt_ctx_ = nullptr; return; } - stream_->time_base = AVRational{1, fps_}; + if (using_encoded_meta_) { + stream_->time_base = AVRational{1, 90000}; + const int fps_eff = (meta && meta->codec && meta->codec->fps > 0) ? meta->codec->fps : fps_; + stream_->avg_frame_rate = AVRational{std::max(1, fps_eff), 1}; + stream_->r_frame_rate = stream_->avg_frame_rate; + } else { + stream_->time_base = AVRational{1, fps_}; + } stream_->codecpar->codec_type = AVMEDIA_TYPE_VIDEO; stream_->codecpar->codec_id = codec_id; - stream_->codecpar->width = frame->width; - stream_->codecpar->height = frame->height; + stream_->codecpar->width = (using_encoded_meta_ && meta && meta->codec && meta->codec->width > 0) + ? meta->codec->width + : frame->width; + stream_->codecpar->height = (using_encoded_meta_ && meta && meta->codec && meta->codec->height > 0) + ? meta->codec->height + : frame->height; stream_->codecpar->format = AV_PIX_FMT_YUV420P; #if HAS_MPP - if (!mpp_encoder_) { - mpp_encoder_ = std::make_unique(); - if (!mpp_encoder_->Init(frame->width, frame->height, frame->format, - codec_, fps_, bitrate_kbps_)) { - mpp_encoder_.reset(); + if (!using_encoded_meta_) { + if (!mpp_encoder_) { + mpp_encoder_ = std::make_unique(); + if (!mpp_encoder_->Init(frame->width, frame->height, frame->format, + codec_, fps_, bitrate_kbps_)) { + mpp_encoder_.reset(); + } + } + + if (mpp_encoder_ && !mpp_encoder_->Header().empty()) { + const auto& hdr = mpp_encoder_->Header(); + stream_->codecpar->extradata_size = static_cast(hdr.size()); + stream_->codecpar->extradata = static_cast( + av_mallocz(hdr.size() + AV_INPUT_BUFFER_PADDING_SIZE)); + std::memcpy(stream_->codecpar->extradata, hdr.data(), hdr.size()); } } - - if (mpp_encoder_ && !mpp_encoder_->Header().empty()) { - const auto& hdr = mpp_encoder_->Header(); - stream_->codecpar->extradata_size = static_cast(hdr.size()); - stream_->codecpar->extradata = static_cast( - av_mallocz(hdr.size() + AV_INPUT_BUFFER_PADDING_SIZE)); - std::memcpy(stream_->codecpar->extradata, hdr.data(), hdr.size()); - } #endif + if (using_encoded_meta_ && meta && meta->codec && !meta->codec->extradata.empty()) { + std::vector ex = meta->codec->extradata; + const bool is_mp4 = (format_ != "ts"); + + // MP4 prefers AVCC/HVCC. Do a best-effort conversion for H264 AnnexB extradata. + if (is_mp4 && codec_id == AV_CODEC_ID_H264) { + if (!(ex.size() >= 1 && ex[0] == 1) && HasAnnexBStartCode(ex.data(), ex.size())) { + std::vector avcc; + if (BuildAvccFromAnnexB(ex, avcc)) { + ex = std::move(avcc); + } + } + } + + stream_->codecpar->extradata_size = static_cast(ex.size()); + stream_->codecpar->extradata = static_cast( + av_mallocz(ex.size() + AV_INPUT_BUFFER_PADDING_SIZE)); + std::memcpy(stream_->codecpar->extradata, ex.data(), ex.size()); + } + if (!(fmt_ctx_->oformat->flags & AVFMT_NOFILE)) { if (avio_open(&fmt_ctx_->pb, current_path_.c_str(), AVIO_FLAG_WRITE) < 0) { avformat_free_context(fmt_ctx_); @@ -450,6 +607,52 @@ private: } void WriteFrame(FramePtr frame) { +#if HAS_FFMPEG + if (!file_open_ || !fmt_ctx_ || !stream_) return; + + if (using_encoded_meta_) { + auto meta = TryGetEncodedMeta(frame); + if (!meta || !meta->codec) return; + + const bool is_mp4 = (format_ != "ts"); + std::vector sample = meta->pkt.data; + + if (is_mp4 && stream_->codecpar->codec_id == AV_CODEC_ID_H264 && + HasAnnexBStartCode(sample.data(), sample.size())) { + sample = AnnexBToLengthPrefixed(sample.data(), sample.size()); + } + if (sample.empty()) return; + + const int fps_eff = (meta->codec->fps > 0) ? meta->codec->fps : fps_; + const int64_t frame_dur = fps_eff > 0 ? (90000 / std::max(1, fps_eff)) : 0; + + int64_t rel_ms = meta->pkt.pts_ms - segment_base_pts_ms_; + if (rel_ms < 0) rel_ms = 0; + int64_t pts90k = rel_ms * 90; + if (last_pts90k_ >= 0 && pts90k <= last_pts90k_) { + pts90k = last_pts90k_ + (frame_dur > 0 ? frame_dur : 1); + } + last_pts90k_ = pts90k; + + AVPacket* pkt = av_packet_alloc(); + if (!pkt) return; + if (av_new_packet(pkt, static_cast(sample.size())) < 0) { + av_packet_free(&pkt); + return; + } + std::memcpy(pkt->data, sample.data(), sample.size()); + pkt->stream_index = stream_->index; + pkt->flags = meta->pkt.key ? AV_PKT_FLAG_KEY : 0; + pkt->pts = pts90k; + pkt->dts = pts90k; + pkt->duration = frame_dur; + (void)av_interleaved_write_frame(fmt_ctx_, pkt); + av_packet_free(&pkt); + ++segment_frames_; + return; + } +#endif + #if HAS_FFMPEG && HAS_MPP if (!file_open_ || !fmt_ctx_ || !mpp_encoder_) return; @@ -476,8 +679,7 @@ private: av_interleaved_write_frame(fmt_ctx_, pkt); av_packet_free(&pkt); ++segment_frames_; -#elif HAS_FFMPEG - // Software encoding fallback would go here +#else (void)frame; #endif } @@ -494,6 +696,9 @@ private: fmt_ctx_ = nullptr; stream_ = nullptr; file_open_ = false; + using_encoded_meta_ = false; + segment_base_pts_ms_ = 0; + last_pts90k_ = -1; std::cout << "[storage] closed: " << current_path_ << " (" << segment_frames_ << " frames)\n"; @@ -509,6 +714,7 @@ private: std::string filename_pattern_; int fps_ = 25; int bitrate_kbps_ = 2000; + bool reuse_encoded_meta_ = true; std::mutex mu_; @@ -520,6 +726,9 @@ private: std::string current_path_; std::chrono::steady_clock::time_point segment_start_; int64_t segment_frames_ = 0; + bool using_encoded_meta_ = false; + int64_t segment_base_pts_ms_ = 0; + int64_t last_pts90k_ = -1; #if HAS_FFMPEG AVFormatContext* fmt_ctx_ = nullptr; diff --git a/src/ai_scheduler.cpp b/src/ai_scheduler.cpp index bdf0646..d45f112 100644 --- a/src/ai_scheduler.cpp +++ b/src/ai_scheduler.cpp @@ -1,5 +1,6 @@ #include "ai_scheduler.h" +#include #include #include @@ -10,6 +11,30 @@ 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: @@ -83,98 +108,118 @@ ModelHandle AiScheduler::LoadModel(const std::string& model_path, std::string& e size_t model_size = file.tellg(); file.seekg(0, std::ios::beg); - auto ctx = std::make_shared(); - ctx->model_data.resize(model_size); - ctx->path = model_path; - - if (!file.read(reinterpret_cast(ctx->model_data.data()), model_size)) { + auto model_data = std::make_shared>(); + model_data->resize(model_size); + if (!file.read(reinterpret_cast(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); + auto group = std::make_shared(); + group->path = model_path; + + const int pool_size = DefaultContextPoolSize(); + group->contexts.reserve(static_cast(pool_size)); + + for (int i = 0; i < pool_size; ++i) { + auto ctx = std::make_shared(); + 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; } - // Prefer using all NPU cores (RK3588). - // This does not change the shared-context behavior; it only hints RKNN runtime scheduling. - ret = rknn_set_core_mask(ctx->ctx, RKNN_NPU_CORE_0_1_2); - if (ret < 0) { - LogWarn("[AiScheduler] rknn_set_core_mask failed: " + std::to_string(ret)); - } - - // 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)); - } - - // Preallocate output buffers for borrowed inference. - ctx->output_buffers.resize(ctx->n_output); - for (uint32_t i = 0; i < ctx->n_output; ++i) { - uint32_t out_sz = ctx->output_attrs[i].size; - if (out_sz == 0 && ctx->output_attrs[i].size_with_stride > 0) { - out_sz = ctx->output_attrs[i].size_with_stride; - } - if (out_sz == 0 && ctx->output_attrs[i].n_elems > 0) { - out_sz = ctx->output_attrs[i].n_elems * TensorTypeSizeBytes(ctx->output_attrs[i].type); - } - if (out_sz > 0) { - ctx->output_buffers[i].resize(out_sz); - } else { - ctx->output_buffers[i].clear(); - } - } - - // 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 lock(models_mutex_); - models_by_handle_[handle] = ctx; - models_by_path_[model_path] = ctx; + 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) + - ", input=" + std::to_string(ctx->input_w) + "x" + std::to_string(ctx->input_h) + - "x" + std::to_string(ctx->input_c) + - ", outputs=" + std::to_string(ctx->n_output) + ")"); + ", 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 @@ -205,16 +250,19 @@ void AiScheduler::UnloadModel(ModelHandle handle) { bool AiScheduler::GetModelInfo(ModelHandle handle, ModelInfo& info) const { #if defined(RK3588_ENABLE_RKNN) - std::shared_ptr ctx; + std::shared_ptr group; { std::lock_guard lock(models_mutex_); auto it = models_by_handle_.find(handle); if (it == models_by_handle_.end() || !it->second) { return false; } - ctx = it->second; + 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; @@ -233,7 +281,7 @@ InferResult AiScheduler::Infer(ModelHandle handle, const InferInput& input) { InferResult result; #if defined(RK3588_ENABLE_RKNN) - std::shared_ptr ctx; + std::shared_ptr group; { std::lock_guard lock(models_mutex_); auto it = models_by_handle_.find(handle); @@ -242,7 +290,20 @@ InferResult AiScheduler::Infer(ModelHandle handle, const InferInput& input) { total_errors_.fetch_add(1); return result; } - ctx = it->second; + 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. @@ -332,7 +393,7 @@ AiScheduler::BorrowedInferResult AiScheduler::InferBorrowed(ModelHandle handle, BorrowedInferResult result; #if defined(RK3588_ENABLE_RKNN) - std::shared_ptr ctx; + std::shared_ptr group; { std::lock_guard lock(models_mutex_); auto it = models_by_handle_.find(handle); @@ -341,7 +402,20 @@ AiScheduler::BorrowedInferResult AiScheduler::InferBorrowed(ModelHandle handle, total_errors_.fetch_add(1); return result; } - ctx = it->second; + 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) {