[OpenMP] Eliminate the ThreadStates array in favor of indirection

[jdoerfert]

Summary

If we have thread states, the program is going to be rather slow. If we
don't we want to avoid wasting shared memory. This patch introduces a
slight penalty (malloc + indirection) for the slow path and reduces
resource usage for the fast path.

[jdoerfert]

Use assumption more aggressively

[tianshilei1992]

openmp/libomptarget/DeviceRTL/include/State.h

112

Does it work if we just have ThreadStateTy *ThreadStates?

openmp/libomptarget/DeviceRTL/src/State.cpp

270

Is it a valid state when it reaches to this point? If not, I think it'd better to have an assertion here.

[jdoerfert]

openmp/libomptarget/DeviceRTL/include/State.h

112

We could, yes. Requires to rewrite the code more :(

openmp/libomptarget/DeviceRTL/src/State.cpp

270

State is valid, yes. I guess we can assert load returns non zero.

[jdoerfert]

Add assertion

openmp/libomptarget/DeviceRTL/include/State.h

112

Not without rewriting a lot. We use the pointer to build a list of nested thread states. Otherwise we would need to set it up to build the list and have the head in the outermost array, or traverse the list on a lookup, neither seems worth it now.

[tianshilei1992]

LG

[jdoerfert]

Fix some issues

openmp/libomptarget/DeviceRTL/include/State.h

struct ThreadStateTy {
   }
 
   void init(ThreadStateTy *PreviousTS) {
     ICVState = PreviousTS ? PreviousTS->ICVState : TeamState.ICVState;
     PreviousThreadState = PreviousTS;
   }
 };
 
-extern ThreadStateTy *ThreadStates[mapping::MaxThreadsPerTeam];
+extern ThreadStateTy **ThreadStates;
 #pragma omp allocate(ThreadStates) allocator(omp_pteam_mem_alloc)
 
 /// Initialize the state machinery. Must be called by all threads.
 void init(bool IsSPMD);
 
 /// TODO
 enum ValueKind {
   VK_NThreads,

openmp/libomptarget/DeviceRTL/include/Types.h

 using uint8_t = unsigned char;
 using int16_t = short;
 using uint16_t = unsigned short;
 using int32_t = int;
 using uint32_t = unsigned int;
 using int64_t = long;
 using uint64_t = unsigned long;
 using size_t = decltype(sizeof(char));
+// TODO: Properly implement this
+using intptr_t = int64_t;
+using uintptr_t = uint64_t;
 
 static_assert(sizeof(int8_t) == 1, "type size mismatch");
 static_assert(sizeof(uint8_t) == 1, "type size mismatch");
 static_assert(sizeof(int16_t) == 2, "type size mismatch");
 static_assert(sizeof(uint16_t) == 2, "type size mismatch");
 static_assert(sizeof(int32_t) == 4, "type size mismatch");
 static_assert(sizeof(uint32_t) == 4, "type size mismatch");
 static_assert(sizeof(int64_t) == 8, "type size mismatch");

openmp/libomptarget/DeviceRTL/src/State.cpp

 //===------ State.cpp - OpenMP State & ICV interface ------------- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 //===----------------------------------------------------------------------===//
 
 #include "State.h"
 #include "Configuration.h"
 #include "Debug.h"
 #include "Interface.h"
+#include "Mapping.h"
 #include "Synchronization.h"
 #include "Types.h"
 #include "Utils.h"
 
 using namespace _OMP;
 
 #pragma omp begin declare target device_type(nohost)
 

 
 void state::TeamStateTy::assertEqual(TeamStateTy &Other) const {
   ICVState.assertEqual(Other.ICVState);
   ASSERT(ParallelTeamSize == Other.ParallelTeamSize);
   ASSERT(HasThreadState == Other.HasThreadState);
 }
 
 state::TeamStateTy SHARED(_OMP::state::TeamState);
-
-__attribute__((loader_uninitialized))
-state::ThreadStateTy *_OMP::state::ThreadStates[mapping::MaxThreadsPerTeam];
-#pragma omp allocate(_OMP::state::ThreadStates) allocator(omp_pteam_mem_alloc)
+state::ThreadStateTy **SHARED(_OMP::state::ThreadStates);
 
 namespace {
 
 int returnValIfLevelIsActive(int Level, int Val, int DefaultVal,
                              int OutOfBoundsVal = -1) {
   if (Level == 0)
     return DefaultVal;
   int LevelVar = omp_get_level();
   if (OMP_UNLIKELY(Level < 0 || Level > LevelVar))
     return OutOfBoundsVal;
   int ActiveLevel = icv::ActiveLevel;
   if (OMP_UNLIKELY(Level != ActiveLevel))
     return DefaultVal;
   return Val;
 }
 
 } // namespace
 
 void state::init(bool IsSPMD) {
   SharedMemorySmartStack.init(IsSPMD);
   if (mapping::isInitialThreadInLevel0(IsSPMD)) {
     TeamState.init(IsSPMD);
     DebugEntryRAII::init();
+    ThreadStates = nullptr;
   }
-
-  ThreadStates[mapping::getThreadIdInBlock()] = nullptr;
 }
 
 void state::enterDataEnvironment(IdentTy *Ident) {
   ASSERT(config::mayUseThreadStates() &&
          "Thread state modified while explicitly disabled!");
+  if (!config::mayUseThreadStates())
+    return;
 
   unsigned TId = mapping::getThreadIdInBlock();
   ThreadStateTy *NewThreadState =
       static_cast<ThreadStateTy *>(__kmpc_alloc_shared(sizeof(ThreadStateTy)));
+  uintptr_t *ThreadStatesBitsPtr = reinterpret_cast<uintptr_t *>(&ThreadStates);
+  if (!atomic::load(ThreadStatesBitsPtr, atomic::seq_cst)) {
+    uint32_t Bytes = sizeof(ThreadStates[0]) * mapping::getBlockSize();
+    void *ThreadStatesPtr =
+        memory::allocShared(Bytes, "Thread state array allocation");
+    if (!atomic::cas(ThreadStatesBitsPtr, uintptr_t(0),
+                     reinterpret_cast<uintptr_t>(ThreadStatesPtr),
+                     atomic::seq_cst, atomic::seq_cst))
+      memory::freeShared(ThreadStatesPtr, Bytes,
+                         "Thread state array allocated multiple times");
+    ASSERT(atomic::load(ThreadStatesBitsPtr, atomic::seq_cst) &&
+           "Expected valid thread states bit!");
+  }
   NewThreadState->init(ThreadStates[TId]);
   TeamState.HasThreadState = true;
   ThreadStates[TId] = NewThreadState;
 }
 
 void state::exitDataEnvironment() {
   ASSERT(config::mayUseThreadStates() &&
          "Thread state modified while explicitly disabled!");
 
   unsigned TId = mapping::getThreadIdInBlock();
   resetStateForThread(TId);
 }
 
 void state::resetStateForThread(uint32_t TId) {
+  if (!config::mayUseThreadStates())
+    return;
   if (OMP_LIKELY(!TeamState.HasThreadState || !ThreadStates[TId]))
     return;
 
   ThreadStateTy *PreviousThreadState = ThreadStates[TId]->PreviousThreadState;
   __kmpc_free_shared(ThreadStates[TId], sizeof(ThreadStateTy));
   ThreadStates[TId] = PreviousThreadState;
 }
 
 void state::runAndCheckState(void(Func(void))) {
   TeamStateTy OldTeamState = TeamState;
   OldTeamState.assertEqual(TeamState);
 
   Func();
 
   OldTeamState.assertEqual(TeamState);
 }
 
 void state::assumeInitialState(bool IsSPMD) {
   TeamStateTy InitialTeamState;
   InitialTeamState.init(IsSPMD);
   InitialTeamState.assertEqual(TeamState);
-  ASSERT(!ThreadStates[mapping::getThreadIdInBlock()]);
   ASSERT(mapping::isSPMDMode() == IsSPMD);
 }
 
 extern "C" {
 void omp_set_dynamic(int V) {}
 
 int omp_get_dynamic(void) { return 0; }