[gcc r14-1026] libgomp: Honor OpenMP's nteams-var ICV as upper limit on num teams [PR109875]

[gcc r14-1026] libgomp: Honor OpenMP's nteams-var ICV as upper limit on num teams [PR109875]

[Tobias Burnus]

https://gcc.gnu.org/g:ad0f80d945cc36fbb60fd1e04d90681d4302de8b

commit r14-1026-gad0f80d945cc36fbb60fd1e04d90681d4302de8b
Author: Tobias Burnus <tobias@codesourcery.com>
Date:   Sun May 21 20:36:19 2023 +0200

    libgomp: Honor OpenMP's nteams-var ICV as upper limit on num teams [PR109875]
    
    The nteams-var ICV exists per device and can be set either via the routine
    omp_set_num_teams or as environment variable (OMP_NUM_TEAMS with optional
    _ALL/_DEV/_DEV_<num> suffix); it is default-initialized to zero. The number
    of teams created is described under the num_teams clause. If the clause is
    absent, the number of teams is implementation defined but at least
    one team must exist and, if nteams-var is positive, at most nteams-var
    teams may exist.
    
    The latter condition was not honored in a target region before this
    commit, such that too many teams were created.
    
    Already before this commit, both the num_teams([lower:]upper) clause
    (on the host and in target regions) and, only on the host, the nteams-var
    ICV were honored. And as only one teams is created for host fallback,
    unless the clause specifies otherwise, the nteams-var ICV was and is
    effectively honored.
    
    libgomp/ChangeLog:
    
            PR libgomp/109875
            * config/gcn/target.c (GOMP_teams4): Honor nteams-var ICV.
            * config/nvptx/target.c (GOMP_teams4): Likewise.
            * testsuite/libgomp.c-c++-common/teams-nteams-icv-1.c: New test.
            * testsuite/libgomp.c-c++-common/teams-nteams-icv-2.c: New test.
            * testsuite/libgomp.c-c++-common/teams-nteams-icv-3.c: New test.
            * testsuite/libgomp.c-c++-common/teams-nteams-icv-4.c: New test.

Diff:
---
 libgomp/config/gcn/target.c                        |   4 +-
 libgomp/config/nvptx/target.c                      |   4 +-
 .../libgomp.c-c++-common/teams-nteams-icv-1.c      | 198 +++++++++++++++++++++
 .../libgomp.c-c++-common/teams-nteams-icv-2.c      |   8 +
 .../libgomp.c-c++-common/teams-nteams-icv-3.c      |   8 +
 .../libgomp.c-c++-common/teams-nteams-icv-4.c      |  14 ++
 6 files changed, 234 insertions(+), 2 deletions(-)

diff --git a/libgomp/config/gcn/target.c b/libgomp/config/gcn/target.c
index c6691fde3c6..ea5eb1ff5ed 100644
--- a/libgomp/config/gcn/target.c
+++ b/libgomp/config/gcn/target.c
@@ -48,7 +48,9 @@ GOMP_teams4 (unsigned int num_teams_lower, unsigned int num_teams_upper,
      multiple times at least for some workgroups.  */
   (void) num_teams_lower;
   if (!num_teams_upper || num_teams_upper >= num_workgroups)
-    num_teams_upper = num_workgroups;
+    num_teams_upper = ((GOMP_ADDITIONAL_ICVS.nteams > 0
+			&& num_workgroups > GOMP_ADDITIONAL_ICVS.nteams)
+		       ? GOMP_ADDITIONAL_ICVS.nteams : num_workgroups);
   else if (workgroup_id >= num_teams_upper)
     return false;
   gomp_num_teams_var = num_teams_upper - 1;
diff --git a/libgomp/config/nvptx/target.c b/libgomp/config/nvptx/target.c
index f102d7d02d9..125d92a2ea9 100644
--- a/libgomp/config/nvptx/target.c
+++ b/libgomp/config/nvptx/target.c
@@ -55,7 +55,9 @@ GOMP_teams4 (unsigned int num_teams_lower, unsigned int num_teams_upper,
 	= thread_limit > INT_MAX ? UINT_MAX : thread_limit;
     }
   if (!num_teams_upper)
-    num_teams_upper = num_blocks;
+    num_teams_upper = ((GOMP_ADDITIONAL_ICVS.nteams > 0
+			&& num_blocks > GOMP_ADDITIONAL_ICVS.nteams)
+		       ? GOMP_ADDITIONAL_ICVS.nteams : num_blocks);
   else if (num_blocks < num_teams_lower)
     num_teams_upper = num_teams_lower;
   else if (num_blocks < num_teams_upper)
diff --git a/libgomp/testsuite/libgomp.c-c++-common/teams-nteams-icv-1.c b/libgomp/testsuite/libgomp.c-c++-common/teams-nteams-icv-1.c
new file mode 100644
index 00000000000..c3c21091e97
--- /dev/null
+++ b/libgomp/testsuite/libgomp.c-c++-common/teams-nteams-icv-1.c
@@ -0,0 +1,198 @@
+/* Check that the nteams ICV is honored. */
+/* PR libgomp/109875  */
+
+/*  This base version of testcases is supposed to be run with all
+    OMP_NUM_TEAMS* env vars being unset.
+
+    The variants teams-nteams-icv-{2,3,4}.c test it by setting the
+    various OMP_NUM_TEAMS* env vars and #define MY_... for checking.
+
+    Currently, only <num> 0,1,2 is supported for the envar via #define
+    and with remote execution, dg-set-target-env-var does not work with
+    DejaGNU, hence, gcc/testsuite/lib/gcc-dg.exp marks those tests as
+    UNSUPPORTED. */
+
+#define MY_MAX_DEVICES 3
+
+/* OpenMP currently has:
+   - nteams-var ICV is initialized to 0; one ICV per device
+   - OMP_NUM_TEAMS(_DEV(_<dev-num>)) overrides it
+     OMP_NUM_TEAMS_ALL overrides it
+   - Number of teams is:
+     -> the value specific by num_teams([lower:]upper)
+	with lower := upper if unspecified
+     -> Otherwise, if nteams-var ICV > 0, #teams <= nteams-var ICV
+     -> Otherwise, if nteams-var ICV <= 0, #teams > 1
+ GCC uses 3 as default on the host and 1 for host fallback.
+ For offloading, it is device specific >> 1.  */
+
+#include <omp.h>
+
+int
+main ()
+{
+  int num_teams_env = -1, num_teams_env_dev = -1;
+  int num_teams_env_devs[MY_MAX_DEVICES];
+
+#ifdef MY_OMP_NUM_TEAMS_ALL
+  num_teams_env = num_teams_env_dev = MY_OMP_NUM_TEAMS_ALL;
+#endif
+
+#ifdef MY_OMP_NUM_TEAMS
+  num_teams_env = MY_OMP_NUM_TEAMS;
+#endif
+
+#ifdef MY_OMP_NUM_TEAMS_DEV
+  num_teams_env_dev = MY_OMP_NUM_TEAMS_DEV;
+#endif
+
+#if MY_MAX_DEVICES != 3
+  #error "Currently strictly assuming MY_MAX_DEVICES = 3"
+#endif
+
+#if defined(MY_OMP_NUM_TEAMS_DEV_4) || defined(MY_OMP_NUM_TEAMS_DEV_5)
+  #error "Currently strictly assuming MY_MAX_DEVICES = 3"
+#endif
+
+#ifdef MY_OMP_NUM_TEAMS_DEV_0
+  num_teams_env_devs[0] = MY_OMP_NUM_TEAMS_DEV_0;
+#else
+  num_teams_env_devs[0] = num_teams_env_dev;
+#endif
+
+#ifdef MY_OMP_NUM_TEAMS_DEV_1
+  num_teams_env_devs[1] = MY_OMP_NUM_TEAMS_DEV_1;
+#else
+  num_teams_env_devs[1] = num_teams_env_dev;
+#endif
+
+#ifdef MY_OMP_NUM_TEAMS_DEV_2
+  num_teams_env_devs[2] = MY_OMP_NUM_TEAMS_DEV_2;
+#else
+  num_teams_env_devs[2] = num_teams_env_dev;
+#endif
+
+  /* Check that the number of teams (initial device and in target) is
+     >= 1 and, if omp_get_max_teams() > 0, it does not
+     exceed omp_get_max_teams (). */
+
+  int nteams, num_teams;
+
+  /* Assume that omp_get_max_teams (); returns the ICV, i.e. 0 as default init
+     and not the number of teams that would be run; hence: '>='.  */
+  nteams = omp_get_max_teams ();
+  if (nteams < 0 || (num_teams_env >= 0 && nteams != num_teams_env))
+    __builtin_abort ();
+  num_teams = -1;
+
+  #pragma omp teams
+   if (omp_get_team_num () == 0)
+     num_teams = omp_get_num_teams ();
+  if (num_teams < 1 || (nteams > 0 && num_teams > nteams))
+    __builtin_abort ();
+
+  /* GCC hard codes 3 teams - check for it.  */
+  if (nteams <= 0 && num_teams != 3)
+    __builtin_abort ();
+
+  /* For each device, including host fallback.  */
+  for (int dev = 0; dev <= omp_get_num_devices (); dev++)
+    {
+      int num_teams_icv = num_teams_env_dev;
+      if (dev == omp_get_num_devices ())
+	num_teams_icv = num_teams_env;
+      else if (dev < MY_MAX_DEVICES)
+	num_teams_icv = num_teams_env_devs[dev];
+
+      nteams = -1;
+      #pragma omp target device(dev) map(from: nteams)
+	nteams = omp_get_max_teams ();
+      if (nteams < 0 || (num_teams_icv >= 0 && nteams != num_teams_icv))
+	__builtin_abort ();
+
+      num_teams = -1;
+      #pragma omp target teams device(dev) map(from: num_teams)
+	if (omp_get_team_num () == 0)
+	  num_teams = omp_get_num_teams ();
+
+      if (num_teams < 1 || (nteams > 0 && num_teams > nteams))
+	__builtin_abort ();
+
+      /* GCC hard codes 1 team for host fallback - check for it.  */
+      if (dev == omp_get_num_devices () && num_teams != 1)
+	__builtin_abort ();
+    }
+
+  /* Now set the nteams-var ICV and check that omp_get_max_teams()
+     returns the set value and that the following holds:
+     num_teams >= 1 and num_teams <= nteams-var ICV.
+
+     Additionally, implementation defined, assume:
+     - num_teams == (not '<=') nteams-var ICV, except:
+     - num_teams == 1 for host fallback.  */
+
+  omp_set_num_teams (5);
+
+  nteams = omp_get_max_teams ();
+  if (nteams != 5)
+    __builtin_abort ();
+  num_teams = -1;
+
+  #pragma omp teams
+   if (omp_get_team_num () == 0)
+     num_teams = omp_get_num_teams ();
+  if (num_teams != 5)
+    __builtin_abort ();
+
+  /* For each device, including host fallback.  */
+  for (int dev = 0; dev <= omp_get_num_devices (); dev++)
+    {
+      #pragma omp target device(dev) firstprivate(dev)
+	omp_set_num_teams (7 + dev);
+
+      #pragma omp target device(dev) map(from: nteams)
+	nteams = omp_get_max_teams ();
+      if (nteams != 7 + dev)
+	__builtin_abort ();
+
+      num_teams = -1;
+      #pragma omp target teams device(dev) map(from: num_teams)
+	if (omp_get_team_num () == 0)
+	  num_teams = omp_get_num_teams ();
+
+      if (dev == omp_get_num_devices ())
+	{
+	  if (num_teams != 1)
+	    __builtin_abort ();
+	}
+      else
+	{
+	  if (num_teams != 7 + dev)
+	    __builtin_abort ();
+	}
+    }
+
+  /* Now use the num_teams clause explicitly.  */
+
+  num_teams = -1;
+  #pragma omp teams num_teams(6)
+   if (omp_get_team_num () == 0)
+     num_teams = omp_get_num_teams ();
+  if (num_teams != 6)
+    __builtin_abort ();
+
+  /* For each device, including host fallback.  */
+  for (int dev = 0; dev <= omp_get_num_devices (); dev++)
+    {
+      num_teams = -1;
+      #pragma omp target teams device(dev) map(from: num_teams) num_teams(dev+3)
+	if (omp_get_team_num () == 0)
+	  num_teams = omp_get_num_teams ();
+
+      /* This must match the set value, also with host fallback.  */
+      if (num_teams != 3 + dev)
+	__builtin_abort ();
+    }
+
+  return 0;
+}
diff --git a/libgomp/testsuite/libgomp.c-c++-common/teams-nteams-icv-2.c b/libgomp/testsuite/libgomp.c-c++-common/teams-nteams-icv-2.c
new file mode 100644
index 00000000000..7bd80de065b
--- /dev/null
+++ b/libgomp/testsuite/libgomp.c-c++-common/teams-nteams-icv-2.c
@@ -0,0 +1,8 @@
+/* PR libgomp/109875  */
+/* { dg-set-target-env-var OMP_NUM_TEAMS_ALL 9 } */
+/* { dg-set-target-env-var OMP_NUM_TEAMS_DEV 7 } */
+
+#define MY_OMP_NUM_TEAMS_ALL 9
+#define MY_OMP_NUM_TEAMS_DEV 7
+
+#include "teams-nteams-icv-1.c"
diff --git a/libgomp/testsuite/libgomp.c-c++-common/teams-nteams-icv-3.c b/libgomp/testsuite/libgomp.c-c++-common/teams-nteams-icv-3.c
new file mode 100644
index 00000000000..10a1cdc3503
--- /dev/null
+++ b/libgomp/testsuite/libgomp.c-c++-common/teams-nteams-icv-3.c
@@ -0,0 +1,8 @@
+/* PR libgomp/109875  */
+/* { dg-set-target-env-var OMP_NUM_TEAMS_ALL 7 } */
+/* { dg-set-target-env-var OMP_NUM_TEAMS 8 } */
+
+#define MY_OMP_NUM_TEAMS_ALL 7
+#define MY_OMP_NUM_TEAMS 8
+
+#include "teams-nteams-icv-1.c"
diff --git a/libgomp/testsuite/libgomp.c-c++-common/teams-nteams-icv-4.c b/libgomp/testsuite/libgomp.c-c++-common/teams-nteams-icv-4.c
new file mode 100644
index 00000000000..c5d65774db6
--- /dev/null
+++ b/libgomp/testsuite/libgomp.c-c++-common/teams-nteams-icv-4.c
@@ -0,0 +1,14 @@
+/* PR libgomp/109875  */
+/* { dg-set-target-env-var OMP_NUM_TEAMS_ALL 7 } */
+/* { dg-set-target-env-var OMP_NUM_TEAMS 4 } */
+/* { dg-set-target-env-var OMP_NUM_TEAMS_DEV 8 } */
+/* { dg-set-target-env-var OMP_NUM_TEAMS_DEV_0 5 } */
+/* { dg-set-target-env-var OMP_NUM_TEAMS_DEV_1 11 } */
+
+#define MY_OMP_NUM_TEAMS_ALL 7
+#define MY_OMP_NUM_TEAMS 4
+#define MY_OMP_NUM_TEAMS_DEV 8
+#define MY_OMP_NUM_TEAMS_DEV_0 5
+#define MY_OMP_NUM_TEAMS_DEV_1 11
+
+#include "teams-nteams-icv-1.c"