[OpenMP][libomp] Add support for offline CPUs in Linux
If some CPUs are offline, then make sure they are not included in the
fullMask even if norespect is given to KMP_AFFINITY.
Differential Revision: https://reviews.llvm.org/D112274
GitOrigin-RevId: 1dd797168e9f8cb592c254c642f196dab3447092
diff --git a/runtime/src/kmp.h b/runtime/src/kmp.h
index 46a95c1..ee068ab 100644
--- a/runtime/src/kmp.h
+++ b/runtime/src/kmp.h
@@ -4274,6 +4274,15 @@
}
}
}
+ /// Instead of erroring out, return non-zero when
+ /// unsuccessful fopen() for any reason
+ int try_open(const char *filename, const char *mode) {
+ KMP_ASSERT(!f);
+ f = fopen(filename, mode);
+ if (!f)
+ return errno;
+ return 0;
+ }
/// Set the FILE* object to stdout and output there
/// No open call should happen before this call.
void set_stdout() {
diff --git a/runtime/src/kmp_affinity.cpp b/runtime/src/kmp_affinity.cpp
index 1bec347..9521cc4 100644
--- a/runtime/src/kmp_affinity.cpp
+++ b/runtime/src/kmp_affinity.cpp
@@ -26,6 +26,7 @@
#define HWLOC_GROUP_KIND_INTEL_DIE 104
#define HWLOC_GROUP_KIND_WINDOWS_PROCESSOR_GROUP 220
#endif
+#include <ctype.h>
// The machine topology
kmp_topology_t *__kmp_topology = nullptr;
@@ -1030,7 +1031,67 @@
return buf;
}
-void __kmp_affinity_entire_machine_mask(kmp_affin_mask_t *mask) {
+// Return (possibly empty) affinity mask representing the offline CPUs
+// Caller must free the mask
+kmp_affin_mask_t *__kmp_affinity_get_offline_cpus() {
+ kmp_affin_mask_t *offline;
+ KMP_CPU_ALLOC(offline);
+ KMP_CPU_ZERO(offline);
+#if KMP_OS_LINUX
+ int n, begin_cpu, end_cpu;
+ kmp_safe_raii_file_t offline_file;
+ auto skip_ws = [](FILE *f) {
+ int c;
+ do {
+ c = fgetc(f);
+ } while (isspace(c));
+ if (c != EOF)
+ ungetc(c, f);
+ };
+ // File contains CSV of integer ranges representing the offline CPUs
+ // e.g., 1,2,4-7,9,11-15
+ int status = offline_file.try_open("/sys/devices/system/cpu/offline", "r");
+ if (status != 0)
+ return offline;
+ while (!feof(offline_file)) {
+ skip_ws(offline_file);
+ n = fscanf(offline_file, "%d", &begin_cpu);
+ if (n != 1)
+ break;
+ skip_ws(offline_file);
+ int c = fgetc(offline_file);
+ if (c == EOF || c == ',') {
+ // Just single CPU
+ end_cpu = begin_cpu;
+ } else if (c == '-') {
+ // Range of CPUs
+ skip_ws(offline_file);
+ n = fscanf(offline_file, "%d", &end_cpu);
+ if (n != 1)
+ break;
+ skip_ws(offline_file);
+ c = fgetc(offline_file); // skip ','
+ } else {
+ // Syntax problem
+ break;
+ }
+ // Ensure a valid range of CPUs
+ if (begin_cpu < 0 || begin_cpu >= __kmp_xproc || end_cpu < 0 ||
+ end_cpu >= __kmp_xproc || begin_cpu > end_cpu) {
+ continue;
+ }
+ // Insert [begin_cpu, end_cpu] into offline mask
+ for (int cpu = begin_cpu; cpu <= end_cpu; ++cpu) {
+ KMP_CPU_SET(cpu, offline);
+ }
+ }
+#endif
+ return offline;
+}
+
+// Return the number of available procs
+int __kmp_affinity_entire_machine_mask(kmp_affin_mask_t *mask) {
+ int avail_proc = 0;
KMP_CPU_ZERO(mask);
#if KMP_GROUP_AFFINITY
@@ -1043,6 +1104,7 @@
int num = __kmp_GetActiveProcessorCount(group);
for (i = 0; i < num; i++) {
KMP_CPU_SET(i + group * (CHAR_BIT * sizeof(DWORD_PTR)), mask);
+ avail_proc++;
}
}
} else
@@ -1051,10 +1113,18 @@
{
int proc;
+ kmp_affin_mask_t *offline_cpus = __kmp_affinity_get_offline_cpus();
for (proc = 0; proc < __kmp_xproc; proc++) {
+ // Skip offline CPUs
+ if (KMP_CPU_ISSET(proc, offline_cpus))
+ continue;
KMP_CPU_SET(proc, mask);
+ avail_proc++;
}
+ KMP_CPU_FREE(offline_cpus);
}
+
+ return avail_proc;
}
// All of the __kmp_affinity_create_*_map() routines should allocate the
@@ -3561,8 +3631,8 @@
__kmp_affin_fullMask);
KMP_INFORM(InitOSProcSetNotRespect, "KMP_AFFINITY", buf);
}
- __kmp_affinity_entire_machine_mask(__kmp_affin_fullMask);
- __kmp_avail_proc = __kmp_xproc;
+ __kmp_avail_proc =
+ __kmp_affinity_entire_machine_mask(__kmp_affin_fullMask);
#if KMP_OS_WINDOWS
// Set the process affinity mask since threads' affinity
// masks must be subset of process mask in Windows* OS
diff --git a/runtime/src/z_Linux_util.cpp b/runtime/src/z_Linux_util.cpp
index 0580a4e..5cd6ad6 100644
--- a/runtime/src/z_Linux_util.cpp
+++ b/runtime/src/z_Linux_util.cpp
@@ -1801,8 +1801,12 @@
int r = 0;
-#if KMP_OS_LINUX || KMP_OS_DRAGONFLY || KMP_OS_FREEBSD || KMP_OS_NETBSD || \
- KMP_OS_OPENBSD || KMP_OS_HURD
+#if KMP_OS_LINUX
+
+ __kmp_type_convert(sysconf(_SC_NPROCESSORS_CONF), &(r));
+
+#elif KMP_OS_DRAGONFLY || KMP_OS_FREEBSD || KMP_OS_NETBSD || KMP_OS_OPENBSD || \
+ KMP_OS_HURD
__kmp_type_convert(sysconf(_SC_NPROCESSORS_ONLN), &(r));
