[Openmp-commits] [openmp] 57b4c52 - [OpenMP][FIX] Eliminate race on the IsSPMD global

[Johannes Doerfert via Openmp-commits]

Author: Johannes Doerfert
Date: 2022-02-16T14:44:20-06:00
New Revision: 57b4c5267b7293b1990d8418477b24732ba0468b

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

LOG: [OpenMP][FIX] Eliminate race on the IsSPMD global

The `IsSPMD` global can only be read by threads other than the main
thread *after* initialization is complete. To allow usage of
`mapping::getBlockSize` before initialization is done, we can pass the
`IsSPMD` state explicitly. This is similar to other APIs that take
`IsSPMD` explicitly to avoid such a race, e.g.,
`mapping::isInitialThreadInLevel0(IsSPMD)`

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

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 4f65d28da513f..36cfae7c5efa4 100644
--- a/openmp/libomptarget/DeviceRTL/include/Mapping.h
+++ b/openmp/libomptarget/DeviceRTL/include/Mapping.h
@@ -79,7 +79,12 @@ uint32_t getNumberOfWarpsInBlock();
 uint32_t getBlockId();
 
 /// Return the block size, thus number of threads in the block.
+///
+/// Note: The version taking \p IsSPMD mode explicitly can be used during the
+/// initialization of the target region, that is before `mapping::isSPMDMode()`
+/// can be called by any thread other than the main one.
 uint32_t getBlockSize();
+uint32_t getBlockSize(bool IsSPMD);
 
 /// Return the number of blocks in the kernel.
 uint32_t getNumberOfBlocks();

diff  --git a/openmp/libomptarget/DeviceRTL/src/Kernel.cpp b/openmp/libomptarget/DeviceRTL/src/Kernel.cpp
index 65b554b729731..8b7a8a2495c45 100644
--- a/openmp/libomptarget/DeviceRTL/src/Kernel.cpp
+++ b/openmp/libomptarget/DeviceRTL/src/Kernel.cpp
@@ -100,7 +100,7 @@ int32_t __kmpc_target_init(IdentTy *Ident, int8_t Mode,
   // doing any work.  mapping::getBlockSize() does not include any of the main
   // thread's warp, so none of its threads can ever be active worker threads.
   if (UseGenericStateMachine &&
-      mapping::getThreadIdInBlock() < mapping::getBlockSize())
+      mapping::getThreadIdInBlock() < mapping::getBlockSize(IsSPMD))
     genericStateMachine(Ident);
 
   return mapping::getThreadIdInBlock();

diff  --git a/openmp/libomptarget/DeviceRTL/src/Mapping.cpp b/openmp/libomptarget/DeviceRTL/src/Mapping.cpp
index 29e9b127290d6..23615cb2f4ed1 100644
--- a/openmp/libomptarget/DeviceRTL/src/Mapping.cpp
+++ b/openmp/libomptarget/DeviceRTL/src/Mapping.cpp
@@ -212,11 +212,14 @@ uint32_t mapping::getThreadIdInBlock() {
 
 uint32_t mapping::getWarpSize() { return impl::getWarpSize(); }
 
-uint32_t mapping::getBlockSize() {
+uint32_t mapping::getBlockSize(bool IsSPMD) {
   uint32_t BlockSize = mapping::getNumberOfProcessorElements() -
-                       (!mapping::isSPMDMode() * impl::getWarpSize());
+                       (!IsSPMD * impl::getWarpSize());
   return BlockSize;
 }
+uint32_t mapping::getBlockSize() {
+  return mapping::getBlockSize(mapping::isSPMDMode());
+}
 
 uint32_t mapping::getKernelSize() { return impl::getKernelSize(); }
 

diff  --git a/openmp/libomptarget/DeviceRTL/src/State.cpp b/openmp/libomptarget/DeviceRTL/src/State.cpp
index 800176eb5eda5..a04f5cccb1738 100644
--- a/openmp/libomptarget/DeviceRTL/src/State.cpp
+++ b/openmp/libomptarget/DeviceRTL/src/State.cpp
@@ -236,7 +236,7 @@ struct TeamStateTy {
 TeamStateTy SHARED(TeamState);
 
 void TeamStateTy::init(bool IsSPMD) {
-  ICVState.NThreadsVar = mapping::getBlockSize();
+  ICVState.NThreadsVar = mapping::getBlockSize(IsSPMD);
   ICVState.LevelVar = 0;
   ICVState.ActiveLevelVar = 0;
   ICVState.MaxActiveLevelsVar = 1;