lavu/cpu.c: Add av_force_cpu_count() to override auto-detection.

2 files changed, 19 insertions, 0 deletions

diff --git a/libavutil/cpu.c b/libavutil/cpu.c
index 8960415d00..eae1485f36 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)
 {
@@ -186,6 +187,7 @@ int av_cpu_count(void)
     static atomic_int printed = ATOMIC_VAR_INIT(0);
 
     int nb_cpus = 1;
+    int count   = 0;
 #if HAVE_WINRT
     SYSTEM_INFO sysinfo;
 #endif
@@ -224,9 +226,20 @@ int av_cpu_count(void)
     if (!atomic_exchange_explicit(&printed, 1, memory_order_relaxed))
         av_log(NULL, AV_LOG_DEBUG, "detected %d logical cores\n", nb_cpus);
 
+    count = atomic_load_explicit(&cpu_count, memory_order_relaxed);
+
+    if (count > 0) {
+        nb_cpus = count;
+        av_log(NULL, AV_LOG_DEBUG, "overriding to %d logical cores\n", nb_cpus);
+    }
+
     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 b555422dae..c069076439 100644
--- a/libavutil/cpu.h
+++ b/libavutil/cpu.h
@@ -99,6 +99,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.
  *
  * Note that this is affected by the build configuration and the CPU flags mask,