[gcc/devel/omp/gcc-11] openmp: Add omp_set_num_teams, omp_get_max_teams, omp_[gs]et_teams_thread_limit

[gcc/devel/omp/gcc-11] openmp: Add omp_set_num_teams, omp_get_max_teams, omp_[gs]et_teams_thread_limit

[Tobias Burnus]

https://gcc.gnu.org/g:6665382c0efd940e7d983c9244764a4bcebdb6ff

commit 6665382c0efd940e7d983c9244764a4bcebdb6ff
Author: Jakub Jelinek <jakub@redhat.com>
Date:   Mon Oct 11 13:05:38 2021 +0200

    openmp: Add omp_set_num_teams, omp_get_max_teams, omp_[gs]et_teams_thread_limit
    
    OpenMP 5.1 adds env vars and functions to set and query new ICVs used
    as fallback if thread_limit or num_teams clauses aren't specified on
    teams construct.
    
    The following patch implements those, though further work will be needed:
    1) OpenMP 5.1 also changed the num_teams clause, so that it can specify
       both lower and upper limit for how many teams should be created and
       changed the meaning when only one expression is provided, instead of
       num_teams(expr) in 5.0 meaning num_teams(1:expr) in 5.1, it now means
       num_teams(expr:expr), i.e. while previously we could create 1 to expr
       teams, in 5.1 we have some low limit by default equal to the single
       expression provided and may not create fewer teams.
       For host teams (which we don't currently implement efficiently for
       NUMA hosts) we trivially satisfy it now by always honoring what the
       user asked for, but for the offloading teams I think we'll need to
       rethink the APIs; currently teams construct is just a call that returns
       and possibly lowers the number of teams; and whenever possible we try
       to evaluate num_teams/thread_limit already on the target construct
       and the GOMP_teams call just sets the number of teams to the minimum
       of provided and requested teams; for some cases e.g. where target
       is not combined with teams and num_teams expression calls some functions
       etc., we need to call those functions in the target region and so it is
       late to figure number of teams, but also hw could just limit what it
       is willing to create; in that case I'm afraid we need to run the target
       body multiple times and arrange for omp_get_team_num () returning the
       right values
    2) we need to finally implement the NUMA handling for GOMP_teams_reg
    3) I now realize I haven't added some testcase coverage, will do that
       incrementally
    4) libgomp.texi needs updates for these new APIs, but also others like
       the allocator
    
    2021-10-11  Jakub Jelinek  <jakub@redhat.com>
    
    gcc/
            * omp-low.c (omp_runtime_api_call): Handle omp_get_max_teams,
            omp_[sg]et_teams_thread_limit and omp_set_num_teams.
    libgomp/
            * omp.h.in (omp_set_num_teams, omp_get_max_teams,
            omp_set_teams_thread_limit, omp_get_teams_thread_limit): Declare.
            * omp_lib.f90.in (omp_set_num_teams, omp_get_max_teams,
            omp_set_teams_thread_limit, omp_get_teams_thread_limit): Declare.
            * omp_lib.h.in (omp_set_num_teams, omp_get_max_teams,
            omp_set_teams_thread_limit, omp_get_teams_thread_limit): Declare.
            * libgomp.h (gomp_nteams_var, gomp_teams_thread_limit_var): Declare.
            * libgomp.map (OMP_5.1): Export omp_get_max_teams{,_},
            omp_get_teams_thread_limit{,_}, omp_set_num_teams{,_,_8_} and
            omp_set_teams_thread_limit{,_,_8_}.
            * icv.c (omp_set_num_teams, omp_get_max_teams,
            omp_set_teams_thread_limit, omp_get_teams_thread_limit): New
            functions.
            * env.c (gomp_nteams_var, gomp_teams_thread_limit_var): Define.
            (omp_display_env): Print OMP_NUM_TEAMS and OMP_TEAMS_THREAD_LIMIT.
            (initialize_env): Handle OMP_NUM_TEAMS and OMP_TEAMS_THREAD_LIMIT env
            vars.
            * teams.c (GOMP_teams_reg): If thread_limit is not specified, use
            gomp_teams_thread_limit_var as fallback if not zero.  If num_teams
            is not specified, use gomp_nteams_var.
            * fortran.c (omp_set_num_teams, omp_get_max_teams,
            omp_set_teams_thread_limit, omp_get_teams_thread_limit): Add
            ialias_redirect.
            (omp_set_num_teams_, omp_set_num_teams_8_, omp_get_max_teams_,
            omp_set_teams_thread_limit_, omp_set_teams_thread_limit_8_,
            omp_get_teams_thread_limit_): New functions.
    
    (cherry picked from commit 07dd3bcda17f97cf5476c3d6f2f2501c1e0712e6)

Diff:
---
 gcc/ChangeLog.omp      |  8 ++++++++
 gcc/omp-low.c          |  6 +++++-
 libgomp/ChangeLog.omp  | 32 ++++++++++++++++++++++++++++++++
 libgomp/env.c          |  7 +++++++
 libgomp/fortran.c      | 40 ++++++++++++++++++++++++++++++++++++++++
 libgomp/icv.c          | 30 ++++++++++++++++++++++++++++++
 libgomp/libgomp.h      |  2 ++
 libgomp/libgomp.map    | 10 ++++++++++
 libgomp/omp.h.in       |  5 +++++
 libgomp/omp_lib.f90.in | 30 ++++++++++++++++++++++++++++++
 libgomp/omp_lib.h.in   |  4 ++++
 libgomp/teams.c        |  4 +++-
 12 files changed, 176 insertions(+), 2 deletions(-)

diff --git a/gcc/ChangeLog.omp b/gcc/ChangeLog.omp
index c69454167c0..a8e0ff15b3f 100644
--- a/gcc/ChangeLog.omp
+++ b/gcc/ChangeLog.omp
@@ -1,3 +1,11 @@
+2021-10-11  Tobias Burnus  <tobias@codesourcery.com>
+
+	Backported from master:
+	2021-10-11  Jakub Jelinek  <jakub@redhat.com>
+
+	* omp-low.c (omp_runtime_api_call): Handle omp_get_max_teams,
+	omp_[sg]et_teams_thread_limit and omp_set_num_teams.
+
 2021-10-01  Tobias Burnus  <tobias@codesourcery.com>
 
 	Backported from master:
diff --git a/gcc/omp-low.c b/gcc/omp-low.c
index 0fc6df6ed41..a2d225ae39e 100644
--- a/gcc/omp-low.c
+++ b/gcc/omp-low.c
@@ -4252,6 +4252,7 @@ omp_runtime_api_call (const_tree fndecl)
       "get_level",
       "get_max_active_levels",
       "get_max_task_priority",
+      "get_max_teams",
       "get_max_threads",
       "get_nested",
       "get_num_devices",
@@ -4264,6 +4265,7 @@ omp_runtime_api_call (const_tree fndecl)
       "get_proc_bind",
       "get_supported_active_levels",
       "get_team_num",
+      "get_teams_thread_limit",
       "get_thread_limit",
       "get_thread_num",
       "get_wtick",
@@ -4297,8 +4299,10 @@ omp_runtime_api_call (const_tree fndecl)
       "set_dynamic",
       "set_max_active_levels",
       "set_nested",
+      "set_num_teams",
       "set_num_threads",
-      "set_schedule"
+      "set_schedule",
+      "set_teams_thread_limit"
     };
 
   int mode = 0;
diff --git a/libgomp/ChangeLog.omp b/libgomp/ChangeLog.omp
index 45693c407a0..dda93a86fcc 100644
--- a/libgomp/ChangeLog.omp
+++ b/libgomp/ChangeLog.omp
@@ -1,3 +1,35 @@
+2021-10-11  Tobias Burnus  <tobias@codesourcery.com>
+
+	Backported from master:
+	2021-10-11  Jakub Jelinek  <jakub@redhat.com>
+
+	* omp.h.in (omp_set_num_teams, omp_get_max_teams,
+	omp_set_teams_thread_limit, omp_get_teams_thread_limit): Declare.
+	* omp_lib.f90.in (omp_set_num_teams, omp_get_max_teams,
+	omp_set_teams_thread_limit, omp_get_teams_thread_limit): Declare.
+	* omp_lib.h.in (omp_set_num_teams, omp_get_max_teams,
+	omp_set_teams_thread_limit, omp_get_teams_thread_limit): Declare.
+	* libgomp.h (gomp_nteams_var, gomp_teams_thread_limit_var): Declare.
+	* libgomp.map (OMP_5.1): Export omp_get_max_teams{,_},
+	omp_get_teams_thread_limit{,_}, omp_set_num_teams{,_,_8_} and
+	omp_set_teams_thread_limit{,_,_8_}.
+	* icv.c (omp_set_num_teams, omp_get_max_teams,
+	omp_set_teams_thread_limit, omp_get_teams_thread_limit): New
+	functions.
+	* env.c (gomp_nteams_var, gomp_teams_thread_limit_var): Define.
+	(omp_display_env): Print OMP_NUM_TEAMS and OMP_TEAMS_THREAD_LIMIT.
+	(initialize_env): Handle OMP_NUM_TEAMS and OMP_TEAMS_THREAD_LIMIT env
+	vars.
+	* teams.c (GOMP_teams_reg): If thread_limit is not specified, use
+	gomp_teams_thread_limit_var as fallback if not zero.  If num_teams
+	is not specified, use gomp_nteams_var.
+	* fortran.c (omp_set_num_teams, omp_get_max_teams,
+	omp_set_teams_thread_limit, omp_get_teams_thread_limit): Add
+	ialias_redirect.
+	(omp_set_num_teams_, omp_set_num_teams_8_, omp_get_max_teams_,
+	omp_set_teams_thread_limit_, omp_set_teams_thread_limit_8_,
+	omp_get_teams_thread_limit_): New functions.
+
 2021-10-09  Tobias Burnus  <tobias@codesourcery.com>
 
 	Backported from master:
diff --git a/libgomp/env.c b/libgomp/env.c
index 69ce1d24ee0..de45c25d540 100644
--- a/libgomp/env.c
+++ b/libgomp/env.c
@@ -90,6 +90,8 @@ unsigned long gomp_places_list_len;
 uintptr_t gomp_def_allocator = omp_default_mem_alloc;
 int gomp_debug_var;
 unsigned int gomp_num_teams_var;
+int gomp_nteams_var;
+int gomp_teams_thread_limit_var;
 bool gomp_display_affinity_var;
 char *gomp_affinity_format_var = "level %L thread %i affinity %A";
 size_t gomp_affinity_format_len;
@@ -1319,6 +1321,9 @@ omp_display_env (int verbose)
 	   gomp_global_icv.thread_limit_var);
   fprintf (stderr, "  OMP_MAX_ACTIVE_LEVELS = '%u'\n",
 	   gomp_global_icv.max_active_levels_var);
+  fprintf (stderr, "  OMP_NUM_TEAMS = '%u'\n", gomp_nteams_var);
+  fprintf (stderr, "  OMP_TEAMS_THREAD_LIMIT = '%u'\n",
+	   gomp_teams_thread_limit_var);
 
   fprintf (stderr, "  OMP_CANCELLATION = '%s'\n",
 	   gomp_cancel_var ? "TRUE" : "FALSE");
@@ -1453,6 +1458,8 @@ initialize_env (void)
 				 &gomp_nthreads_var_list,
 				 &gomp_nthreads_var_list_len))
     gomp_global_icv.nthreads_var = gomp_available_cpus;
+  parse_int ("OMP_NUM_TEAMS", &gomp_nteams_var, false);
+  parse_int ("OMP_TEAMS_THREAD_LIMIT", &gomp_teams_thread_limit_var, false);
   bool ignore = false;
   if (parse_bind_var ("OMP_PROC_BIND",
 		      &gomp_global_icv.bind_var,
diff --git a/libgomp/fortran.c b/libgomp/fortran.c
index d171c7992eb..14b5c8d51dd 100644
--- a/libgomp/fortran.c
+++ b/libgomp/fortran.c
@@ -67,6 +67,10 @@ ialias_redirect (omp_get_thread_limit)
 ialias_redirect (omp_set_max_active_levels)
 ialias_redirect (omp_get_max_active_levels)
 ialias_redirect (omp_get_supported_active_levels)
+ialias_redirect (omp_set_num_teams)
+ialias_redirect (omp_get_max_teams)
+ialias_redirect (omp_set_teams_thread_limit)
+ialias_redirect (omp_get_teams_thread_limit)
 ialias_redirect (omp_get_level)
 ialias_redirect (omp_get_ancestor_thread_num)
 ialias_redirect (omp_get_team_size)
@@ -478,6 +482,42 @@ omp_in_final_ (void)
   return omp_in_final ();
 }
 
+void
+omp_set_num_teams_ (const int32_t *num_teams)
+{
+  omp_set_num_teams (*num_teams);
+}
+
+void
+omp_set_num_teams_8_ (const int64_t *num_teams)
+{
+  omp_set_max_active_levels (TO_INT (*num_teams));
+}
+
+int32_t
+omp_get_max_teams_ (void)
+{
+  return omp_get_max_teams ();
+}
+
+void
+omp_set_teams_thread_limit_ (const int32_t *thread_limit)
+{
+  omp_set_teams_thread_limit (*thread_limit);
+}
+
+void
+omp_set_teams_thread_limit_8_ (const int64_t *thread_limit)
+{
+  omp_set_teams_thread_limit (TO_INT (*thread_limit));
+}
+
+int32_t
+omp_get_teams_thread_limit_ (void)
+{
+  return omp_get_teams_thread_limit ();
+}
+
 int32_t
 omp_get_cancellation_ (void)
 {
diff --git a/libgomp/icv.c b/libgomp/icv.c
index 4b4892fb35e..b7fad42e5c1 100644
--- a/libgomp/icv.c
+++ b/libgomp/icv.c
@@ -148,6 +148,32 @@ omp_get_supported_active_levels (void)
   return gomp_supported_active_levels;
 }
 
+void
+omp_set_num_teams (int num_teams)
+{
+  if (num_teams >= 0)
+    gomp_nteams_var = num_teams;
+}
+
+int
+omp_get_max_teams (void)
+{
+  return gomp_nteams_var;
+}
+
+void
+omp_set_teams_thread_limit (int thread_limit)
+{
+  if (thread_limit >= 0)
+    gomp_teams_thread_limit_var = thread_limit;
+}
+
+int
+omp_get_teams_thread_limit (void)
+{
+  return gomp_teams_thread_limit_var;
+}
+
 int
 omp_get_cancellation (void)
 {
@@ -248,6 +274,10 @@ ialias (omp_get_thread_limit)
 ialias (omp_set_max_active_levels)
 ialias (omp_get_max_active_levels)
 ialias (omp_get_supported_active_levels)
+ialias (omp_set_num_teams)
+ialias (omp_get_max_teams)
+ialias (omp_set_teams_thread_limit)
+ialias (omp_get_teams_thread_limit)
 ialias (omp_get_cancellation)
 ialias (omp_get_proc_bind)
 ialias (omp_get_max_task_priority)
diff --git a/libgomp/libgomp.h b/libgomp/libgomp.h
index 0c8085e350c..f7745a239ec 100644
--- a/libgomp/libgomp.h
+++ b/libgomp/libgomp.h
@@ -458,6 +458,8 @@ 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_nteams_var;
+extern int gomp_teams_thread_limit_var;
 extern int gomp_debug_var;
 extern bool gomp_display_affinity_var;
 extern char *gomp_affinity_format_var;
diff --git a/libgomp/libgomp.map b/libgomp/libgomp.map
index 058d904e24a..5aff705fa8c 100644
--- a/libgomp/libgomp.map
+++ b/libgomp/libgomp.map
@@ -214,6 +214,16 @@ OMP_5.1 {
 	omp_display_env;
 	omp_display_env_;
 	omp_display_env_8_;
+	omp_set_num_teams;
+	omp_set_num_teams_;
+	omp_set_num_teams_8_;
+	omp_get_max_teams;
+	omp_get_max_teams_;
+	omp_set_teams_thread_limit;
+	omp_set_teams_thread_limit_;
+	omp_set_teams_thread_limit_8_;
+	omp_get_teams_thread_limit;
+	omp_get_teams_thread_limit_;
 } OMP_5.0.2;
 
 GOMP_1.0 {
diff --git a/libgomp/omp.h.in b/libgomp/omp.h.in
index d75ee13c16d..e0177dca6fb 100644
--- a/libgomp/omp.h.in
+++ b/libgomp/omp.h.in
@@ -261,6 +261,11 @@ extern int omp_get_max_task_priority (void) __GOMP_NOTHROW;
 
 extern void omp_fulfill_event (omp_event_handle_t) __GOMP_NOTHROW;
 
+extern void omp_set_num_teams (int) __GOMP_NOTHROW;
+extern int omp_get_max_teams (void) __GOMP_NOTHROW;
+extern void omp_set_teams_thread_limit (int) __GOMP_NOTHROW;
+extern int omp_get_teams_thread_limit (void) __GOMP_NOTHROW;
+
 extern void *omp_target_alloc (__SIZE_TYPE__, int) __GOMP_NOTHROW;
 extern void omp_target_free (void *, int) __GOMP_NOTHROW;
 extern int omp_target_is_present (const void *, int) __GOMP_NOTHROW;
diff --git a/libgomp/omp_lib.f90.in b/libgomp/omp_lib.f90.in
index 1063eee0c94..973b87be80a 100644
--- a/libgomp/omp_lib.f90.in
+++ b/libgomp/omp_lib.f90.in
@@ -564,6 +564,36 @@
           end function omp_get_max_task_priority
         end interface
 
+        interface omp_set_num_teams
+          subroutine omp_set_num_teams (num_teams)
+            integer (4), intent (in) :: num_teams
+          end subroutine omp_set_num_teams
+          subroutine omp_set_num_teams_8 (num_teams)
+            integer (8), intent (in) :: num_teams
+          end subroutine omp_set_num_teams_8
+        end interface
+
+        interface
+          function omp_get_max_teams ()
+            integer (4) :: omp_get_max_teams
+          end function omp_get_max_teams
+        end interface
+
+        interface omp_set_teams_thread_limit
+          subroutine omp_set_teams_thread_limit (thread_limit)
+            integer (4), intent (in) :: thread_limit
+          end subroutine omp_set_teams_thread_limit
+          subroutine omp_set_teams_thread_limit_8 (thread_limit)
+            integer (8), intent (in) :: thread_limit
+          end subroutine omp_set_teams_thread_limit_8
+        end interface
+
+        interface
+          function omp_get_teams_thread_limit ()
+            integer (4) :: omp_get_teams_thread_limit
+          end function omp_get_teams_thread_limit
+        end interface
+
         interface
           subroutine omp_fulfill_event (event)
             use omp_lib_kinds
diff --git a/libgomp/omp_lib.h.in b/libgomp/omp_lib.h.in
index f40321c479b..d5c02605dd7 100644
--- a/libgomp/omp_lib.h.in
+++ b/libgomp/omp_lib.h.in
@@ -252,6 +252,10 @@
       external omp_get_max_task_priority
       integer(4) omp_get_max_task_priority
 
+      external omp_set_num_teams, omp_set_teams_thread_limit
+      external omp_get_max_teams, omp_get_teams_thread_limit
+      integer(4) omp_get_max_teams, omp_get_teams_thread_limit
+
       external omp_fulfill_event
 
       external omp_set_affinity_format, omp_get_affinity_format
diff --git a/libgomp/teams.c b/libgomp/teams.c
index 8bac3dba6e4..9409f8ee6aa 100644
--- a/libgomp/teams.c
+++ b/libgomp/teams.c
@@ -37,6 +37,8 @@ GOMP_teams_reg (void (*fn) (void *), void *data, unsigned int num_teams,
   (void) flags;
   (void) num_teams;
   unsigned old_thread_limit_var = 0;
+  if (thread_limit == 0)
+    thread_limit = gomp_teams_thread_limit_var;
   if (thread_limit)
     {
       struct gomp_task_icv *icv = gomp_icv (true);
@@ -45,7 +47,7 @@ GOMP_teams_reg (void (*fn) (void *), void *data, unsigned int num_teams,
 	= thread_limit > INT_MAX ? UINT_MAX : thread_limit;
     }
   if (num_teams == 0)
-    num_teams = 3;
+    num_teams = gomp_nteams_var ? gomp_nteams_var : 3;
   gomp_num_teams = num_teams;
   for (gomp_team_num = 0; gomp_team_num < num_teams; gomp_team_num++)
     fn (data);