[FFmpeg-devel] [PATCH] fftools/cmdutils.c: Add cmd line option to override detection of cpu count

[Thilo Borgmann]

Hi,

add an option to override auto-detection of cpu count. Defaults to auto-detection, of course.

-Thilo
-------------- next part --------------
From 38612f3e1339354dbaa6be4f36072320ff71c707 Mon Sep 17 00:00:00 2001
From: Thilo Borgmann <thilo.borgmann at mail.de>
Date: Sat, 5 Jun 2021 14:26:23 +0200
Subject: [PATCH] fftools/cmdutils.c: Add cmd line option to override detection
 of cpu count

Suggested-By: ffmpeg at fb.com
---
 doc/fftools-common-opts.texi |  7 +++++++
 fftools/cmdutils.c           | 27 +++++++++++++++++++++++++++
 fftools/cmdutils.h           |  7 +++++++
 libavutil/cpu.c              | 14 ++++++++++++++
 libavutil/cpu.h              |  6 ++++++
 5 files changed, 61 insertions(+)

diff --git a/doc/fftools-common-opts.texi b/doc/fftools-common-opts.texi
index 5260ecb8f3..7643dd8396 100644
--- a/doc/fftools-common-opts.texi
+++ b/doc/fftools-common-opts.texi
@@ -353,6 +353,13 @@ Possible flags for this option are:
 @end table
 @end table
 
+ at item -cpucount @var{count} (@emph{global})
+Override detection of CPU count. This option is intended
+for testing. Do not use it unless you know what you're doing.
+ at example
+ffmpeg -cpucount 2
+ at end example
+
 @item -max_alloc @var{bytes}
 Set the maximum size limit for allocating a block on the heap by ffmpeg's
 family of malloc functions. Exercise @strong{extreme caution} when using
diff --git a/fftools/cmdutils.c b/fftools/cmdutils.c
index 4eb68d2201..583a706e5d 100644
--- a/fftools/cmdutils.c
+++ b/fftools/cmdutils.c
@@ -72,6 +72,7 @@ AVDictionary *format_opts, *codec_opts, *resample_opts;
 static FILE *report_file;
 static int report_file_level = AV_LOG_DEBUG;
 int hide_banner = 0;
+int cpu_count = -1;
 
 enum show_muxdemuxers {
     SHOW_DEFAULT,
@@ -866,6 +867,32 @@ int opt_cpuflags(void *optctx, const char *opt, const char *arg)
     return 0;
 }
 
+int opt_cpucount(void *optctx, const char *opt, const char *arg)
+{
+    int ret;
+    int count;
+
+    static const AVOption opts[] = {
+        {"count", NULL, 0, AV_OPT_TYPE_INT, { .i64 = -1}, -1, INT_MAX, NULL},
+        {NULL},
+    };
+    static const AVClass class = {
+        .class_name = "cpucount",
+        .item_name  = av_default_item_name,
+        .option     = opts,
+        .version    = LIBAVUTIL_VERSION_INT,
+    };
+    const AVClass *pclass = &class;
+
+    ret = av_opt_eval_int(&pclass, opts, arg, &count);
+
+    if (!ret) {
+        av_force_cpu_count(count);
+    }
+
+    return ret;
+}
+
 int opt_loglevel(void *optctx, const char *opt, const char *arg)
 {
     const struct { const char *name; int level; } log_levels[] = {
diff --git a/fftools/cmdutils.h b/fftools/cmdutils.h
index 1917510589..29a45dd0be 100644
--- a/fftools/cmdutils.h
+++ b/fftools/cmdutils.h
@@ -50,6 +50,7 @@ extern AVDictionary *sws_dict;
 extern AVDictionary *swr_opts;
 extern AVDictionary *format_opts, *codec_opts, *resample_opts;
 extern int hide_banner;
+extern int cpu_count;
 
 /**
  * Register a program-specific cleanup routine.
@@ -88,6 +89,11 @@ void log_callback_help(void* ptr, int level, const char* fmt, va_list vl);
  */
 int opt_cpuflags(void *optctx, const char *opt, const char *arg);
 
+/**
+ * Override the cpucount.
+ */
+int opt_cpucount(void *optctx, const char *opt, const char *arg);
+
 /**
  * Fallback for options that are not explicitly handled, these will be
  * parsed through AVOptions.
@@ -239,6 +245,7 @@ void show_help_options(const OptionDef *options, const char *msg, int req_flags,
     { "report",      0,                    { .func_arg = opt_report },       "generate a report" },                     \
     { "max_alloc",   HAS_ARG,              { .func_arg = opt_max_alloc },    "set maximum size of a single allocated block", "bytes" }, \
     { "cpuflags",    HAS_ARG | OPT_EXPERT, { .func_arg = opt_cpuflags },     "force specific cpu flags", "flags" },     \
+    { "cpucount",    HAS_ARG | OPT_EXPERT, { .func_arg = opt_cpucount },     "force specific cpu count", "count" },     \
     { "hide_banner", OPT_BOOL | OPT_EXPERT, {&hide_banner},     "do not show program banner", "hide_banner" },          \
     CMDUTILS_COMMON_OPTIONS_AVDEVICE                                                                                    \
 
diff --git a/libavutil/cpu.c b/libavutil/cpu.c
index 52f6b9a3bf..ccd5b5adac 100644
--- a/libavutil/cpu.c
+++ b/libavutil/cpu.c
@@ -48,6 +48,7 @@
 #endif
 
 static atomic_int cpu_flags = ATOMIC_VAR_INIT(-1);
+static atomic_int cpu_count = ATOMIC_VAR_INIT(-1);
 
 static int get_cpu_flags(void)
 {
@@ -306,14 +307,27 @@ int av_cpu_count(void)
     nb_cpus = sysinfo.dwNumberOfProcessors;
 #endif
 
+    int count = atomic_load_explicit(&cpu_count, memory_order_relaxed);
+
     if (!printed) {
         av_log(NULL, AV_LOG_DEBUG, "detected %d logical cores\n", nb_cpus);
+        if (count > 0) {
+            av_log(NULL, AV_LOG_DEBUG, "overriding to %d logical cores\n", count);
+        }
         printed = 1;
     }
 
+    if (count > 0) {
+        nb_cpus = count;
+    }
+
     return nb_cpus;
 }
 
+void av_force_cpu_count(int count){
+    atomic_store_explicit(&cpu_count, count, memory_order_relaxed);
+}
+
 size_t av_cpu_max_align(void)
 {
     if (ARCH_MIPS)
diff --git a/libavutil/cpu.h b/libavutil/cpu.h
index 83099dd969..9c682ad50b 100644
--- a/libavutil/cpu.h
+++ b/libavutil/cpu.h
@@ -119,6 +119,12 @@ int av_parse_cpu_caps(unsigned *flags, const char *s);
  */
 int av_cpu_count(void);
 
+/**
+ * Overrides cpu count detection and forces the specified count.
+ * Count < 1 disables forcing of specific count.
+ */
+void av_force_cpu_count(int count);
+
 /**
  * Get the maximum data alignment that may be required by FFmpeg.
  *
-- 
2.20.1 (Apple Git-117)