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

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.

Reviewed By: tianshilei1992

Differential Revision: https://reviews.llvm.org/D112874
diff --git a/openmp/libomptarget/DeviceRTL/src/Mapping.cpp b/openmp/libomptarget/DeviceRTL/src/Mapping.cpp
index 9bd26c8..be937d1 100644
--- a/openmp/libomptarget/DeviceRTL/src/Mapping.cpp
+++ b/openmp/libomptarget/DeviceRTL/src/Mapping.cpp
@@ -164,20 +164,30 @@
 } // namespace impl
 } // namespace _OMP
 
-bool mapping::isMainThreadInGenericMode(bool IsSPMD) {
-  if (IsSPMD || icv::Level)
-    return false;
-
-  // Check if this is the last warp in the block.
+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.
+  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();
@@ -220,7 +230,7 @@
 static int SHARED(IsSPMDMode);
 
 void mapping::init(bool IsSPMD) {
-  if (!mapping::getThreadIdInBlock())
+  if (mapping::isInitialThreadInLevel0(IsSPMD))
     IsSPMDMode = IsSPMD;
 }