safesight-edge/src/ai_scheduler.cpp

#include "ai_scheduler.h"

#include <algorithm>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <limits>

#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;
    }
}

uint64_t GetEnvU64(const char* name, uint64_t default_value) {
    if (!name) return default_value;
    const char* v = std::getenv(name);
    if (!v || !*v) return default_value;
    try {
        return static_cast<uint64_t>(std::stoull(v));
    } catch (...) {
        return default_value;
    }
}

uint64_t DefaultMaxModelBytes() {
    // Guardrail against tellg() failures or accidentally huge model files.
    // Override with env: RK3588_MODEL_MAX_BYTES.
    const uint64_t def = 512ull * 1024ull * 1024ull;
    const uint64_t v = GetEnvU64("RK3588_MODEL_MAX_BYTES", def);
    // Clamp to a sane range.
    return std::min<uint64_t>(std::max<uint64_t>(v, 4ull * 1024ull * 1024ull), 4ull * 1024ull * 1024ull * 1024ull);
}

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);
}

}  // 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;
    }

    const std::streampos end_pos = file.tellg();
    if (end_pos <= 0) {
        err = "Failed to read model size: " + model_path;
        return kInvalidModelHandle;
    }
    const uint64_t model_size_u64 = static_cast<uint64_t>(end_pos);
    const uint64_t max_bytes = DefaultMaxModelBytes();
    if (model_size_u64 > max_bytes) {
        err = "Model file too large (" + std::to_string(model_size_u64) + " bytes, max=" +
              std::to_string(max_bytes) + "): " + model_path;
        return kInvalidModelHandle;
    }
    if (model_size_u64 > static_cast<uint64_t>(std::numeric_limits<std::streamsize>::max())) {
        err = "Model file too large for this build: " + model_path;
        return kInvalidModelHandle;
    }

    const size_t model_size = static_cast<size_t>(model_size_u64);
    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()), static_cast<std::streamsize>(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.
        {
            rknn_core_mask 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));
            LogInfo("[ai_scheduler] output[" + std::to_string(j) + "] type=" + 
                     std::to_string(ctx->output_attrs[j].type) + " qnt_type=" + 
                     std::to_string(ctx->output_attrs[j].qnt_type) + " zp=" + 
                     std::to_string(ctx->output_attrs[j].zp) + " scale=" + 
                     std::to_string(ctx->output_attrs[j].scale));
        }

        ctx->output_buffers.resize(ctx->n_output);
        for (uint32_t j = 0; j < ctx->n_output; ++j) {
            // FP32 output when want_float=1: 4 bytes per element
            uint32_t out_sz = ctx->output_attrs[j].n_elems * sizeof(float);
            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;
    }

    if (ctx->n_input != 1) {
        result.error = "Model expects " + std::to_string(ctx->n_input) + " inputs; current pipeline provides 1";
        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;
    }

    struct InputMemGuard {
        rknn_context ctx = 0;
        rknn_tensor_mem* mem = nullptr;
        ~InputMemGuard() {
            if (ctx && mem) {
                rknn_destroy_mem(ctx, mem);
                mem = nullptr;
            }
        }
    };
    InputMemGuard input_mem{ctx->ctx, nullptr};
    bool used_io_mem = false;

    // Best-effort RKNN zero-copy input via DMA-BUF.
    if (input.dma_fd >= 0) {
        void* virt_base = const_cast<void*>(input.data);
        const int32_t offset = input.dma_offset;
        if (offset != 0) {
            virt_base = static_cast<uint8_t*>(virt_base) - offset;
        }
        input_mem.mem = rknn_create_mem_from_fd(ctx->ctx, input.dma_fd,
                                               virt_base,
                                               static_cast<uint32_t>(input.size),
                                               offset);
        if (input_mem.mem) {
            const rknn_tensor_attr model_attr = ctx->input_attrs.empty() ? rknn_tensor_attr{} : ctx->input_attrs[0];
            const uint32_t required_size = (model_attr.size_with_stride > 0)
                                               ? model_attr.size_with_stride
                                               : model_attr.size;
            const bool fmt_match = (ctx->input_attrs.empty()) ? false
                                  : ((input.is_nhwc && model_attr.fmt == RKNN_TENSOR_NHWC) ||
                                     (!input.is_nhwc && model_attr.fmt == RKNN_TENSOR_NCHW));
            const bool type_match = (ctx->input_attrs.empty()) ? false : (model_attr.type == input.type);
            const bool can_passthrough = fmt_match && type_match && required_size > 0 && input.size >= required_size;

            rknn_tensor_attr attr = model_attr;
            attr.index = 0;
            attr.pass_through = can_passthrough ? 1 : 0;
            if (!can_passthrough) {
                // Allow RKNN driver to convert/pack to model input when formats/types differ.
                attr.type = input.type;
                attr.fmt = input.is_nhwc ? RKNN_TENSOR_NHWC : RKNN_TENSOR_NCHW;
            }

            const int mem_ret = rknn_set_io_mem(ctx->ctx, input_mem.mem, &attr);
            if (mem_ret == 0) {
                used_io_mem = true;
            } else {
                // Fallback to normal inputs_set path.
                rknn_destroy_mem(ctx->ctx, input_mem.mem);
                input_mem.mem = nullptr;
            }
        }
    }

    if (!used_io_mem) {
        // Setup input (legacy copy path).
        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, 1, inputs);
        if (ret < 0) {
            result.error = "rknn_inputs_set failed: " + std::to_string(ret);
            total_errors_.fetch_add(1);
            return result;
        }
    }

    // Run inference
    int 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 INT8 quantized output, manual dequantize
    }

    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 (ctx->n_input != 1) {
        result.error = "Model expects " + std::to_string(ctx->n_input) + " inputs; current pipeline provides 1";
        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;

    struct InputMemGuard {
        rknn_context ctx = 0;
        rknn_tensor_mem* mem = nullptr;
        ~InputMemGuard() {
            if (ctx && mem) {
                rknn_destroy_mem(ctx, mem);
                mem = nullptr;
            }
        }
    };
    InputMemGuard input_mem{ctx->ctx, nullptr};
    bool used_io_mem = false;

    // Best-effort RKNN zero-copy input via DMA-BUF.
    if (input.dma_fd >= 0) {
        void* virt_base = const_cast<void*>(input.data);
        const int32_t offset = input.dma_offset;
        if (offset != 0) {
            virt_base = static_cast<uint8_t*>(virt_base) - offset;
        }
        input_mem.mem = rknn_create_mem_from_fd(ctx->ctx, input.dma_fd,
                                               virt_base,
                                               static_cast<uint32_t>(input.size),
                                               offset);
        if (input_mem.mem) {
            const rknn_tensor_attr model_attr = ctx->input_attrs.empty() ? rknn_tensor_attr{} : ctx->input_attrs[0];
            const uint32_t required_size = (model_attr.size_with_stride > 0)
                                               ? model_attr.size_with_stride
                                               : model_attr.size;
            const bool fmt_match = (ctx->input_attrs.empty()) ? false
                                  : ((input.is_nhwc && model_attr.fmt == RKNN_TENSOR_NHWC) ||
                                     (!input.is_nhwc && model_attr.fmt == RKNN_TENSOR_NCHW));
            const bool type_match = (ctx->input_attrs.empty()) ? false : (model_attr.type == input.type);
            const bool can_passthrough = fmt_match && type_match && required_size > 0 && input.size >= required_size;

            rknn_tensor_attr attr = model_attr;
            attr.index = 0;
            attr.pass_through = can_passthrough ? 1 : 0;
            if (!can_passthrough) {
                attr.type = input.type;
                attr.fmt = input.is_nhwc ? RKNN_TENSOR_NHWC : RKNN_TENSOR_NCHW;
            }
            const int mem_ret = rknn_set_io_mem(ctx->ctx, input_mem.mem, &attr);
            if (mem_ret == 0) {
                used_io_mem = true;
            } else {
                rknn_destroy_mem(ctx->ctx, input_mem.mem);
                input_mem.mem = nullptr;
            }
        }
    }

    if (!used_io_mem) {
        // Setup input (legacy copy path).
        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, 1, inputs);
        if (ret < 0) {
            result.error = "rknn_inputs_set failed: " + std::to_string(ret);
            total_errors_.fetch_add(1);
            return result;
        }
    }

    int 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 = 1;  // Request FP32 output for direct use
        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);
        // When want_float=1, RKNN outputs FP32
        out.type = RKNN_TENSOR_FLOAT32;
        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