[OpenMP][libomp] Fix disabled affinity

Fix setting affinity type and topology method when affinity is disabled
and fix places that were not taking into account that affinity can be
explicitly disabled by putting proper KMP_AFFINITY_CAPABLE() check.

Differential Revision: https://reviews.llvm.org/D137176

4 files changed, 46 insertions, 2 deletions

diff --git a/openmp/runtime/src/kmp.h b/openmp/runtime/src/kmp.h
index 6f7da88f66ac..8a2bcedb4c9a 100644
--- a/openmp/runtime/src/kmp.h
+++ b/openmp/runtime/src/kmp.h
@@ -3655,6 +3655,8 @@ static inline void __kmp_assign_root_init_mask() {
   }
 }
 static inline void __kmp_reset_root_init_mask(int gtid) {
+  if (!KMP_AFFINITY_CAPABLE())
+    return;
   kmp_info_t *th = __kmp_threads[gtid];
   kmp_root_t *r = th->th.th_root;
   if (r->r.r_uber_thread == th && r->r.r_affinity_assigned) {
diff --git a/openmp/runtime/src/kmp_affinity.cpp b/openmp/runtime/src/kmp_affinity.cpp
index e9d0b99f6417..43bf79403f44 100644
--- a/openmp/runtime/src/kmp_affinity.cpp
+++ b/openmp/runtime/src/kmp_affinity.cpp
@@ -675,7 +675,11 @@ void kmp_topology_t::print(const char *env_var) const {
   kmp_hw_t print_types[KMP_HW_LAST + 2];
 
   // Num Available Threads
-  KMP_INFORM(AvailableOSProc, env_var, num_hw_threads);
+  if (num_hw_threads) {
+    KMP_INFORM(AvailableOSProc, env_var, num_hw_threads);
+  } else {
+    KMP_INFORM(AvailableOSProc, env_var, __kmp_xproc);
+  }
 
   // Uniform or not
   if (is_uniform()) {
@@ -3062,7 +3066,8 @@ static bool __kmp_affinity_create_cpuinfo_map(int *line,
       }
 
       // Skip this proc if it is not included in the machine model.
-      if (!KMP_CPU_ISSET(threadInfo[num_avail][osIdIndex],
+      if (KMP_AFFINITY_CAPABLE() &&
+          !KMP_CPU_ISSET(threadInfo[num_avail][osIdIndex],
                          __kmp_affin_fullMask)) {
         INIT_PROC_INFO(threadInfo[num_avail]);
         continue;
@@ -4525,6 +4530,9 @@ void __kmp_affinity_uninitialize(void) {
     *affinity = KMP_AFFINITY_INIT(affinity->env_var);
   }
   if (__kmp_affin_origMask != NULL) {
+    if (KMP_AFFINITY_CAPABLE()) {
+      __kmp_set_system_affinity(__kmp_affin_origMask, FALSE);
+    }
     KMP_CPU_FREE(__kmp_affin_origMask);
     __kmp_affin_origMask = NULL;
   }
diff --git a/openmp/runtime/src/kmp_settings.cpp b/openmp/runtime/src/kmp_settings.cpp
index 080f4015b6e0..88eb150fc0c2 100644
--- a/openmp/runtime/src/kmp_settings.cpp
+++ b/openmp/runtime/src/kmp_settings.cpp
@@ -6296,6 +6296,15 @@ void __kmp_env_initialize(char const *string) {
           __kmp_affinity_top_method = affinity_top_method_all;
         }
       }
+    } else {
+      // If affinity is disabled, then still need to assign topology method
+      // to attempt machine detection and affinity types
+      if (__kmp_affinity_top_method == affinity_top_method_default)
+        __kmp_affinity_top_method = affinity_top_method_all;
+      if (__kmp_affinity.type == affinity_default)
+        __kmp_affinity.type = affinity_disabled;
+      if (__kmp_hh_affinity.type == affinity_default)
+        __kmp_hh_affinity.type = affinity_disabled;
     }
 
 #ifdef KMP_DEBUG
diff --git a/openmp/runtime/test/affinity/disabled.c b/openmp/runtime/test/affinity/disabled.c
new file mode 100644
index 000000000000..18261010a096
--- /dev/null
+++ b/openmp/runtime/test/affinity/disabled.c
@@ -0,0 +1,25 @@
+// RUN: %libomp-compile
+// RUN: env KMP_AFFINITY=disabled %libomp-run
+// RUN: env KMP_AFFINITY=disabled,reset %libomp-run
+// REQUIRES: affinity
+#include <stdio.h>
+#include <stdlib.h>
+#include <omp.h>
+
+int main() {
+  int nthreads, correct_value;;
+  int a = 0;
+  #pragma omp parallel reduction(+: a)
+  {
+    a += omp_get_thread_num();
+    #pragma omp single
+    nthreads = omp_get_num_threads();
+  }
+  correct_value = nthreads * (nthreads - 1) / 2;
+  if (a != correct_value) {
+    printf("Incorrect value: %d should be %d\n", a, correct_value);
+    return EXIT_FAILURE;
+  }
+  return EXIT_SUCCESS;
+}
+