[flang][OpenMP] Use OmpDirectiveSpecification in DISPATCH (#148008)

kparzysz · web-flow · commit 61a9d2c22dfe · 2025-07-11T07:28:54.000-05:00
Dispatch is the last construct (after ATOMIC and ALLOCATORS) where the
associated block requires a specific form.
Using OmpDirectiveSpecification for the begin and the optional end
directives will make the structure of all block directives more uniform.
diff --git a/flang/include/flang/Parser/dump-parse-tree.h b/flang/include/flang/Parser/dump-parse-tree.h
@@ -710,8 +710,6 @@ class ParseTreeDumper {
   NODE(parser, OpenMPDepobjConstruct)
   NODE(parser, OpenMPUtilityConstruct)
   NODE(parser, OpenMPDispatchConstruct)
-  NODE(parser, OmpDispatchDirective)
-  NODE(parser, OmpEndDispatchDirective)
   NODE(parser, OpenMPFlushConstruct)
   NODE(parser, OpenMPLoopConstruct)
   NODE(parser, OpenMPExecutableAllocate)
diff --git a/flang/include/flang/Parser/parse-tree.h b/flang/include/flang/Parser/parse-tree.h
@@ -4939,19 +4939,11 @@ struct OpenMPDepobjConstruct {
 //                    nocontext-clause |
 //                    novariants-clause |
 //                    nowait-clause
-struct OmpDispatchDirective {
-  TUPLE_CLASS_BOILERPLATE(OmpDispatchDirective);
-  CharBlock source;
-  std::tuple<Verbatim, OmpClauseList> t;
-};
-
-EMPTY_CLASS(OmpEndDispatchDirective);
-
 struct OpenMPDispatchConstruct {
   TUPLE_CLASS_BOILERPLATE(OpenMPDispatchConstruct);
   CharBlock source;
-  std::tuple<OmpDispatchDirective, Block,
-      std::optional<OmpEndDispatchDirective>>
+  std::tuple<OmpDirectiveSpecification, Block,
+      std::optional<OmpDirectiveSpecification>>
       t;
 };
 
diff --git a/flang/lib/Parser/openmp-parsers.cpp b/flang/lib/Parser/openmp-parsers.cpp
@@ -1302,6 +1302,32 @@ TYPE_PARSER(sourced( //
     construct<OpenMPAllocatorsConstruct>(
         "ALLOCATORS"_tok >= OmpAllocatorsConstructParser{})))
 
+struct OmpDispatchConstructParser {
+  using resultType = OpenMPDispatchConstruct;
+
+  std::optional<resultType> Parse(ParseState &state) const {
+    auto dirSpec{Parser<OmpDirectiveSpecification>{}.Parse(state)};
+    if (!dirSpec || dirSpec->DirId() != llvm::omp::Directive::OMPD_dispatch) {
+      return std::nullopt;
+    }
+
+    // This should be a function call. That will be checked in semantics.
+    Block block;
+    if (auto stmt{attempt(Parser<ExecutionPartConstruct>{}).Parse(state)}) {
+      block.emplace_back(std::move(*stmt));
+    }
+    // Allow empty block. Check for this in semantics.
+
+    auto end{OmpEndDirectiveParser{llvm::omp::Directive::OMPD_dispatch}};
+    return OpenMPDispatchConstruct{
+        std::move(*dirSpec), std::move(block), *maybe(end).Parse(state)};
+  }
+};
+
+TYPE_PARSER(sourced( //
+    construct<OpenMPDispatchConstruct>(
+        "DISPATCH"_tok >= OmpDispatchConstructParser{})))
+
 // Parser for an arbitrary OpenMP ATOMIC construct.
 //
 // Depending on circumstances, an ATOMIC construct applies to one or more
@@ -1631,16 +1657,6 @@ TYPE_PARSER(sourced(construct<OmpCriticalDirective>(verbatim("CRITICAL"_tok),
 TYPE_PARSER(construct<OpenMPCriticalConstruct>(
     Parser<OmpCriticalDirective>{}, block, Parser<OmpEndCriticalDirective>{}))
 
-TYPE_PARSER(sourced(construct<OmpDispatchDirective>(
-    verbatim("DISPATCH"_tok), Parser<OmpClauseList>{})))
-
-TYPE_PARSER(
-    construct<OmpEndDispatchDirective>(startOmpLine >> "END DISPATCH"_tok))
-
-TYPE_PARSER(sourced(construct<OpenMPDispatchConstruct>(
-    Parser<OmpDispatchDirective>{} / endOmpLine, block,
-    maybe(Parser<OmpEndDispatchDirective>{} / endOmpLine))))
-
 // 2.11.3 Executable Allocate directive
 TYPE_PARSER(
     sourced(construct<OpenMPExecutableAllocate>(verbatim("ALLOCATE"_tok),
diff --git a/flang/lib/Parser/unparse.cpp b/flang/lib/Parser/unparse.cpp
@@ -2758,6 +2758,9 @@ class UnparseVisitor {
     Put("\n");
     EndOpenMP();
   }
+  void Unparse(const OpenMPDispatchConstruct &x) { //
+    UnparseBlockConstruct(x);
+  }
   void Unparse(const OpenMPRequiresConstruct &y) {
     BeginOpenMP();
     Word("!$OMP REQUIRES ");
@@ -2777,15 +2780,6 @@ class UnparseVisitor {
     Walk(x.v);
     return false;
   }
-  void Unparse(const OmpDispatchDirective &x) {
-    Word("!$OMP DISPATCH");
-    Walk(x.t);
-    Put("\n");
-  }
-  void Unparse(const OmpEndDispatchDirective &) {
-    Word("!$OMP END DISPATCH");
-    Put("\n");
-  }
   void Unparse(const OmpErrorDirective &x) {
     Word("!$OMP ERROR ");
     Walk(x.t);
diff --git a/flang/lib/Semantics/check-omp-structure.cpp b/flang/lib/Semantics/check-omp-structure.cpp
@@ -509,8 +509,8 @@ template <typename Checker> struct DirectiveSpellingVisitor {
     checker_(std::get<parser::Verbatim>(x.t).source, Directive::OMPD_allocate);
     return false;
   }
-  bool Pre(const parser::OmpDispatchDirective &x) {
-    checker_(std::get<parser::Verbatim>(x.t).source, Directive::OMPD_dispatch);
+  bool Pre(const parser::OpenMPDispatchConstruct &x) {
+    checker_(GetDirName(x.t).source, Directive::OMPD_dispatch);
     return false;
   }
   bool Pre(const parser::OmpErrorDirective &x) {
@@ -1595,28 +1595,31 @@ void OmpStructureChecker::Enter(const parser::OmpErrorDirective &x) {
 }
 
 void OmpStructureChecker::Enter(const parser::OpenMPDispatchConstruct &x) {
-  PushContextAndClauseSets(x.source, llvm::omp::Directive::OMPD_dispatch);
+  auto &dirSpec{std::get<parser::OmpDirectiveSpecification>(x.t)};
   const auto &block{std::get<parser::Block>(x.t)};
-  if (block.empty() || block.size() > 1) {
+  PushContextAndClauseSets(
+      dirSpec.DirName().source, llvm::omp::Directive::OMPD_dispatch);
+
+  if (block.empty()) {
     context_.Say(x.source,
-        "The DISPATCH construct is empty or contains more than one statement"_err_en_US);
+        "The DISPATCH construct should contain a single function or subroutine call"_err_en_US);
     return;
   }
 
-  auto it{block.begin()};
   bool passChecks{false};
-  if (const parser::AssignmentStmt *
-      assignStmt{parser::Unwrap<parser::AssignmentStmt>(*it)}) {
+  omp::SourcedActionStmt action{omp::GetActionStmt(block)};
+  if (const auto *assignStmt{
+          parser::Unwrap<parser::AssignmentStmt>(*action.stmt)}) {
     if (parser::Unwrap<parser::FunctionReference>(assignStmt->t)) {
       passChecks = true;
     }
-  } else if (parser::Unwrap<parser::CallStmt>(*it)) {
+  } else if (parser::Unwrap<parser::CallStmt>(*action.stmt)) {
     passChecks = true;
   }
 
   if (!passChecks) {
-    context_.Say(x.source,
-        "The DISPATCH construct does not contain a SUBROUTINE or FUNCTION"_err_en_US);
+    context_.Say(action.source,
+        "The body of the DISPATCH construct should be a function or a subroutine call"_err_en_US);
   }
 }
 
diff --git a/flang/test/Parser/OpenMP/dispatch.f90 b/flang/test/Parser/OpenMP/dispatch.f90
@@ -1,4 +1,4 @@
-! RUN: %flang_fc1 -fopenmp -fdebug-dump-parse-tree %s | FileCheck %s
+! RUN: %flang_fc1 -fopenmp -fdebug-dump-parse-tree %s | FileCheck %s --check-prefix=PARSE-TREE
 ! RUN: %flang_fc1 -fopenmp -fdebug-unparse %s | FileCheck %s --check-prefix="UNPARSE"
 
 integer function func(a, b, c)
@@ -12,40 +12,57 @@ subroutine sub(x)
   integer :: r
   type(c_ptr) :: x
   integer :: a = 14, b = 7, c = 21
+
 !UNPARSE: !$OMP DISPATCH DEVICE(3_4) NOWAIT NOCONTEXT(.false._4) NOVARIANTS(.true._4)
-!CHECK: | | ExecutionPartConstruct -> ExecutableConstruct -> OpenMPConstruct -> OpenMPDispatchConstruct
-!CHECK-NEXT: | | | OmpDispatchDirective
-!CHECK: | | | | OmpClauseList -> OmpClause -> Device -> OmpDeviceClause
-!CHECK-NEXT: | | | | | Scalar -> Integer -> Expr = '3_4'
-!CHECK-NEXT: | | | | | | LiteralConstant -> IntLiteralConstant = '3'
-!CHECK-NEXT: | | | | OmpClause -> Nowait
-!CHECK-NEXT: | | | | OmpClause -> Nocontext -> Scalar -> Logical -> Expr = '.false._4'
-!CHECK-NEXT: | | | | | LiteralConstant -> LogicalLiteralConstant
-!CHECK-NEXT: | | | | | | bool = 'false'
-!CHECK-NEXT: | | | | OmpClause -> Novariants -> Scalar -> Logical -> Expr = '.true._4'
-!CHECK-NEXT: | | | | | EQ
-!CHECK-NEXT: | | | | | | Expr = '1_4'
-!CHECK-NEXT: | | | | | | | LiteralConstant -> IntLiteralConstant = '1'
-!CHECK-NEXT: | | | | | | Expr = '1_4'
-!CHECK-NEXT: | | | | | | | LiteralConstant -> IntLiteralConstant = '1'
-!CHECK-NEXT: | | | Block
- 
+!UNPARSE:   r=func(a,b,c)
+!UNPARSE: !$OMP END DISPATCH
+
+!PARSE-TREE: ExecutionPartConstruct -> ExecutableConstruct -> OpenMPConstruct -> OpenMPDispatchConstruct
+!PARSE-TREE: | OmpDirectiveSpecification
+!PARSE-TREE: | | OmpDirectiveName -> llvm::omp::Directive = dispatch
+!PARSE-TREE: | | OmpClauseList -> OmpClause -> Device -> OmpDeviceClause
+!PARSE-TREE: | | | Scalar -> Integer -> Expr = '3_4'
+!PARSE-TREE: | | | | LiteralConstant -> IntLiteralConstant = '3'
+!PARSE-TREE: | | OmpClause -> Nowait
+!PARSE-TREE: | | OmpClause -> Nocontext -> Scalar -> Logical -> Expr = '.false._4'
+!PARSE-TREE: | | | LiteralConstant -> LogicalLiteralConstant
+!PARSE-TREE: | | | | bool = 'false'
+!PARSE-TREE: | | OmpClause -> Novariants -> Scalar -> Logical -> Expr = '.true._4'
+!PARSE-TREE: | | | EQ
+!PARSE-TREE: | | | | Expr = '1_4'
+!PARSE-TREE: | | | | | LiteralConstant -> IntLiteralConstant = '1'
+!PARSE-TREE: | | | | Expr = '1_4'
+!PARSE-TREE: | | | | | LiteralConstant -> IntLiteralConstant = '1'
+!PARSE-TREE: | | Flags = None
+!PARSE-TREE: | Block
+!PARSE-TREE: | | ExecutionPartConstruct -> ExecutableConstruct -> ActionStmt -> AssignmentStmt
+![...]
+!PARSE-TREE: | OmpDirectiveSpecification
+!PARSE-TREE: | | OmpDirectiveName -> llvm::omp::Directive = dispatch
+!PARSE-TREE: | | OmpClauseList ->
+!PARSE-TREE: | | Flags = None
+
   !$omp dispatch device(3) nowait nocontext(.false.) novariants(1.eq.1)
   r = func(a, b, c)
-!UNPARSE: !$OMP END DISPATCH
-!CHECK: | | | OmpEndDispatchDirective
   !$omp end dispatch
 
 !! Test the "no end dispatch" option.
-!UNPARSE: !$OMP DISPATCH  DEVICE(3_4) IS_DEVICE_PTR(x)
-!CHECK: | | ExecutionPartConstruct -> ExecutableConstruct -> OpenMPConstruct -> OpenMPDispatchConstruct
-!CHECK-NEXT: | | | OmpDispatchDirective
-!CHECK: | | | | OmpClause -> IsDevicePtr ->  OmpObjectList -> OmpObject -> Designator -> DataRef -> Name = 'x'  
+!UNPARSE: !$OMP DISPATCH DEVICE(3_4) IS_DEVICE_PTR(x)
+!UNPARSE:   r=func(a+1_4,b+2_4,c+3_4)
+
+!PARSE-TREE: ExecutionPartConstruct -> ExecutableConstruct -> OpenMPConstruct -> OpenMPDispatchConstruct
+!PARSE-TREE: | OmpDirectiveSpecification
+!PARSE-TREE: | | OmpDirectiveName -> llvm::omp::Directive = dispatch
+!PARSE-TREE: | | OmpClauseList -> OmpClause -> Device -> OmpDeviceClause
+!PARSE-TREE: | | | Scalar -> Integer -> Expr = '3_4'
+!PARSE-TREE: | | | | LiteralConstant -> IntLiteralConstant = '3'
+!PARSE-TREE: | | OmpClause -> IsDevicePtr -> OmpObjectList -> OmpObject -> Designator -> DataRef -> Name = 'x'
+!PARSE-TREE: | | Flags = None
+!PARSE-TREE: | Block
+!PARSE-TREE: | | ExecutionPartConstruct -> ExecutableConstruct -> ActionStmt -> AssignmentStmt
+!PARSE-TREE-NOT: OmpDirectiveSpecification
+
   !$omp dispatch device(3) is_device_ptr(x)
   r = func(a+1, b+2, c+3)
-!CHECK-NOT: | | | OmpEndDispatchDirective
 
 end subroutine sub
-
-
-
diff --git a/flang/test/Semantics/OpenMP/dispatch.f90 b/flang/test/Semantics/OpenMP/dispatch.f90
@@ -1,24 +1,20 @@
-! RUN: %python %S/../test_errors.py %s %flang -fopenmp
+! RUN: %python %S/../test_errors.py %s %flang -fopenmp -fopenmp-version=52
 
 subroutine sb1
   integer :: r
   r = 1
-  !ERROR: The DISPATCH construct does not contain a SUBROUTINE or FUNCTION
   !$omp dispatch nowait
+!ERROR: The body of the DISPATCH construct should be a function or a subroutine call
   print *,r
 end subroutine
+
 subroutine sb2
-  integer :: r
-!ERROR: The DISPATCH construct is empty or contains more than one statement
+!ERROR: The DISPATCH construct should contain a single function or subroutine call
   !$omp dispatch
-  call foo()
-  r = bar()
   !$omp end dispatch
-contains
-  subroutine foo
-  end subroutine foo
-  function bar
-    integer :: bar
-    bar = 2
-  end function
+end subroutine
+
+subroutine sb3
+!ERROR: The DISPATCH construct should contain a single function or subroutine call
+  !$omp dispatch
 end subroutine
