[OpenMP] Fix balanced affinity so thread's private affinity mask is updated

Balanced affinity only updated the thread's affinity with the operating system.
This change also has the thread's private mask reflect that change as well so
that any API that probes the thread's affinity mask will report the correct
mask value.

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

llvm-svn: 343142

diff --git a/openmp/runtime/src/kmp.h b/openmp/runtime/src/kmp.h
index 34c722a..ce8fddf 100644 (file)
--- a/openmp/runtime/src/kmp.h
+++ b/openmp/runtime/src/kmp.h
@@ -3362,7 +3362,7 @@ extern int __kmp_aux_get_affinity_max_proc();
 extern int __kmp_aux_set_affinity_mask_proc(int proc, void **mask);
 extern int __kmp_aux_unset_affinity_mask_proc(int proc, void **mask);
 extern int __kmp_aux_get_affinity_mask_proc(int proc, void **mask);
-extern void __kmp_balanced_affinity(int tid, int team_size);
+extern void __kmp_balanced_affinity(kmp_info_t *th, int team_size);
 #if KMP_OS_LINUX
 extern int kmp_set_thread_affinity_mask_initial(void);
 #endif

diff --git a/openmp/runtime/src/kmp_affinity.cpp b/openmp/runtime/src/kmp_affinity.cpp
index bda80ea..a9a21cf 100644 (file)
--- a/openmp/runtime/src/kmp_affinity.cpp
+++ b/openmp/runtime/src/kmp_affinity.cpp
@@ -5012,8 +5012,10 @@ int __kmp_aux_get_affinity_mask_proc(int proc, void **mask) {
 }
 
 // Dynamic affinity settings - Affinity balanced
-void __kmp_balanced_affinity(int tid, int nthreads) {
+void __kmp_balanced_affinity(kmp_info_t *th, int nthreads) {
+  KMP_DEBUG_ASSERT(th);
   bool fine_gran = true;
+  int tid = th->th.th_info.ds.ds_tid;
 
   switch (__kmp_affinity_gran) {
   case affinity_gran_fine:
@@ -5061,8 +5063,7 @@ void __kmp_balanced_affinity(int tid, int nthreads) {
     KMP_DEBUG_ASSERT2(KMP_AFFINITY_CAPABLE(),
                       "Illegal set affinity operation when not capable");
 
-    kmp_affin_mask_t *mask;
-    KMP_CPU_ALLOC_ON_STACK(mask);
+    kmp_affin_mask_t *mask = th->th.th_affin_mask;
     KMP_CPU_ZERO(mask);
 
     if (fine_gran) {
@@ -5082,11 +5083,9 @@ void __kmp_balanced_affinity(int tid, int nthreads) {
                  __kmp_gettid(), tid, buf);
     }
     __kmp_set_system_affinity(mask, TRUE);
-    KMP_CPU_FREE_FROM_STACK(mask);
   } else { // Non-uniform topology
 
-    kmp_affin_mask_t *mask;
-    KMP_CPU_ALLOC_ON_STACK(mask);
+    kmp_affin_mask_t *mask = th->th.th_affin_mask;
     KMP_CPU_ZERO(mask);
 
     int core_level = __kmp_affinity_find_core_level(
@@ -5250,7 +5249,6 @@ void __kmp_balanced_affinity(int tid, int nthreads) {
                  __kmp_gettid(), tid, buf);
     }
     __kmp_set_system_affinity(mask, TRUE);
-    KMP_CPU_FREE_FROM_STACK(mask);
   }
 }
 

diff --git a/openmp/runtime/src/kmp_barrier.cpp b/openmp/runtime/src/kmp_barrier.cpp
index e3925cb..2b78b54 100644 (file)
--- a/openmp/runtime/src/kmp_barrier.cpp
+++ b/openmp/runtime/src/kmp_barrier.cpp
@@ -1970,7 +1970,7 @@ void __kmp_fork_barrier(int gtid, int tid) {
 #if KMP_AFFINITY_SUPPORTED
     // Call dynamic affinity settings
     if (__kmp_affinity_type == affinity_balanced && team->t.t_size_changed) {
-      __kmp_balanced_affinity(tid, team->t.t_nproc);
+      __kmp_balanced_affinity(this_thr, team->t.t_nproc);
     }
 #endif // KMP_AFFINITY_SUPPORTED
 #if OMP_40_ENABLED && KMP_AFFINITY_SUPPORTED