From 70c9726e06f02e81da8204d0a769ad133dd93384 Mon Sep 17 00:00:00 2001
From: Ian Romanick <ian.d.romanick@intel.com>
Date: Mon, 7 Jun 2021 13:06:22 -0700
Subject: [PATCH] util: Consider CPU affinity when detecting number of CPUs
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

A similar path can be used on at least FreeBSD using cpuset_getaffinity.
This is how Ninja determines the number of available CPUs on that
platform.  See the GetProcessorCount function in util.cc:

https://github.com/ninja-build/ninja/blob/master/src/util.cc

v2: Fix counting the number of available CPUs.  The CPU_COUNT API does
not work the way I thought it did. :face_palm: Noticed by Marek.

Reviewed-by: Adam Jackson <ajax@redhat.com> [v1]
Reviewed-by: Pierre-Eric Pelloux-Prayer <pierre-eric.pelloux-prayer@amd.com> [v1]
Reviewed-by: Marek Olšák <marek.olsak@amd.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/11228>
---
 meson.build             |  7 +++++++
 src/util/u_cpu_detect.c | 32 +++++++++++++++++++++++++++-----
 2 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/meson.build b/meson.build
index 156396a936d..af12fb7d02a 100644
--- a/meson.build
+++ b/meson.build
@@ -1288,6 +1288,13 @@ if (cc.has_header_symbol('sys/mkdev.h', 'major') and
   pre_args += '-DMAJOR_IN_MKDEV'
 endif
 
+if cc.check_header('sched.h')
+  pre_args += '-DHAS_SCHED_H'
+  if cc.has_function('sched_getaffinity')
+    pre_args += '-DHAS_SCHED_GETAFFINITY'
+  endif
+endif
+
 if not ['linux'].contains(host_machine.system())
   # Deprecated on Linux and requires <sys/types.h> on FreeBSD and OpenBSD
   if cc.check_header('sys/sysctl.h', prefix : '#include <sys/types.h>')
diff --git a/src/util/u_cpu_detect.c b/src/util/u_cpu_detect.c
index e0ed1a4188b..e8664a66933 100644
--- a/src/util/u_cpu_detect.c
+++ b/src/util/u_cpu_detect.c
@@ -86,6 +86,9 @@
 #endif
 #endif
 
+#if defined(HAS_SCHED_H)
+#include <sched.h>
+#endif
 
 DEBUG_GET_ONCE_BOOL_OPTION(dump_cpu, "GALLIUM_DUMP_CPU", false)
 
@@ -573,12 +576,29 @@ util_cpu_detect_once(void)
       available_cpus = MAX2(1, system_info.dwNumberOfProcessors);
    }
 #elif defined(PIPE_OS_UNIX)
+#  if defined(HAS_SCHED_GETAFFINITY)
+   {
+      /* sched_setaffinity() can be used to further restrict the number of
+       * CPUs on which the process can run.  Use sched_getaffinity() to
+       * determine the true number of available CPUs.
+       *
+       * FIXME: The Linux manual page for sched_getaffinity describes how this
+       * simple implementation will fail with > 1024 CPUs, and we'll fall back
+       * to the _SC_NPROCESSORS_ONLN path.  Support for > 1024 CPUs can be
+       * added to this path once someone has such a system for testing.
+       */
+      cpu_set_t affin;
+      if (sched_getaffinity(getpid(), sizeof(affin), &affin) == 0)
+         available_cpus = CPU_COUNT(&affin);
+   }
+#  endif
+
    /* Linux, FreeBSD, DragonFly, and Mac OS X should have
     * _SC_NOPROCESSORS_ONLN.  NetBSD and OpenBSD should have HW_NCPUONLINE.
     * This is what FFmpeg uses on those platforms.
     */
 #  if defined(PIPE_OS_BSD) && defined(HW_NCPUONLINE)
-   {
+   if (available_cpus == 0) {
       const int mib[] = { CTL_HW, HW_NCPUONLINE };
       int ncpu;
       int len = sizeof(ncpu);
@@ -587,11 +607,13 @@ util_cpu_detect_once(void)
       available_cpus = ncpu;
    }
 #  elif defined(_SC_NPROCESSORS_ONLN)
-   available_cpus = sysconf(_SC_NPROCESSORS_ONLN);
-   if (available_cpus == ~0)
-      available_cpus = 1;
+   if (available_cpus == 0) {
+      available_cpus = sysconf(_SC_NPROCESSORS_ONLN);
+      if (available_cpus == ~0)
+         available_cpus = 1;
+   }
 #  elif defined(PIPE_OS_BSD)
-   {
+   if (available_cpus == 0) {
       const int mib[] = { CTL_HW, HW_NCPU };
       int ncpu;
       int len = sizeof(ncpu);