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

jpeyton52 · jpeyton52 · commit e525f0d4e224 · 2018-09-26T20:43:23.000Z
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
@@ -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
@@ -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
@@ -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

Original file line number	Diff line number	Diff line change
`@@ -1970,7 +1970,7 @@ void __kmp_fork_barrier(int gtid, int tid) {`
`1970`	`1970`	`#if KMP_AFFINITY_SUPPORTED`
`1971`	`1971`	`// Call dynamic affinity settings`
`1972`	`1972`	`if (__kmp_affinity_type == affinity_balanced && team->t.t_size_changed) {`
`1973`		`- __kmp_balanced_affinity(tid, team->t.t_nproc);`
	`1973`	`+ __kmp_balanced_affinity(this_thr, team->t.t_nproc);`
`1974`	`1974`	`}`
`1975`	`1975`	`#endif // KMP_AFFINITY_SUPPORTED`
`1976`	`1976`	`#if OMP_40_ENABLED && KMP_AFFINITY_SUPPORTED`