[PATCH] selftests: powerpc: Fix online CPU selection

[Harish <harish@linux.ibm.com>]

On systems with large number of cpus, test fails trying to set
affinity by calling sched_setaffinity() with smaller size for
cpuset. This patch fixes it by making sure that the size of
allocated cpu set is dependent on the number of CPUs as reported
by get_nprocs().

Reported-by: Shirisha Ganta <shiganta@in.ibm.com>
Signed-off-by: Harish <harish@linux.ibm.com>
Signed-off-by: Sandipan Das <sandipan@linux.ibm.com>
---
 .../powerpc/benchmarks/context_switch.c        | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/tools/testing/selftests/powerpc/benchmarks/context_switch.c b/tools/testing/selftests/powerpc/benchmarks/context_switch.c
index a2e8c9da7fa5..de6c49d6f88f 100644
--- a/tools/testing/selftests/powerpc/benchmarks/context_switch.c
+++ b/tools/testing/selftests/powerpc/benchmarks/context_switch.c
@@ -19,6 +19,7 @@
 #include <limits.h>
 #include <sys/time.h>
 #include <sys/syscall.h>
+#include <sys/sysinfo.h>
 #include <sys/types.h>
 #include <sys/shm.h>
 #include <linux/futex.h>
@@ -104,8 +105,9 @@ static void start_thread_on(void *(*fn)(void *), void *arg, unsigned long cpu)
 
 static void start_process_on(void *(*fn)(void *), void *arg, unsigned long cpu)
 {
-	int pid;
-	cpu_set_t cpuset;
+	int pid, ncpus;
+	cpu_set_t *cpuset;
+	size_t size;
 
 	pid = fork();
 	if (pid == -1) {
@@ -116,12 +118,16 @@ static void start_process_on(void *(*fn)(void *), void *arg, unsigned long cpu)
 	if (pid)
 		return;
 
-	CPU_ZERO(&cpuset);
-	CPU_SET(cpu, &cpuset);
+	size = CPU_ALLOC_SIZE(ncpus);
+	ncpus = get_nprocs();
+	cpuset = CPU_ALLOC(ncpus);
+	CPU_ZERO_S(size, cpuset);
+	CPU_SET_S(cpu, size, cpuset);
 
-	if (sched_setaffinity(0, sizeof(cpuset), &cpuset)) {
+	if (sched_setaffinity(0, size, cpuset)) {
 		perror("sched_setaffinity");
-		exit(1);
+		CPU_FREE(cpuset);
+		exit(-1);
 	}
 
 	fn(arg);
-- 
2.24.1