[OPENMP][NVPTX]Improved lock/critical constructs.

Summary: Improved support for critical constructs + omp_..._lock... constructs.

Reviewers: gtbercea, kkwli0, caomhin

Subscribers: guansong, jfb, openmp-commits

Differential Revision: https://reviews.llvm.org/D54766

3 files changed, 11 insertions, 24 deletions

diff --git a/openmp/libomptarget/deviceRTLs/nvptx/src/critical.cu b/openmp/libomptarget/deviceRTLs/nvptx/src/critical.cu
index fef81013e6c..2b92b9a1ca8 100644
--- a/openmp/libomptarget/deviceRTLs/nvptx/src/critical.cu
+++ b/openmp/libomptarget/deviceRTLs/nvptx/src/critical.cu
@@ -15,18 +15,14 @@
 
 #include "omptarget-nvptx.h"
 
-EXTERN
-void __kmpc_critical(kmp_Indent *loc, int32_t global_tid,
-                     kmp_CriticalName *lck) {
+EXTERN void __kmpc_critical(kmp_Indent *loc, int32_t global_tid,
+                            kmp_CriticalName *lck) {
   PRINT0(LD_IO, "call to kmpc_critical()\n");
-  omptarget_nvptx_TeamDescr &teamDescr = getMyTeamDescriptor();
-  omp_set_lock(teamDescr.CriticalLock());
+  omp_set_lock((omp_lock_t *)lck);
 }
 
-EXTERN
-void __kmpc_end_critical(kmp_Indent *loc, int32_t global_tid,
-                         kmp_CriticalName *lck) {
+EXTERN void __kmpc_end_critical(kmp_Indent *loc, int32_t global_tid,
+                                kmp_CriticalName *lck) {
   PRINT0(LD_IO, "call to kmpc_end_critical()\n");
-  omptarget_nvptx_TeamDescr &teamDescr = getMyTeamDescriptor();
-  omp_unset_lock(teamDescr.CriticalLock());
+  omp_unset_lock((omp_lock_t *)lck);
 }
diff --git a/openmp/libomptarget/deviceRTLs/nvptx/src/libcall.cu b/openmp/libomptarget/deviceRTLs/nvptx/src/libcall.cu
index ea9225db9b9..3e8f248f11e 100644
--- a/openmp/libomptarget/deviceRTLs/nvptx/src/libcall.cu
+++ b/openmp/libomptarget/deviceRTLs/nvptx/src/libcall.cu
@@ -404,23 +404,21 @@ EXTERN int omp_get_max_task_priority(void) {
 #define SET 1
 
 EXTERN void omp_init_lock(omp_lock_t *lock) {
-  *lock = UNSET;
+  omp_unset_lock(lock);
   PRINT0(LD_IO, "call omp_init_lock()\n");
 }
 
 EXTERN void omp_destroy_lock(omp_lock_t *lock) {
+  omp_unset_lock(lock);
   PRINT0(LD_IO, "call omp_destroy_lock()\n");
 }
 
 EXTERN void omp_set_lock(omp_lock_t *lock) {
   // int atomicCAS(int* address, int compare, int val);
   // (old == compare ? val : old)
-  int compare = UNSET;
-  int val = SET;
 
   // TODO: not sure spinning is a good idea here..
-  while (atomicCAS(lock, compare, val) != UNSET) {
-
+  while (atomicCAS(lock, UNSET, SET) != UNSET) {
     clock_t start = clock();
     clock_t now;
     for (;;) {
@@ -436,9 +434,7 @@ EXTERN void omp_set_lock(omp_lock_t *lock) {
 }
 
 EXTERN void omp_unset_lock(omp_lock_t *lock) {
-  int compare = SET;
-  int val = UNSET;
-  int old = atomicCAS(lock, compare, val);
+  (void)atomicExch(lock, UNSET);
 
   PRINT0(LD_IO, "call omp_unset_lock()\n");
 }
@@ -446,10 +442,7 @@ EXTERN void omp_unset_lock(omp_lock_t *lock) {
 EXTERN int omp_test_lock(omp_lock_t *lock) {
   // int atomicCAS(int* address, int compare, int val);
   // (old == compare ? val : old)
-  int compare = UNSET;
-  int val = SET;
-
-  int ret = atomicCAS(lock, compare, val);
+  int ret = atomicAdd(lock, 0);
 
   PRINT(LD_IO, "call omp_test_lock() return %d\n", ret);
 
diff --git a/openmp/libomptarget/deviceRTLs/nvptx/src/omptarget-nvptx.h b/openmp/libomptarget/deviceRTLs/nvptx/src/omptarget-nvptx.h
index 2a6de28d274..b63feaefbdd 100644
--- a/openmp/libomptarget/deviceRTLs/nvptx/src/omptarget-nvptx.h
+++ b/openmp/libomptarget/deviceRTLs/nvptx/src/omptarget-nvptx.h
@@ -251,7 +251,6 @@ public:
   INLINE omptarget_nvptx_WorkDescr &WorkDescr() {
     return workDescrForActiveParallel;
   }
-  INLINE omp_lock_t *CriticalLock() { return &criticalLock; }
   INLINE uint64_t *getLastprivateIterBuffer() { return &lastprivateIterBuffer; }
 
   // init
@@ -303,7 +302,6 @@ private:
       levelZeroTaskDescr; // icv for team master initial thread
   omptarget_nvptx_WorkDescr
       workDescrForActiveParallel; // one, ONLY for the active par
-  omp_lock_t criticalLock;
   uint64_t lastprivateIterBuffer;
 
   __align__(16)