[flang][OpenMP] Added semantic checks for atomic capture construct

[NimishMishra]

Summary

This patch adds semantic checks for atomic capture construct.

Depends on D127620

[NimishMishra]

There are still some semantic checks left. I will cover them in a separate patch

[shraiysh]

Two minor comments at the moment. Will review it later this week.

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

1498

Can we please start adding some documentation to the function right above it, like here

1806–1821

nit: wrap this loop in a lambda inside this function and use twice instead of duplication.

To avoid higher cyclomatic complexity, use early return in the lambda function when the if conditions become nullptr. What happens if none of the clauses in the list match the if conditions? Is it possible to have one clause match some conditions and not the rest? If not, then we should assert instead of if and if yes, should their else parts be ignored? (I think the innermost if must always be true if control reaches there, but please feel free to correct me).

[NimishMishra]

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

1498

This would certainly be nice. If others agree too, we can go ahead and document the functions.

[kiranchandramohan]

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

1498

+1 for more documentation, particularly for special support functions.

[peixin]

Can you split this patch? The small patch with single function is easier to review. If there are multiple checks with similar style, you can put them together.

1. check "capture statement is of the form v = x if atomic construct is read or capture"

check "write statement is of the form x = expr if atomic construct is write"

2. check "x must not have the ALLOCATABLE attribute."

check "operator must refer to the intrinsic operator and not to a user-defined operator"
check "hint-expression is a constant expression that evaluates to a scalar value with kind omp_sync_hint_kind and a value that is a valid synchronization hint."

Or even split point 2 into three small patches.

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

1498

+1

But I prefer not to add the comments like this The comments for op, converter, loc, eval are useless. My preference is to use function name and argument name comment themselves. If they cannot comment itself such as with some special support, it's better to add comments for them.

[NimishMishra]

In D126612#3550182, @peixin wrote:

Can you split this patch? The small patch with single function is easier to review. If there are multiple checks with similar style, you can put them together.

Sure. Working on splitting this patch's contents across multiple smaller patches

[NimishMishra]

flang/include/flang/Parser/parse-tree.h

3715

I have an idea here that could potentially ease lowering of atomic capture construct. To lower such constructs, we need to know what the two statements are before we can create operations for them.

My idea was if we could introduce attributes as these in the parse-tree node itself which shall be filled in during semantics. Then during lowering, we just look at Stmt1Type and Stmt2Type and create required operations for the two statements of capture construct.

The only caveat is that the node is const at the point where OmpStructureChecker is instantiated. Any ideas on how to go about this problem will greatly help.

[NimishMishra]

The design of this solution could be improved. Following whatever reviews come in https://reviews.llvm.org/D127272, we can use the same basic structural checks from that patch here to understand what the statements 1 and 2 of the atomic capture construct actually are (among read/write/update statements). That would remove dependence on global boolean flags, which don't look good design wise.

flang/include/flang/Parser/parse-tree.h

 struct OmpAtomicCapture {
   TUPLE_CLASS_BOILERPLATE(OmpAtomicCapture);
   CharBlock source;
   WRAPPER_CLASS(Stmt1, Statement<AssignmentStmt>);
   WRAPPER_CLASS(Stmt2, Statement<AssignmentStmt>);
   std::tuple<OmpAtomicClauseList, Verbatim, OmpAtomicClauseList, Stmt1, Stmt2,
       OmpEndAtomic>
       t;
+  enum class StmtType {
+    CaptureStmt,
+    WriteStmt,
+    UpdateStmt
+  } Stmt1Type,
+      Stmt2Type;
 };
 
 // ATOMIC
 struct OmpAtomic {
   TUPLE_CLASS_BOILERPLATE(OmpAtomic);
   CharBlock source;
   std::tuple<Verbatim, OmpAtomicClauseList, Statement<AssignmentStmt>,
       std::optional<OmpEndAtomic>>

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

private:
   void CheckDoWhile(const parser::OpenMPLoopConstruct &x);
   void CheckCycleConstraints(const parser::OpenMPLoopConstruct &x);
   template <typename T, typename D> bool IsOperatorValid(const T &, const D &);
   void CheckAtomicMemoryOrderClause(
       const parser::OmpAtomicClauseList *, const parser::OmpAtomicClauseList *);
   void CheckAtomicUpdateStmt(const parser::AssignmentStmt &);
   void CheckAtomicCaptureStmt(const parser::AssignmentStmt &);
   void CheckAtomicWriteStmt(const parser::AssignmentStmt &);
+  void CheckAtomicCaptureConstruct(const parser::OmpAtomicCapture &);
   void CheckAtomicConstructStructure(const parser::OpenMPAtomicConstruct &);
   void CheckDistLinear(const parser::OpenMPLoopConstruct &x);
   void CheckSIMDNest(const parser::OpenMPConstruct &x);
   void CheckTargetNest(const parser::OpenMPConstruct &x);
   void CheckCancellationNest(
       const parser::CharBlock &source, const parser::OmpCancelType::Type &type);
   std::int64_t GetOrdCollapseLevel(const parser::OpenMPLoopConstruct &x);
   bool CheckReductionOperators(const parser::OmpClause::Reduction &);
inline void ErrIfNonScalarAssignmentStmt(
       const parser::Variable &, const parser::Expr &);
   enum directiveNestType {
     SIMDNest,
     TargetBlockOnlyTeams,
     TargetNest,
     LastType
   };
   int directiveNest_[LastType + 1] = {0};
+
+  // Flags used exclusively for atomic capture semantic checks
+  bool shouldSuppressOutput = false; // check semantics without reporting output
+  bool isValidConstruct; // if the construct is valid or not
 };
 } // namespace Fortran::semantics
 #endif // FORTRAN_SEMANTICS_CHECK_OMP_STRUCTURE_H_

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

 // EndBlockDirective.
 void OmpStructureChecker::Leave(const parser::OmpEndBlockDirective &x) {
   if ((GetContext().directive == llvm::omp::Directive::OMPD_end_single) ||
       (GetContext().directive == llvm::omp::Directive::OMPD_end_workshare)) {
     dirContext_.pop_back();
   }
 }
 
 inline void OmpStructureChecker::ErrIfAllocatableVariable(
     const parser::Variable &var) {
   // Err out if the given symbol has
   // ALLOCATABLE attribute
   if (const auto *e{GetExpr(context_, var)})
     for (const Symbol &symbol : evaluate::CollectSymbols(*e))
       if (IsAllocatable(symbol)) {
         const auto &designator =
             std::get<common::Indirection<parser::Designator>>(var.u);
         const auto *dataRef =
             std::get_if<Fortran::parser::DataRef>(&designator.value().u);
         const Fortran::parser::Name *name =
             dataRef ? std::get_if<Fortran::parser::Name>(&dataRef->u) : nullptr;
-        if (name)
+        if (name) {
+          isValidConstruct = false;
+          if (!shouldSuppressOutput)
             context_.Say(name->source,
                 "%s must not have ALLOCATABLE "
                 "attribute"_err_en_US,
                 name->ToString());
         }
       }
+}
 
 inline void OmpStructureChecker::ErrIfLHSAndRHSSymbolsMatch(
     const parser::Variable &var, const parser::Expr &expr) {
   // Err out if the symbol on the LHS is also used on the RHS of the assignment
   // statement
   const auto *e{GetExpr(context_, expr)};
   const auto *v{GetExpr(context_, var)};
   if (e && v) {
     const Symbol &varSymbol = evaluate::GetSymbolVector(*v).front();
     for (const Symbol &symbol : evaluate::GetSymbolVector(*e)) {
       if (varSymbol == symbol) {
+        isValidConstruct = false;
+        if (!shouldSuppressOutput)
           context_.Say(expr.source,
               "RHS expression "
               "on atomic assignment statement"
               " cannot access '%s'"_err_en_US,
               var.GetSource().ToString());
       }
     }
   }
 }
 
 inline void OmpStructureChecker::ErrIfNonScalarAssignmentStmt(
     const parser::Variable &var, const parser::Expr &expr) {
   // Err out if either the variable on the LHS or the expression on the RHS of
   // the assignment statement are non-scalar (i.e. have rank > 0)
   const auto *e{GetExpr(context_, expr)};
   const auto *v{GetExpr(context_, var)};
   if (e && v) {
-    if (e->Rank() != 0)
+    if (e->Rank() != 0) {
+      isValidConstruct = false;
+      if (!shouldSuppressOutput)
         context_.Say(expr.source,
             "Expected scalar expression "
             "on the RHS of atomic assignment "
             "statement"_err_en_US);
-    if (v->Rank() != 0)
+    }
+    if (v->Rank() != 0) {
+      isValidConstruct = false;
+      if (!shouldSuppressOutput)
         context_.Say(var.GetSource(),
             "Expected scalar variable "
             "on the LHS of atomic assignment "
             "statement"_err_en_US);
     }
   }
+}
 
 template <typename T, typename D>
 bool OmpStructureChecker::IsOperatorValid(const T &node, const D &variable) {
   using AllowedBinaryOperators =
       std::variant<parser::Expr::Add, parser::Expr::Multiply,
           parser::Expr::Subtract, parser::Expr::Divide, parser::Expr::AND,
           parser::Expr::OR, parser::Expr::EQV, parser::Expr::NEQV>;
   using BinaryOperators = std::variant<parser::Expr::Add,
       parser::Expr::Multiply, parser::Expr::Subtract, parser::Expr::Divide,
       parser::Expr::AND, parser::Expr::OR, parser::Expr::EQV,
       parser::Expr::NEQV, parser::Expr::Power, parser::Expr::Concat,
       parser::Expr::LT, parser::Expr::LE, parser::Expr::EQ, parser::Expr::NE,
       parser::Expr::GE, parser::Expr::GT>;
 
   if constexpr (common::HasMember<T, BinaryOperators>) {
     const auto &variableName{variable.GetSource().ToString()};
     const auto &exprLeft{std::get<0>(node.t)};
     const auto &exprRight{std::get<1>(node.t)};
     if ((exprLeft.value().source.ToString() != variableName) &&
         (exprRight.value().source.ToString() != variableName)) {
+      isValidConstruct = false;
+      if (!shouldSuppressOutput)
         context_.Say(variable.GetSource(),
             "Atomic update statement should be of form "
             "`%s = %s operator expr` OR `%s = expr operator %s`"_err_en_US,
             variableName, variableName, variableName, variableName);
     }
     return common::HasMember<T, AllowedBinaryOperators>;
   }
   return true;
 }
 
 void OmpStructureChecker::CheckAtomicCaptureStmt(
     const parser::AssignmentStmt &assignmentStmt) {
   const auto &var{std::get<parser::Variable>(assignmentStmt.t)};
   const auto &expr{std::get<parser::Expr>(assignmentStmt.t)};
+  isValidConstruct = true;
   common::visit(
       common::visitors{
           [&](const common::Indirection<parser::Designator> &designator) {
             const auto *dataRef =
                 std::get_if<Fortran::parser::DataRef>(&designator.value().u);
             const auto *name = dataRef
                 ? std::get_if<Fortran::parser::Name>(&dataRef->u)
                 : nullptr;
-            if (name && IsAllocatable(*name->symbol))
+            if (name && IsAllocatable(*name->symbol)) {
+              isValidConstruct = false;
+              if (!shouldSuppressOutput)
                 context_.Say(name->source,
                     "%s must not have ALLOCATABLE "
                     "attribute"_err_en_US,
                     name->ToString());
+            }
           },
           [&](const auto &) {
             // Anything other than a `parser::Designator` is not allowed
+            isValidConstruct = false;
+            if (!shouldSuppressOutput)
               context_.Say(expr.source,
                   "Expected scalar variable "
                   "of intrinsic type on RHS of atomic "
                   "assignment statement"_err_en_US);
           }},
       expr.u);
   ErrIfLHSAndRHSSymbolsMatch(var, expr);
   ErrIfNonScalarAssignmentStmt(var, expr);
 }
 
 void OmpStructureChecker::CheckAtomicWriteStmt(
     const parser::AssignmentStmt &assignmentStmt) {
   const auto &var{std::get<parser::Variable>(assignmentStmt.t)};
   const auto &expr{std::get<parser::Expr>(assignmentStmt.t)};
+  isValidConstruct = true;
   ErrIfAllocatableVariable(var);
   ErrIfLHSAndRHSSymbolsMatch(var, expr);
   ErrIfNonScalarAssignmentStmt(var, expr);
 }
 
 void OmpStructureChecker::CheckAtomicUpdateStmt(
     const parser::AssignmentStmt &assignment) {
   const auto &expr{std::get<parser::Expr>(assignment.t)};
   const auto &var{std::get<parser::Variable>(assignment.t)};
+  isValidConstruct = true;
   common::visit(
       common::visitors{
           [&](const common::Indirection<parser::FunctionReference> &x) {
             const auto &procedureDesignator{
                 std::get<parser::ProcedureDesignator>(x.value().v.t)};
             const parser::Name *name{
                 std::get_if<parser::Name>(&procedureDesignator.u)};
             if (name &&
                 !(name->source == "max" || name->source == "min" ||
                     name->source == "iand" || name->source == "ior" ||
                     name->source == "ieor")) {
+              isValidConstruct = false;
+              if (!shouldSuppressOutput)
                 context_.Say(expr.source,
                     "Invalid intrinsic procedure name in "
                     "OpenMP ATOMIC (UPDATE) statement"_err_en_US);
             } else if (name) {
               bool foundMatch{false};
               if (auto varDesignatorIndirection =
                       std::get_if<Fortran::common::Indirection<
                           Fortran::parser::Designator>>(&var.u)) {
                 const auto &varDesignator = varDesignatorIndirection->value();
                 if (const auto *dataRef = std::get_if<Fortran::parser::DataRef>(
                         &varDesignator.u)) {
                   if (const auto *name =
                           std::get_if<Fortran::parser::Name>(&dataRef->u)) {
                     const auto &varSymbol = *name->symbol;
                     if (const auto *e{GetExpr(context_, expr)}) {
                       for (const Symbol &symbol :
                           evaluate::CollectSymbols(*e)) {
                         if (symbol == varSymbol) {
                           foundMatch = true;
                           break;
                         }
                       }
                     }
                   }
                 }
               }
               if (!foundMatch) {
+                isValidConstruct = false;
+                if (!shouldSuppressOutput)
                   context_.Say(expr.source,
                       "Atomic update variable '%s' not found in the "
                       "argument list of intrinsic procedure"_err_en_US,
                       var.GetSource().ToString());
               }
             }
           },
           [&](const auto &x) {
             if (!IsOperatorValid(x, var)) {
+              isValidConstruct = false;
+              if (!shouldSuppressOutput)
                 context_.Say(expr.source,
                     "Invalid operator in OpenMP ATOMIC (UPDATE) statement"_err_en_US);
             }
           },
       },
       expr.u);
   ErrIfAllocatableVariable(var);
 }
 
+void OmpStructureChecker::CheckAtomicCaptureConstruct(
+    const parser::OmpAtomicCapture &atomicCaptureConstruct) {
+  const parser::AssignmentStmt &stmt1 =
+      std::get<3>(atomicCaptureConstruct.t).v.statement;
+  const parser::AssignmentStmt &stmt2 =
+      std::get<4>(atomicCaptureConstruct.t).v.statement;
+  const auto &stmt1Expr{std::get<parser::Expr>(stmt1.t)};
+  const auto &stmt1Var{std::get<parser::Variable>(stmt1.t)};
+  const auto &stmt2Expr{std::get<parser::Expr>(stmt2.t)};
+  const auto &stmt2Var{std::get<parser::Variable>(stmt2.t)};
+
+  enum class StmtType {
+    Undefined,
+    CaptureStmt,
+    WriteStmt,
+    UpdateStmt
+  } Stmt1Type,
+      Stmt2Type;
+
+  shouldSuppressOutput = true;
+
+  // Statement 1 checks
+  CheckAtomicCaptureStmt(stmt1);
+  Stmt1Type = isValidConstruct ? StmtType::CaptureStmt : StmtType::Undefined;
+  CheckAtomicUpdateStmt(stmt1);
+  Stmt1Type = (isValidConstruct && Stmt1Type == StmtType::Undefined)
+      ? StmtType::UpdateStmt
+      : Stmt1Type;
+  // Statement 2 checks if statement 1 is capture-stmt
+  if (Stmt1Type == StmtType::CaptureStmt) {
+    // Check for [capture-stmt, update-stmt]
+    CheckAtomicUpdateStmt(stmt2);
+    Stmt2Type = isValidConstruct ? StmtType::UpdateStmt : StmtType::Undefined;
+
+    // Check for [capture-stmt, write-stmt]
+    CheckAtomicWriteStmt(stmt2);
+    Stmt2Type = (isValidConstruct && Stmt2Type == StmtType::Undefined)
+        ? StmtType::WriteStmt
+        : Stmt2Type;
+  }
+  // Statement 2 checks if statement 1 is update-stmt
+  else if (Stmt1Type == StmtType::UpdateStmt) {
+    // Check for [update-stmt, capture-stmt]
+    CheckAtomicCaptureStmt(stmt2);
+    Stmt2Type = isValidConstruct ? StmtType::CaptureStmt : StmtType::Undefined;
+  }
+  shouldSuppressOutput = false;
+
+  if (Stmt1Type == StmtType::Undefined || Stmt2Type == StmtType::Undefined)
+    context_.Say(stmt1Expr.source,
+        "Invalid atomic capture construct statements. "
+        "Expected one of [update-stmt, capture-stmt], "
+        "[capture-stmt, update-stmt], or [capture-stmt, write-stmt]"_err_en_US);
+
+  else if (Stmt1Type == StmtType::CaptureStmt) {
+    // Variable captured in stmt1 should be assigned in stmt2
+    const auto *e{GetExpr(context_, stmt1Expr)};
+    const auto *v{GetExpr(context_, stmt2Var)};
+    if (e && v) {
+      const Symbol &stmt2VarSymbol = evaluate::GetSymbolVector(*v).front();
+      const Symbol &stmt1ExprSymbol = evaluate::GetSymbolVector(*e).front();
+      if (stmt2VarSymbol != stmt1ExprSymbol)
+        context_.Say(stmt1Expr.source,
+            "Captured variable %s "
+            "expected to be assigned in statement 2 of "
+            "atomic capture construct"_err_en_US,
+            stmt1ExprSymbol.name().ToString());
+    }
+  } else if (Stmt1Type == StmtType::UpdateStmt) {
+    // Variable updated in stmt1 should be captured in stmt2
+    const auto *e{GetExpr(context_, stmt2Expr)};
+    const auto *v{GetExpr(context_, stmt1Var)};
+    if (e && v) {
+      const Symbol &stmt1VarSymbol = evaluate::GetSymbolVector(*v).front();
+      const Symbol &stmt2ExprSymbol = evaluate::GetSymbolVector(*e).front();
+      if (stmt1VarSymbol != stmt2ExprSymbol)
+        context_.Say(stmt1Var.GetSource(),
+            "Updated variable %s "
+            "expected to be captured in statement 2 of "
+            "atomic capture construct"_err_en_US,
+            stmt1Var.GetSource().ToString());
+    }
+  }
+}
+
 void OmpStructureChecker::CheckAtomicMemoryOrderClause(
     const parser::OmpAtomicClauseList *leftHandClauseList,
     const parser::OmpAtomicClauseList *rightHandClauseList) {
   int numMemoryOrderClause = 0;
   auto checkForValidMemoryOrderClause =
       [&](const parser::OmpAtomicClauseList *clauseList) {
         for (const auto &clause : clauseList->v) {
           if (std::get_if<Fortran::parser::OmpMemoryOrderClause>(&clause.u)) {
             numMemoryOrderClause++;
             if (numMemoryOrderClause > 1) {
               context_.Say(clause.source,
                   "More than one memory order clause not allowed on "
                   "OpenMP Atomic construct"_err_en_US);
               return;
             }
           }
         }
       };
   if (leftHandClauseList) {
     checkForValidMemoryOrderClause(leftHandClauseList);
   }
   if (rightHandClauseList) {
     checkForValidMemoryOrderClause(rightHandClauseList);
   }
 }
 
 void OmpStructureChecker::Enter(const parser::OpenMPAtomicConstruct &x) {
   common::visit(
       common::visitors{
           [&](const parser::OmpAtomic &atomicConstruct) {
             const auto &dir{std::get<parser::Verbatim>(atomicConstruct.t)};
common::visit(
                 &std::get<0>(atomicWrite.t), &std::get<2>(atomicWrite.t));
             CheckHintClause<const parser::OmpAtomicClauseList>(
                 &std::get<0>(atomicWrite.t), &std::get<2>(atomicWrite.t));
             CheckAtomicWriteStmt(
                 std::get<parser::Statement<parser::AssignmentStmt>>(
                     atomicWrite.t)
                     .statement);
           },
-          [&](const auto &atomicConstruct) {
-            const auto &dir{std::get<parser::Verbatim>(atomicConstruct.t)};
+          [&](const parser::OmpAtomicCapture &atomicCapture) {
+            const auto &dir{std::get<parser::Verbatim>(atomicCapture.t)};
             PushContextAndClauseSets(
                 dir.source, llvm::omp::Directive::OMPD_atomic);
-            CheckAtomicMemoryOrderClause(&std::get<0>(atomicConstruct.t),
-                &std::get<2>(atomicConstruct.t));
+            CheckAtomicMemoryOrderClause(
+                &std::get<0>(atomicCapture.t), &std::get<2>(atomicCapture.t));
             CheckHintClause<const parser::OmpAtomicClauseList>(
-                &std::get<0>(atomicConstruct.t),
-                &std::get<2>(atomicConstruct.t));
+                &std::get<0>(atomicCapture.t), &std::get<2>(atomicCapture.t));
+            CheckAtomicCaptureConstruct(atomicCapture);
           },
       },
       x.u);
 }
 
 void OmpStructureChecker::Leave(const parser::OpenMPAtomicConstruct &) {
   dirContext_.pop_back();
 }

flang/test/Semantics/OpenMP/omp-atomic-assignment-stmt.f90

 ! RUN: %python %S/../test_errors.py %s %flang_fc1 -fopenmp
 ! Semantic checks for various assignments related to atomic constructs
 
 program sample
     use omp_lib
-    integer :: x, v
+    integer :: x, v, b
     integer :: y(10)
     integer, allocatable :: k
     integer a(10)
     type sample_type
         integer :: y
         integer :: m
     endtype
     type(sample_type) :: z
program sample
     !$omp atomic write
     !ERROR: No intrinsic or user-defined ASSIGNMENT(=) matches scalar INTEGER(4) and rank 1 array of INTEGER(4)
     !ERROR: Expected scalar expression on the RHS of atomic assignment statement
         x = a
 
     !$omp atomic write
     !ERROR: Expected scalar variable on the LHS of atomic assignment statement
         a = x
+
+    !$omp atomic capture
+       v = x
+       x = x + 1
+    !$omp end atomic
+
+    !$omp atomic release capture
+    !ERROR: Invalid atomic capture construct statements. Expected one of [update-stmt, capture-stmt], [capture-stmt, update-stmt], or [capture-stmt, write-stmt]
+        v = x
+        x = b + (x * 1)
+    !$omp end atomic
+
+    !$omp atomic capture hint(0)
+        v = x
+        x = 1
+    !$omp end atomic
+
+    !$omp atomic capture
+    !ERROR: Invalid atomic capture construct statements. Expected one of [update-stmt, capture-stmt], [capture-stmt, update-stmt], or [capture-stmt, write-stmt]
+        v = 1
+        x = 4
+    !$omp end atomic
+
+    !$omp atomic capture
+    !ERROR: Captured variable x expected to be assigned in statement 2 of atomic capture construct
+        v = x
+        b = b + 1
+    !$omp end atomic
+
+    !$omp atomic capture
+    !ERROR: Captured variable x expected to be assigned in statement 2 of atomic capture construct
+        v = x
+        b = 10
+    !$omp end atomic
+
+    !$omp atomic capture
+    !ERROR: Updated variable x expected to be captured in statement 2 of atomic capture construct
+        x = x + 10
+        v = b
+    !$omp end atomic
 end program