[OpenMP] Introduce aligned synchronization into the new device RT

[jdoerfert]

Summary

We will later use the fact that a barrier is aligned to reason about
thread divergence. For now we introduce the assumption and some more
documentation.

[tianshilei1992]

openmp/libomptarget/DeviceRTL/src/Synchronization.cpp

186

We don't have to use a named barrier for aligned barrier, do we?

[jdoerfert]

openmp/libomptarget/DeviceRTL/src/Synchronization.cpp

186

I guess not.

[jdoerfert]

use syncthreads

[JonChesterfield]

openmp/libomptarget/DeviceRTL/src/Parallelism.cpp

97

Thank you for the comments here

[tianshilei1992]

LGTM

openmp/libomptarget/DeviceRTL/include/Synchronization.h

 void init(bool IsSPMD);
 
 /// Synchronize all threads in a warp identified by \p Mask.
 void warp(LaneMaskTy Mask);
 
 /// Synchronize all threads in a block.
 void threads();
 
+/// Synchronizing threads is allowed even if they all hit different instances of
+/// `synchronize::threads()`. However, `synchronize::threadsAligned()` is more
+/// restrictive in that it requires all threads to hit the same instance. The
+/// noinline is removed by the openmp-opt pass and helps to preserve the
+/// information till then.
+///{
+#pragma omp begin assumes ext_aligned_barrier
+
+/// Synchronize all threads in a block, they are are reaching the same
+/// instruction (hence all threads in the block are "aligned").
+__attribute__((noinline)) void threadsAligned();
+
+#pragma omp end assumes
+///}
+
 } // namespace synchronize
 
 namespace fence {
 
 /// Memory fence with \p Ordering semantics for the team.
 void team(int Ordering);
 
 /// Memory fence with \p Ordering semantics for the contention group.

openmp/libomptarget/DeviceRTL/src/Kernel.cpp

 ///
 /// \param Ident               Source location identification, can be NULL.
 ///
 int32_t __kmpc_target_init(IdentTy *Ident, int8_t Mode,
                            bool UseGenericStateMachine, bool) {
   const bool IsSPMD = Mode & OMP_TGT_EXEC_MODE_SPMD;
   if (IsSPMD) {
     inititializeRuntime(/* IsSPMD */ true);
-    synchronize::threads();
+    synchronize::threadsAligned();
   } else {
     inititializeRuntime(/* IsSPMD */ false);
     // No need to wait since only the main threads will execute user
     // code and workers will run into a barrier right away.
   }
 
   if (IsSPMD) {
     state::assumeInitialState(IsSPMD);

openmp/libomptarget/DeviceRTL/src/Parallelism.cpp

if (OMP_UNLIKELY(!if_expr || icv::Level)) {
     ++icv::Level;
     invokeMicrotask(TId, 0, fn, args, nargs);
     state::exitDataEnvironment();
     return;
   }
 
   uint32_t NumThreads = determineNumberOfThreads(num_threads);
   if (mapping::isSPMDMode()) {
-    synchronize::threads();
+    // Avoid the race between the read of the `icv::Level` above and the write
+    // below by synchronizing all threads here.
+    synchronize::threadsAligned();
     {
+      // Note that the order here is important. `icv::Level` has to be updated
+      // last or the other updates will cause a thread specific state to be
+      // created.
       state::ValueRAII ParallelTeamSizeRAII(state::ParallelTeamSize, NumThreads,
                                             1u, TId == 0);
       state::ValueRAII ActiveLevelRAII(icv::ActiveLevel, 1u, 0u, TId == 0);
       state::ValueRAII LevelRAII(icv::Level, 1u, 0u, TId == 0);
-      synchronize::threads();
+
+      // Synchronize all threads after the main thread (TId == 0) set up the
+      // team state properly.
+      synchronize::threadsAligned();
+
+      ASSERT(state::ParallelTeamSize == NumThreads);
+      ASSERT(icv::ActiveLevel == 1u);
+      ASSERT(icv::Level == 1u);
 
       if (TId < NumThreads)
         invokeMicrotask(TId, 0, fn, args, nargs);
-      synchronize::threads();
+
+      // Synchronize all threads at the end of a parallel region.
+      synchronize::threadsAligned();
     }
+
+    ASSERT(state::ParallelTeamSize == 1u);
+    ASSERT(icv::ActiveLevel == 0u);
+    ASSERT(icv::Level == 0u);
     return;
   }
 
   // We do *not* create a new data environment because all threads in the team
   // that are active are now running this parallel region. They share the
   // TeamState, which has an increase level-var and potentially active-level
   // set, but they do not have individual ThreadStates yet. If they ever
   // modify the ICVs beyond this point a ThreadStates will be allocated.
void __kmpc_parallel_51(IdentTy *ident, int32_t, int32_t if_expr,
   if (nargs) {
     __kmpc_begin_sharing_variables(&GlobalArgs, nargs);
 #pragma unroll
     for (int I = 0; I < nargs; I++)
       GlobalArgs[I] = args[I];
   }
 
   {
+    // Note that the order here is important. `icv::Level` has to be updated
+    // last or the other updates will cause a thread specific state to be
+    // created.
     state::ValueRAII ParallelTeamSizeRAII(state::ParallelTeamSize, NumThreads,
                                           1u, true);
     state::ValueRAII ParallelRegionFnRAII(state::ParallelRegionFn, wrapper_fn,
                                           (void *)nullptr, true);
     state::ValueRAII ActiveLevelRAII(icv::ActiveLevel, 1u, 0u, true);
     state::ValueRAII LevelRAII(icv::Level, 1u, 0u, true);
 
     // Master signals work to activate workers.

openmp/libomptarget/DeviceRTL/src/State.cpp

 
 /// Fallback implementations are missing to trigger a link time error.
 /// Implementations for new devices, including the host, should go into a
 /// dedicated begin/end declare variant.
 ///
 ///{
 
 extern "C" {
-void *malloc(uint64_t Size);
-void free(void *Ptr);
+__attribute__((leaf)) void *malloc(uint64_t Size);
+__attribute__((leaf)) void free(void *Ptr);
 }
 
 ///}
 
 /// AMDGCN implementations of the shuffle sync idiom.
 ///
 ///{
 #pragma omp begin declare variant match(device = {arch(amdgcn)})

openmp/libomptarget/DeviceRTL/src/Synchronization.cpp

 }
 
 void syncWarp(__kmpc_impl_lanemask_t) {
   // AMDGCN doesn't need to sync threads in a warp
 }
 
 void syncThreads() { __builtin_amdgcn_s_barrier(); }
 
+void syncThreadsAligned() { syncThreads(); }
+
 void fenceTeam(int Ordering) { __builtin_amdgcn_fence(Ordering, "workgroup"); }
 
 void fenceKernel(int Ordering) { __builtin_amdgcn_fence(Ordering, "agent"); }
 
 void fenceSystem(int Ordering) { __builtin_amdgcn_fence(Ordering, ""); }
 
 #pragma omp end declare variant
 ///}

 
 void syncWarp(__kmpc_impl_lanemask_t Mask) { __nvvm_bar_warp_sync(Mask); }
 
 void syncThreads() {
   constexpr int BarrierNo = 8;
   asm volatile("barrier.sync %0;" : : "r"(BarrierNo) : "memory");
 }
 
+void syncThreadsAligned() { __syncthreads(); }
+
 constexpr uint32_t OMP_SPIN = 1000;
 constexpr uint32_t UNSET = 0;
 constexpr uint32_t SET = 1;
 
 // TODO: This seems to hide a bug in the declare variant handling. If it is
 // called before it is defined
 //       here the overload won't happen. Investigate lalter!
 void unsetLock(omp_lock_t *Lock) {
   (void)atomicExchange((uint32_t *)Lock, UNSET, __ATOMIC_SEQ_CST);
void synchronize::init(bool IsSPMD) {
   if (!IsSPMD)
     impl::namedBarrierInit();
 }
 
 void synchronize::warp(LaneMaskTy Mask) { impl::syncWarp(Mask); }
 
 void synchronize::threads() { impl::syncThreads(); }
 
+void synchronize::threadsAligned() { impl::syncThreadsAligned(); }
+
 void fence::team(int Ordering) { impl::fenceTeam(Ordering); }
 
 void fence::kernel(int Ordering) { impl::fenceKernel(Ordering); }
 
 void fence::system(int Ordering) { impl::fenceSystem(Ordering); }
 
 uint32_t atomic::load(uint32_t *Addr, int Ordering) {
   return impl::atomicLoad(Addr, Ordering);
 }
 
 void atomic::store(uint32_t *Addr, uint32_t V, int Ordering) {
-   impl::atomicStore(Addr, V, Ordering);
 }
 
 uint32_t atomic::inc(uint32_t *Addr, uint32_t V, int Ordering) {
   return impl::atomicInc(Addr, V, Ordering);
 }
 
 uint32_t atomic::add(uint32_t *Addr, uint32_t V, int Ordering) {
   return impl::atomicAdd(Addr, V, Ordering);
void __kmpc_barrier(IdentTy *Loc, int32_t TId) {
   if (mapping::isSPMDMode())
     return __kmpc_barrier_simple_spmd(Loc, TId);
 
   impl::namedBarrier();
 }
 
 __attribute__((noinline)) void __kmpc_barrier_simple_spmd(IdentTy *Loc,
                                                           int32_t TId) {
-  synchronize::threads();
+  synchronize::threadsAligned();
 }
 
 int32_t __kmpc_master(IdentTy *Loc, int32_t TId) {
   return omp_get_team_num() == 0;
 }
 
 void __kmpc_end_master(IdentTy *Loc, int32_t TId) {}