[OpenMP] Allow affinity to re-detect for child processes

Current atfork() handler for child processes does not reset
the affinity masks array which prevents users from setting their own
affinity in child processes.

Differential Revision: https://reviews.llvm.org/D99218

openmp/runtime/src/z_Linux_util.cpp

@@ -1303,6 +1303,8 @@ static void __kmp_atfork_child(void) {
   if (__kmp_nested_proc_bind.bind_types != NULL) {
     __kmp_nested_proc_bind.bind_types[0] = proc_bind_false;
   }
+  __kmp_affinity_masks = NULL;
+  __kmp_affinity_num_masks = 0;
 #endif // KMP_AFFINITY_SUPPORTED
 
 #if KMP_USE_MONITOR

openmp/runtime/test/affinity/libomp_test_affinity.h

@@ -0,0 +1,131 @@
+#ifndef LIBOMP_TEST_AFFINITY_H
+#define LIBOMP_TEST_AFFINITY_H
+
+#ifndef _GNU_SOURCE
+#define _GNU_SOURCE
+#endif
+#include <sched.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+typedef struct affinity_mask_t {
+  size_t setsize;
+  cpu_set_t *set;
+} affinity_mask_t;
+
+#define AFFINITY_MAX_CPUS (32 * 64)
+
+// Operating system affinity mask API
+static void affinity_mask_zero(affinity_mask_t *mask) {
+  CPU_ZERO_S(mask->setsize, mask->set);
+}
+
+static affinity_mask_t *affinity_mask_alloc() {
+  size_t setsize = CPU_ALLOC_SIZE(AFFINITY_MAX_CPUS);
+  cpu_set_t *set = CPU_ALLOC(AFFINITY_MAX_CPUS);
+  affinity_mask_t *retval = (affinity_mask_t *)malloc(sizeof(affinity_mask_t));
+  retval->setsize = setsize;
+  retval->set = set;
+  affinity_mask_zero(retval);
+  return retval;
+}
+
+static void affinity_mask_free(affinity_mask_t *mask) { CPU_FREE(mask->set); }
+
+static void affinity_mask_copy(affinity_mask_t *dest,
+                               const affinity_mask_t *src) {
+  memcpy(dest->set, src->set, dest->setsize);
+}
+
+static void affinity_mask_set(affinity_mask_t *mask, int cpu) {
+  CPU_SET_S(cpu, mask->setsize, mask->set);
+}
+
+static void affinity_mask_clr(affinity_mask_t *mask, int cpu) {
+  CPU_CLR_S(cpu, mask->setsize, mask->set);
+}
+
+static int affinity_mask_isset(const affinity_mask_t *mask, int cpu) {
+  return CPU_ISSET_S(cpu, mask->setsize, mask->set);
+}
+
+static int affinity_mask_count(const affinity_mask_t *mask) {
+  return CPU_COUNT_S(mask->setsize, mask->set);
+}
+
+static int affinity_mask_equal(const affinity_mask_t *mask1,
+                               const affinity_mask_t *mask2) {
+  return CPU_EQUAL_S(mask1->setsize, mask1->set, mask2->set);
+}
+
+static void get_thread_affinity(affinity_mask_t *mask) {
+  if (sched_getaffinity(0, mask->setsize, mask->set) != 0) {
+    perror("sched_getaffinity()");
+    exit(EXIT_FAILURE);
+  }
+}
+
+static void set_thread_affinity(const affinity_mask_t *mask) {
+  if (sched_setaffinity(0, mask->setsize, mask->set) != 0) {
+    perror("sched_setaffinity()");
+    exit(EXIT_FAILURE);
+  }
+}
+
+static void affinity_update_snprintf_values(char **ptr, size_t *remaining,
+                                            size_t n, size_t *retval) {
+  if (n > *remaining && *remaining > 0) {
+    *ptr += *remaining;
+    *remaining = 0;
+  } else {
+    *ptr += n;
+    *remaining -= n;
+  }
+  *retval += n;
+}
+
+static size_t affinity_mask_snprintf(char *buf, size_t bufsize,
+                                     const affinity_mask_t *mask) {
+  int cpu, need_comma, begin, end;
+  size_t n;
+  char *ptr = buf;
+  size_t remaining = bufsize;
+  size_t retval = 0;
+
+  n = snprintf(ptr, remaining, "%c", '{');
+  affinity_update_snprintf_values(&ptr, &remaining, n, &retval);
+
+  need_comma = 0;
+  for (cpu = 0; cpu < AFFINITY_MAX_CPUS; cpu++) {
+    if (!affinity_mask_isset(mask, cpu))
+      continue;
+    if (need_comma) {
+      n = snprintf(ptr, remaining, "%c", ',');
+      affinity_update_snprintf_values(&ptr, &remaining, n, &retval);
+    }
+    begin = cpu;
+    // Find end of range (inclusive end)
+    for (end = begin + 1; end < AFFINITY_MAX_CPUS; ++end) {
+      if (!affinity_mask_isset(mask, end))
+        break;
+    }
+    end--;
+
+    if (end - begin >= 2) {
+      n = snprintf(ptr, remaining, "%d-%d", begin, end);
+      affinity_update_snprintf_values(&ptr, &remaining, n, &retval);
+    } else if (end - begin == 1) {
+      n = snprintf(ptr, remaining, "%d,%d", begin, end);
+      affinity_update_snprintf_values(&ptr, &remaining, n, &retval);
+    } else if (end - begin == 0) {
+      n = snprintf(ptr, remaining, "%d", begin);
+      affinity_update_snprintf_values(&ptr, &remaining, n, &retval);
+    }
+    need_comma = 1;
+    cpu = end;
+  }
+  n = snprintf(ptr, remaining, "%c", '}');
+  affinity_update_snprintf_values(&ptr, &remaining, n, &retval);
+  return retval;
+}
+#endif

openmp/runtime/test/affinity/redetect.c

@@ -0,0 +1,101 @@
+// RUN: %libomp-compile
+// RUN: env KMP_AFFINITY=none %libomp-run
+// REQUIRES: linux
+
+// Check if forked child process resets affinity properly by restricting
+// child's affinity to a subset of the parent and then checking it after
+// a parallel region
+
+#define _GNU_SOURCE
+#include "libomp_test_affinity.h"
+#include <omp.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <sys/wait.h>
+#include <sys/types.h>
+#include <unistd.h>
+
+// Set the affinity mask of the calling thread to a proper subset of the
+// original affinity mask, specifically, one processor less.
+void set_subset_affinity(affinity_mask_t *mask) {
+  int cpu;
+  affinity_mask_t *original_mask = affinity_mask_alloc();
+  affinity_mask_copy(original_mask, mask);
+  // Find first processor to clear for subset mask
+  for (cpu = 0; cpu <= AFFINITY_MAX_CPUS; ++cpu) {
+    if (affinity_mask_isset(original_mask, cpu)) {
+      affinity_mask_clr(mask, cpu);
+      break;
+    }
+  }
+  affinity_mask_free(original_mask);
+  set_thread_affinity(mask);
+}
+
+int main(int argc, char **argv) {
+  char buf[1024] = {0};
+  char *other_buf;
+  size_t n;
+  int child_exit_status, exit_status;
+  affinity_mask_t *mask = affinity_mask_alloc();
+  get_thread_affinity(mask);
+  n = affinity_mask_snprintf(buf, sizeof(buf), mask);
+  printf("Orignal Mask: %s\n", buf);
+
+  if (affinity_mask_count(mask) == 1) {
+    printf("Only one processor in affinity mask, skipping test.\n");
+    exit(EXIT_SUCCESS);
+  }
+
+  #pragma omp parallel
+  {
+    #pragma omp single
+    printf("Hello! Thread %d executed single region in parent process\n",
+           omp_get_thread_num());
+  }
+
+  pid_t pid = fork();
+  if (pid < 0) {
+    perror("fork()");
+    exit(EXIT_FAILURE);
+  }
+
+  if (pid == 0) {
+    // Let child set a new initial mask
+    set_subset_affinity(mask);
+    #pragma omp parallel
+    {
+      #pragma omp single
+      printf("Hello! Thread %d executed single region in child process\n",
+             omp_get_thread_num());
+    }
+    affinity_mask_t *new_mask = affinity_mask_alloc();
+    get_thread_affinity(new_mask);
+    if (!affinity_mask_equal(mask, new_mask)) {
+      affinity_mask_snprintf(buf, sizeof(buf), mask);
+      fprintf(stderr, "Original Mask = %s\n", buf);
+      affinity_mask_snprintf(buf, sizeof(buf), new_mask);
+      fprintf(stderr, "New Mask = %s\n", buf);
+      affinity_mask_free(new_mask);
+      fprintf(stderr, "Child affinity mask did not reset properly\n");
+      exit(EXIT_FAILURE);
+    }
+    affinity_mask_free(new_mask);
+    exit_status = EXIT_SUCCESS;
+  } else {
+    pid_t child_pid = pid;
+    pid = wait(&child_exit_status);
+    if (pid == -1) {
+      perror("wait()");
+      exit(EXIT_FAILURE);
+    }
+    if (WIFEXITED(child_exit_status)) {
+      exit_status = WEXITSTATUS(child_exit_status);
+    } else {
+      exit_status = EXIT_FAILURE;
+    }
+  }
+
+  affinity_mask_free(mask);
+  return exit_status;
+}