Improvements to process affinity mask setting

[tlwilmar]

Summary

A couple improvements:

• Add ability to limit fullMask size when KMP_HW_SUBSET limits resources.
• Make KMP_HW_SUBSET work for affinity_none, and only limit fullMask in this case.

Patch by Andrey Churbanov.

[AndreyChurbanov]

LGTM

[jlpeyton]

LGTM

openmp/trunk/runtime/src/kmp_affinity.cpp

 //
 kmp_affin_mask_t *__kmp_affin_fullMask = NULL;
 
 static int nCoresPerPkg, nPackages;
 static int __kmp_nThreadsPerCore;
 #ifndef KMP_DFLT_NTH_CORES
 static int __kmp_ncores;
 #endif
+static int *__kmp_pu_os_idx = NULL;
 
 //
 // __kmp_affinity_uniform_topology() doesn't work when called from
 // places which support arbitrarily many levels in the machine topology
 // map, i.e. the non-default cases in __kmp_affinity_create_cpuinfo_map()
 // __kmp_affinity_create_x2apicid_map().
 //
 inline static bool
if (! KMP_AFFINITY_CAPABLE())
         KMP_CPU_FREE(oldMask);
         return 0;
     }
 
     //
     // Allocate the data structure to be returned.
     //
     AddrUnsPair *retval = (AddrUnsPair *)__kmp_allocate(sizeof(AddrUnsPair) * __kmp_avail_proc);
+    __kmp_pu_os_idx = (int*)__kmp_allocate(sizeof(int) * __kmp_avail_proc);
 
     //
     // When affinity is off, this routine will still be called to set
     // __kmp_ncores, as well as __kmp_nThreadsPerCore,
     // nCoresPerPkg, & nPackages.  Make sure all these vars are set
     // correctly, and return if affinity is not enabled.
     //
 
for(socket = hwloc_get_obj_by_type(__kmp_hwloc_topology, HWLOC_OBJ_SOCKET, 0);
                 if(! KMP_CPU_ISSET(pu->os_index, __kmp_affin_fullMask))
                     continue;         // skip inactive (inaccessible) unit
                 KA_TRACE(20, ("Hwloc inserting %d (%d) %d (%d) %d (%d) into address2os\n",
                     socket->os_index, socket->logical_index, core->os_index, core->logical_index, pu->os_index,pu->logical_index));
                 addr.labels[0] = socket_identifier; // package
                 addr.labels[1] = core_identifier; // core
                 addr.labels[2] = pu_identifier; // pu
                 retval[nActiveThreads] = AddrUnsPair(addr, pu->os_index);
+                __kmp_pu_os_idx[nActiveThreads] = pu->os_index; // keep os index for each active pu
                 nActiveThreads++;
                 ++num_active_threads; // count active threads per core
             }
             if (num_active_threads) { // were there any active threads on the core?
                 ++__kmp_ncores;       // count total active cores
                 ++num_active_cores;   // count active cores per socket
                 if (num_active_threads > __kmp_nThreadsPerCore)
                     __kmp_nThreadsPerCore = num_active_threads; // calc maximum
if (__kmp_affinity_verbose) {
         } else {
             KMP_INFORM(InitOSProcSetNotRespect, "KMP_AFFINITY", buf);
         }
         KMP_INFORM(AvailableOSProc, "KMP_AFFINITY", __kmp_avail_proc);
         KMP_INFORM(Uniform, "KMP_AFFINITY");
         KMP_INFORM(Topology, "KMP_AFFINITY", nPackages, nCoresPerPkg,
           __kmp_nThreadsPerCore, __kmp_ncores);
     }
+    KMP_DEBUG_ASSERT(__kmp_pu_os_idx == NULL);
+    __kmp_pu_os_idx = (int*)__kmp_allocate(sizeof(int) * __kmp_avail_proc);
     if (__kmp_affinity_type == affinity_none) {
+        int avail_ct = 0;
+        unsigned int i;
+        KMP_CPU_SET_ITERATE(i, __kmp_affin_fullMask) {
+            if (! KMP_CPU_ISSET(i, __kmp_affin_fullMask))
+                continue;
+            __kmp_pu_os_idx[avail_ct++] = i; // suppose indices are flat
+        }
         return 0;
     }
 
     //
     // Contruct the data structure to be returned.
     //
     *address2os = (AddrUnsPair*)
       __kmp_allocate(sizeof(**address2os) * __kmp_avail_proc);
     int avail_ct = 0;
     unsigned int i;
     KMP_CPU_SET_ITERATE(i, __kmp_affin_fullMask) {
         //
         // Skip this proc if it is not included in the machine model.
         //
         if (! KMP_CPU_ISSET(i, __kmp_affin_fullMask)) {
             continue;
         }
-
+        __kmp_pu_os_idx[avail_ct] = i; // suppose indices are flat
         Address addr(1);
         addr.labels[0] = i;
         (*address2os)[avail_ct++] = AddrUnsPair(addr,i);
     }
     if (__kmp_affinity_verbose) {
         KMP_INFORM(OSProcToPackage, "KMP_AFFINITY");
     }
 
if ((! KMP_AFFINITY_CAPABLE()) || (__kmp_get_proc_group(__kmp_affin_fullMask) >= 0)) {
         return -1;
     }
 
     //
     // Contruct the data structure to be returned.
     //
     *address2os = (AddrUnsPair*)
       __kmp_allocate(sizeof(**address2os) * __kmp_avail_proc);
+    KMP_DEBUG_ASSERT(__kmp_pu_os_idx == NULL);
+    __kmp_pu_os_idx = (int*)__kmp_allocate(sizeof(int) * __kmp_avail_proc);
     int avail_ct = 0;
     int i;
     KMP_CPU_SET_ITERATE(i, __kmp_affin_fullMask) {
         //
         // Skip this proc if it is not included in the machine model.
         //
         if (! KMP_CPU_ISSET(i, __kmp_affin_fullMask)) {
             continue;
         }
-
+        __kmp_pu_os_idx[avail_ct] = i;  // suppose indices are flat
         Address addr(2);
         addr.labels[0] = i / (CHAR_BIT * sizeof(DWORD_PTR));
         addr.labels[1] = i % (CHAR_BIT * sizeof(DWORD_PTR));
         (*address2os)[avail_ct++] = AddrUnsPair(addr,i);
 
         if (__kmp_affinity_verbose) {
             KMP_INFORM(AffOSProcToGroup, "KMP_AFFINITY", i, addr.labels[0],
               addr.labels[1]);
if (__kmp_affinity_verbose) {
             KMP_INFORM(Uniform, "KMP_AFFINITY");
         } else {
             KMP_INFORM(NonUniform, "KMP_AFFINITY");
         }
         KMP_INFORM(Topology, "KMP_AFFINITY", nPackages, nCoresPerPkg,
           __kmp_nThreadsPerCore, __kmp_ncores);
 
     }
-
+    KMP_DEBUG_ASSERT(__kmp_pu_os_idx == NULL);
+    KMP_DEBUG_ASSERT(nApics == __kmp_avail_proc);
+    __kmp_pu_os_idx = (int*)__kmp_allocate(sizeof(int) * __kmp_avail_proc);
+    for (i = 0; i < nApics; ++i) {
+        __kmp_pu_os_idx[i] = threadInfo[i].osId;
+    }
     if (__kmp_affinity_type == affinity_none) {
         __kmp_free(threadInfo);
         KMP_CPU_FREE(oldMask);
         return 0;
     }
 
     //
     // Now that we've determined the number of packages, the number of cores
if (__kmp_affinity_verbose) {
         for (level = 1; level <= pkgLevel; level++) {
             __kmp_str_buf_print(&buf, " x %d", maxCt[level]);
         }
         KMP_INFORM(TopologyExtra, "KMP_AFFINITY", buf.str, nCoresPerPkg,
           __kmp_nThreadsPerCore, __kmp_ncores);
 
         __kmp_str_buf_free(&buf);
     }
-
+    KMP_DEBUG_ASSERT(__kmp_pu_os_idx == NULL);
+    KMP_DEBUG_ASSERT(nApics == __kmp_avail_proc);
+    __kmp_pu_os_idx = (int*)__kmp_allocate(sizeof(int) * __kmp_avail_proc);
+    for (proc = 0; (int)proc < nApics; ++proc) {
+        __kmp_pu_os_idx[proc] = retval[proc].second;
+    }
     if (__kmp_affinity_type == affinity_none) {
         __kmp_free(last);
         __kmp_free(maxCt);
         __kmp_free(counts);
         __kmp_free(totals);
         __kmp_free(retval);
         KMP_CPU_FREE(oldMask);
         return 0;
# if KMP_MIC && REDUCE_TEAM_SIZE
     //
     if ((__kmp_dflt_team_nth == 0) && (teamSize > 0)) {
         __kmp_dflt_team_nth = teamSize;
         KA_TRACE(20, ("__kmp_affinity_create_cpuinfo_map: setting __kmp_dflt_team_nth = %d\n",
           __kmp_dflt_team_nth));
     }
 # endif // KMP_MIC && REDUCE_TEAM_SIZE
 
+    KMP_DEBUG_ASSERT(__kmp_pu_os_idx == NULL);
+    KMP_DEBUG_ASSERT(num_avail == __kmp_avail_proc);
+    __kmp_pu_os_idx = (int*)__kmp_allocate(sizeof(int) * __kmp_avail_proc);
+    for (i = 0; i < num_avail; ++i) { // fill the os indices
+        __kmp_pu_os_idx[i] = threadInfo[i][osIdIndex];
+    }
+
     if (__kmp_affinity_type == affinity_none) {
         __kmp_free(lastId);
         __kmp_free(totals);
         __kmp_free(maxCt);
         __kmp_free(counts);
         CLEANUP_THREAD_INFO;
         return 0;
     }

 # endif /* OMP_40_ENABLED */
 
 #undef ADD_MASK
 #undef ADD_MASK_OSID
 
 static void
 __kmp_apply_thread_places(AddrUnsPair **pAddr, int depth)
 {
+    int i, j, k, n_old = 0, n_new = 0, proc_num = 0;
     if (__kmp_place_num_sockets == 0 &&
         __kmp_place_num_cores == 0 &&
         __kmp_place_num_threads_per_core == 0 )
-        return;   // no topology limiting actions requested, exit
+        goto _exit;   // no topology limiting actions requested, exit
     if (__kmp_place_num_sockets == 0)
         __kmp_place_num_sockets = nPackages;    // use all available sockets
     if (__kmp_place_num_cores == 0)
         __kmp_place_num_cores = nCoresPerPkg;   // use all available cores
     if (__kmp_place_num_threads_per_core == 0 ||
         __kmp_place_num_threads_per_core > __kmp_nThreadsPerCore)
         __kmp_place_num_threads_per_core = __kmp_nThreadsPerCore; // use all HW contexts
 
     if ( !__kmp_affinity_uniform_topology() ) {
         KMP_WARNING( AffHWSubsetNonUniform );
-        return; // don't support non-uniform topology
+        goto _exit; // don't support non-uniform topology
     }
-    if ( depth != 3 ) {
+    if ( depth > 3 ) {
         KMP_WARNING( AffHWSubsetNonThreeLevel );
-        return; // don't support not-3-level topology
+        goto _exit; // don't support not-3-level topology
     }
     if (__kmp_place_socket_offset + __kmp_place_num_sockets > nPackages) {
         KMP_WARNING(AffHWSubsetManySockets);
-        return;
+        goto _exit;
     }
     if ( __kmp_place_core_offset + __kmp_place_num_cores > nCoresPerPkg ) {
         KMP_WARNING( AffHWSubsetManyCores );
-        return;
+        goto _exit;
     }
 
-    AddrUnsPair *newAddr = (AddrUnsPair *)__kmp_allocate( sizeof(AddrUnsPair) *
+    AddrUnsPair *newAddr;
+    if (pAddr) // pAddr is NULL in case of affinity_none
+        newAddr = (AddrUnsPair *)__kmp_allocate( sizeof(AddrUnsPair) *
             __kmp_place_num_sockets * __kmp_place_num_cores * __kmp_place_num_threads_per_core);
 
-    int i, j, k, n_old = 0, n_new = 0;
-    for (i = 0; i < nPackages; ++i)
+    for (i = 0; i < nPackages; ++i) {
         if (i < __kmp_place_socket_offset ||
-            i >= __kmp_place_socket_offset + __kmp_place_num_sockets)
+            i >= __kmp_place_socket_offset + __kmp_place_num_sockets) {
             n_old += nCoresPerPkg * __kmp_nThreadsPerCore; // skip not-requested socket
-        else
-            for (j = 0; j < nCoresPerPkg; ++j) // walk through requested socket
+            if (__kmp_pu_os_idx != NULL) {
+                for (j = 0; j < nCoresPerPkg; ++j) { // walk through skipped socket
+                    for (k = 0; k < __kmp_nThreadsPerCore; ++k) {
+                        KMP_CPU_CLR(__kmp_pu_os_idx[proc_num], __kmp_affin_fullMask);
+                        ++proc_num;
+                    }
+                }
+            }
+        } else {
+            for (j = 0; j < nCoresPerPkg; ++j) { // walk through requested socket
                 if (j < __kmp_place_core_offset ||
-                    j >= __kmp_place_core_offset + __kmp_place_num_cores)
+                    j >= __kmp_place_core_offset + __kmp_place_num_cores) {
                     n_old += __kmp_nThreadsPerCore; // skip not-requested core
-                else
+                    if (__kmp_pu_os_idx != NULL) {
+                        for (k = 0; k < __kmp_nThreadsPerCore; ++k) { // walk through skipped core
+                            KMP_CPU_CLR(__kmp_pu_os_idx[proc_num], __kmp_affin_fullMask);
+                            ++proc_num;
+                        }
+                    }
+                } else {
                     for (k = 0; k < __kmp_nThreadsPerCore; ++k) { // walk through requested core
                         if (k < __kmp_place_num_threads_per_core) {
+                            if (pAddr)
                                 newAddr[n_new] = (*pAddr)[n_old]; // collect requested thread's data
                             n_new++;
+                        } else {
+                            if (__kmp_pu_os_idx != NULL)
+                                KMP_CPU_CLR(__kmp_pu_os_idx[proc_num], __kmp_affin_fullMask);
                         }
                         n_old++;
-                    }
+                        ++proc_num;
+                    }
+                }
+            }
+        }
     KMP_DEBUG_ASSERT(n_old == nPackages * nCoresPerPkg * __kmp_nThreadsPerCore);
     KMP_DEBUG_ASSERT(n_new == __kmp_place_num_sockets * __kmp_place_num_cores *
                      __kmp_place_num_threads_per_core);
 
     nPackages = __kmp_place_num_sockets;                      // correct nPackages
     nCoresPerPkg = __kmp_place_num_cores;                     // correct nCoresPerPkg
     __kmp_nThreadsPerCore = __kmp_place_num_threads_per_core; // correct __kmp_nThreadsPerCore
     __kmp_avail_proc = n_new;                                 // correct avail_proc
     __kmp_ncores = nPackages * __kmp_place_num_cores;         // correct ncores
 
+    if (pAddr) {
         __kmp_free( *pAddr );
         *pAddr = newAddr;      // replace old topology with new one
     }
+_exit:
+    if (__kmp_pu_os_idx != NULL) {
+        __kmp_free(__kmp_pu_os_idx);
+        __kmp_pu_os_idx = NULL;
+    }
+}
 
 
 static AddrUnsPair *address2os = NULL;
 static int           * procarr = NULL;
 static int     __kmp_aff_depth = 0;
 
+#define KMP_EXIT_AFF_NONE                             \
+    KMP_ASSERT(__kmp_affinity_type == affinity_none); \
+    KMP_ASSERT(address2os == NULL);                   \
+    __kmp_apply_thread_places(NULL, 0);               \
+    return;
+
 static void
 __kmp_aux_affinity_initialize(void)
 {
     if (__kmp_affinity_masks != NULL) {
         KMP_ASSERT(__kmp_affin_fullMask != NULL);
         return;
     }
 

 # if KMP_USE_HWLOC
         if (depth < 0) {
             if (__kmp_affinity_verbose) {
                 KMP_INFORM(AffUsingHwloc, "KMP_AFFINITY");
             }
             if(!__kmp_hwloc_error) {
                 depth = __kmp_affinity_create_hwloc_map(&address2os, &msg_id);
                 if (depth == 0) {
-                    KMP_ASSERT(__kmp_affinity_type == affinity_none);
-                    KMP_ASSERT(address2os == NULL);
-                    return;
+                    KMP_EXIT_AFF_NONE;
                 } else if(depth < 0 && __kmp_affinity_verbose) {
                     KMP_INFORM(AffIgnoringHwloc, "KMP_AFFINITY");
                 }
             } else if(__kmp_affinity_verbose) {
                 KMP_INFORM(AffIgnoringHwloc, "KMP_AFFINITY");
             }
         }
 # endif
 
 # if KMP_ARCH_X86 || KMP_ARCH_X86_64
 
         if (depth < 0) {
             if (__kmp_affinity_verbose) {
                 KMP_INFORM(AffInfoStr, "KMP_AFFINITY", KMP_I18N_STR(Decodingx2APIC));
             }
 
             file_name = NULL;
             depth = __kmp_affinity_create_x2apicid_map(&address2os, &msg_id);
             if (depth == 0) {
-                KMP_ASSERT(__kmp_affinity_type == affinity_none);
-                KMP_ASSERT(address2os == NULL);
-                return;
+                KMP_EXIT_AFF_NONE;
             }
 
             if (depth < 0) {
                 if (__kmp_affinity_verbose) {
                     if (msg_id != kmp_i18n_null) {
                         KMP_INFORM(AffInfoStrStr, "KMP_AFFINITY", __kmp_i18n_catgets(msg_id),
                           KMP_I18N_STR(DecodingLegacyAPIC));
                     }
                     else {
                         KMP_INFORM(AffInfoStr, "KMP_AFFINITY", KMP_I18N_STR(DecodingLegacyAPIC));
                     }
                 }
 
                 file_name = NULL;
                 depth = __kmp_affinity_create_apicid_map(&address2os, &msg_id);
                 if (depth == 0) {
-                    KMP_ASSERT(__kmp_affinity_type == affinity_none);
-                    KMP_ASSERT(address2os == NULL);
-                    return;
+                    KMP_EXIT_AFF_NONE;
                 }
             }
         }
 
 # endif /* KMP_ARCH_X86 || KMP_ARCH_X86_64 */
 
 # if KMP_OS_LINUX
 
# if KMP_OS_LINUX
             if (f == NULL) {
                 msg_id = kmp_i18n_str_CantOpenCpuinfo;
             }
             else {
                 file_name = "/proc/cpuinfo";
                 depth = __kmp_affinity_create_cpuinfo_map(&address2os, &line, &msg_id, f);
                 fclose(f);
                 if (depth == 0) {
-                    KMP_ASSERT(__kmp_affinity_type == affinity_none);
-                    KMP_ASSERT(address2os == NULL);
-                    return;
+                    KMP_EXIT_AFF_NONE;
                 }
             }
         }
 
 # endif /* KMP_OS_LINUX */
 
 # if KMP_GROUP_AFFINITY
 
# endif /* KMP_GROUP_AFFINITY */
                     KMP_INFORM(UsingFlatOSFileLine, file_name, line, __kmp_i18n_catgets(msg_id));
                 }
             }
             // FIXME - print msg if msg_id = kmp_i18n_null ???
 
             file_name = "";
             depth = __kmp_affinity_create_flat_map(&address2os, &msg_id);
             if (depth == 0) {
-                KMP_ASSERT(__kmp_affinity_type == affinity_none);
-                KMP_ASSERT(address2os == NULL);
-                return;
+                KMP_EXIT_AFF_NONE;
             }
             KMP_ASSERT(depth > 0);
             KMP_ASSERT(address2os != NULL);
         }
     }
 
     //
     // If the user has specified that a paricular topology discovery method
     // is to be used, then we abort if that method fails.  The exception is
     // group affinity, which might have been implicitly set.
     //
 
 # if KMP_ARCH_X86 || KMP_ARCH_X86_64
 
     else if (__kmp_affinity_top_method == affinity_top_method_x2apicid) {
         if (__kmp_affinity_verbose) {
             KMP_INFORM(AffInfoStr, "KMP_AFFINITY",
               KMP_I18N_STR(Decodingx2APIC));
         }
 
         depth = __kmp_affinity_create_x2apicid_map(&address2os, &msg_id);
         if (depth == 0) {
-            KMP_ASSERT(__kmp_affinity_type == affinity_none);
-            KMP_ASSERT(address2os == NULL);
-            return;
+            KMP_EXIT_AFF_NONE;
         }
         if (depth < 0) {
             KMP_ASSERT(msg_id != kmp_i18n_null);
             KMP_FATAL(MsgExiting, __kmp_i18n_catgets(msg_id));
         }
     }
     else if (__kmp_affinity_top_method == affinity_top_method_apicid) {
         if (__kmp_affinity_verbose) {
             KMP_INFORM(AffInfoStr, "KMP_AFFINITY",
               KMP_I18N_STR(DecodingLegacyAPIC));
         }
 
         depth = __kmp_affinity_create_apicid_map(&address2os, &msg_id);
         if (depth == 0) {
-            KMP_ASSERT(__kmp_affinity_type == affinity_none);
-            KMP_ASSERT(address2os == NULL);
-            return;
+            KMP_EXIT_AFF_NONE;
         }
         if (depth < 0) {
             KMP_ASSERT(msg_id != kmp_i18n_null);
             KMP_FATAL(MsgExiting, __kmp_i18n_catgets(msg_id));
         }
     }
 
 # endif /* KMP_ARCH_X86 || KMP_ARCH_X86_64 */
else if (__kmp_affinity_top_method == affinity_top_method_cpuinfo) {
                 KMP_FATAL(FileLineMsgExiting, filename, line, __kmp_i18n_catgets(msg_id));
             }
             else {
                 KMP_FATAL(FileMsgExiting, filename, __kmp_i18n_catgets(msg_id));
             }
         }
         if (__kmp_affinity_type == affinity_none) {
             KMP_ASSERT(depth == 0);
-            KMP_ASSERT(address2os == NULL);
-            return;
+            KMP_EXIT_AFF_NONE;
         }
     }
 
 # if KMP_GROUP_AFFINITY
 
     else if (__kmp_affinity_top_method == affinity_top_method_group) {
         if (__kmp_affinity_verbose) {
             KMP_INFORM(AffWindowsProcGroupMap, "KMP_AFFINITY");
# endif /* KMP_GROUP_AFFINITY */
 
     else if (__kmp_affinity_top_method == affinity_top_method_flat) {
         if (__kmp_affinity_verbose) {
             KMP_INFORM(AffUsingFlatOS, "KMP_AFFINITY");
         }
 
         depth = __kmp_affinity_create_flat_map(&address2os, &msg_id);
         if (depth == 0) {
-            KMP_ASSERT(__kmp_affinity_type == affinity_none);
-            KMP_ASSERT(address2os == NULL);
-            return;
+            KMP_EXIT_AFF_NONE;
         }
         // should not fail
         KMP_ASSERT(depth > 0);
         KMP_ASSERT(address2os != NULL);
     }
 
 # if KMP_USE_HWLOC
     else if (__kmp_affinity_top_method == affinity_top_method_hwloc) {
         if (__kmp_affinity_verbose) {
             KMP_INFORM(AffUsingHwloc, "KMP_AFFINITY");
         }
         depth = __kmp_affinity_create_hwloc_map(&address2os, &msg_id);
         if (depth == 0) {
-            KMP_ASSERT(__kmp_affinity_type == affinity_none);
-            KMP_ASSERT(address2os == NULL);
-            return;
+            KMP_EXIT_AFF_NONE;
         }
     }
 # endif // KMP_USE_HWLOC
 
     if (address2os == NULL) {
         if (KMP_AFFINITY_CAPABLE()
           && (__kmp_affinity_verbose || (__kmp_affinity_warnings
           && (__kmp_affinity_type != affinity_none)))) {
# endif
 
         default:
         KMP_ASSERT2(0, "Unexpected affinity setting");
     }
 
     __kmp_free(osId2Mask);
     machine_hierarchy.init(address2os, __kmp_avail_proc);
 }
+#undef KMP_EXIT_AFF_NONE
 
 
 void
 __kmp_affinity_initialize(void)
 {
     //
     // Much of the code above was written assumming that if a machine was not
     // affinity capable, then __kmp_affinity_type == affinity_none.  We now