|
16 | 16 | #include <string.h>
|
17 | 17 | #include <sys/ioctl.h>
|
18 | 18 | #include <sys/stat.h>
|
| 19 | +#include <sys/sysinfo.h> |
19 | 20 | #include <sys/types.h>
|
20 | 21 | #include <sys/utsname.h>
|
21 | 22 | #include <unistd.h>
|
@@ -88,28 +89,40 @@ void *get_auxv_entry(int type)
|
88 | 89 |
|
89 | 90 | int pick_online_cpu(void)
|
90 | 91 | {
|
91 |
| - cpu_set_t mask; |
92 |
| - int cpu; |
| 92 | + int ncpus, cpu = -1; |
| 93 | + cpu_set_t *mask; |
| 94 | + size_t size; |
| 95 | + |
| 96 | + ncpus = get_nprocs_conf(); |
| 97 | + size = CPU_ALLOC_SIZE(ncpus); |
| 98 | + mask = CPU_ALLOC(ncpus); |
| 99 | + if (!mask) { |
| 100 | + perror("malloc"); |
| 101 | + return -1; |
| 102 | + } |
93 | 103 |
|
94 |
| - CPU_ZERO(&mask); |
| 104 | + CPU_ZERO_S(size, mask); |
95 | 105 |
|
96 |
| - if (sched_getaffinity(0, sizeof(mask), &mask)) { |
| 106 | + if (sched_getaffinity(0, size, mask)) { |
97 | 107 | perror("sched_getaffinity");
|
98 |
| - return -1; |
| 108 | + goto done; |
99 | 109 | }
|
100 | 110 |
|
101 | 111 | /* We prefer a primary thread, but skip 0 */
|
102 |
| - for (cpu = 8; cpu < CPU_SETSIZE; cpu += 8) |
103 |
| - if (CPU_ISSET(cpu, &mask)) |
104 |
| - return cpu; |
| 112 | + for (cpu = 8; cpu < ncpus; cpu += 8) |
| 113 | + if (CPU_ISSET_S(cpu, size, mask)) |
| 114 | + goto done; |
105 | 115 |
|
106 | 116 | /* Search for anything, but in reverse */
|
107 |
| - for (cpu = CPU_SETSIZE - 1; cpu >= 0; cpu--) |
108 |
| - if (CPU_ISSET(cpu, &mask)) |
109 |
| - return cpu; |
| 117 | + for (cpu = ncpus - 1; cpu >= 0; cpu--) |
| 118 | + if (CPU_ISSET_S(cpu, size, mask)) |
| 119 | + goto done; |
110 | 120 |
|
111 | 121 | printf("No cpus in affinity mask?!\n");
|
112 |
| - return -1; |
| 122 | + |
| 123 | +done: |
| 124 | + CPU_FREE(mask); |
| 125 | + return cpu; |
113 | 126 | }
|
114 | 127 |
|
115 | 128 | bool is_ppc64le(void)
|
|
0 commit comments