Signed-off-by: Raja Rathour <imraja729@gmail.com> This patch implements asynchronous model execution for the LibTorch backend in FFmpeg's DNN module. Key changes: - Integrated a worker thread and a pending queue to handle inference. - Prevents the main filter thread from blocking during model execution. - Aligns LibTorch backend behavior with the existing OpenVINO async implementation. - Improves overall throughput for deep learning filters using LibTorch. The implementation has been tested with various torch models to ensure stability and correct frame synchronization. --- libavfilter/dnn/dnn_backend_torch.cpp | 115 +++++++++++++++++++++++++- 1 file changed, 111 insertions(+), 4 deletions(-) diff --git a/libavfilter/dnn/dnn_backend_torch.cpp b/libavfilter/dnn/dnn_backend_torch.cpp index 2e4326d9d4..ad81aff8da 100644 --- a/libavfilter/dnn/dnn_backend_torch.cpp +++ b/libavfilter/dnn/dnn_backend_torch.cpp @@ -25,6 +25,10 @@ #include <torch/torch.h> #include <torch/script.h> +#include <thread> +#include <mutex> +#include <condition_variable> +#include <atomic> extern "C" { #include "dnn_io_proc.h" @@ -39,9 +43,16 @@ typedef struct THModel { DNNModel model; DnnContext *ctx; torch::jit::Module *jit_model; - SafeQueue *request_queue; + SafeQueue *request_queue; // Holds available/idle request slots Queue *task_queue; Queue *lltask_queue; + + // --- Async Support --- + SafeQueue *pending_queue; // Holds requests waiting for inference + std::thread *worker_thread; // The background worker + std::mutex *mutex; // Protects the condition variable + std::condition_variable *cond; // Wakes up worker when new task arrives + std::atomic<bool> worker_stop; // Flag to stop the thread } THModel; typedef struct THInferRequest { @@ -119,6 +130,32 @@ static void dnn_free_model_th(DNNModel **model) return; th_model = (THModel *) (*model); + + // --- Stop and Join Worker Thread --- + if (th_model->worker_thread) { + { + std::lock_guard<std::mutex> lock(*th_model->mutex); + th_model->worker_stop = true; + } + th_model->cond->notify_all(); + + if (th_model->worker_thread->joinable()) { + th_model->worker_thread->join(); + } + delete th_model->worker_thread; + delete th_model->mutex; + delete th_model->cond; + } + + if (th_model->pending_queue) { + // Clear remaining items (if any) + while (ff_safe_queue_size(th_model->pending_queue) != 0) { + ff_safe_queue_pop_front(th_model->pending_queue); + } + ff_safe_queue_destroy(th_model->pending_queue); + } + // ----------------------------------- + while (ff_safe_queue_size(th_model->request_queue) != 0) { THRequestItem *item = (THRequestItem *)ff_safe_queue_pop_front(th_model->request_queue); destroy_request_item(&item); @@ -318,6 +355,41 @@ err: } } +// --- Worker Thread Function --- +static void th_worker_thread(THModel *th_model) { + while (true) { + THRequestItem *request = NULL; + + { + // Acquire lock to check condition + std::unique_lock<std::mutex> lock(*th_model->mutex); + + // Wait until: We are told to stop OR there is work in the queue + th_model->cond->wait(lock, [&]{ + return th_model->worker_stop || ff_safe_queue_size(th_model->pending_queue) > 0; + }); + + // If stopped and no work left, exit + if (th_model->worker_stop && ff_safe_queue_size(th_model->pending_queue) == 0) { + break; + } + + // Get work + request = (THRequestItem *)ff_safe_queue_pop_front(th_model->pending_queue); + } + + // Process work (Lock released so we don't block submission) + if (request) { + int ret = th_start_inference(request); + if (ret != 0) { + av_log(th_model->ctx, AV_LOG_ERROR, "Async inference failed\n"); + } + // Always callback to clean up and notify FFmpeg + infer_completion_callback(request); + } + } +} + static int execute_model_th(THRequestItem *request, Queue *lltask_queue) { THModel *th_model = NULL; @@ -343,9 +415,24 @@ static int execute_model_th(THRequestItem *request, Queue *lltask_queue) if ( ret != 0) { goto err; } + + // --- EXECUTION LOGIC (ASYNC vs SYNC) --- if (task->async) { - avpriv_report_missing_feature(th_model->ctx, "LibTorch async"); + // 1. Acquire lock + std::lock_guard<std::mutex> lock(*th_model->mutex); + + // 2. Push to pending queue + if (ff_safe_queue_push_back(th_model->pending_queue, request) < 0) { + return AVERROR(ENOMEM); + } + + // 3. Wake up worker + th_model->cond->notify_one(); + + // 4. Return immediately (Success) + return 0; } else { + // Synchronous fallback ret = th_start_inference((void *)(request)); if (ret != 0) { goto err; @@ -484,6 +571,25 @@ static DNNModel *dnn_load_model_th(DnnContext *ctx, DNNFunctionType func_type, A goto fail; } + // --- INITIALIZE ASYNC QUEUE AND THREAD --- + th_model->pending_queue = ff_safe_queue_create(); + if (!th_model->pending_queue) { + av_log(ctx, AV_LOG_ERROR, "Failed to create pending queue\n"); + goto fail; + } + + try { + th_model->mutex = new std::mutex(); + th_model->cond = new std::condition_variable(); + th_model->worker_stop = false; + + // Start worker thread + th_model->worker_thread = new std::thread(th_worker_thread, th_model); + } catch (...) { + av_log(ctx, AV_LOG_ERROR, "Failed to initialize worker thread or mutexes\n"); + goto fail; + } + model->get_input = &get_input_th; model->get_output = &get_output_th; model->filter_ctx = filter_ctx; @@ -519,7 +625,8 @@ static int dnn_execute_model_th(const DNNModel *model, DNNExecBaseParams *exec_p return AVERROR(ENOMEM); } - ret = ff_dnn_fill_task(task, exec_params, th_model, 0, 1); + // Set 'async' flag based on context (ctx->async) instead of hardcoded 0 + ret = ff_dnn_fill_task(task, exec_params, th_model, ctx->async, 1); if (ret != 0) { av_freep(&task); av_log(ctx, AV_LOG_ERROR, "unable to fill task.\n"); @@ -580,4 +687,4 @@ extern const DNNModule ff_dnn_backend_torch = { .get_result = dnn_get_result_th, .flush = dnn_flush_th, .free_model = dnn_free_model_th, -}; +}; \ No newline at end of file -- 2.48.1