pythongh-109649: Add affinity parameter to os.cpu_count()

Implement cpu_count(affinity=True) with sched_getaffinity() on Unix
and GetProcessAffinityMask() on Windows.

Changes:

* Fix test_posix.test_sched_getaffinity(): restore the old CPU mask
  when the test completes!
* Doc: Specify that os.cpu_count() counts *logicial* CPUs and mention
  that Linux cgroups are ignored.
* _Py_popcount32() uses UINT32_C() for M1, M2 and M4 constants.
* Add _Py_popcount64(). Add tests on _Py_popcount64().

Author: vstinner

Doc/library/os.rst

@@ -5153,14 +5153,19 @@ Miscellaneous System Information
    .. availability:: Unix.
 
 
-.. function:: cpu_count()
+.. function:: cpu_count(*, affinity=False)
 
-   Return the number of CPUs in the system. Returns ``None`` if undetermined.
+   Return the number of logical CPUs in the system. Returns ``None`` if
+   undetermined.
 
-   This number is not equivalent to the number of CPUs the current process can
-   use.  The number of usable CPUs can be obtained with
-   ``len(os.sched_getaffinity(0))``
+   If *affinity* is true, return the number of logical CPUs the current process
+   can use. On error, raise an exception, instead of returning ``None``.
 
+   Linux control groups, *cgroups*, are not taken in account to get the number
+   of logical CPUs.
+
+   .. versionchanged:: 3.13
+      Add *affinity* parameter.
 
    .. versionadded:: 3.4
 

Doc/whatsnew/3.13.rst

@@ -163,6 +163,13 @@ opcode
   documented or exposed through ``dis``, and were not intended to be
   used externally.
 
+os
+--
+
+* Add *affinity* parameter to :func:`os.cpu_count` to get the number of CPUs
+  the current process can use.
+  (Contributed by Victor Stinner in :gh:`109649`.)
+
 pathlib
 -------
 

Include/internal/pycore_bitutils.h

@@ -110,18 +110,19 @@ _Py_popcount32(uint32_t x)
     // 32-bit SWAR (SIMD Within A Register) popcount
 
     // Binary: 0 1 0 1 ...
-    const uint32_t M1 = 0x55555555;
+    const uint32_t M1 = UINT32_C(0x55555555);
     // Binary: 00 11 00 11. ..
-    const uint32_t M2 = 0x33333333;
+    const uint32_t M2 = UINT32_C(0x33333333);
     // Binary: 0000 1111 0000 1111 ...
-    const uint32_t M4 = 0x0F0F0F0F;
+    const uint32_t M4 = UINT32_C(0x0F0F0F0F);
 
     // Put count of each 2 bits into those 2 bits
     x = x - ((x >> 1) & M1);
     // Put count of each 4 bits into those 4 bits
     x = (x & M2) + ((x >> 2) & M2);
     // Put count of each 8 bits into those 8 bits
     x = (x + (x >> 4)) & M4;
+
     // Sum of the 4 byte counts.
     // Take care when considering changes to the next line. Portability and
     // correctness are delicate here, thanks to C's "integer promotions" (C99
@@ -140,6 +141,45 @@ _Py_popcount32(uint32_t x)
 }
 
 
+static inline int
+_Py_popcount64(uint64_t x)
+{
+#if (defined(__clang__) || defined(__GNUC__))
+
+#if SIZEOF_INT >= 8
+    Py_BUILD_ASSERT(sizeof(x) <= sizeof(unsigned int));
+    return __builtin_popcount(x);
+#else
+    // The C standard guarantees that unsigned long will always be big enough
+    // to hold a uint32_t value without losing information.
+    Py_BUILD_ASSERT(sizeof(x) <= sizeof(unsigned long));
+    return __builtin_popcountl(x);
+#endif
+
+#else
+    // 64-bit SWAR (SIMD Within A Register) popcount
+
+    // Binary: 0 1 0 1 ...
+    const uint64_t M1 = UINT64_C(0x5555555555555555);
+    // Binary: 00 11 00 11. ..
+    const uint64_t M2 = UINT64_C(0x3333333333333333);
+    // Binary: 0000 1111 0000 1111 ...
+    const uint64_t M4 = UINT64_C(0xF0F0F0F0F0F0F0F);
+
+    // Put count of each 2 bits into those 2 bits
+    x = x - ((x >> 1) & M1);
+    // Put count of each 4 bits into those 4 bits
+    x = (x & M2) + ((x >> 2) & M2);
+    // Put count of each 8 bits into those 8 bits
+    x = (x + (x >> 4)) & M4;
+
+    // Sum of the 8 byte counts.
+    // See _Py_popcount32() for the rationale on the final cast.
+    return (uint64_t)(x * UINT64_C(0x0101010101010101)) >> 56;
+#endif
+}
+
+
 // Return the index of the most significant 1 bit in 'x'. This is the smallest
 // integer k such that x < 2**k. Equivalent to floor(log2(x)) + 1 for x != 0.
 static inline int

Include/internal/pycore_global_objects_fini_generated.h

@@ -273,6 +273,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(aclose)
         STRUCT_FOR_ID(add)
         STRUCT_FOR_ID(add_done_callback)
+        STRUCT_FOR_ID(affinity)
         STRUCT_FOR_ID(after_in_child)
         STRUCT_FOR_ID(after_in_parent)
         STRUCT_FOR_ID(aggregate_class)

Include/internal/pycore_global_strings.h

@@ -3996,14 +3996,43 @@ def test_oserror_filename(self):
                     self.fail(f"No exception thrown by {func}")
 
 class CPUCountTests(unittest.TestCase):
+    def check_cpu_count(self, cpus):
+        if cpus is None:
+            # None is valid: os.cpu_count() cannot determine
+            # the number of CPUs
+            return
+
+        self.assertIsInstance(cpus, int)
+        self.assertGreater(cpus, 0)
+
     def test_cpu_count(self):
         cpus = os.cpu_count()
-        if cpus is not None:
-            self.assertIsInstance(cpus, int)
-            self.assertGreater(cpus, 0)
-        else:
+        self.check_cpu_count(cpus)
+
+    def test_cpu_count_affinity(self):
+        cpus = os.cpu_count(affinity=True)
+        self.check_cpu_count(cpus)
+
+    @unittest.skipUnless(hasattr(os, 'sched_setaffinity'),
+                         "don't have sched affinity support")
+    def test_cpu_count_affinity_setaffinity(self):
+        ncpu = os.cpu_count()
+        if ncpu is None:
             self.skipTest("Could not determine the number of CPUs")
 
+        # Disable one CPU
+        mask = os.sched_getaffinity(0)
+        if len(mask) <= 1:
+            self.skipTest(f"sched_getaffinity() returns less than "
+                          f"2 CPUs: {sorted(mask)}")
+        self.addCleanup(os.sched_setaffinity, 0, list(mask))
+        mask.pop()
+        os.sched_setaffinity(0, mask)
+
+        # test cpu_count(affinity=True)
+        affinity = os.cpu_count(affinity=True)
+        self.assertEqual(affinity, ncpu - 1)
+
 
 # FD inheritance check is only useful for systems with process support.
 @support.requires_subprocess()

Include/internal/pycore_runtime_init_generated.h

@@ -1205,6 +1205,7 @@ def test_sched_getaffinity(self):
     @requires_sched_affinity
     def test_sched_setaffinity(self):
         mask = posix.sched_getaffinity(0)
+        self.addCleanup(posix.sched_setaffinity, 0, list(mask))
         if len(mask) > 1:
             # Empty masks are forbidden
             mask.pop()

Include/internal/pycore_unicodeobject_generated.h

@@ -0,0 +1,2 @@
+Add *affinity* parameter to :func:`os.cpu_count` to get the number of CPUs the
+current process can use. Patch by Victor Stinner.