Add fallback to os.cpu_count() if os.sched_getaffinity() is unavailable (e.g. on Darwin)

bibikar · oleksandr-pavlyk · commit fabb5358465e · 2020-02-04T13:03:46.000-06:00
diff --git a/fft_bench.py b/fft_bench.py
@@ -40,13 +40,19 @@ def valid_dtype(dtype_str):
     return dtype
 
 
+if hasattr(os, 'sched_getaffinity'):
+    default_num_threads = len(os.sched_getaffinity(0))
+else:
+    default_num_threads = os.cpu_count()
+
+
 # Parse args
 parser = argparse.ArgumentParser(description='Benchmark FFT using NumPy and '
                                  'SciPy.')
 
 fft_group = parser.add_argument_group(title='FFT problem arguments')
 fft_group.add_argument('-t', '--threads', '--num-threads', '--core-number',
-                       type=int, default=len(os.sched_getaffinity(0)),
+                       type=int, default=default_num_threads,
                        help='Number of threads to use for FFT computation. '
                        '%(prog)s will attempt to use mkl-service to get/set '
                        'number of threads globally, and will also try to '
diff --git a/perf.py b/perf.py
@@ -63,7 +63,10 @@ def set_threads(num_threads=None, verbose=False):
             print(f'TAG: WARNING: np.__mkl_version__ = {np.__mkl_version__}, '
                   f'but mkl-service module was not found. Number of threads '
                   f'is likely inaccurate!')
-            return len(os.sched_getaffinity(0)), 'len(os.sched_getaffinity(0))'
+            if hasattr(os, 'sched_getaffinity'):
+                return len(os.sched_getaffinity(0)), 'len(os.sched_getaffinity(0))'
+            else:
+                return os.cpu_count(), 'os.cpu_count()'
         else:
             # no MKL, so assume not threaded
             return 1, 'guessing'
