[libomptarget][nfc] Refactor amdgpu partial barrier to simplify adding a second one

[JonChesterfield]

Summary

[libomptarget][nfc] Refactor amdgpu partial barrier to simplify adding a second one

D101976 would require a second barrier instance. This NFC to amdgpu makes it
simpler to add one (an extra global, one more line in init). Also renames the
current barrier to L0.

[jdoerfert]

LG

openmp/libomptarget/deviceRTLs/amdgcn/src/target_impl.hip

EXTERN double __kmpc_impl_get_wtime() {
   return 0;
 }
 
 // Warp vote function
 EXTERN __kmpc_impl_lanemask_t __kmpc_impl_activemask() {
   return __builtin_amdgcn_read_exec();
 }
 
-uint32_t __kmpc_L1_Barrier [[clang::loader_uninitialized]];
-#pragma allocate(__kmpc_L1_Barrier) allocator(omp_pteam_mem_alloc)
-
-EXTERN void __kmpc_impl_target_init() {
-  // Don't have global ctors, and shared memory is not zero init
-  __atomic_store_n(&__kmpc_L1_Barrier, 0u, __ATOMIC_RELEASE);
-}
-
-EXTERN void __kmpc_impl_named_sync(uint32_t num_threads) {
+static void pteam_mem_barrier(uint32_t num_threads, uint32_t * barrier_state)
+{
   __atomic_thread_fence(__ATOMIC_ACQUIRE);
 
   uint32_t num_waves = num_threads / WARPSIZE;
 
   // Partial barrier implementation for amdgcn.
   // Uses two 16 bit unsigned counters. One for the number of waves to have
   // reached the barrier, and one to count how many times the barrier has been
   // passed. These are packed in a single atomically accessed 32 bit integer.
   // Low bits for the number of waves, assumed zero before this call.
   // High bits to count the number of times the barrier has been passed.
 
   // precondition: num_waves != 0;
   // invariant: num_waves * WARPSIZE == num_threads;
   // precondition: num_waves < 0xffffu;
 
   // Increment the low 16 bits once, using the lowest active thread.
   uint64_t lowestActiveThread = __kmpc_impl_ffs(__kmpc_impl_activemask()) - 1;
   bool isLowest = GetLaneId() == lowestActiveThread;
 
   if (isLowest) {
-    uint32_t load = __atomic_fetch_add(&__kmpc_L1_Barrier, 1,
+    uint32_t load = __atomic_fetch_add(barrier_state, 1,
                                        __ATOMIC_RELAXED); // commutative
 
     // Record the number of times the barrier has been passed
     uint32_t generation = load & 0xffff0000u;
 
     if ((load & 0x0000ffffu) == (num_waves - 1)) {
       // Reached num_waves in low bits so this is the last wave.
       // Set low bits to zero and increment high bits
       load += 0x00010000u; // wrap is safe
       load &= 0xffff0000u; // because bits zeroed second
 
       // Reset the wave counter and release the waiting waves
-      __atomic_store_n(&__kmpc_L1_Barrier, load, __ATOMIC_RELAXED);
+      __atomic_store_n(barrier_state, load, __ATOMIC_RELAXED);
     } else {
       // more waves still to go, spin until generation counter changes
       do {
         __builtin_amdgcn_s_sleep(0);
-        load = __atomic_load_n(&__kmpc_L1_Barrier, __ATOMIC_RELAXED);
+        load = __atomic_load_n(barrier_state, __ATOMIC_RELAXED);
       } while ((load & 0xffff0000u) == generation);
     }
   }
   __atomic_thread_fence(__ATOMIC_RELEASE);
 }
 
+uint32_t __kmpc_L0_Barrier [[clang::loader_uninitialized]];
+#pragma allocate(__kmpc_L0_Barrier) allocator(omp_pteam_mem_alloc)
+
+EXTERN void __kmpc_impl_target_init() {
+  // Don't have global ctors, and shared memory is not zero init
+  __atomic_store_n(&__kmpc_L0_Barrier, 0u, __ATOMIC_RELEASE);
+}
+
+EXTERN void __kmpc_impl_named_sync(uint32_t num_threads) {
+  pteam_mem_barrier(num_threads, &__kmpc_L0_Barrier);
+}
+
 namespace {
 uint32_t get_grid_dim(uint32_t n, uint16_t d) {
   uint32_t q = n / d;
   return q + (n > q * d);
 }
 uint32_t get_workgroup_dim(uint32_t group_id, uint32_t grid_size,
                            uint16_t group_size) {
   uint32_t r = grid_size - group_id * group_size;