[OpenMP][AIX] Implement __kmp_get_load_balance() for AIX (#91520)

AIX has the `/proc` filesystem where `/proc/<pid>/lwp/<tid>/lwpsinfo` has
the thread state in binary, similar to Linux's
`/proc/<pid>/task/<tid>/stat` where the state is in ASCII. However, the
definition of state info `R` in `lwpsinfo` is `runnable`. In Linux,
state `R` means the thread is `running`. Therefore, `lwpsinfo` is not
ideal for our purpose of getting the current load of the system. This
patch uses `perfstat_cpu()` in AIX system library `libperfstat.a` to
obtain the number of threads current running on logical CPUs.

2 files changed, 80 insertions, 12 deletions

diff --git a/openmp/runtime/CMakeLists.txt b/openmp/runtime/CMakeLists.txt
index 57ed54b..bcae02e 100644
--- a/openmp/runtime/CMakeLists.txt
+++ b/openmp/runtime/CMakeLists.txt
@@ -132,10 +132,13 @@ set(LIBOMP_ASMFLAGS "" CACHE STRING
   "Appended user specified assembler flags.")
 set(LIBOMP_LDFLAGS "" CACHE STRING
   "Appended user specified linker flags.")
-if("${LIBOMP_ARCH}" STREQUAL "ppc" AND ${CMAKE_SYSTEM_NAME} MATCHES "AIX")
-  # PPC (32-bit) on AIX needs libatomic for __atomic_load_8, etc.
-  set(LIBOMP_LIBFLAGS "-latomic" CACHE STRING
+if(${CMAKE_SYSTEM_NAME} MATCHES "AIX")
+  set(LIBOMP_LIBFLAGS "-lperfstat" CACHE STRING
     "Appended user specified linked libs flags. (e.g., -lm)")
+  if("${LIBOMP_ARCH}" STREQUAL "ppc")
+    # PPC (32-bit) on AIX needs libatomic for __atomic_load_8, etc.
+    set(LIBOMP_LIBFLAGS "${LIBOMP_LIBFLAGS} -latomic")
+  endif()
 else()
   set(LIBOMP_LIBFLAGS "" CACHE STRING
     "Appended user specified linked libs flags. (e.g., -lm)")
diff --git a/openmp/runtime/src/z_Linux_util.cpp b/openmp/runtime/src/z_Linux_util.cpp
index affb577..7c90740a 100644
--- a/openmp/runtime/src/z_Linux_util.cpp
+++ b/openmp/runtime/src/z_Linux_util.cpp
@@ -31,6 +31,7 @@
 #include <sys/resource.h>
 #if KMP_OS_AIX
 #include <sys/ldr.h>
+#include <libperfstat.h>
 #else
 #include <sys/syscall.h>
 #endif
@@ -2427,6 +2428,79 @@ int __kmp_get_load_balance(int max) {
   return ret_avg;
 }
 
+#elif KMP_OS_AIX
+
+// The function returns number of running (not sleeping) threads, or -1 in case
+// of error.
+int __kmp_get_load_balance(int max) {
+
+  static int glb_running_threads = 0; // Saved count of the running threads for
+                                      // the thread balance algorithm.
+  static double glb_call_time = 0; // Thread balance algorithm call time.
+  int running_threads = 0; // Number of running threads in the system.
+
+  double call_time = 0.0;
+
+  __kmp_elapsed(&call_time);
+
+  if (glb_call_time &&
+      (call_time - glb_call_time < __kmp_load_balance_interval))
+    return glb_running_threads;
+
+  glb_call_time = call_time;
+
+  if (max <= 0) {
+    max = INT_MAX;
+  }
+
+  // Check how many perfstat_cpu_t structures are available.
+  int logical_cpus = perfstat_cpu(NULL, NULL, sizeof(perfstat_cpu_t), 0);
+  if (logical_cpus <= 0) {
+    glb_call_time = -1;
+    return -1;
+  }
+
+  perfstat_cpu_t *cpu_stat = (perfstat_cpu_t *)KMP_INTERNAL_MALLOC(
+      logical_cpus * sizeof(perfstat_cpu_t));
+  if (cpu_stat == NULL) {
+    glb_call_time = -1;
+    return -1;
+  }
+
+  // Set first CPU as the name of the first logical CPU for which the info is
+  // desired.
+  perfstat_id_t first_cpu_name;
+  strcpy(first_cpu_name.name, FIRST_CPU);
+
+  // Get the stat info of logical CPUs.
+  int rc = perfstat_cpu(&first_cpu_name, cpu_stat, sizeof(perfstat_cpu_t),
+                        logical_cpus);
+  KMP_DEBUG_ASSERT(rc == logical_cpus);
+  if (rc <= 0) {
+    KMP_INTERNAL_FREE(cpu_stat);
+    glb_call_time = -1;
+    return -1;
+  }
+  for (int i = 0; i < logical_cpus; ++i) {
+    running_threads += cpu_stat[i].runque;
+    if (running_threads >= max)
+      break;
+  }
+
+  // There _might_ be a timing hole where the thread executing this
+  // code gets skipped in the load balance, and running_threads is 0.
+  // Assert in the debug builds only!!!
+  KMP_DEBUG_ASSERT(running_threads > 0);
+  if (running_threads <= 0)
+    running_threads = 1;
+
+  KMP_INTERNAL_FREE(cpu_stat);
+
+  glb_running_threads = running_threads;
+
+  return running_threads;
+}
+
 #else // Linux* OS
 
 // The function returns number of running (not sleeping) threads, or -1 in case
@@ -2498,14 +2572,9 @@ int __kmp_get_load_balance(int max) {
 
   proc_entry = readdir(proc_dir);
   while (proc_entry != NULL) {
-#if KMP_OS_AIX
-    // Proc entry name starts with a digit. Assume it is a  process' directory.
-    if (isdigit(proc_entry->d_name[0])) {
-#else
     // Proc entry is a directory and name starts with a digit. Assume it is a
     // process' directory.
     if (proc_entry->d_type == DT_DIR && isdigit(proc_entry->d_name[0])) {
-#endif
 
 #ifdef KMP_DEBUG
       ++total_processes;
@@ -2549,11 +2618,7 @@ int __kmp_get_load_balance(int max) {
         task_entry = readdir(task_dir);
         while (task_entry != NULL) {
           // It is a directory and name starts with a digit.
-#if KMP_OS_AIX
-          if (isdigit(task_entry->d_name[0])) {
-#else
           if (proc_entry->d_type == DT_DIR && isdigit(task_entry->d_name[0])) {
-#endif
 
             // Construct complete stat file path. Easiest way would be:
             //  __kmp_str_buf_print( & stat_path, "%s/%s/stat", task_path.str,