[OpenMP] Extend omp teams to permit nested omp atomic

[jdenny]

Summary

OpenMP 5.2, sec. 10.2 "teams Construct", p. 232, L9-12 restricts what
regions can be strictly nested within a teams construct. This patch
relaxes Clang's enforcement of this restriction in the case of nested
atomic constructs unless -fno-openmp-extensions is specified.
Cases like the following then seem to work fine with no additional
implementation changes:

#pragma omp target teams map(tofrom:x)
#pragma omp atomic update
x++;

[ABataev]

clang/lib/Sema/SemaOpenMP.cpp

4983

Can you add a check for 5.2?

[jdenny]

clang/lib/Sema/SemaOpenMP.cpp

4983

This is an extension relative not only to 5.2 and but also to previous versions. I quoted 5.2 just because it's the most recent.

Given that, I'm not sure why this should check for 5.2. Am I misunderstanding your request?

[ABataev]

clang/lib/Sema/SemaOpenMP.cpp

4983

Just looks like it is allowed in 5.2. Am I missing something?

[jdenny]

clang/lib/Sema/SemaOpenMP.cpp

4983

I'm not aware of any change in the 5.2 spec that would permit this. The text I cited in the review summary does not include atomic in the list of regions "that may be strictly nested inside the teams region":

distribute regions, including any distribute regions arising from composite constructs,
parallel regions, including any parallel regions arising from combined constructs, loop
regions, omp_get_num_teams() regions, and omp_get_team_num() regions are the
only OpenMP regions that may be strictly nested inside the teams region.

[ABataev]

LG

[jdenny]

Fixed a bug in the new tests.

[jdenny]

Thanks. Will try to push tomorrow.

clang/lib/Sema/SemaOpenMP.cpp

if (!NestingProhibited &&
         (ParentRegion == OMPD_teams || ParentRegion == OMPD_target_teams)) {
       // OpenMP [5.1, 2.22, Nesting of Regions]
       // distribute, distribute simd, distribute parallel worksharing-loop,
       // distribute parallel worksharing-loop SIMD, loop, parallel regions,
       // including any parallel regions arising from combined constructs,
       // omp_get_num_teams() regions, and omp_get_team_num() regions are the
       // only OpenMP regions that may be strictly nested inside the teams
       // region.
+      //
+      // As an extension, we permit atomic within teams as well.
       NestingProhibited = !isOpenMPParallelDirective(CurrentRegion) &&
                           !isOpenMPDistributeDirective(CurrentRegion) &&
-                          CurrentRegion != OMPD_loop;
+                          CurrentRegion != OMPD_loop &&
+                          !(SemaRef.getLangOpts().OpenMPExtensions &&
+                            CurrentRegion == OMPD_atomic);
       Recommend = ShouldBeInParallelRegion;
     }
     if (!NestingProhibited && CurrentRegion == OMPD_loop) {
       // OpenMP [5.1, 2.11.7, loop Construct, Restrictions]
       // If the bind clause is present on the loop construct and binding is
       // teams then the corresponding loop region must be strictly nested inside
       // a teams region.
       NestingProhibited = BindKind == OMPC_BIND_teams &&

clang/test/OpenMP/nesting_of_regions.cpp

-// RUN: %clang_cc1 -fsyntax-only -fopenmp -fopenmp-version=45 -verify=expected,omp45,omp45warn %s
-// RUN: %clang_cc1 -fsyntax-only -fopenmp -verify=expected,omp50 %s
-// RUN: %clang_cc1 -fsyntax-only -fopenmp -fopenmp-version=45 -verify=expected,omp45 -Wno-openmp %s
-// RUN: %clang_cc1 -fsyntax-only -fopenmp -fopenmp-version=45 -verify=expected,omp45 -Wno-source-uses-openmp %s
+// RUN: %clang_cc1 -fsyntax-only -fopenmp -fopenmp-version=45 -fno-openmp-extensions -verify=expected,omp45,omp45warn,omp %s
+// RUN: %clang_cc1 -fsyntax-only -fopenmp -fno-openmp-extensions -verify=expected,omp50,omp %s
+// RUN: %clang_cc1 -fsyntax-only -fopenmp -fopenmp-extensions -verify=expected,omp50 %s
+// RUN: %clang_cc1 -fsyntax-only -fopenmp -fopenmp-version=45 -verify=expected,omp45,omp -fno-openmp-extensions -Wno-openmp %s
+// RUN: %clang_cc1 -fsyntax-only -fopenmp -fopenmp-version=45 -verify=expected,omp45,omp -fno-openmp-extensions -Wno-source-uses-openmp %s
 
-// RUN: %clang_cc1 -fsyntax-only -fopenmp-simd -fopenmp-version=45 -verify=expected,omp45,omp45warn %s
-// RUN: %clang_cc1 -fsyntax-only -fopenmp-simd -verify=expected,omp50 %s
+// RUN: %clang_cc1 -fsyntax-only -fopenmp-simd -fopenmp-version=45 -fno-openmp-extensions -verify=expected,omp45,omp45warn,omp %s
+// RUN: %clang_cc1 -fsyntax-only -fopenmp-simd -verify=expected,omp50,omp -fno-openmp-extensions %s
 // SIMD-ONLY0-NOT: {{__kmpc|__tgt}}
 
 void bar();
 
 template <class T>
 void foo() {
   T a = T();
 // PARALLEL DIRECTIVE

 #pragma omp teams
   {
 #pragma omp ordered // expected-error {{region cannot be closely nested inside 'teams' region; perhaps you forget to enclose 'omp ordered' directive into a for or a parallel for region with 'ordered' clause?}}
     bar();
   }
 #pragma omp target
 #pragma omp teams
   {
-#pragma omp atomic // expected-error {{region cannot be closely nested inside 'teams' region; perhaps you forget to enclose 'omp atomic' directive into a parallel region?}}
+#pragma omp atomic // omp-error {{region cannot be closely nested inside 'teams' region; perhaps you forget to enclose 'omp atomic' directive into a parallel region?}}
     ++a;
   }
 #pragma omp target
 #pragma omp teams
   {
 #pragma omp target // expected-error {{region cannot be nested inside 'target' region}}
     ++a;
   }
#pragma omp flush // expected-error {{region cannot be closely nested inside 'target teams' region; perhaps you forget to enclose 'omp flush' directive into a parallel region?}}
   }
 #pragma omp target teams
   {
 #pragma omp ordered // expected-error {{region cannot be closely nested inside 'target teams' region; perhaps you forget to enclose 'omp ordered' directive into a for or a parallel for region with 'ordered' clause?}}
     bar();
   }
 #pragma omp target teams
   {
-#pragma omp atomic // expected-error {{region cannot be closely nested inside 'target teams' region; perhaps you forget to enclose 'omp atomic' directive into a parallel region?}}
+#pragma omp atomic // omp-error {{region cannot be closely nested inside 'target teams' region; perhaps you forget to enclose 'omp atomic' directive into a parallel region?}}
     ++a;
   }
 #pragma omp target teams
   {
 #pragma omp target // expected-error {{region cannot be nested inside 'target teams' region}}
     ++a;
   }
 #pragma omp target teams

 #pragma omp teams
   {
 #pragma omp ordered // expected-error {{region cannot be closely nested inside 'teams' region; perhaps you forget to enclose 'omp ordered' directive into a for or a parallel for region with 'ordered' clause?}}
     bar();
   }
 #pragma omp target
 #pragma omp teams
   {
-#pragma omp atomic // expected-error {{region cannot be closely nested inside 'teams' region; perhaps you forget to enclose 'omp atomic' directive into a parallel region?}}
+#pragma omp atomic // omp-error {{region cannot be closely nested inside 'teams' region; perhaps you forget to enclose 'omp atomic' directive into a parallel region?}}
     ++a;
   }
 #pragma omp target
 #pragma omp teams
   {
 #pragma omp target // expected-error {{region cannot be nested inside 'target' region}}
     ++a;
   }
#pragma omp flush // expected-error {{region cannot be closely nested inside 'target teams' region; perhaps you forget to enclose 'omp flush' directive into a parallel region?}}
   }
 #pragma omp target teams
   {
 #pragma omp ordered // expected-error {{region cannot be closely nested inside 'target teams' region; perhaps you forget to enclose 'omp ordered' directive into a for or a parallel for region with 'ordered' clause?}}
     bar();
   }
 #pragma omp target teams
   {
-#pragma omp atomic // expected-error {{region cannot be closely nested inside 'target teams' region; perhaps you forget to enclose 'omp atomic' directive into a parallel region?}}
+#pragma omp atomic // omp-error {{region cannot be closely nested inside 'target teams' region; perhaps you forget to enclose 'omp atomic' directive into a parallel region?}}
     ++a;
   }
 #pragma omp target teams
   {
 #pragma omp target // expected-error {{region cannot be nested inside 'target teams' region}}
     ++a;
   }
 #pragma omp target teams

openmp/libomptarget/test/offloading/target-teams-atomic.c

+// Check that omp atomic is permitted and behaves when strictly nested within
+// omp target teams.  This is an extension to OpenMP 5.2 and is enabled by
+// default.
+
+// RUN: %libomptarget-compile-run-and-check-generic
+
+#include <omp.h>
+#include <stdbool.h>
+#include <stdio.h>
+#include <string.h>
+
+// High parallelism increases our chances of detecting a lack of atomicity.
+#define NUM_TEAMS_TRY 256
+
+int main() {
+  //      CHECK: update: num_teams=[[#NUM_TEAMS:]]{{$}}
+  // CHECK-NEXT: update: x=[[#NUM_TEAMS]]{{$}}
+  int x = 0;
+  int numTeams;
+  #pragma omp target teams num_teams(NUM_TEAMS_TRY) map(tofrom:x, numTeams)
+  {
+    #pragma omp atomic update
+    ++x;
+    if (omp_get_team_num() == 0)
+      numTeams = omp_get_num_teams();
+  }
+  printf("update: num_teams=%d\n", numTeams);
+  printf("update: x=%d\n", x);
+
+  // CHECK-NEXT: capture: x=[[#NUM_TEAMS]]{{$}}
+  // CHECK-NEXT: capture: xCapturedCount=[[#NUM_TEAMS]]{{$}}
+  bool xCaptured[numTeams];
+  memset(xCaptured, 0, sizeof xCaptured);
+  x = 0;
+  #pragma omp target teams num_teams(NUM_TEAMS_TRY) map(tofrom:x, numTeams)
+  {
+    int v;
+    #pragma omp atomic capture
+    v = x++;
+    xCaptured[v] = true;
+  }
+  printf("capture: x=%d\n", x);
+  int xCapturedCount = 0;
+  for (int i = 0; i < numTeams; ++i) {
+    if (xCaptured[i])
+      ++xCapturedCount;
+  }
+  printf("capture: xCapturedCount=%d\n", xCapturedCount);
+  return 0;
+}

openmp/runtime/test/teams/teams-atomic.c

+// Check that omp atomic is permitted and behaves when strictly nested within
+// omp teams.  This is an extension to OpenMP 5.2 and is enabled by default.
+
+// RUN: %libomp-compile-and-run | FileCheck %s
+
+#include <omp.h>
+#include <stdbool.h>
+#include <stdio.h>
+#include <string.h>
+
+// High parallelism increases our chances of detecting a lack of atomicity.
+#define NUM_TEAMS_TRY 256
+
+int main() {
+  //      CHECK: update: num_teams=[[#NUM_TEAMS:]]{{$}}
+  // CHECK-NEXT: update: x=[[#NUM_TEAMS]]{{$}}
+  int x = 0;
+  int numTeams;
+  #pragma omp teams num_teams(NUM_TEAMS_TRY)
+  {
+    #pragma omp atomic update
+    ++x;
+    if (omp_get_team_num() == 0)
+      numTeams = omp_get_num_teams();
+  }
+  printf("update: num_teams=%d\n", numTeams);
+  printf("update: x=%d\n", x);
+
+  // CHECK-NEXT: capture: x=[[#NUM_TEAMS]]{{$}}
+  // CHECK-NEXT: capture: xCapturedCount=[[#NUM_TEAMS]]{{$}}
+  bool xCaptured[numTeams];
+  memset(xCaptured, 0, sizeof xCaptured);
+  x = 0;
+  #pragma omp teams num_teams(NUM_TEAMS_TRY)
+  {
+    int v;
+    #pragma omp atomic capture
+    v = x++;
+    xCaptured[v] = true;
+  }
+  printf("capture: x=%d\n", x);
+  int xCapturedCount = 0;
+  for (int i = 0; i < numTeams; ++i) {
+    if (xCaptured[i])
+      ++xCapturedCount;
+  }
+  printf("capture: xCapturedCount=%d\n", xCapturedCount);
+  return 0;
+}