[FFmpeg-devel] [PATCH v2 FFmpeg 5/20] libavfilter/dnn: libtorch add CUDA suppoort

[m.kaindl0208 at gmail.com]

Signed-off-by: MaximilianKaindl <m.kaindl0208 at gmail.com>
---
 libavfilter/dnn/dnn_backend_torch.cpp | 37 +++++++++++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/libavfilter/dnn/dnn_backend_torch.cpp b/libavfilter/dnn/dnn_backend_torch.cpp
index 2e4326d9d4..062821949d 100644
--- a/libavfilter/dnn/dnn_backend_torch.cpp
+++ b/libavfilter/dnn/dnn_backend_torch.cpp
@@ -34,6 +34,10 @@ extern "C" {
 #include "queue.h"
 #include "safe_queue.h"
 }
+#if (CONFIG_LIBTORCH_CUDA == 1)
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAStream.h>
+#endif

 typedef struct THModel {
     DNNModel model;
@@ -435,7 +439,40 @@ static DNNModel *dnn_load_model_th(DnnContext *ctx, DNNFunctionType func_type, A
             av_log(ctx, AV_LOG_ERROR, "No XPU device found\n");
             goto fail;
         }
+#if (CONFIG_LIBTORCH_CUDA == 0)
         at::detail::getXPUHooks().initXPU();
+#else
+        at::detail::getXPUHooks().init();
+    } else if (device.is_cuda()) {
+        if (!torch::cuda::is_available()) {
+            av_log(ctx, AV_LOG_ERROR, "CUDA is not available!\n");
+            goto fail;
+        }
+        // Initialize CUDA
+        try {
+            int device_idx = 0;
+            const char *device_num = strstr(device_name, ":");
+            if (device_num) {
+                char *endptr = NULL;
+                device_idx = strtol(device_num + 1, &endptr, 10);
+                if (*endptr != '\0' && !isspace(*endptr)) {
+                    av_log(ctx, AV_LOG_ERROR, "Invalid device number format: %s\n", device_num + 1);
+                    goto fail;
+                }
+            }
+            if (device_idx >= static_cast<int>(torch::cuda::device_count())) {
+                av_log(ctx, AV_LOG_ERROR, "Requested CUDA device %d but only %ld devices available\n", device_idx,
+                        torch::cuda::device_count());
+                goto fail;
+            }
+            c10::cuda::set_device(device_idx);
+            c10::cuda::setCurrentCUDAStream(c10::cuda::getDefaultCUDAStream());
+            torch::cuda::synchronize();
+        } catch (const c10::Error &e) {
+            av_log(ctx, AV_LOG_ERROR, "CUDA initialization failed: %s\n", e.what());
+            goto fail;
+        }
+#endif
     } else if (!device.is_cpu()) {
         av_log(ctx, AV_LOG_ERROR, "Not supported device:\"%s\"\n", device_name);
         goto fail;
--
2.34.1