[flang][openmp] Add parser/semantic support for workdistribute #154377

skc7 · 2025-08-19T16:28:43Z

This PR adds workdistribute parser and semantic support in flang.

This is part 2 of workdistribute frontend support.

Part1 : #154376 [OpenMP] Add workdistribute construct in openMP dialect and in llvm frontend
Part2 : #154377 [flang][openmp] Add parser/semantic support for workdistribute
Part3 : #154378 [flang][openmp] Add Lowering to omp mlir for workdistribute construct

llvmbot · 2025-08-20T09:26:08Z

@llvm/pr-subscribers-flang-parser

@llvm/pr-subscribers-flang-semantics

Author: Chaitanya (skc7)

Changes

This PR adds workdistribute parser and semantic support in flang.

This is part 2 of workdistribute frontend support.

Part1 : #154376 [OpenMP] Add workdistribute construct in openMP dialect and in llvm frontend
Part2 : #154377 [flang][openmp] Add parser/semantic support for workdistribute
Part3 : #154378 [flang][openmp] Add Lowering to omp mlir for workdistribute construct

Full diff: https://github.com/llvm/llvm-project/pull/154377.diff

9 Files Affected:

(modified) flang/include/flang/Semantics/openmp-directive-sets.h (+7)
(modified) flang/lib/Parser/openmp-parsers.cpp (+5-1)
(modified) flang/lib/Semantics/check-omp-structure.cpp (+95)
(modified) flang/lib/Semantics/check-omp-structure.h (+1)
(modified) flang/lib/Semantics/resolve-directives.cpp (+7-1)
(added) flang/test/Parser/OpenMP/workdistribute.f90 (+27)
(added) flang/test/Semantics/OpenMP/workdistribute01.f90 (+16)
(added) flang/test/Semantics/OpenMP/workdistribute02.f90 (+34)
(added) flang/test/Semantics/OpenMP/workdistribute03.f90 (+34)

diff --git a/flang/include/flang/Semantics/openmp-directive-sets.h b/flang/include/flang/Semantics/openmp-directive-sets.h
index cc66cc833e8b7..01e8481e05721 100644
--- a/flang/include/flang/Semantics/openmp-directive-sets.h
+++ b/flang/include/flang/Semantics/openmp-directive-sets.h
@@ -143,6 +143,7 @@ static const OmpDirectiveSet topTargetSet{
     Directive::OMPD_target_teams_distribute_parallel_do_simd,
     Directive::OMPD_target_teams_distribute_simd,
     Directive::OMPD_target_teams_loop,
+    Directive::OMPD_target_teams_workdistribute,
 };
 
 static const OmpDirectiveSet allTargetSet{topTargetSet};
@@ -172,6 +173,7 @@ static const OmpDirectiveSet topTeamsSet{
     Directive::OMPD_teams_distribute_parallel_do_simd,
     Directive::OMPD_teams_distribute_simd,
     Directive::OMPD_teams_loop,
+    Directive::OMPD_teams_workdistribute,
 };
 
 static const OmpDirectiveSet bottomTeamsSet{
@@ -187,6 +189,7 @@ static const OmpDirectiveSet allTeamsSet{
         Directive::OMPD_target_teams_distribute_parallel_do_simd,
         Directive::OMPD_target_teams_distribute_simd,
         Directive::OMPD_target_teams_loop,
+        Directive::OMPD_target_teams_workdistribute,
     } | topTeamsSet,
 };
 
@@ -230,6 +233,9 @@ static const OmpDirectiveSet blockConstructSet{
     Directive::OMPD_taskgroup,
     Directive::OMPD_teams,
     Directive::OMPD_workshare,
+    Directive::OMPD_target_teams_workdistribute,
+    Directive::OMPD_teams_workdistribute,
+    Directive::OMPD_workdistribute,
 };
 
 static const OmpDirectiveSet loopConstructSet{
@@ -376,6 +382,7 @@ static const OmpDirectiveSet nestedReduceWorkshareAllowedSet{
 };
 
 static const OmpDirectiveSet nestedTeamsAllowedSet{
+    Directive::OMPD_workdistribute,
     Directive::OMPD_distribute,
     Directive::OMPD_distribute_parallel_do,
     Directive::OMPD_distribute_parallel_do_simd,
diff --git a/flang/lib/Parser/openmp-parsers.cpp b/flang/lib/Parser/openmp-parsers.cpp
index 56cee4ab38e9b..51b49a591b02f 100644
--- a/flang/lib/Parser/openmp-parsers.cpp
+++ b/flang/lib/Parser/openmp-parsers.cpp
@@ -1870,11 +1870,15 @@ TYPE_PARSER( //
     MakeBlockConstruct(llvm::omp::Directive::OMPD_target_data) ||
     MakeBlockConstruct(llvm::omp::Directive::OMPD_target_parallel) ||
     MakeBlockConstruct(llvm::omp::Directive::OMPD_target_teams) ||
+    MakeBlockConstruct(
+        llvm::omp::Directive::OMPD_target_teams_workdistribute) ||
     MakeBlockConstruct(llvm::omp::Directive::OMPD_target) ||
     MakeBlockConstruct(llvm::omp::Directive::OMPD_task) ||
     MakeBlockConstruct(llvm::omp::Directive::OMPD_taskgroup) ||
     MakeBlockConstruct(llvm::omp::Directive::OMPD_teams) ||
-    MakeBlockConstruct(llvm::omp::Directive::OMPD_workshare))
+    MakeBlockConstruct(llvm::omp::Directive::OMPD_teams_workdistribute) ||
+    MakeBlockConstruct(llvm::omp::Directive::OMPD_workshare) ||
+    MakeBlockConstruct(llvm::omp::Directive::OMPD_workdistribute))
 #undef MakeBlockConstruct
 
 // OMP SECTIONS Directive
diff --git a/flang/lib/Semantics/check-omp-structure.cpp b/flang/lib/Semantics/check-omp-structure.cpp
index 2b36b085ae08d..4c4e17c39c03a 100644
--- a/flang/lib/Semantics/check-omp-structure.cpp
+++ b/flang/lib/Semantics/check-omp-structure.cpp
@@ -141,6 +141,67 @@ class OmpWorkshareBlockChecker {
   parser::CharBlock source_;
 };
 
+// 'OmpWorkdistributeBlockChecker' is used to check the validity of the
+// assignment statements and the expressions enclosed in an OpenMP
+// workdistribute construct
+class OmpWorkdistributeBlockChecker {
+public:
+  OmpWorkdistributeBlockChecker(
+      SemanticsContext &context, parser::CharBlock source)
+      : context_{context}, source_{source} {}
+
+  template <typename T> bool Pre(const T &) { return true; }
+  template <typename T> void Post(const T &) {}
+
+  bool Pre(const parser::AssignmentStmt &assignment) {
+    const auto &var{std::get<parser::Variable>(assignment.t)};
+    const auto &expr{std::get<parser::Expr>(assignment.t)};
+    const auto *lhs{GetExpr(context_, var)};
+    const auto *rhs{GetExpr(context_, expr)};
+    if (lhs && rhs) {
+      Tristate isDefined{semantics::IsDefinedAssignment(
+          lhs->GetType(), lhs->Rank(), rhs->GetType(), rhs->Rank())};
+      if (isDefined == Tristate::Yes) {
+        context_.Say(expr.source,
+            "Defined assignment statement is not "
+            "allowed in a WORKDISTRIBUTE construct"_err_en_US);
+      }
+    }
+    return true;
+  }
+
+  bool Pre(const parser::Expr &expr) {
+    if (const auto *e{GetExpr(context_, expr)}) {
+      for (const Symbol &symbol : evaluate::CollectSymbols(*e)) {
+        const Symbol &root{GetAssociationRoot(symbol)};
+        if (IsFunction(root)) {
+          std::string attrs{""};
+          if (!IsElementalProcedure(root)) {
+            attrs = " non-ELEMENTAL";
+          }
+          if (root.attrs().test(Attr::IMPURE)) {
+            if (attrs != "") {
+              attrs = "," + attrs;
+            }
+            attrs = " IMPURE" + attrs;
+          }
+          if (attrs != "") {
+            context_.Say(expr.source,
+                "User defined%s function '%s' is not allowed in a "
+                "WORKDISTRIBUTE construct"_err_en_US,
+                attrs, root.name());
+          }
+        }
+      }
+    }
+    return false;
+  }
+
+private:
+  SemanticsContext &context_;
+  parser::CharBlock source_;
+};
+
 // `OmpUnitedTaskDesignatorChecker` is used to check if the designator
 // can appear within the TASK construct
 class OmpUnitedTaskDesignatorChecker {
@@ -813,6 +874,13 @@ void OmpStructureChecker::Enter(const parser::OpenMPBlockConstruct &x) {
           "TARGET construct with nested TEAMS region contains statements or "
           "directives outside of the TEAMS construct"_err_en_US);
     }
+    if (GetContext().directive == llvm::omp::Directive::OMPD_workdistribute &&
+        GetContextParent().directive != llvm::omp::Directive::OMPD_teams) {
+      context_.Say(x.BeginDir().DirName().source,
+          "%s region can only be strictly nested within the "
+          "teams region"_err_en_US,
+          ContextDirectiveAsFortran());
+    }
   }
 
   CheckNoBranching(block, beginSpec.DirId(), beginSpec.source);
@@ -896,6 +964,17 @@ void OmpStructureChecker::Enter(const parser::OpenMPBlockConstruct &x) {
     HasInvalidWorksharingNesting(
         beginSpec.source, llvm::omp::nestedWorkshareErrSet);
     break;
+  case llvm::omp::OMPD_workdistribute:
+    if (!CurrentDirectiveIsNested()) {
+      context_.Say(beginSpec.source,
+          "A workdistribute region must be nested inside teams region only."_err_en_US);
+    }
+    CheckWorkdistributeBlockStmts(block, beginSpec.source);
+    break;
+  case llvm::omp::OMPD_teams_workdistribute:
+  case llvm::omp::OMPD_target_teams_workdistribute:
+    CheckWorkdistributeBlockStmts(block, beginSpec.source);
+    break;
   case llvm::omp::Directive::OMPD_scope:
   case llvm::omp::Directive::OMPD_single:
     // TODO: This check needs to be extended while implementing nesting of
@@ -4497,6 +4576,22 @@ void OmpStructureChecker::CheckWorkshareBlockStmts(
   }
 }
 
+void OmpStructureChecker::CheckWorkdistributeBlockStmts(
+    const parser::Block &block, parser::CharBlock source) {
+  OmpWorkdistributeBlockChecker ompWorkdistributeBlockChecker{context_, source};
+
+  for (auto it{block.begin()}; it != block.end(); ++it) {
+    if (parser::Unwrap<parser::AssignmentStmt>(*it)) {
+      parser::Walk(*it, ompWorkdistributeBlockChecker);
+    } else {
+      context_.Say(source,
+          "The structured block in a WORKDISTRIBUTE construct may consist of "
+          "only "
+          "SCALAR or ARRAY assignments"_err_en_US);
+    }
+  }
+}
+
 void OmpStructureChecker::CheckIfContiguous(const parser::OmpObject &object) {
   if (auto contig{IsContiguous(context_, object)}; contig && !*contig) {
     const parser::Name *name{GetObjectName(object)};
diff --git a/flang/lib/Semantics/check-omp-structure.h b/flang/lib/Semantics/check-omp-structure.h
index b548a455ee077..08519c4836ed8 100644
--- a/flang/lib/Semantics/check-omp-structure.h
+++ b/flang/lib/Semantics/check-omp-structure.h
@@ -245,6 +245,7 @@ class OmpStructureChecker
       llvmOmpClause clause, const parser::OmpObjectList &ompObjectList);
   bool CheckTargetBlockOnlyTeams(const parser::Block &);
   void CheckWorkshareBlockStmts(const parser::Block &, parser::CharBlock);
+  void CheckWorkdistributeBlockStmts(const parser::Block &, parser::CharBlock);
 
   void CheckIteratorRange(const parser::OmpIteratorSpecifier &x);
   void CheckIteratorModifier(const parser::OmpIterator &x);
diff --git a/flang/lib/Semantics/resolve-directives.cpp b/flang/lib/Semantics/resolve-directives.cpp
index 6a4660c9882ab..069d78554e9c2 100644
--- a/flang/lib/Semantics/resolve-directives.cpp
+++ b/flang/lib/Semantics/resolve-directives.cpp
@@ -1736,10 +1736,13 @@ bool OmpAttributeVisitor::Pre(const parser::OpenMPBlockConstruct &x) {
   case llvm::omp::Directive::OMPD_task:
   case llvm::omp::Directive::OMPD_taskgroup:
   case llvm::omp::Directive::OMPD_teams:
+  case llvm::omp::Directive::OMPD_workdistribute:
   case llvm::omp::Directive::OMPD_workshare:
   case llvm::omp::Directive::OMPD_parallel_workshare:
   case llvm::omp::Directive::OMPD_target_teams:
+  case llvm::omp::Directive::OMPD_target_teams_workdistribute:
   case llvm::omp::Directive::OMPD_target_parallel:
+  case llvm::omp::Directive::OMPD_teams_workdistribute:
     PushContext(dirSpec.source, dirId);
     break;
   default:
@@ -1769,9 +1772,12 @@ void OmpAttributeVisitor::Post(const parser::OpenMPBlockConstruct &x) {
   case llvm::omp::Directive::OMPD_target:
   case llvm::omp::Directive::OMPD_task:
   case llvm::omp::Directive::OMPD_teams:
+  case llvm::omp::Directive::OMPD_workdistribute:
   case llvm::omp::Directive::OMPD_parallel_workshare:
   case llvm::omp::Directive::OMPD_target_teams:
-  case llvm::omp::Directive::OMPD_target_parallel: {
+  case llvm::omp::Directive::OMPD_target_parallel:
+  case llvm::omp::Directive::OMPD_target_teams_workdistribute:
+  case llvm::omp::Directive::OMPD_teams_workdistribute: {
     bool hasPrivate;
     for (const auto *allocName : allocateNames_) {
       hasPrivate = false;
diff --git a/flang/test/Parser/OpenMP/workdistribute.f90 b/flang/test/Parser/OpenMP/workdistribute.f90
new file mode 100644
index 0000000000000..874538d5f0296
--- /dev/null
+++ b/flang/test/Parser/OpenMP/workdistribute.f90
@@ -0,0 +1,27 @@
+!RUN: %flang_fc1 -fdebug-unparse -fopenmp %s | FileCheck --ignore-case --check-prefix="UNPARSE" %s
+!RUN: %flang_fc1 -fdebug-dump-parse-tree -fopenmp -fopenmp-version=61 %s | FileCheck --check-prefix="PARSE-TREE" %s
+
+!UNPARSE: SUBROUTINE teams_workdistribute
+!UNPARSE:  USE :: iso_fortran_env
+!UNPARSE:  REAL(KIND=4_4) a
+!UNPARSE:  REAL(KIND=4_4), DIMENSION(10_4) :: x
+!UNPARSE:  REAL(KIND=4_4), DIMENSION(10_4) :: y
+!UNPARSE: !$OMP TEAMS WORKDISTRIBUTE
+!UNPARSE:   y=a*x+y
+!UNPARSE: !$OMP END TEAMS WORKDISTRIBUTE
+!UNPARSE: END SUBROUTINE teams_workdistribute
+
+!PARSE-TREE: | | | OmpBeginDirective
+!PARSE-TREE: | | | | OmpDirectiveName -> llvm::omp::Directive = teams workdistribute
+!PARSE-TREE: | | | OmpEndDirective
+!PARSE-TREE: | | | | OmpDirectiveName -> llvm::omp::Directive = teams workdistribute
+
+subroutine teams_workdistribute()
+  use iso_fortran_env
+  real(kind=real32) :: a
+  real(kind=real32), dimension(10) :: x
+  real(kind=real32), dimension(10) :: y
+  !$omp teams workdistribute
+  y = a * x + y
+  !$omp end teams workdistribute
+end subroutine teams_workdistribute
diff --git a/flang/test/Semantics/OpenMP/workdistribute01.f90 b/flang/test/Semantics/OpenMP/workdistribute01.f90
new file mode 100644
index 0000000000000..76ddbc74ea4ec
--- /dev/null
+++ b/flang/test/Semantics/OpenMP/workdistribute01.f90
@@ -0,0 +1,16 @@
+! RUN: %python %S/../test_errors.py %s %flang -fopenmp
+! OpenMP Version 6.0
+! workdistribute Construct
+! Invalid do construct inside !$omp workdistribute
+
+subroutine workdistribute()
+  integer n, i
+  !ERROR: A workdistribute region must be nested inside teams region only.
+  !ERROR: The structured block in a WORKDISTRIBUTE construct may consist of only SCALAR or ARRAY assignments
+  !$omp workdistribute
+  do i = 1, n
+    print *, "omp workdistribute"
+  end do
+  !$omp end workdistribute
+
+end subroutine workdistribute
diff --git a/flang/test/Semantics/OpenMP/workdistribute02.f90 b/flang/test/Semantics/OpenMP/workdistribute02.f90
new file mode 100644
index 0000000000000..ad2cde2c3daf0
--- /dev/null
+++ b/flang/test/Semantics/OpenMP/workdistribute02.f90
@@ -0,0 +1,34 @@
+! RUN: %python %S/../test_errors.py %s %flang -fopenmp
+! OpenMP Version 6.0
+! workdistribute Construct
+! The !omp workdistribute construct must not contain any user defined
+! function calls unless the function is ELEMENTAL.
+
+module my_mod
+  contains
+  integer function my_func()
+    my_func = 10
+  end function my_func
+
+  impure integer function impure_my_func()
+    impure_my_func = 20
+  end function impure_my_func
+
+  impure elemental integer function impure_ele_my_func()
+    impure_ele_my_func = 20
+  end function impure_ele_my_func
+end module my_mod
+
+subroutine workdistribute(aa, bb, cc, n)
+  use my_mod
+  integer n
+  real aa(n), bb(n), cc(n)
+  !$omp teams
+  !$omp workdistribute
+  !ERROR: User defined non-ELEMENTAL function 'my_func' is not allowed in a WORKDISTRIBUTE construct
+  aa = my_func()
+  aa = bb * cc
+  !$omp end workdistribute
+  !$omp end teams
+
+end subroutine workdistribute
diff --git a/flang/test/Semantics/OpenMP/workdistribute03.f90 b/flang/test/Semantics/OpenMP/workdistribute03.f90
new file mode 100644
index 0000000000000..eac28cb39c47f
--- /dev/null
+++ b/flang/test/Semantics/OpenMP/workdistribute03.f90
@@ -0,0 +1,34 @@
+! RUN: %python %S/../test_errors.py %s %flang -fopenmp
+! OpenMP Version 6.0
+! workdistribute Construct
+! All array assignments, scalar assignments, and masked array assignments
+! must be intrinsic assignments.
+
+module defined_assign
+  interface assignment(=)
+    module procedure work_assign
+  end interface
+
+  contains
+    subroutine work_assign(a,b)
+      integer, intent(out) :: a
+      logical, intent(in) :: b(:)
+    end subroutine work_assign
+end module defined_assign
+
+program omp_workdistribute
+  use defined_assign
+
+  integer :: a, aa(10), bb(10)
+  logical :: l(10)
+  l = .TRUE.
+
+  !$omp teams
+  !$omp workdistribute
+  !ERROR: Defined assignment statement is not allowed in a WORKDISTRIBUTE construct
+  a = l
+  aa = bb
+  !$omp end workdistribute
+  !$omp end teams
+
+end program omp_workdistribute

llvmbot · 2025-08-20T09:26:08Z

@llvm/pr-subscribers-flang-openmp

Author: Chaitanya (skc7)

Changes

This PR adds workdistribute parser and semantic support in flang.

This is part 2 of workdistribute frontend support.

Part1 : #154376 [OpenMP] Add workdistribute construct in openMP dialect and in llvm frontend
Part2 : #154377 [flang][openmp] Add parser/semantic support for workdistribute
Part3 : #154378 [flang][openmp] Add Lowering to omp mlir for workdistribute construct

Full diff: https://github.com/llvm/llvm-project/pull/154377.diff

9 Files Affected:

(modified) flang/include/flang/Semantics/openmp-directive-sets.h (+7)
(modified) flang/lib/Parser/openmp-parsers.cpp (+5-1)
(modified) flang/lib/Semantics/check-omp-structure.cpp (+95)
(modified) flang/lib/Semantics/check-omp-structure.h (+1)
(modified) flang/lib/Semantics/resolve-directives.cpp (+7-1)
(added) flang/test/Parser/OpenMP/workdistribute.f90 (+27)
(added) flang/test/Semantics/OpenMP/workdistribute01.f90 (+16)
(added) flang/test/Semantics/OpenMP/workdistribute02.f90 (+34)
(added) flang/test/Semantics/OpenMP/workdistribute03.f90 (+34)

diff --git a/flang/include/flang/Semantics/openmp-directive-sets.h b/flang/include/flang/Semantics/openmp-directive-sets.h
index cc66cc833e8b7..01e8481e05721 100644
--- a/flang/include/flang/Semantics/openmp-directive-sets.h
+++ b/flang/include/flang/Semantics/openmp-directive-sets.h
@@ -143,6 +143,7 @@ static const OmpDirectiveSet topTargetSet{
     Directive::OMPD_target_teams_distribute_parallel_do_simd,
     Directive::OMPD_target_teams_distribute_simd,
     Directive::OMPD_target_teams_loop,
+    Directive::OMPD_target_teams_workdistribute,
 };
 
 static const OmpDirectiveSet allTargetSet{topTargetSet};
@@ -172,6 +173,7 @@ static const OmpDirectiveSet topTeamsSet{
     Directive::OMPD_teams_distribute_parallel_do_simd,
     Directive::OMPD_teams_distribute_simd,
     Directive::OMPD_teams_loop,
+    Directive::OMPD_teams_workdistribute,
 };
 
 static const OmpDirectiveSet bottomTeamsSet{
@@ -187,6 +189,7 @@ static const OmpDirectiveSet allTeamsSet{
         Directive::OMPD_target_teams_distribute_parallel_do_simd,
         Directive::OMPD_target_teams_distribute_simd,
         Directive::OMPD_target_teams_loop,
+        Directive::OMPD_target_teams_workdistribute,
     } | topTeamsSet,
 };
 
@@ -230,6 +233,9 @@ static const OmpDirectiveSet blockConstructSet{
     Directive::OMPD_taskgroup,
     Directive::OMPD_teams,
     Directive::OMPD_workshare,
+    Directive::OMPD_target_teams_workdistribute,
+    Directive::OMPD_teams_workdistribute,
+    Directive::OMPD_workdistribute,
 };
 
 static const OmpDirectiveSet loopConstructSet{
@@ -376,6 +382,7 @@ static const OmpDirectiveSet nestedReduceWorkshareAllowedSet{
 };
 
 static const OmpDirectiveSet nestedTeamsAllowedSet{
+    Directive::OMPD_workdistribute,
     Directive::OMPD_distribute,
     Directive::OMPD_distribute_parallel_do,
     Directive::OMPD_distribute_parallel_do_simd,
diff --git a/flang/lib/Parser/openmp-parsers.cpp b/flang/lib/Parser/openmp-parsers.cpp
index 56cee4ab38e9b..51b49a591b02f 100644
--- a/flang/lib/Parser/openmp-parsers.cpp
+++ b/flang/lib/Parser/openmp-parsers.cpp
@@ -1870,11 +1870,15 @@ TYPE_PARSER( //
     MakeBlockConstruct(llvm::omp::Directive::OMPD_target_data) ||
     MakeBlockConstruct(llvm::omp::Directive::OMPD_target_parallel) ||
     MakeBlockConstruct(llvm::omp::Directive::OMPD_target_teams) ||
+    MakeBlockConstruct(
+        llvm::omp::Directive::OMPD_target_teams_workdistribute) ||
     MakeBlockConstruct(llvm::omp::Directive::OMPD_target) ||
     MakeBlockConstruct(llvm::omp::Directive::OMPD_task) ||
     MakeBlockConstruct(llvm::omp::Directive::OMPD_taskgroup) ||
     MakeBlockConstruct(llvm::omp::Directive::OMPD_teams) ||
-    MakeBlockConstruct(llvm::omp::Directive::OMPD_workshare))
+    MakeBlockConstruct(llvm::omp::Directive::OMPD_teams_workdistribute) ||
+    MakeBlockConstruct(llvm::omp::Directive::OMPD_workshare) ||
+    MakeBlockConstruct(llvm::omp::Directive::OMPD_workdistribute))
 #undef MakeBlockConstruct
 
 // OMP SECTIONS Directive
diff --git a/flang/lib/Semantics/check-omp-structure.cpp b/flang/lib/Semantics/check-omp-structure.cpp
index 2b36b085ae08d..4c4e17c39c03a 100644
--- a/flang/lib/Semantics/check-omp-structure.cpp
+++ b/flang/lib/Semantics/check-omp-structure.cpp
@@ -141,6 +141,67 @@ class OmpWorkshareBlockChecker {
   parser::CharBlock source_;
 };
 
+// 'OmpWorkdistributeBlockChecker' is used to check the validity of the
+// assignment statements and the expressions enclosed in an OpenMP
+// workdistribute construct
+class OmpWorkdistributeBlockChecker {
+public:
+  OmpWorkdistributeBlockChecker(
+      SemanticsContext &context, parser::CharBlock source)
+      : context_{context}, source_{source} {}
+
+  template <typename T> bool Pre(const T &) { return true; }
+  template <typename T> void Post(const T &) {}
+
+  bool Pre(const parser::AssignmentStmt &assignment) {
+    const auto &var{std::get<parser::Variable>(assignment.t)};
+    const auto &expr{std::get<parser::Expr>(assignment.t)};
+    const auto *lhs{GetExpr(context_, var)};
+    const auto *rhs{GetExpr(context_, expr)};
+    if (lhs && rhs) {
+      Tristate isDefined{semantics::IsDefinedAssignment(
+          lhs->GetType(), lhs->Rank(), rhs->GetType(), rhs->Rank())};
+      if (isDefined == Tristate::Yes) {
+        context_.Say(expr.source,
+            "Defined assignment statement is not "
+            "allowed in a WORKDISTRIBUTE construct"_err_en_US);
+      }
+    }
+    return true;
+  }
+
+  bool Pre(const parser::Expr &expr) {
+    if (const auto *e{GetExpr(context_, expr)}) {
+      for (const Symbol &symbol : evaluate::CollectSymbols(*e)) {
+        const Symbol &root{GetAssociationRoot(symbol)};
+        if (IsFunction(root)) {
+          std::string attrs{""};
+          if (!IsElementalProcedure(root)) {
+            attrs = " non-ELEMENTAL";
+          }
+          if (root.attrs().test(Attr::IMPURE)) {
+            if (attrs != "") {
+              attrs = "," + attrs;
+            }
+            attrs = " IMPURE" + attrs;
+          }
+          if (attrs != "") {
+            context_.Say(expr.source,
+                "User defined%s function '%s' is not allowed in a "
+                "WORKDISTRIBUTE construct"_err_en_US,
+                attrs, root.name());
+          }
+        }
+      }
+    }
+    return false;
+  }
+
+private:
+  SemanticsContext &context_;
+  parser::CharBlock source_;
+};
+
 // `OmpUnitedTaskDesignatorChecker` is used to check if the designator
 // can appear within the TASK construct
 class OmpUnitedTaskDesignatorChecker {
@@ -813,6 +874,13 @@ void OmpStructureChecker::Enter(const parser::OpenMPBlockConstruct &x) {
           "TARGET construct with nested TEAMS region contains statements or "
           "directives outside of the TEAMS construct"_err_en_US);
     }
+    if (GetContext().directive == llvm::omp::Directive::OMPD_workdistribute &&
+        GetContextParent().directive != llvm::omp::Directive::OMPD_teams) {
+      context_.Say(x.BeginDir().DirName().source,
+          "%s region can only be strictly nested within the "
+          "teams region"_err_en_US,
+          ContextDirectiveAsFortran());
+    }
   }
 
   CheckNoBranching(block, beginSpec.DirId(), beginSpec.source);
@@ -896,6 +964,17 @@ void OmpStructureChecker::Enter(const parser::OpenMPBlockConstruct &x) {
     HasInvalidWorksharingNesting(
         beginSpec.source, llvm::omp::nestedWorkshareErrSet);
     break;
+  case llvm::omp::OMPD_workdistribute:
+    if (!CurrentDirectiveIsNested()) {
+      context_.Say(beginSpec.source,
+          "A workdistribute region must be nested inside teams region only."_err_en_US);
+    }
+    CheckWorkdistributeBlockStmts(block, beginSpec.source);
+    break;
+  case llvm::omp::OMPD_teams_workdistribute:
+  case llvm::omp::OMPD_target_teams_workdistribute:
+    CheckWorkdistributeBlockStmts(block, beginSpec.source);
+    break;
   case llvm::omp::Directive::OMPD_scope:
   case llvm::omp::Directive::OMPD_single:
     // TODO: This check needs to be extended while implementing nesting of
@@ -4497,6 +4576,22 @@ void OmpStructureChecker::CheckWorkshareBlockStmts(
   }
 }
 
+void OmpStructureChecker::CheckWorkdistributeBlockStmts(
+    const parser::Block &block, parser::CharBlock source) {
+  OmpWorkdistributeBlockChecker ompWorkdistributeBlockChecker{context_, source};
+
+  for (auto it{block.begin()}; it != block.end(); ++it) {
+    if (parser::Unwrap<parser::AssignmentStmt>(*it)) {
+      parser::Walk(*it, ompWorkdistributeBlockChecker);
+    } else {
+      context_.Say(source,
+          "The structured block in a WORKDISTRIBUTE construct may consist of "
+          "only "
+          "SCALAR or ARRAY assignments"_err_en_US);
+    }
+  }
+}
+
 void OmpStructureChecker::CheckIfContiguous(const parser::OmpObject &object) {
   if (auto contig{IsContiguous(context_, object)}; contig && !*contig) {
     const parser::Name *name{GetObjectName(object)};
diff --git a/flang/lib/Semantics/check-omp-structure.h b/flang/lib/Semantics/check-omp-structure.h
index b548a455ee077..08519c4836ed8 100644
--- a/flang/lib/Semantics/check-omp-structure.h
+++ b/flang/lib/Semantics/check-omp-structure.h
@@ -245,6 +245,7 @@ class OmpStructureChecker
       llvmOmpClause clause, const parser::OmpObjectList &ompObjectList);
   bool CheckTargetBlockOnlyTeams(const parser::Block &);
   void CheckWorkshareBlockStmts(const parser::Block &, parser::CharBlock);
+  void CheckWorkdistributeBlockStmts(const parser::Block &, parser::CharBlock);
 
   void CheckIteratorRange(const parser::OmpIteratorSpecifier &x);
   void CheckIteratorModifier(const parser::OmpIterator &x);
diff --git a/flang/lib/Semantics/resolve-directives.cpp b/flang/lib/Semantics/resolve-directives.cpp
index 6a4660c9882ab..069d78554e9c2 100644
--- a/flang/lib/Semantics/resolve-directives.cpp
+++ b/flang/lib/Semantics/resolve-directives.cpp
@@ -1736,10 +1736,13 @@ bool OmpAttributeVisitor::Pre(const parser::OpenMPBlockConstruct &x) {
   case llvm::omp::Directive::OMPD_task:
   case llvm::omp::Directive::OMPD_taskgroup:
   case llvm::omp::Directive::OMPD_teams:
+  case llvm::omp::Directive::OMPD_workdistribute:
   case llvm::omp::Directive::OMPD_workshare:
   case llvm::omp::Directive::OMPD_parallel_workshare:
   case llvm::omp::Directive::OMPD_target_teams:
+  case llvm::omp::Directive::OMPD_target_teams_workdistribute:
   case llvm::omp::Directive::OMPD_target_parallel:
+  case llvm::omp::Directive::OMPD_teams_workdistribute:
     PushContext(dirSpec.source, dirId);
     break;
   default:
@@ -1769,9 +1772,12 @@ void OmpAttributeVisitor::Post(const parser::OpenMPBlockConstruct &x) {
   case llvm::omp::Directive::OMPD_target:
   case llvm::omp::Directive::OMPD_task:
   case llvm::omp::Directive::OMPD_teams:
+  case llvm::omp::Directive::OMPD_workdistribute:
   case llvm::omp::Directive::OMPD_parallel_workshare:
   case llvm::omp::Directive::OMPD_target_teams:
-  case llvm::omp::Directive::OMPD_target_parallel: {
+  case llvm::omp::Directive::OMPD_target_parallel:
+  case llvm::omp::Directive::OMPD_target_teams_workdistribute:
+  case llvm::omp::Directive::OMPD_teams_workdistribute: {
     bool hasPrivate;
     for (const auto *allocName : allocateNames_) {
       hasPrivate = false;
diff --git a/flang/test/Parser/OpenMP/workdistribute.f90 b/flang/test/Parser/OpenMP/workdistribute.f90
new file mode 100644
index 0000000000000..874538d5f0296
--- /dev/null
+++ b/flang/test/Parser/OpenMP/workdistribute.f90
@@ -0,0 +1,27 @@
+!RUN: %flang_fc1 -fdebug-unparse -fopenmp %s | FileCheck --ignore-case --check-prefix="UNPARSE" %s
+!RUN: %flang_fc1 -fdebug-dump-parse-tree -fopenmp -fopenmp-version=61 %s | FileCheck --check-prefix="PARSE-TREE" %s
+
+!UNPARSE: SUBROUTINE teams_workdistribute
+!UNPARSE:  USE :: iso_fortran_env
+!UNPARSE:  REAL(KIND=4_4) a
+!UNPARSE:  REAL(KIND=4_4), DIMENSION(10_4) :: x
+!UNPARSE:  REAL(KIND=4_4), DIMENSION(10_4) :: y
+!UNPARSE: !$OMP TEAMS WORKDISTRIBUTE
+!UNPARSE:   y=a*x+y
+!UNPARSE: !$OMP END TEAMS WORKDISTRIBUTE
+!UNPARSE: END SUBROUTINE teams_workdistribute
+
+!PARSE-TREE: | | | OmpBeginDirective
+!PARSE-TREE: | | | | OmpDirectiveName -> llvm::omp::Directive = teams workdistribute
+!PARSE-TREE: | | | OmpEndDirective
+!PARSE-TREE: | | | | OmpDirectiveName -> llvm::omp::Directive = teams workdistribute
+
+subroutine teams_workdistribute()
+  use iso_fortran_env
+  real(kind=real32) :: a
+  real(kind=real32), dimension(10) :: x
+  real(kind=real32), dimension(10) :: y
+  !$omp teams workdistribute
+  y = a * x + y
+  !$omp end teams workdistribute
+end subroutine teams_workdistribute
diff --git a/flang/test/Semantics/OpenMP/workdistribute01.f90 b/flang/test/Semantics/OpenMP/workdistribute01.f90
new file mode 100644
index 0000000000000..76ddbc74ea4ec
--- /dev/null
+++ b/flang/test/Semantics/OpenMP/workdistribute01.f90
@@ -0,0 +1,16 @@
+! RUN: %python %S/../test_errors.py %s %flang -fopenmp
+! OpenMP Version 6.0
+! workdistribute Construct
+! Invalid do construct inside !$omp workdistribute
+
+subroutine workdistribute()
+  integer n, i
+  !ERROR: A workdistribute region must be nested inside teams region only.
+  !ERROR: The structured block in a WORKDISTRIBUTE construct may consist of only SCALAR or ARRAY assignments
+  !$omp workdistribute
+  do i = 1, n
+    print *, "omp workdistribute"
+  end do
+  !$omp end workdistribute
+
+end subroutine workdistribute
diff --git a/flang/test/Semantics/OpenMP/workdistribute02.f90 b/flang/test/Semantics/OpenMP/workdistribute02.f90
new file mode 100644
index 0000000000000..ad2cde2c3daf0
--- /dev/null
+++ b/flang/test/Semantics/OpenMP/workdistribute02.f90
@@ -0,0 +1,34 @@
+! RUN: %python %S/../test_errors.py %s %flang -fopenmp
+! OpenMP Version 6.0
+! workdistribute Construct
+! The !omp workdistribute construct must not contain any user defined
+! function calls unless the function is ELEMENTAL.
+
+module my_mod
+  contains
+  integer function my_func()
+    my_func = 10
+  end function my_func
+
+  impure integer function impure_my_func()
+    impure_my_func = 20
+  end function impure_my_func
+
+  impure elemental integer function impure_ele_my_func()
+    impure_ele_my_func = 20
+  end function impure_ele_my_func
+end module my_mod
+
+subroutine workdistribute(aa, bb, cc, n)
+  use my_mod
+  integer n
+  real aa(n), bb(n), cc(n)
+  !$omp teams
+  !$omp workdistribute
+  !ERROR: User defined non-ELEMENTAL function 'my_func' is not allowed in a WORKDISTRIBUTE construct
+  aa = my_func()
+  aa = bb * cc
+  !$omp end workdistribute
+  !$omp end teams
+
+end subroutine workdistribute
diff --git a/flang/test/Semantics/OpenMP/workdistribute03.f90 b/flang/test/Semantics/OpenMP/workdistribute03.f90
new file mode 100644
index 0000000000000..eac28cb39c47f
--- /dev/null
+++ b/flang/test/Semantics/OpenMP/workdistribute03.f90
@@ -0,0 +1,34 @@
+! RUN: %python %S/../test_errors.py %s %flang -fopenmp
+! OpenMP Version 6.0
+! workdistribute Construct
+! All array assignments, scalar assignments, and masked array assignments
+! must be intrinsic assignments.
+
+module defined_assign
+  interface assignment(=)
+    module procedure work_assign
+  end interface
+
+  contains
+    subroutine work_assign(a,b)
+      integer, intent(out) :: a
+      logical, intent(in) :: b(:)
+    end subroutine work_assign
+end module defined_assign
+
+program omp_workdistribute
+  use defined_assign
+
+  integer :: a, aa(10), bb(10)
+  logical :: l(10)
+  l = .TRUE.
+
+  !$omp teams
+  !$omp workdistribute
+  !ERROR: Defined assignment statement is not allowed in a WORKDISTRIBUTE construct
+  a = l
+  aa = bb
+  !$omp end workdistribute
+  !$omp end teams
+
+end program omp_workdistribute

mjklemm

Please slightly adjust the commented error message and put the workdistribute under a version check for OpenMP 6.0.

mjklemm · 2025-08-20T15:43:16Z

flang/lib/Semantics/check-omp-structure.cpp

+          "The structured block in a WORKDISTRIBUTE construct may consist of "
+          "only "
+          "SCALAR or ARRAY assignments"_err_en_US);


Suggested change

"The structured block in a WORKDISTRIBUTE construct may consist of "

"only "

"SCALAR or ARRAY assignments"_err_en_US);

"The structured block in a `WORKDISTRIBUTE` construct may consist of "

"only scalar or array assignments"_err_en_US);

Same thing about concatenating...

Fixed all the error strings concatenations. Also added check for version 6.0 for WORKDISTRIBUTE.

kparzysz · 2025-08-20T16:21:41Z

flang/lib/Semantics/check-omp-structure.cpp

+      if (isDefined == Tristate::Yes) {
+        context_.Say(expr.source,
+            "Defined assignment statement is not "
+            "allowed in a WORKDISTRIBUTE construct"_err_en_US);


Please concatenate the parts of error messages. This makes it easier to grep sources for a message.

kparzysz · 2025-08-20T16:22:08Z

flang/lib/Semantics/check-omp-structure.cpp

+          if (attrs != "") {
+            context_.Say(expr.source,
+                "User defined%s function '%s' is not allowed in a "
+                "WORKDISTRIBUTE construct"_err_en_US,


kparzysz · 2025-08-20T16:23:09Z

flang/lib/Semantics/check-omp-structure.cpp

+  }
+
+  bool Pre(const parser::Expr &expr) {
+    if (const auto *e{GetExpr(context_, expr)}) {


Consider an early return if e is nullptr.

fixed it in latest patch.

kparzysz · 2025-08-20T16:29:10Z

flang/lib/Semantics/check-omp-structure.cpp

+      for (const Symbol &symbol : evaluate::CollectSymbols(*e)) {
+        const Symbol &root{GetAssociationRoot(symbol)};
+        if (IsFunction(root)) {
+          std::string attrs{""};


The easiest thing would be to declare a vector of strings, add individual pieces, and then use join with ", " as a separator (see llvm/include/llvm/ADT/StringExtras.h).

Thanks for the feedback. Fixed it.

kparzysz · 2025-08-20T16:29:48Z

flang/lib/Semantics/check-omp-structure.cpp

+        GetContextParent().directive != llvm::omp::Directive::OMPD_teams) {
+      context_.Say(x.BeginDir().DirName().source,
+          "%s region can only be strictly nested within the "
+          "teams region"_err_en_US,


Please concat message.

kparzysz · 2025-08-20T16:30:36Z

flang/lib/Semantics/check-omp-structure.cpp

+          "The structured block in a WORKDISTRIBUTE construct may consist of "
+          "only "
+          "SCALAR or ARRAY assignments"_err_en_US);


Same thing about concatenating...

ergawy · 2025-08-21T05:17:30Z

flang/lib/Semantics/check-omp-structure.cpp

+          "%s region can only be strictly nested within the "
+          "teams region"_err_en_US,


nit: I think we capitalize construct keywords and do not use the in fron of them. Just to be consistent with the rest of the error messages.

Suggested change

"%s region can only be strictly nested within the "

"teams region"_err_en_US,

"%s region can only be strictly nested within TEAMS region"_err_en_US,

ergawy · 2025-08-21T05:21:08Z

flang/lib/Semantics/check-omp-structure.cpp

+  case llvm::omp::OMPD_workdistribute:
+    if (!CurrentDirectiveIsNested()) {
+      context_.Say(beginSpec.source,
+          "A workdistribute region must be nested inside teams region only."_err_en_US);


nit: capitalization of construct names here as well.

ergawy

Thanks Chaitanya. Just a few small comments.

kparzysz

LGTM with a comment

kparzysz · 2025-08-22T16:54:55Z

flang/lib/Semantics/check-omp-structure.cpp

+    context_.Say(source,
+        "WORKDISTRIBUTE construct is only supported from openMP 6.0"_err_en_US);


We usually generate messages in the form of "xyz is not allowed in %s, %s", with the string arguments being ThisVersion(version) and TryVersion(sinceVersion). ThisVersion(v) expands to "OpenMP v.v", TryVersion(v) expands to "try -fopenmp-version=v".
See https://github.com/llvm/llvm-project/blob/main/flang/lib/Semantics/check-omp-structure.cpp#L205-L207 for an example.

Didn't know this. Thanks for feedback. Updated in latest patch.

This was referenced Aug 19, 2025

[OpenMP] Add workdistribute construct in openMP dialect and in llvm frontend #154376

Merged

[flang][openmp] Add Lowering to omp mlir for workdistribute construct #154378

Merged

skc7 force-pushed the users/skc7/flang_wd_frontend_ps branch from e30b5cc to 4442fce Compare August 20, 2025 07:06

skc7 requested review from ergawy, ivanradanov, mjklemm and kparzysz August 20, 2025 09:25

skc7 marked this pull request as ready for review August 20, 2025 09:25

llvmbot added flang Flang issues not falling into any other category flang:openmp flang:semantics flang:parser labels Aug 20, 2025

skc7 mentioned this pull request Aug 20, 2025

[flang] Add support for workdistribute construct in flang frontend #146029

Closed

mjklemm requested changes Aug 20, 2025

View reviewed changes

kparzysz reviewed Aug 20, 2025

View reviewed changes

ergawy reviewed Aug 21, 2025

View reviewed changes

skc7 requested review from ergawy, kparzysz and mjklemm August 22, 2025 15:25

kparzysz approved these changes Aug 22, 2025

View reviewed changes

Base automatically changed from users/skc7/flang_wd_frontend_mlir to main August 25, 2025 06:50

skc7 added 4 commits August 25, 2025 12:25

[flang][openmp] Add parser/semantic support for workdistribute

d7da0b7

Fix error messages.

7cf35e2

Fix comment in test

f00d196

Fix version comment

d36bd6a

skc7 force-pushed the users/skc7/flang_wd_frontend_ps branch from cda9510 to d36bd6a Compare August 25, 2025 06:58

skc7 merged commit 374db67 into main Aug 25, 2025
9 checks passed

skc7 deleted the users/skc7/flang_wd_frontend_ps branch August 25, 2025 13:07

		"%s region can only be strictly nested within the "
		"teams region"_err_en_US,

		context_.Say(source,
		"WORKDISTRIBUTE construct is only supported from openMP 6.0"_err_en_US);

[flang][openmp] Add parser/semantic support for workdistribute #154377

[flang][openmp] Add parser/semantic support for workdistribute #154377

Uh oh!

Conversation

skc7 commented Aug 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Aug 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Aug 20, 2025

Uh oh!

mjklemm left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ergawy left a comment

Choose a reason for hiding this comment

Uh oh!

kparzysz left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

skc7 commented Aug 19, 2025 •

edited

Loading

llvmbot commented Aug 20, 2025 •

edited

Loading