[OpenMP][libomp] Hold old __kmp_threads arrays until library shutdown

[jlpeyton]

Summary

When many nested teams are formed, __kmp_threads may be reallocated
to accommodate new threads. This reallocation causes a data
race when another existing team's thread simultaneously references
__kmp_threads. This patch keeps the old thread arrays around until library
shutdown so these lingering references can complete without issue and
access to __kmp_threads remains a simple array reference.

Fixes: https://github.com/llvm/llvm-project/issues/54708

[tlwilmar]

LGTM

openmp/runtime/src/kmp.h

typedef union KMP_ALIGN_CACHE kmp_root {
   double r_align; /* use worst case alignment */
   char r_pad[KMP_PAD(kmp_base_root_t, CACHE_LINE)];
 } kmp_root_t;
 
 struct fortran_inx_info {
   kmp_int32 data;
 };
 
+// This list type exists to hold old __kmp_threads arrays so that
+// old references to them may complete while reallocation takes place when
+// expanding the array. The items in this list are kept alive until library
+// shutdown.
+typedef struct kmp_old_threads_list_t {
+  kmp_info_t **threads;
+  struct kmp_old_threads_list_t *next;
+} kmp_old_threads_list_t;
+
 /* ------------------------------------------------------------------------ */
 
 extern int __kmp_settings;
 extern int __kmp_duplicate_library_ok;
 #if USE_ITT_BUILD
 extern int __kmp_forkjoin_frames;
 extern int __kmp_forkjoin_frames_mode;
 #endif

 extern int __kmp_nteams;
 extern int __kmp_teams_thread_limit;
 
 /* ------------------------------------------------------------------------- */
 
 /* the following are protected by the fork/join lock */
 /* write: lock  read: anytime */
 extern kmp_info_t **__kmp_threads; /* Descriptors for the threads */
+/* Holds old arrays of __kmp_threads until library shutdown */
+extern kmp_old_threads_list_t *__kmp_old_threads_list;
 /* read/write: lock */
 extern volatile kmp_team_t *__kmp_team_pool;
 extern volatile kmp_info_t *__kmp_thread_pool;
 extern kmp_info_t *__kmp_thread_pool_insert_pt;
 
 // total num threads reachable from some root thread including all root threads
 extern volatile int __kmp_nth;
 /* total number of threads reachable from some root thread including all root

openmp/runtime/src/kmp_global.cpp

 /* STATE mostly syncronized with global lock */
 /* data written to rarely by primary threads, read often by workers */
 /* TODO: None of this global padding stuff works consistently because the order
    of declaration is not necessarily correlated to storage order. To fix this,
    all the important globals must be put in a big structure instead. */
 KMP_ALIGN_CACHE
 kmp_info_t **__kmp_threads = NULL;
 kmp_root_t **__kmp_root = NULL;
+kmp_old_threads_list_t *__kmp_old_threads_list = NULL;
 
 /* data read/written to often by primary threads */
 KMP_ALIGN_CACHE
 volatile int __kmp_nth = 0;
 volatile int __kmp_all_nth = 0;
 volatile kmp_info_t *__kmp_thread_pool = NULL;
 volatile kmp_team_t *__kmp_team_pool = NULL;
 

openmp/runtime/src/kmp_runtime.cpp

#endif
   newThreads = (kmp_info_t **)__kmp_allocate(
       (sizeof(kmp_info_t *) + sizeof(kmp_root_t *)) * newCapacity + CACHE_LINE);
   newRoot =
       (kmp_root_t **)((char *)newThreads + sizeof(kmp_info_t *) * newCapacity);
   KMP_MEMCPY(newThreads, __kmp_threads,
              __kmp_threads_capacity * sizeof(kmp_info_t *));
   KMP_MEMCPY(newRoot, __kmp_root,
              __kmp_threads_capacity * sizeof(kmp_root_t *));
+  // Put old __kmp_threads array on a list. Any ongoing references to the old
+  // list will be valid. This list is cleaned up at library shutdown.
+  kmp_old_threads_list_t *node =
+      (kmp_old_threads_list_t *)__kmp_allocate(sizeof(kmp_old_threads_list_t));
+  node->threads = __kmp_threads;
+  node->next = __kmp_old_threads_list;
+  __kmp_old_threads_list = node;
 
-  kmp_info_t **temp_threads = __kmp_threads;
   *(kmp_info_t * *volatile *)&__kmp_threads = newThreads;
   *(kmp_root_t * *volatile *)&__kmp_root = newRoot;
-  __kmp_free(temp_threads);
   added += newCapacity - __kmp_threads_capacity;
   *(volatile int *)&__kmp_threads_capacity = newCapacity;
 
   if (newCapacity > __kmp_tp_capacity) {
     __kmp_acquire_bootstrap_lock(&__kmp_tp_cached_lock);
     if (__kmp_tp_cached && newCapacity > __kmp_tp_capacity) {
       __kmp_threadprivate_resize_cache(newCapacity);
     } else { // increase __kmp_tp_capacity to correspond with kmp_threads size
#endif /* KMP_AFFINITY_SUPPORTED */
   }
   __kmp_free(__kmp_threads);
   // __kmp_threads and __kmp_root were allocated at once, as single block, so
   // there is no need in freeing __kmp_root.
   __kmp_threads = NULL;
   __kmp_root = NULL;
   __kmp_threads_capacity = 0;
 
+  // Free old __kmp_threads arrays if they exist.
+  kmp_old_threads_list_t *ptr = __kmp_old_threads_list;
+  while (ptr) {
+    kmp_old_threads_list_t *next = ptr->next;
+    __kmp_free(ptr->threads);
+    __kmp_free(ptr);
+    ptr = next;
+  }
+
 #if KMP_USE_DYNAMIC_LOCK
   __kmp_cleanup_indirect_user_locks();
 #else
   __kmp_cleanup_user_locks();
 #endif
 #if OMPD_SUPPORT
   if (ompd_state) {
     __kmp_free(ompd_env_block);