[gomp-nvptx 12/13] libgomp: handle multiple teams on NVPTX

[Alexander Monakov]

        * config/nvptx/icv-device.c (omp_get_num_teams): Update.
        (omp_get_team_num): Ditto.
        * config/nvptx/target.c (GOMP_teams): Update.
        * config/nvptx/team.c (nvptx_thrs): Place in shared memory.
        * icv.c (gomp_num_teams_var): Define.
        * libgomp.h (gomp_num_teams_var): Declare.
        (nvptx_thrs): Place in shared memory.
---
 libgomp/ChangeLog.gomp-nvptx      | 10 ++++++++++
 libgomp/config/nvptx/icv-device.c |  8 ++++----
 libgomp/config/nvptx/target.c     | 13 ++++++++++++-
 libgomp/config/nvptx/team.c       |  2 +-
 libgomp/icv.c                     |  1 +
 libgomp/libgomp.h                 |  3 ++-
 6 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/libgomp/config/nvptx/icv-device.c 
b/libgomp/config/nvptx/icv-device.c
index 0e5fef0..bd11002 100644
--- a/libgomp/config/nvptx/icv-device.c
+++ b/libgomp/config/nvptx/icv-device.c
@@ -47,15 +47,15 @@ omp_get_num_devices (void)
 int
 omp_get_num_teams (void)
 {
-  /* FORNOW.  */
-  return 1;
+  return gomp_num_teams_var + 1;
 }
 
 int
 omp_get_team_num (void)
 {
-  /* FORNOW.  */
-  return 0;
+  int ctaid;
+  asm ("mov.u32 %0, %%ctaid.x;" : "=r" (ctaid));
+  return ctaid;
 }
 
 int
diff --git a/libgomp/config/nvptx/target.c b/libgomp/config/nvptx/target.c
index ad36013..9f34ae8 100644
--- a/libgomp/config/nvptx/target.c
+++ b/libgomp/config/nvptx/target.c
@@ -35,5 +35,16 @@ GOMP_teams (unsigned int num_teams, unsigned int 
thread_limit)
       icv->thread_limit_var
        = thread_limit > INT_MAX ? UINT_MAX : thread_limit;
     }
-  (void) num_teams;
+  unsigned int num_blocks, block_id;
+  asm ("mov.u32 %0, %%nctaid.x;" : "=r" (num_blocks));
+  asm ("mov.u32 %0, %%ctaid.x;" : "=r" (block_id));
+  if (!num_teams || num_teams >= num_blocks)
+    num_teams = num_blocks;
+  else if (block_id >= num_teams)
+    {
+      gomp_free_thread (nvptx_thrs);
+      free (nvptx_thrs);
+      asm ("exit;");
+    }
+  gomp_num_teams_var = num_teams - 1;
 }
diff --git a/libgomp/config/nvptx/team.c b/libgomp/config/nvptx/team.c
index c18517a..909f296 100644
--- a/libgomp/config/nvptx/team.c
+++ b/libgomp/config/nvptx/team.c
@@ -29,7 +29,7 @@
 #include "libgomp.h"
 #include <stdlib.h>
 
-struct gomp_thread *nvptx_thrs;
+struct gomp_thread *nvptx_thrs __attribute__((shared));
 
 static void gomp_thread_start (struct gomp_thread_pool *);
 
diff --git a/libgomp/icv.c b/libgomp/icv.c
index aa79423..18e35e5 100644
--- a/libgomp/icv.c
+++ b/libgomp/icv.c
@@ -56,6 +56,7 @@ unsigned long gomp_bind_var_list_len;
 void **gomp_places_list;
 unsigned long gomp_places_list_len;
 int gomp_debug_var;
+unsigned int gomp_num_teams_var;
 char *goacc_device_type;
 int goacc_device_num;
 
diff --git a/libgomp/libgomp.h b/libgomp/libgomp.h
index 1d137f1..0ef2a05 100644
--- a/libgomp/libgomp.h
+++ b/libgomp/libgomp.h
@@ -363,6 +363,7 @@ extern char *gomp_bind_var_list;
 extern unsigned long gomp_bind_var_list_len;
 extern void **gomp_places_list;
 extern unsigned long gomp_places_list_len;
+extern unsigned int gomp_num_teams_var;
 extern int gomp_debug_var;
 extern int goacc_device_num;
 extern char *goacc_device_type;
@@ -648,7 +649,7 @@ enum gomp_cancel_kind
 /* ... and here is that TLS data.  */
 
 #if defined __nvptx__
-extern struct gomp_thread *nvptx_thrs;
+extern struct gomp_thread *nvptx_thrs __attribute__((shared));
 static inline struct gomp_thread *gomp_thread (void)
 {
   int tid;