[Openmp-commits] [openmp] ccb5d27 - [OpenMP][FIX] Avoid a race between initialization and first state reads

[Johannes Doerfert via Openmp-commits]

Author: Johannes Doerfert
Date: 2021-11-02T23:21:49-05:00
New Revision: ccb5d2726a8bc0092857f9e69963e872d78cd4d8

URL: https://github.com/llvm/llvm-project/commit/ccb5d2726a8bc0092857f9e69963e872d78cd4d8
DIFF: https://github.com/llvm/llvm-project/commit/ccb5d2726a8bc0092857f9e69963e872d78cd4d8.diff

LOG: [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

Added: 
    

Modified: 
    openmp/libomptarget/DeviceRTL/include/Mapping.h
    openmp/libomptarget/DeviceRTL/src/Kernel.cpp
    openmp/libomptarget/DeviceRTL/src/Mapping.cpp
    openmp/libomptarget/DeviceRTL/src/State.cpp

Removed: 
    


################################################################################
diff  --git a/openmp/libomptarget/DeviceRTL/include/Mapping.h b/openmp/libomptarget/DeviceRTL/include/Mapping.h
index a3193f3575640..4f65d28da513f 100644
--- a/openmp/libomptarget/DeviceRTL/include/Mapping.h
+++ b/openmp/libomptarget/DeviceRTL/include/Mapping.h
@@ -34,9 +34,19 @@ bool isSPMDMode();
 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();

diff  --git a/openmp/libomptarget/DeviceRTL/src/Kernel.cpp b/openmp/libomptarget/DeviceRTL/src/Kernel.cpp
index 94bf432acb9a9..bf3d4ca24090b 100644
--- a/openmp/libomptarget/DeviceRTL/src/Kernel.cpp
+++ b/openmp/libomptarget/DeviceRTL/src/Kernel.cpp
@@ -83,7 +83,7 @@ int32_t __kmpc_target_init(IdentTy *Ident, int8_t Mode,
     return -1;
   }
 
-  if (mapping::isMainThreadInGenericMode(IsSPMD))
+  if (mapping::isInitialThreadInLevel0(IsSPMD))
     return -1;
 
   if (UseGenericStateMachine)

diff  --git a/openmp/libomptarget/DeviceRTL/src/Mapping.cpp b/openmp/libomptarget/DeviceRTL/src/Mapping.cpp
index 9bd26c80636ef..be937d1ca69e5 100644
--- a/openmp/libomptarget/DeviceRTL/src/Mapping.cpp
+++ b/openmp/libomptarget/DeviceRTL/src/Mapping.cpp
@@ -164,20 +164,30 @@ 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();
@@ -220,7 +230,7 @@ uint32_t mapping::getNumberOfWarpsInBlock() {
 static int SHARED(IsSPMDMode);
 
 void mapping::init(bool IsSPMD) {
-  if (!mapping::getThreadIdInBlock())
+  if (mapping::isInitialThreadInLevel0(IsSPMD))
     IsSPMDMode = IsSPMD;
 }
 

diff  --git a/openmp/libomptarget/DeviceRTL/src/State.cpp b/openmp/libomptarget/DeviceRTL/src/State.cpp
index 40d10a65817ea..a16fa1b1a0fac 100644
--- a/openmp/libomptarget/DeviceRTL/src/State.cpp
+++ b/openmp/libomptarget/DeviceRTL/src/State.cpp
@@ -366,7 +366,7 @@ void *&state::lookupPtr(ValueKind Kind, bool IsReadonly) {
 
 void state::init(bool IsSPMD) {
   SharedMemorySmartStack.init(IsSPMD);
-  if (!mapping::getThreadIdInBlock())
+  if (mapping::isInitialThreadInLevel0(IsSPMD))
     TeamState.init(IsSPMD);
 
   ThreadStates[mapping::getThreadIdInBlock()] = nullptr;