Fix balanced affinity so thread's private affinity mask is updated

[jlpeyton]

Summary

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.

[hbae]

LGTM.

openmp/trunk/runtime/src/kmp.h

 #endif
 extern void __kmp_affinity_determine_capable(const char *env_var);
 extern int __kmp_aux_set_affinity(void **mask);
 extern int __kmp_aux_get_affinity(void **mask);
 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
 #endif /* KMP_AFFINITY_SUPPORTED */
 
 extern void __kmp_cleanup_hierarchy();
 extern void __kmp_get_hierarchy(kmp_uint32 nproc, kmp_bstate_t *thr_bar);
 

openmp/trunk/runtime/src/kmp_affinity.cpp

int __kmp_aux_get_affinity_mask_proc(int proc, void **mask) {
   if (!KMP_CPU_ISSET(proc, __kmp_affin_fullMask)) {
     return 0;
   }
 
   return KMP_CPU_ISSET(proc, (kmp_affin_mask_t *)(*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:
   case affinity_gran_thread:
     break;
   case affinity_gran_core:
     if (__kmp_nThreadsPerCore > 1) {
       fine_gran = false;
if (__kmp_affinity_uniform_topology()) {
     } else { // tid >= big_nth
       coreID = (tid - big_cores) / chunk;
       threadID = ((tid - big_cores) % chunk) % __kmp_nth_per_core;
     }
 
     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) {
       int osID = address2os[coreID * __kmp_nth_per_core + threadID].second;
       KMP_CPU_SET(osID, mask);
     } else {
       for (int i = 0; i < __kmp_nth_per_core; i++) {
         int osID;
         osID = address2os[coreID * __kmp_nth_per_core + i].second;
         KMP_CPU_SET(osID, mask);
       }
     }
     if (__kmp_affinity_verbose) {
       char buf[KMP_AFFIN_MASK_PRINT_LEN];
       __kmp_affinity_print_mask(buf, KMP_AFFIN_MASK_PRINT_LEN, mask);
       KMP_INFORM(BoundToOSProcSet, "KMP_AFFINITY", (kmp_int32)getpid(),
                  __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(
         address2os, __kmp_avail_proc, __kmp_aff_depth - 1);
     int ncores = __kmp_affinity_compute_ncores(address2os, __kmp_avail_proc,
                                                __kmp_aff_depth - 1, core_level);
     int nth_per_core = __kmp_affinity_max_proc_per_core(
         address2os, __kmp_avail_proc, __kmp_aff_depth - 1, core_level);
if (__kmp_affinity_uniform_topology()) {
 
     if (__kmp_affinity_verbose) {
       char buf[KMP_AFFIN_MASK_PRINT_LEN];
       __kmp_affinity_print_mask(buf, KMP_AFFIN_MASK_PRINT_LEN, mask);
       KMP_INFORM(BoundToOSProcSet, "KMP_AFFINITY", (kmp_int32)getpid(),
                  __kmp_gettid(), tid, buf);
     }
     __kmp_set_system_affinity(mask, TRUE);
-    KMP_CPU_FREE_FROM_STACK(mask);
   }
 }
 
 #if KMP_OS_LINUX
 // We don't need this entry for Windows because
 // there is GetProcessAffinityMask() api
 //
 // The intended usage is indicated by these steps:

openmp/trunk/runtime/src/kmp_barrier.cpp

 
 #if OMP_40_ENABLED && KMP_AFFINITY_SUPPORTED
   kmp_proc_bind_t proc_bind = team->t.t_proc_bind;
   if (proc_bind == proc_bind_intel) {
 #endif
 #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
   } else if (proc_bind != proc_bind_false) {
     if (this_thr->th.th_new_place == this_thr->th.th_current_place) {
       KA_TRACE(100, ("__kmp_fork_barrier: T#%d already in correct place %d\n",
                      __kmp_gtid_from_thread(this_thr),
                      this_thr->th.th_current_place));