[OpenMP][FIX] Avoid a race between initialization and first state reads

[jdoerfert]

Summary

When we pick state 0 to initialize state but thread N is going to be the
"main thread", in generic mode, we would require extra synchronization.
Instead, we should pick the main thread to initialize state in generic
mode and any thread in SPMD mode.

[tianshilei1992]

openmp/libomptarget/DeviceRTL/include/Mapping.h

46

nit: in OpenMP term, it is "initial thread".

[jdoerfert]

openmp/libomptarget/DeviceRTL/include/Mapping.h

46

But in generic mode the initial thread, or main thread (thread 0), changes between parallel level 0 and 1. we used to use the one for level 1, now the one for level 0. As this is really asking "should I use this thread for initialization" and not "is this the initial thread (in some unspecified level)" I think the name fits.

[jdoerfert]

openmp/libomptarget/DeviceRTL/include/Mapping.h

46

I'm rewording.

[jdoerfert]

Rename, add helper, and avoid another race

openmp/libomptarget/DeviceRTL/include/Mapping.h

 
 /// Return true if the kernel is executed in SPMD mode.
 bool isSPMDMode();
 
 /// Return true if the kernel is executed in generic mode.
 bool isGenericMode();
 
 /// Return true if the executing thread is the main thread in generic mode.
+/// These functions will lookup state and it is required that that is OK for the
+/// thread and location. See also `isInitialThreadInLevel0` for a stateless
+/// alternative for certain situations, e.g. during initialization.
 bool isMainThreadInGenericMode();
 bool isMainThreadInGenericMode(bool IsSPMD);
 
+/// Return true if this thread is the initial thread in parallel level 0.
+///
+/// The thread for which this returns true should be used for single threaded
+/// initialization tasks. We pick a special thread to ensure there are no
+/// races between the initialization and the first read of initialized state.
+bool isInitialThreadInLevel0(bool IsSPMD);
+
 /// Return true if the executing thread has the lowest Id of the active threads
 /// in the warp.
 bool isLeaderInWarp();
 
 /// Return a mask describing all active threads in the warp.
 LaneMaskTy activemask();
 
 /// Return a mask describing all threads with a smaller Id in the warp.

openmp/libomptarget/DeviceRTL/src/Kernel.cpp

if (IsSPMD) {
     // code and workers will run into a barrier right away.
   }
 
   if (IsSPMD) {
     state::assumeInitialState(IsSPMD);
     return -1;
   }
 
-  if (mapping::isMainThreadInGenericMode(IsSPMD))
+  if (mapping::isInitialThreadInLevel0(IsSPMD))
     return -1;
 
   if (UseGenericStateMachine)
     genericStateMachine(Ident);
 
   return mapping::getThreadIdInBlock();
 }
 

openmp/libomptarget/DeviceRTL/src/Mapping.cpp

 #pragma omp end declare variant
 ///}
 
 uint32_t getWarpSize() { return getGridValue().GV_Warp_Size; }
 
 } // namespace impl
 } // namespace _OMP
 
+static bool isInLastWarp() {
+  uint32_t MainTId = (mapping::getNumberOfProcessorElements() - 1) &
+                     ~(mapping::getWarpSize() - 1);
+  return mapping::getThreadIdInBlock() == MainTId;
+}
+
 bool mapping::isMainThreadInGenericMode(bool IsSPMD) {
   if (IsSPMD || icv::Level)
     return false;
 
   // Check if this is the last warp in the block.
-  uint32_t MainTId = (mapping::getNumberOfProcessorElements() - 1) &
-                     ~(mapping::getWarpSize() - 1);
-  return mapping::getThreadIdInBlock() == MainTId;
+  return isInLastWarp();
 }
 
 bool mapping::isMainThreadInGenericMode() {
   return mapping::isMainThreadInGenericMode(mapping::isSPMDMode());
 }
 
+bool mapping::isInitialThreadInLevel0(bool IsSPMD) {
+  if (IsSPMD)
+    return mapping::getThreadIdInBlock() == 0;
+  return isInLastWarp();
+}
+
 bool mapping::isLeaderInWarp() {
   __kmpc_impl_lanemask_t Active = mapping::activemask();
   __kmpc_impl_lanemask_t LaneMaskLT = mapping::lanemaskLT();
   return utils::popc(Active & LaneMaskLT) == 0;
 }
 
 LaneMaskTy mapping::activemask() { return impl::activemask(); }
 

 }
 
 /// Execution mode
 ///
 ///{
 static int SHARED(IsSPMDMode);
 
 void mapping::init(bool IsSPMD) {
-  if (!mapping::getThreadIdInBlock())
+  if (mapping::isInitialThreadInLevel0(IsSPMD))
     IsSPMDMode = IsSPMD;
 }
 
 bool mapping::isSPMDMode() { return IsSPMDMode; }
 
 bool mapping::isGenericMode() { return !isSPMDMode(); }
 ///}
 

openmp/libomptarget/DeviceRTL/src/State.cpp

void *&state::lookupPtr(ValueKind Kind, bool IsReadonly) {
   default:
     break;
   }
   __builtin_unreachable();
 }
 
 void state::init(bool IsSPMD) {
   SharedMemorySmartStack.init(IsSPMD);
-  if (!mapping::getThreadIdInBlock())
+  if (mapping::isInitialThreadInLevel0(IsSPMD))
     TeamState.init(IsSPMD);
 
   ThreadStates[mapping::getThreadIdInBlock()] = nullptr;
 }
 
 void state::enterDataEnvironment() {
   unsigned TId = mapping::getThreadIdInBlock();
   ThreadStateTy *NewThreadState =