[OpenMP] Added memory scope to atomic::inc API and used the device scope in reduction.

[dhruvachak]

Summary

With https://reviews.llvm.org/D137524, memory scope and ordering
attributes are being used to generate the required instructions for
atomic inc/dec on AMDGPU. This patch adds the memory scope attribute to
the atomic::inc API and uses the device scope in reduction. Without
the device scope in atomic_inc, the default system scope leads to
unnecessary L2 write-backs/invalidates.

[dhruvachak]

The memscope types added are kept consistent with the attributes in OpenMP spec 6.0 TR11. I think the system, device, and team scopes should be sufficient for the OpenMP use cases. Additional scopes are supported by the AMDGPU backend but they haven't been added here.

[arsenm]

openmp/libomptarget/DeviceRTL/include/Synchronization.h

29

Why does this need a new subset abstraction for scopes? I could understand a complete enum around all the names target scopes

openmp/libomptarget/DeviceRTL/src/Reduction.cpp

226

I doubt this is target specific

[dhruvachak]

openmp/libomptarget/DeviceRTL/include/Synchronization.h

29

I tried to keep the memscope arch-independent, hence the additional abstraction. I kept the same memscopes as what the upcoming OpenMP spec is going to have. This way, the call to atomic::inc in DeviceRTL (e.g. in reduction) does not have to be arch-dependent. In other words, I did not want to pass in "agent" directly from the reduction code for amdgpu, as an example.

openmp/libomptarget/DeviceRTL/src/Reduction.cpp

226

Well, I don't know whether it is required for NVPTX for example, so I kept it around.

As an aside, I think the system fence is too conservative, but again I kept it unchanged for archs other than amdgpu.

[jdoerfert]

openmp/libomptarget/DeviceRTL/include/Synchronization.h

33

Can we out a namespace or sth around this, atomic::device and atomic::all are confusing.
What about atomic::scope::ABC?

openmp/libomptarget/DeviceRTL/src/Synchronization.cpp

207

The above is expanded below.

213–226

[jdoerfert]

openmp/libomptarget/DeviceRTL/src/Reduction.cpp

226

As an aside, I think the system fence is too conservative, but again I kept it unchanged for archs other than amdgpu.

1. This is unrelated.
2. Could you elaborate on why you think the change is sound (on AMDGPUs)?

[dhruvachak]

openmp/libomptarget/DeviceRTL/src/Reduction.cpp

226

1. Agreed this is unrelated. I can separate this out if you like.
2. With reference to https://llvm.org/docs/AMDGPUUsage.html#amdgpu-amdhsa-memory-model-gfx90a (gfx90a as an example), the guide points out the required instructions. And that's what we get with this patch for the atomic inc:

s_waitcnt vmcnt(0) lgkmcnt(0)
global_atomic_inc v0, v1, v0, s[2:3] glc
s_waitcnt vmcnt(0)
buffer_wbinvl1_vol

If we bring back the system fence, here's the resulting sequence (fence followed by the atomic inc):

buffer_wbl2
s_waitcnt vmcnt(0) lgkmcnt(0)
buffer_invl2
buffer_wbinvl1_vol
s_waitcnt vmcnt(0) lgkmcnt(0)
global_atomic_inc v0, v1, v0, s[2:3] glc
s_waitcnt vmcnt(0)
buffer_wbinvl1_vol

According to the guide, the L2 instructions are required for coherence between different agents. For the reduction use case, we need coherence within a GPU which is provided by buffer_wbinvl1_vol following the atomic inc. So we can drop the L2 instructions here. The duplicate waitcnt and the buffer_wbinvl1_vol can be dropped too.

In other words, as long as the ordering and scope are correct on the atomic operation, we should not need an additional fence. Perhaps it was required before D137524.

[dhruvachak]

Addressed feedback. Removed the fence change, used the suggested macro, and used the enum scope for clarity.

[dhruvachak]

openmp/libomptarget/DeviceRTL/src/Reduction.cpp

226

I removed this unrelated change, will post a separate patch for it.

openmp/libomptarget/DeviceRTL/src/Synchronization.cpp

213–226

Used the suggested macro. Thanks @jdoerfert

openmp/libomptarget/DeviceRTL/include/Synchronization.h

 enum OrderingTy {
   relaxed = __ATOMIC_RELAXED,
   aquire = __ATOMIC_ACQUIRE,
   release = __ATOMIC_RELEASE,
   acq_rel = __ATOMIC_ACQ_REL,
   seq_cst = __ATOMIC_SEQ_CST,
 };
 
+enum MemScopeTy {
+  all,    // All threads on all devices
+  device, // All threads on the device
+  cgroup  // All threads in the contention group, e.g. the team
+};
+
 /// Atomically increment \p *Addr and wrap at \p V with \p Ordering semantics.
-uint32_t inc(uint32_t *Addr, uint32_t V, OrderingTy Ordering);
+uint32_t inc(uint32_t *Addr, uint32_t V, OrderingTy Ordering,
+             MemScopeTy MemScope = MemScopeTy::all);
 
 /// Atomically perform <op> on \p V and \p *Addr with \p Ordering semantics. The
 /// result is stored in \p *Addr;
 /// {
 
 #define ATOMIC_COMMON_OP(TY)                                                   \
   TY add(TY *Addr, TY V, OrderingTy Ordering);                                 \
   TY mul(TY *Addr, TY V, OrderingTy Ordering);                                 \

openmp/libomptarget/DeviceRTL/src/Reduction.cpp

int32_t __kmpc_nvptx_teams_reduce_nowait_v2(
 
   if (IsMaster) {
     int ModBockId = TeamId % num_of_records;
     if (TeamId < num_of_records) {
       lgcpyFct(GlobalBuffer, ModBockId, reduce_data);
     } else
       lgredFct(GlobalBuffer, ModBockId, reduce_data);
 
+#ifndef __AMDGCN__
     fence::system(atomic::seq_cst);
+#endif
 
     // Increment team counter.
     // This counter is incremented by all teams in the current
     // BUFFER_SIZE chunk.
-    ChunkTeamCount = atomic::inc(&Cnt, num_of_records - 1u, atomic::seq_cst);
+    ChunkTeamCount =
+        atomic::inc(&Cnt, num_of_records - 1u, atomic::seq_cst, atomic::device);
   }
   // Synchronize
   if (mapping::isSPMDMode())
     __kmpc_barrier(Loc, TId);
 
   // reduce_data is global or shared so before being reduced within the
   // warp we need to bring it in local memory:
   // local_reduce_data = reduce_data[i]

openmp/libomptarget/DeviceRTL/src/Synchronization.cpp

 using namespace ompx;
 
 namespace impl {
 
 /// Atomics
 ///
 ///{
 /// NOTE: This function needs to be implemented by every target.
-uint32_t atomicInc(uint32_t *Address, uint32_t Val,
-                   atomic::OrderingTy Ordering);
+uint32_t atomicInc(uint32_t *Address, uint32_t Val, atomic::OrderingTy Ordering,
+                   atomic::MemScopeTy MemScope);
 
 template <typename Ty>
 Ty atomicAdd(Ty *Address, Ty Val, atomic::OrderingTy Ordering) {
   return __atomic_fetch_add(Address, Val, Ordering);
 }
 
 template <typename Ty>
 Ty atomicMul(Ty *Address, Ty V, atomic::OrderingTy Ordering) {
uint32_t atomicExchange(uint32_t *Address, uint32_t Val,
                         atomic::OrderingTy Ordering) {
   uint32_t R;
   __atomic_exchange(Address, &Val, &R, Ordering);
   return R;
 }
 ///}
 
 // Forward declarations defined to be defined for AMDGCN and NVPTX.
-uint32_t atomicInc(uint32_t *A, uint32_t V, atomic::OrderingTy Ordering);
+uint32_t atomicInc(uint32_t *A, uint32_t V, atomic::OrderingTy Ordering,
+                   atomic::MemScopeTy MemScope);
 void namedBarrierInit();
 void namedBarrier();
 void fenceTeam(atomic::OrderingTy Ordering);
 void fenceKernel(atomic::OrderingTy Ordering);
 void fenceSystem(atomic::OrderingTy Ordering);
 void syncWarp(__kmpc_impl_lanemask_t);
 void syncThreads(atomic::OrderingTy Ordering);
 void syncThreadsAligned(atomic::OrderingTy Ordering) { syncThreads(Ordering); }
 void unsetLock(omp_lock_t *);
 int testLock(omp_lock_t *);
 void initLock(omp_lock_t *);
 void destroyLock(omp_lock_t *);
 void setLock(omp_lock_t *);
 void unsetCriticalLock(omp_lock_t *);
 void setCriticalLock(omp_lock_t *);
 
 /// AMDGCN Implementation
 ///
 ///{
 #pragma omp begin declare variant match(device = {arch(amdgcn)})
 
-uint32_t atomicInc(uint32_t *A, uint32_t V, atomic::OrderingTy Ordering) {
+uint32_t atomicIncRelaxed(uint32_t *A, uint32_t V,
+                          atomic::MemScopeTy MemScope) {
+  switch (MemScope) {
+  default:
+    __builtin_unreachable();
+  case atomic::all:
+    return __builtin_amdgcn_atomic_inc32(A, V, atomic::relaxed, "");
+  case atomic::device:
+    return __builtin_amdgcn_atomic_inc32(A, V, atomic::relaxed, "agent");
+  case atomic::cgroup:
+    return __builtin_amdgcn_atomic_inc32(A, V, atomic::relaxed, "workgroup");
+  }
+}
+
+uint32_t atomicIncAquire(uint32_t *A, uint32_t V, atomic::MemScopeTy MemScope) {
+  switch (MemScope) {
+  default:
+    __builtin_unreachable();
+  case atomic::all:
+    return __builtin_amdgcn_atomic_inc32(A, V, atomic::aquire, "");
+  case atomic::device:
+    return __builtin_amdgcn_atomic_inc32(A, V, atomic::aquire, "agent");
+  case atomic::cgroup:
+    return __builtin_amdgcn_atomic_inc32(A, V, atomic::aquire, "workgroup");
+  }
+}
+
+uint32_t atomicIncRelease(uint32_t *A, uint32_t V,
+                          atomic::MemScopeTy MemScope) {
+  switch (MemScope) {
+  default:
+    __builtin_unreachable();
+  case atomic::all:
+    return __builtin_amdgcn_atomic_inc32(A, V, atomic::release, "");
+  case atomic::device:
+    return __builtin_amdgcn_atomic_inc32(A, V, atomic::release, "agent");
+  case atomic::cgroup:
+    return __builtin_amdgcn_atomic_inc32(A, V, atomic::release, "workgroup");
+  }
+}
+
+uint32_t atomicIncAcqRel(uint32_t *A, uint32_t V, atomic::MemScopeTy MemScope) {
+  switch (MemScope) {
+  default:
+    __builtin_unreachable();
+  case atomic::all:
+    return __builtin_amdgcn_atomic_inc32(A, V, atomic::acq_rel, "");
+  case atomic::device:
+    return __builtin_amdgcn_atomic_inc32(A, V, atomic::acq_rel, "agent");
+  case atomic::cgroup:
+    return __builtin_amdgcn_atomic_inc32(A, V, atomic::acq_rel, "workgroup");
+  }
+}
+
+uint32_t atomicIncSeqCst(uint32_t *A, uint32_t V, atomic::MemScopeTy MemScope) {
+  switch (MemScope) {
+  default:
+    __builtin_unreachable();
+  case atomic::all:
+    return __builtin_amdgcn_atomic_inc32(A, V, atomic::seq_cst, "");
+  case atomic::device:
+    return __builtin_amdgcn_atomic_inc32(A, V, atomic::seq_cst, "agent");
+  case atomic::cgroup:
+    return __builtin_amdgcn_atomic_inc32(A, V, atomic::seq_cst, "workgroup");
+  }
+}
+
+uint32_t atomicInc(uint32_t *A, uint32_t V, atomic::OrderingTy Ordering,
+                   atomic::MemScopeTy MemScope) {
   // builtin_amdgcn_atomic_inc32 should expand to this switch when
   // passed a runtime value, but does not do so yet. Workaround here.
   switch (Ordering) {
   default:
     __builtin_unreachable();
   case atomic::relaxed:
-    return __builtin_amdgcn_atomic_inc32(A, V, atomic::relaxed, "");
+    return atomicIncRelaxed(A, V, MemScope);
   case atomic::aquire:
-    return __builtin_amdgcn_atomic_inc32(A, V, atomic::aquire, "");
+    return atomicIncAquire(A, V, MemScope);
   case atomic::release:
-    return __builtin_amdgcn_atomic_inc32(A, V, atomic::release, "");
+    return atomicIncRelease(A, V, MemScope);
   case atomic::acq_rel:
-    return __builtin_amdgcn_atomic_inc32(A, V, atomic::acq_rel, "");
+    return atomicIncAcqRel(A, V, MemScope);
   case atomic::seq_cst:
-    return __builtin_amdgcn_atomic_inc32(A, V, atomic::seq_cst, "");
+    return atomicIncSeqCst(A, V, MemScope);
   }
 }
 
 uint32_t SHARED(namedBarrierTracker);
 
 void namedBarrierInit() {
   // Don't have global ctors, and shared memory is not zero init
   atomic::store(&namedBarrierTracker, 0u, atomic::release);
 }

 
 /// NVPTX Implementation
 ///
 ///{
 #pragma omp begin declare variant match(                                       \
         device = {arch(nvptx, nvptx64)},                                       \
             implementation = {extension(match_any)})
 
-uint32_t atomicInc(uint32_t *Address, uint32_t Val,
-                   atomic::OrderingTy Ordering) {
+uint32_t atomicInc(uint32_t *Address, uint32_t Val, atomic::OrderingTy Ordering,
+                   atomic::MemScopeTy MemScope) {
   return __nvvm_atom_inc_gen_ui(Address, Val);
 }
 
 void namedBarrierInit() {}
 
 void namedBarrier() {
   uint32_t NumThreads = omp_get_num_threads();
   ASSERT(NumThreads % 32 == 0);

 ATOMIC_FP_OP(double, int64_t, uint64_t)
 
 #undef ATOMIC_INT_ONLY_OP
 #undef ATOMIC_FP_ONLY_OP
 #undef ATOMIC_COMMON_OP
 #undef ATOMIC_INT_OP
 #undef ATOMIC_FP_OP
 
-uint32_t atomic::inc(uint32_t *Addr, uint32_t V, atomic::OrderingTy Ordering) {
-  return impl::atomicInc(Addr, V, Ordering);
+uint32_t atomic::inc(uint32_t *Addr, uint32_t V, atomic::OrderingTy Ordering,
+                     atomic::MemScopeTy MemScope) {
+  return impl::atomicInc(Addr, V, Ordering, MemScope);
 }
 
 void unsetCriticalLock(omp_lock_t *Lock) { impl::unsetLock(Lock); }
 
 void setCriticalLock(omp_lock_t *Lock) { impl::setLock(Lock); }
 
 extern "C" {
 void __kmpc_ordered(IdentTy *Loc, int32_t TId) { FunctionTracingRAII(); }