[flang][openacc] Enable lowering support for OpenACC atomic operations

[razvanlupusoru]

Summary

Since the OpenACC atomics specification is a subset of OpenMP atomics,
the same lowering implementation can be used. This change extracts out
the necessary pieces from the OpenMP lowering and puts them in a shared
spot. The shared spot is a header file (in the lib directory) so that
each implementation can template specialize directly.

After putting the OpenMP implementation in a common spot, the following
changes were needed to make it work for OpenACC:

• Ensure parsing works correctly by avoiding hardcoded offsets.
• Templatize based on atomic type.
• The checking whether it is OpenMP or OpenACC is done by checking

for OmpAtomicClauseList (OpenACC does not implement this so we just
templatize with void). It was preferable to check this instead of atomic
type because in some cases, like atomic capture, the read/write/update
implementations are called - and we want compile time evaluation of
these conditional parts.

• The memory order and hint are used only for OpenMP.
• Generate acc dialect operations instead of omp dialect operations.

Depends on D158768

[clementval]

Just to be sure, did you build this change with https://reviews.llvm.org/D158294 in mind?

[razvanlupusoru]

In D158772#4618076, @clementval wrote:

Just to be sure, did you build this change with https://reviews.llvm.org/D158294 in mind?

Yes - it will be straightforward to adapt the changes here on top of that as well.

[clementval]

In D158772#4618093, @razvanlupusoru wrote:

In D158772#4618076, @clementval wrote:

Just to be sure, did you build this change with https://reviews.llvm.org/D158294 in mind?

Yes - it will be straightforward to adapt the changes here on top of that as well.

Ok great! This look good for me.

flang/lib/Lower/Bridge.cpp

   }
 
   void genFIR(const Fortran::parser::CompilerDirective &) {
     // TODO
   }
 
   void genFIR(const Fortran::parser::OpenACCConstruct &acc) {
     mlir::OpBuilder::InsertPoint insertPt = builder->saveInsertionPoint();
+    localSymbols.pushScope();
     genOpenACCConstruct(*this, bridge.getSemanticsContext(), getEval(), acc);
     for (Fortran::lower::pft::Evaluation &e : getEval().getNestedEvaluations())
       genFIR(e);
+    localSymbols.popScope();
     builder->restoreInsertionPoint(insertPt);
   }
 
   void genFIR(const Fortran::parser::OpenACCDeclarativeConstruct &accDecl) {
     genOpenACCDeclarativeConstruct(*this, bridge.getSemanticsContext(),
                                    bridge.fctCtx(), getEval(), accDecl,
                                    accRoutineInfos);
     for (Fortran::lower::pft::Evaluation &e : getEval().getNestedEvaluations())

flang/lib/Lower/DirectivesCommon.h

+//===-- Lower/DirectivesCommon.h -- OpenACC/OpenMP lower common -----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Coding style: https://mlir.llvm.org/getting_started/DeveloperGuide/
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef FORTRAN_LOWER_DIRECTIVES_COMMON_H
+#define FORTRAN_LOWER_DIRECTIVES_COMMON_H
+
+#include "flang/Common/idioms.h"
+#include "flang/Lower/Bridge.h"
+#include "flang/Lower/ConvertExpr.h"
+#include "flang/Lower/ConvertVariable.h"
+#include "flang/Lower/OpenACC.h"
+#include "flang/Lower/OpenMP.h"
+#include "flang/Lower/PFTBuilder.h"
+#include "flang/Lower/StatementContext.h"
+#include "flang/Optimizer/Builder/BoxValue.h"
+#include "flang/Optimizer/Builder/FIRBuilder.h"
+#include "flang/Optimizer/Builder/Todo.h"
+#include "flang/Optimizer/HLFIR/HLFIROps.h"
+#include "flang/Parser/parse-tree.h"
+#include "flang/Semantics/openmp-directive-sets.h"
+#include "flang/Semantics/tools.h"
+#include "mlir/Dialect/OpenACC/OpenACC.h"
+#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
+#include "llvm/Frontend/OpenMP/OMPConstants.h"
+#include <type_traits>
+
+namespace Fortran {
+namespace lower {
+static inline bool checkForSingleVariableOnRHS(
+    const Fortran::parser::AssignmentStmt &assignmentStmt) {
+  // Check if the assignment statement has a single variable on the RHS
+  const Fortran::parser::Expr &expr{
+      std::get<Fortran::parser::Expr>(assignmentStmt.t)};
+  const Fortran::common::Indirection<Fortran::parser::Designator> *designator =
+      std::get_if<Fortran::common::Indirection<Fortran::parser::Designator>>(
+          &expr.u);
+  const Fortran::parser::Name *name =
+      designator
+          ? Fortran::semantics::getDesignatorNameIfDataRef(designator->value())
+          : nullptr;
+  return name != nullptr;
+}
+
+static inline bool
+checkForSymbolMatch(const Fortran::parser::AssignmentStmt &assignmentStmt) {
+  // Check if the symbol on the LHS of the assignment statement is present in
+  // the RHS expression
+  const auto &var{std::get<Fortran::parser::Variable>(assignmentStmt.t)};
+  const auto &expr{std::get<Fortran::parser::Expr>(assignmentStmt.t)};
+  const auto *e{Fortran::semantics::GetExpr(expr)};
+  const auto *v{Fortran::semantics::GetExpr(var)};
+  auto varSyms{Fortran::evaluate::GetSymbolVector(*v)};
+  const Fortran::semantics::Symbol &varSymbol{*varSyms.front()};
+  for (const Fortran::semantics::Symbol &symbol :
+       Fortran::evaluate::GetSymbolVector(*e))
+    if (varSymbol == symbol)
+      return true;
+  return false;
+}
+
+static inline void genOmpAtomicHintAndMemoryOrderClauses(
+    Fortran::lower::AbstractConverter &converter,
+    const Fortran::parser::OmpAtomicClauseList &clauseList,
+    mlir::IntegerAttr &hint,
+    mlir::omp::ClauseMemoryOrderKindAttr &memoryOrder) {
+  fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
+  for (const Fortran::parser::OmpAtomicClause &clause : clauseList.v) {
+    if (const auto *ompClause =
+            std::get_if<Fortran::parser::OmpClause>(&clause.u)) {
+      if (const auto *hintClause =
+              std::get_if<Fortran::parser::OmpClause::Hint>(&ompClause->u)) {
+        const auto *expr = Fortran::semantics::GetExpr(hintClause->v);
+        uint64_t hintExprValue = *Fortran::evaluate::ToInt64(*expr);
+        hint = firOpBuilder.getI64IntegerAttr(hintExprValue);
+      }
+    } else if (const auto *ompMemoryOrderClause =
+                   std::get_if<Fortran::parser::OmpMemoryOrderClause>(
+                       &clause.u)) {
+      if (std::get_if<Fortran::parser::OmpClause::Acquire>(
+              &ompMemoryOrderClause->v.u)) {
+        memoryOrder = mlir::omp::ClauseMemoryOrderKindAttr::get(
+            firOpBuilder.getContext(),
+            mlir::omp::ClauseMemoryOrderKind::Acquire);
+      } else if (std::get_if<Fortran::parser::OmpClause::Relaxed>(
+                     &ompMemoryOrderClause->v.u)) {
+        memoryOrder = mlir::omp::ClauseMemoryOrderKindAttr::get(
+            firOpBuilder.getContext(),
+            mlir::omp::ClauseMemoryOrderKind::Relaxed);
+      } else if (std::get_if<Fortran::parser::OmpClause::SeqCst>(
+                     &ompMemoryOrderClause->v.u)) {
+        memoryOrder = mlir::omp::ClauseMemoryOrderKindAttr::get(
+            firOpBuilder.getContext(),
+            mlir::omp::ClauseMemoryOrderKind::Seq_cst);
+      } else if (std::get_if<Fortran::parser::OmpClause::Release>(
+                     &ompMemoryOrderClause->v.u)) {
+        memoryOrder = mlir::omp::ClauseMemoryOrderKindAttr::get(
+            firOpBuilder.getContext(),
+            mlir::omp::ClauseMemoryOrderKind::Release);
+      }
+    }
+  }
+}
+
+template <typename AtomicListT>
+static inline void genOmpAccAtomicCaptureStatement(
+    Fortran::lower::AbstractConverter &converter,
+    Fortran::lower::pft::Evaluation &eval, mlir::Value fromAddress,
+    mlir::Value toAddress, const AtomicListT *leftHandClauseList,
+    const AtomicListT *rightHandClauseList, mlir::Type elementType) {
+  // Generate `atomic.read` operation for atomic assigment statements
+  fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
+  mlir::Location currentLocation = converter.getCurrentLocation();
+
+  if constexpr (std::is_same<AtomicListT,
+                             Fortran::parser::OmpAtomicClauseList>()) {
+    // If no hint clause is specified, the effect is as if
+    // hint(omp_sync_hint_none) had been specified.
+    mlir::IntegerAttr hint = nullptr;
+
+    mlir::omp::ClauseMemoryOrderKindAttr memoryOrder = nullptr;
+    if (leftHandClauseList)
+      genOmpAtomicHintAndMemoryOrderClauses(converter, *leftHandClauseList,
+                                            hint, memoryOrder);
+    if (rightHandClauseList)
+      genOmpAtomicHintAndMemoryOrderClauses(converter, *rightHandClauseList,
+                                            hint, memoryOrder);
+    firOpBuilder.create<mlir::omp::AtomicReadOp>(
+        currentLocation, fromAddress, toAddress,
+        mlir::TypeAttr::get(elementType), hint, memoryOrder);
+  } else {
+    firOpBuilder.create<mlir::acc::AtomicReadOp>(
+        currentLocation, fromAddress, toAddress,
+        mlir::TypeAttr::get(elementType));
+  }
+}
+
+template <typename AtomicListT>
+static inline void
+genOmpAccAtomicWriteStatement(Fortran::lower::AbstractConverter &converter,
+                              Fortran::lower::pft::Evaluation &eval,
+                              mlir::Value lhsAddr, mlir::Value rhsExpr,
+                              const AtomicListT *leftHandClauseList,
+                              const AtomicListT *rightHandClauseList,
+                              mlir::Value *evaluatedExprValue = nullptr) {
+  // Generate `atomic.write` operation for atomic assignment statements
+  fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
+  mlir::Location currentLocation = converter.getCurrentLocation();
+
+  if constexpr (std::is_same<AtomicListT,
+                             Fortran::parser::OmpAtomicClauseList>()) {
+    // If no hint clause is specified, the effect is as if
+    // hint(omp_sync_hint_none) had been specified.
+    mlir::IntegerAttr hint = nullptr;
+    mlir::omp::ClauseMemoryOrderKindAttr memoryOrder = nullptr;
+    if (leftHandClauseList)
+      genOmpAtomicHintAndMemoryOrderClauses(converter, *leftHandClauseList,
+                                            hint, memoryOrder);
+    if (rightHandClauseList)
+      genOmpAtomicHintAndMemoryOrderClauses(converter, *rightHandClauseList,
+                                            hint, memoryOrder);
+    firOpBuilder.create<mlir::omp::AtomicWriteOp>(currentLocation, lhsAddr,
+                                                  rhsExpr, hint, memoryOrder);
+  } else {
+    firOpBuilder.create<mlir::acc::AtomicWriteOp>(currentLocation, lhsAddr,
+                                                  rhsExpr);
+  }
+}
+
+template <typename AtomicListT>
+static inline void genOmpAccAtomicUpdateStatement(
+    Fortran::lower::AbstractConverter &converter,
+    Fortran::lower::pft::Evaluation &eval, mlir::Value lhsAddr,
+    mlir::Type varType, const Fortran::parser::Variable &assignmentStmtVariable,
+    const Fortran::parser::Expr &assignmentStmtExpr,
+    const AtomicListT *leftHandClauseList,
+    const AtomicListT *rightHandClauseList) {
+  // Generate `atomic.update` operation for atomic assignment statements
+  fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
+  mlir::Location currentLocation = converter.getCurrentLocation();
+  mlir::Operation *atomicUpdateOp = nullptr;
+
+  if constexpr (std::is_same<AtomicListT,
+                             Fortran::parser::OmpAtomicClauseList>()) {
+    // If no hint clause is specified, the effect is as if
+    // hint(omp_sync_hint_none) had been specified.
+    mlir::IntegerAttr hint = nullptr;
+    mlir::omp::ClauseMemoryOrderKindAttr memoryOrder = nullptr;
+
+    if (leftHandClauseList)
+      genOmpAtomicHintAndMemoryOrderClauses(converter, *leftHandClauseList,
+                                            hint, memoryOrder);
+    if (rightHandClauseList)
+      genOmpAtomicHintAndMemoryOrderClauses(converter, *rightHandClauseList,
+                                            hint, memoryOrder);
+    atomicUpdateOp = firOpBuilder.create<mlir::omp::AtomicUpdateOp>(
+        currentLocation, lhsAddr, hint, memoryOrder);
+  } else {
+    atomicUpdateOp = firOpBuilder.create<mlir::acc::AtomicUpdateOp>(
+        currentLocation, lhsAddr);
+  }
+
+  //// Generate body of Atomic Update operation
+  // If an argument for the region is provided then create the block with that
+  // argument. Also update the symbol's address with the argument mlir value.
+  llvm::SmallVector<mlir::Type> varTys = {varType};
+  llvm::SmallVector<mlir::Location> locs = {currentLocation};
+  firOpBuilder.createBlock(&(atomicUpdateOp->getRegion(0)), {}, varTys, locs);
+  mlir::Value val =
+      fir::getBase(atomicUpdateOp->getRegion(0).front().getArgument(0));
+  const auto *varDesignator =
+      std::get_if<Fortran::common::Indirection<Fortran::parser::Designator>>(
+          &assignmentStmtVariable.u);
+  assert(varDesignator && "Variable designator for atomic update assignment "
+                          "statement does not exist");
+  const Fortran::parser::Name *name =
+      Fortran::semantics::getDesignatorNameIfDataRef(varDesignator->value());
+  if (!name)
+    TODO(converter.getCurrentLocation(),
+         "Array references as atomic update variable");
+  assert(name && name->symbol &&
+         "No symbol attached to atomic update variable");
+  converter.bindSymbol(*name->symbol, val);
+  // Set the insert for the terminator operation to go at the end of the
+  // block.
+  mlir::Block &block = atomicUpdateOp->getRegion(0).back();
+  firOpBuilder.setInsertionPointToEnd(&block);
+
+  Fortran::lower::StatementContext stmtCtx;
+  mlir::Value rhsExpr = fir::getBase(converter.genExprValue(
+      *Fortran::semantics::GetExpr(assignmentStmtExpr), stmtCtx));
+  mlir::Value convertResult =
+      firOpBuilder.createConvert(currentLocation, varType, rhsExpr);
+  // Insert the terminator: YieldOp.
+  if constexpr (std::is_same<AtomicListT,
+                             Fortran::parser::OmpAtomicClauseList>()) {
+    firOpBuilder.create<mlir::omp::YieldOp>(currentLocation, convertResult);
+  } else {
+    firOpBuilder.create<mlir::acc::YieldOp>(currentLocation, convertResult);
+  }
+  // Reset the insert point to before the terminator.
+  firOpBuilder.setInsertionPointToStart(&block);
+}
+
+template <typename AtomicT, typename AtomicListT>
+void genOmpAccAtomicWrite(Fortran::lower::AbstractConverter &converter,
+                          Fortran::lower::pft::Evaluation &eval,
+                          const AtomicT &atomicWrite) {
+  const AtomicListT *rightHandClauseList = nullptr;
+  const AtomicListT *leftHandClauseList = nullptr;
+  if constexpr (std::is_same<AtomicListT,
+                             Fortran::parser::OmpAtomicClauseList>()) {
+    // Get the address of atomic read operands.
+    rightHandClauseList = &std::get<2>(atomicWrite.t);
+    leftHandClauseList = &std::get<0>(atomicWrite.t);
+  }
+
+  const Fortran::parser::AssignmentStmt &stmt =
+      std::get<Fortran::parser::Statement<Fortran::parser::AssignmentStmt>>(
+          atomicWrite.t)
+          .statement;
+  const Fortran::evaluate::Assignment &assign = *stmt.typedAssignment->v;
+  Fortran::lower::StatementContext stmtCtx;
+  // Get the value and address of atomic write operands.
+  mlir::Value rhsExpr =
+      fir::getBase(converter.genExprValue(assign.rhs, stmtCtx));
+  mlir::Value lhsAddr =
+      fir::getBase(converter.genExprAddr(assign.lhs, stmtCtx));
+  genOmpAccAtomicWriteStatement(converter, eval, lhsAddr, rhsExpr,
+                                leftHandClauseList, rightHandClauseList);
+}
+
+template <typename AtomicT, typename AtomicListT>
+void genOmpAccAtomicRead(Fortran::lower::AbstractConverter &converter,
+                         Fortran::lower::pft::Evaluation &eval,
+                         const AtomicT &atomicRead) {
+  const AtomicListT *rightHandClauseList = nullptr;
+  const AtomicListT *leftHandClauseList = nullptr;
+  if constexpr (std::is_same<AtomicListT,
+                             Fortran::parser::OmpAtomicClauseList>()) {
+    // Get the address of atomic read operands.
+    rightHandClauseList = &std::get<2>(atomicRead.t);
+    leftHandClauseList = &std::get<0>(atomicRead.t);
+  }
+
+  const auto &assignmentStmtExpr = std::get<Fortran::parser::Expr>(
+      std::get<Fortran::parser::Statement<Fortran::parser::AssignmentStmt>>(
+          atomicRead.t)
+          .statement.t);
+  const auto &assignmentStmtVariable = std::get<Fortran::parser::Variable>(
+      std::get<Fortran::parser::Statement<Fortran::parser::AssignmentStmt>>(
+          atomicRead.t)
+          .statement.t);
+
+  Fortran::lower::StatementContext stmtCtx;
+  const Fortran::semantics::SomeExpr &fromExpr =
+      *Fortran::semantics::GetExpr(assignmentStmtExpr);
+  mlir::Type elementType = converter.genType(fromExpr);
+  mlir::Value fromAddress =
+      fir::getBase(converter.genExprAddr(fromExpr, stmtCtx));
+  mlir::Value toAddress = fir::getBase(converter.genExprAddr(
+      *Fortran::semantics::GetExpr(assignmentStmtVariable), stmtCtx));
+  genOmpAccAtomicCaptureStatement(converter, eval, fromAddress, toAddress,
+                                  leftHandClauseList, rightHandClauseList,
+                                  elementType);
+}
+
+template <typename AtomicT, typename AtomicListT>
+void genOmpAccAtomicUpdate(Fortran::lower::AbstractConverter &converter,
+                           Fortran::lower::pft::Evaluation &eval,
+                           const AtomicT &atomicUpdate) {
+  const AtomicListT *rightHandClauseList = nullptr;
+  const AtomicListT *leftHandClauseList = nullptr;
+  if constexpr (std::is_same<AtomicListT,
+                             Fortran::parser::OmpAtomicClauseList>()) {
+    // Get the address of atomic read operands.
+    rightHandClauseList = &std::get<2>(atomicUpdate.t);
+    leftHandClauseList = &std::get<0>(atomicUpdate.t);
+  }
+
+  const auto &assignmentStmtExpr = std::get<Fortran::parser::Expr>(
+      std::get<Fortran::parser::Statement<Fortran::parser::AssignmentStmt>>(
+          atomicUpdate.t)
+          .statement.t);
+  const auto &assignmentStmtVariable = std::get<Fortran::parser::Variable>(
+      std::get<Fortran::parser::Statement<Fortran::parser::AssignmentStmt>>(
+          atomicUpdate.t)
+          .statement.t);
+
+  Fortran::lower::StatementContext stmtCtx;
+  mlir::Value lhsAddr = fir::getBase(converter.genExprAddr(
+      *Fortran::semantics::GetExpr(assignmentStmtVariable), stmtCtx));
+  mlir::Type varType =
+      fir::getBase(
+          converter.genExprValue(
+              *Fortran::semantics::GetExpr(assignmentStmtVariable), stmtCtx))
+          .getType();
+  genOmpAccAtomicUpdateStatement<AtomicListT>(
+      converter, eval, lhsAddr, varType, assignmentStmtVariable,
+      assignmentStmtExpr, leftHandClauseList, rightHandClauseList);
+}
+
+template <typename AtomicT, typename AtomicListT>
+void genOmpAtomic(Fortran::lower::AbstractConverter &converter,
+                  Fortran::lower::pft::Evaluation &eval,
+                  const AtomicT &atomicConstruct) {
+  const AtomicListT &atomicClauseList =
+      std::get<AtomicListT>(atomicConstruct.t);
+  const auto &assignmentStmtExpr = std::get<Fortran::parser::Expr>(
+      std::get<Fortran::parser::Statement<Fortran::parser::AssignmentStmt>>(
+          atomicConstruct.t)
+          .statement.t);
+  const auto &assignmentStmtVariable = std::get<Fortran::parser::Variable>(
+      std::get<Fortran::parser::Statement<Fortran::parser::AssignmentStmt>>(
+          atomicConstruct.t)
+          .statement.t);
+  Fortran::lower::StatementContext stmtCtx;
+  mlir::Value lhsAddr = fir::getBase(converter.genExprAddr(
+      *Fortran::semantics::GetExpr(assignmentStmtVariable), stmtCtx));
+  mlir::Type varType =
+      fir::getBase(
+          converter.genExprValue(
+              *Fortran::semantics::GetExpr(assignmentStmtVariable), stmtCtx))
+          .getType();
+  // If atomic-clause is not present on the construct, the behaviour is as if
+  // the update clause is specified (for both OpenMP and OpenACC).
+  genOmpAccAtomicUpdateStatement<AtomicListT>(
+      converter, eval, lhsAddr, varType, assignmentStmtVariable,
+      assignmentStmtExpr, &atomicClauseList, nullptr);
+}
+
+template <typename AtomicT, typename AtomicListT>
+void genOmpAccAtomicCapture(Fortran::lower::AbstractConverter &converter,
+                            Fortran::lower::pft::Evaluation &eval,
+                            const AtomicT &atomicCapture) {
+  fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
+  mlir::Location currentLocation = converter.getCurrentLocation();
+
+  mlir::IntegerAttr hint = nullptr;
+  mlir::omp::ClauseMemoryOrderKindAttr memoryOrder = nullptr;
+  if constexpr (std::is_same<AtomicListT,
+                             Fortran::parser::OmpAtomicClauseList>()) {
+    const AtomicListT &rightHandClauseList = std::get<2>(atomicCapture.t);
+    const AtomicListT &leftHandClauseList = std::get<0>(atomicCapture.t);
+    genOmpAtomicHintAndMemoryOrderClauses(converter, leftHandClauseList, hint,
+                                          memoryOrder);
+    genOmpAtomicHintAndMemoryOrderClauses(converter, rightHandClauseList, hint,
+                                          memoryOrder);
+  }
+
+  const Fortran::parser::AssignmentStmt &stmt1 =
+      std::get<typename AtomicT::Stmt1>(atomicCapture.t).v.statement;
+  const auto &stmt1Var{std::get<Fortran::parser::Variable>(stmt1.t)};
+  const auto &stmt1Expr{std::get<Fortran::parser::Expr>(stmt1.t)};
+  const Fortran::parser::AssignmentStmt &stmt2 =
+      std::get<typename AtomicT::Stmt2>(atomicCapture.t).v.statement;
+  const auto &stmt2Var{std::get<Fortran::parser::Variable>(stmt2.t)};
+  const auto &stmt2Expr{std::get<Fortran::parser::Expr>(stmt2.t)};
+
+  // Pre-evaluate expressions to be used in the various operations inside
+  // `atomic.capture` since it is not desirable to have anything other than
+  // a `atomic.read`, `atomic.write`, or `atomic.update` operation
+  // inside `atomic.capture`
+  Fortran::lower::StatementContext stmtCtx;
+  mlir::Value stmt1LHSArg, stmt1RHSArg, stmt2LHSArg, stmt2RHSArg;
+  mlir::Type elementType;
+  // LHS evaluations are common to all combinations of `atomic.capture`
+  stmt1LHSArg = fir::getBase(
+      converter.genExprAddr(*Fortran::semantics::GetExpr(stmt1Var), stmtCtx));
+  stmt2LHSArg = fir::getBase(
+      converter.genExprAddr(*Fortran::semantics::GetExpr(stmt2Var), stmtCtx));
+
+  // Operation specific RHS evaluations
+  if (checkForSingleVariableOnRHS(stmt1)) {
+    // Atomic capture construct is of the form [capture-stmt, update-stmt] or
+    // of the form [capture-stmt, write-stmt]
+    stmt1RHSArg = fir::getBase(converter.genExprAddr(
+        *Fortran::semantics::GetExpr(stmt1Expr), stmtCtx));
+    stmt2RHSArg = fir::getBase(converter.genExprValue(
+        *Fortran::semantics::GetExpr(stmt2Expr), stmtCtx));
+
+  } else {
+    // Atomic capture construct is of the form [update-stmt, capture-stmt]
+    stmt1RHSArg = fir::getBase(converter.genExprValue(
+        *Fortran::semantics::GetExpr(stmt1Expr), stmtCtx));
+    stmt2RHSArg = fir::getBase(converter.genExprAddr(
+        *Fortran::semantics::GetExpr(stmt2Expr), stmtCtx));
+  }
+  // Type information used in generation of `atomic.update` operation
+  mlir::Type stmt1VarType =
+      fir::getBase(converter.genExprValue(
+                       *Fortran::semantics::GetExpr(stmt1Var), stmtCtx))
+          .getType();
+  mlir::Type stmt2VarType =
+      fir::getBase(converter.genExprValue(
+                       *Fortran::semantics::GetExpr(stmt2Var), stmtCtx))
+          .getType();
+
+  mlir::Operation *atomicCaptureOp = nullptr;
+  if constexpr (std::is_same<AtomicListT,
+                             Fortran::parser::OmpAtomicClauseList>()) {
+    atomicCaptureOp = firOpBuilder.create<mlir::omp::AtomicCaptureOp>(
+        currentLocation, hint, memoryOrder);
+  } else {
+    atomicCaptureOp =
+        firOpBuilder.create<mlir::acc::AtomicCaptureOp>(currentLocation);
+  }
+
+  firOpBuilder.createBlock(&(atomicCaptureOp->getRegion(0)));
+  mlir::Block &block = atomicCaptureOp->getRegion(0).back();
+  firOpBuilder.setInsertionPointToStart(&block);
+  if (checkForSingleVariableOnRHS(stmt1)) {
+    if (checkForSymbolMatch(stmt2)) {
+      // Atomic capture construct is of the form [capture-stmt, update-stmt]
+      const Fortran::semantics::SomeExpr &fromExpr =
+          *Fortran::semantics::GetExpr(stmt1Expr);
+      elementType = converter.genType(fromExpr);
+      genOmpAccAtomicCaptureStatement<AtomicListT>(
+          converter, eval, stmt1RHSArg, stmt1LHSArg,
+          /*leftHandClauseList=*/nullptr,
+          /*rightHandClauseList=*/nullptr, elementType);
+      genOmpAccAtomicUpdateStatement<AtomicListT>(
+          converter, eval, stmt1RHSArg, stmt2VarType, stmt2Var, stmt2Expr,
+          /*leftHandClauseList=*/nullptr,
+          /*rightHandClauseList=*/nullptr);
+    } else {
+      // Atomic capture construct is of the form [capture-stmt, write-stmt]
+      const Fortran::semantics::SomeExpr &fromExpr =
+          *Fortran::semantics::GetExpr(stmt1Expr);
+      elementType = converter.genType(fromExpr);
+      genOmpAccAtomicCaptureStatement<AtomicListT>(
+          converter, eval, stmt1RHSArg, stmt1LHSArg,
+          /*leftHandClauseList=*/nullptr,
+          /*rightHandClauseList=*/nullptr, elementType);
+      genOmpAccAtomicWriteStatement<AtomicListT>(
+          converter, eval, stmt1RHSArg, stmt2RHSArg,
+          /*leftHandClauseList=*/nullptr,
+          /*rightHandClauseList=*/nullptr);
+    }
+  } else {
+    // Atomic capture construct is of the form [update-stmt, capture-stmt]
+    firOpBuilder.setInsertionPointToEnd(&block);
+    const Fortran::semantics::SomeExpr &fromExpr =
+        *Fortran::semantics::GetExpr(stmt2Expr);
+    elementType = converter.genType(fromExpr);
+    genOmpAccAtomicCaptureStatement<AtomicListT>(
+        converter, eval, stmt1LHSArg, stmt2LHSArg,
+        /*leftHandClauseList=*/nullptr,
+        /*rightHandClauseList=*/nullptr, elementType);
+    firOpBuilder.setInsertionPointToStart(&block);
+    genOmpAccAtomicUpdateStatement<AtomicListT>(
+        converter, eval, stmt1LHSArg, stmt1VarType, stmt1Var, stmt1Expr,
+        /*leftHandClauseList=*/nullptr,
+        /*rightHandClauseList=*/nullptr);
+  }
+  firOpBuilder.setInsertionPointToEnd(&block);
+  if constexpr (std::is_same<AtomicListT,
+                             Fortran::parser::OmpAtomicClauseList>()) {
+    firOpBuilder.create<mlir::omp::TerminatorOp>(currentLocation);
+  } else {
+    firOpBuilder.create<mlir::acc::TerminatorOp>(currentLocation);
+  }
+  firOpBuilder.setInsertionPointToStart(&block);
+}
+
+} // namespace lower
+} // namespace Fortran
+
+#endif // FORTRAN_LOWER_DIRECTIVES_COMMON_H
 No newline at end of file

flang/lib/Lower/OpenACC.cpp

 //===-- OpenACC.cpp -- OpenACC directive lowering -------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Coding style: https://mlir.llvm.org/getting_started/DeveloperGuide/
 //
 //===----------------------------------------------------------------------===//
 
 #include "flang/Lower/OpenACC.h"
+#include "DirectivesCommon.h"
 #include "flang/Common/idioms.h"
 #include "flang/Lower/Bridge.h"
 #include "flang/Lower/ConvertType.h"
 #include "flang/Lower/Mangler.h"
 #include "flang/Lower/PFTBuilder.h"
 #include "flang/Lower/StatementContext.h"
 #include "flang/Lower/Support/Utils.h"
 #include "flang/Optimizer/Builder/BoxValue.h"
for (auto &mapping : accRoutineInfos) {
     if (!funcOp)
       llvm::report_fatal_error(
           "could not find function to attach OpenACC routine information.");
     attachRoutineInfo(funcOp, mapping.second);
   }
   accRoutineInfos.clear();
 }
 
+static void
+genACC(Fortran::lower::AbstractConverter &converter,
+       Fortran::lower::pft::Evaluation &eval,
+       const Fortran::parser::OpenACCAtomicConstruct &atomicConstruct) {
+  std::visit(
+      Fortran::common::visitors{
+          [&](const Fortran::parser::AccAtomicRead &atomicRead) {
+            Fortran::lower::genOmpAccAtomicRead<Fortran::parser::AccAtomicRead,
+                                                void>(converter, eval,
+                                                      atomicRead);
+          },
+          [&](const Fortran::parser::AccAtomicWrite &atomicWrite) {
+            Fortran::lower::genOmpAccAtomicWrite<
+                Fortran::parser::AccAtomicWrite, void>(converter, eval,
+                                                       atomicWrite);
+          },
+          [&](const Fortran::parser::AccAtomicUpdate &atomicUpdate) {
+            Fortran::lower::genOmpAccAtomicUpdate<
+                Fortran::parser::AccAtomicUpdate, void>(converter, eval,
+                                                        atomicUpdate);
+          },
+          [&](const Fortran::parser::AccAtomicCapture &atomicCapture) {
+            Fortran::lower::genOmpAccAtomicCapture<
+                Fortran::parser::AccAtomicCapture, void>(converter, eval,
+                                                         atomicCapture);
+          },
+      },
+      atomicConstruct.u);
+}
+
 void Fortran::lower::genOpenACCConstruct(
     Fortran::lower::AbstractConverter &converter,
     Fortran::semantics::SemanticsContext &semanticsContext,
     Fortran::lower::pft::Evaluation &eval,
     const Fortran::parser::OpenACCConstruct &accConstruct) {
 
   std::visit(
       common::visitors{
std::visit(
           [&](const Fortran::parser::OpenACCCacheConstruct &cacheConstruct) {
             TODO(converter.genLocation(cacheConstruct.source),
                  "OpenACC Cache construct not lowered yet!");
           },
           [&](const Fortran::parser::OpenACCWaitConstruct &waitConstruct) {
             genACC(converter, eval, waitConstruct);
           },
           [&](const Fortran::parser::OpenACCAtomicConstruct &atomicConstruct) {
-            TODO(converter.genLocation(atomicConstruct.source),
-                 "OpenACC Atomic construct not lowered yet!");
+            genACC(converter, eval, atomicConstruct);
           },
       },
       accConstruct.u);
 }
 
 void Fortran::lower::genOpenACCDeclarativeConstruct(
     Fortran::lower::AbstractConverter &converter,
     Fortran::semantics::SemanticsContext &semanticsContext,

flang/lib/Lower/OpenMP.cpp

 //===-- OpenMP.cpp -- Open MP directive lowering --------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Coding style: https://mlir.llvm.org/getting_started/DeveloperGuide/
 //
 //===----------------------------------------------------------------------===//
 
 #include "flang/Lower/OpenMP.h"
+#include "DirectivesCommon.h"
 #include "flang/Common/idioms.h"
 #include "flang/Lower/Bridge.h"
 #include "flang/Lower/ConvertExpr.h"
 #include "flang/Lower/ConvertVariable.h"
 #include "flang/Lower/PFTBuilder.h"
 #include "flang/Lower/StatementContext.h"
 #include "flang/Optimizer/Builder/BoxValue.h"
 #include "flang/Optimizer/Builder/FIRBuilder.h"
genOMP(Fortran::lower::AbstractConverter &converter,
   genOpWithBody<mlir::omp::SectionsOp>(converter, eval, currentLocation,
                                        /*outerCombined=*/false,
                                        /*clauseList=*/nullptr,
                                        /*reduction_vars=*/mlir::ValueRange(),
                                        /*reductions=*/nullptr, allocateOperands,
                                        allocatorOperands, nowaitClauseOperand);
 }
 
-static bool checkForSingleVariableOnRHS(
-    const Fortran::parser::AssignmentStmt &assignmentStmt) {
-  // Check if the assignment statement has a single variable on the RHS
-  const Fortran::parser::Expr &expr{
-      std::get<Fortran::parser::Expr>(assignmentStmt.t)};
-  const Fortran::common::Indirection<Fortran::parser::Designator> *designator =
-      std::get_if<Fortran::common::Indirection<Fortran::parser::Designator>>(
-          &expr.u);
-  const Fortran::parser::Name *name =
-      designator
-          ? Fortran::semantics::getDesignatorNameIfDataRef(designator->value())
-          : nullptr;
-  return name != nullptr;
-}
-
-static bool
-checkForSymbolMatch(const Fortran::parser::AssignmentStmt &assignmentStmt) {
-  // Check if the symbol on the LHS of the assignment statement is present in
-  // the RHS expression
-  const auto &var{std::get<Fortran::parser::Variable>(assignmentStmt.t)};
-  const auto &expr{std::get<Fortran::parser::Expr>(assignmentStmt.t)};
-  const auto *e{Fortran::semantics::GetExpr(expr)};
-  const auto *v{Fortran::semantics::GetExpr(var)};
-  auto varSyms{Fortran::evaluate::GetSymbolVector(*v)};
-  const Fortran::semantics::Symbol &varSymbol{*varSyms.front()};
-  for (const Fortran::semantics::Symbol &symbol :
-       Fortran::evaluate::GetSymbolVector(*e))
-    if (varSymbol == symbol)
-      return true;
-  return false;
-}
-
-static void genOmpAtomicHintAndMemoryOrderClauses(
-    Fortran::lower::AbstractConverter &converter,
-    const Fortran::parser::OmpAtomicClauseList &clauseList,
-    mlir::IntegerAttr &hint,
-    mlir::omp::ClauseMemoryOrderKindAttr &memoryOrder) {
-  fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
-  for (const Fortran::parser::OmpAtomicClause &clause : clauseList.v) {
-    if (const auto *ompClause =
-            std::get_if<Fortran::parser::OmpClause>(&clause.u)) {
-      if (const auto *hintClause =
-              std::get_if<Fortran::parser::OmpClause::Hint>(&ompClause->u)) {
-        const auto *expr = Fortran::semantics::GetExpr(hintClause->v);
-        uint64_t hintExprValue = *Fortran::evaluate::ToInt64(*expr);
-        hint = firOpBuilder.getI64IntegerAttr(hintExprValue);
-      }
-    } else if (const auto *ompMemoryOrderClause =
-                   std::get_if<Fortran::parser::OmpMemoryOrderClause>(
-                       &clause.u)) {
-      if (std::get_if<Fortran::parser::OmpClause::Acquire>(
-              &ompMemoryOrderClause->v.u)) {
-        memoryOrder = mlir::omp::ClauseMemoryOrderKindAttr::get(
-            firOpBuilder.getContext(),
-            mlir::omp::ClauseMemoryOrderKind::Acquire);
-      } else if (std::get_if<Fortran::parser::OmpClause::Relaxed>(
-                     &ompMemoryOrderClause->v.u)) {
-        memoryOrder = mlir::omp::ClauseMemoryOrderKindAttr::get(
-            firOpBuilder.getContext(),
-            mlir::omp::ClauseMemoryOrderKind::Relaxed);
-      } else if (std::get_if<Fortran::parser::OmpClause::SeqCst>(
-                     &ompMemoryOrderClause->v.u)) {
-        memoryOrder = mlir::omp::ClauseMemoryOrderKindAttr::get(
-            firOpBuilder.getContext(),
-            mlir::omp::ClauseMemoryOrderKind::Seq_cst);
-      } else if (std::get_if<Fortran::parser::OmpClause::Release>(
-                     &ompMemoryOrderClause->v.u)) {
-        memoryOrder = mlir::omp::ClauseMemoryOrderKindAttr::get(
-            firOpBuilder.getContext(),
-            mlir::omp::ClauseMemoryOrderKind::Release);
-      }
-    }
-  }
-}
-
-static void genOmpAtomicCaptureStatement(
-    Fortran::lower::AbstractConverter &converter,
-    Fortran::lower::pft::Evaluation &eval, mlir::Value fromAddress,
-    mlir::Value toAddress,
-    const Fortran::parser::OmpAtomicClauseList *leftHandClauseList,
-    const Fortran::parser::OmpAtomicClauseList *rightHandClauseList,
-    mlir::Type elementType) {
-  // Generate `omp.atomic.read` operation for atomic assigment statements
-  fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
-  mlir::Location currentLocation = converter.getCurrentLocation();
-
-  // If no hint clause is specified, the effect is as if
-  // hint(omp_sync_hint_none) had been specified.
-  mlir::IntegerAttr hint = nullptr;
-
-  mlir::omp::ClauseMemoryOrderKindAttr memoryOrder = nullptr;
-  if (leftHandClauseList)
-    genOmpAtomicHintAndMemoryOrderClauses(converter, *leftHandClauseList, hint,
-                                          memoryOrder);
-  if (rightHandClauseList)
-    genOmpAtomicHintAndMemoryOrderClauses(converter, *rightHandClauseList, hint,
-                                          memoryOrder);
-  firOpBuilder.create<mlir::omp::AtomicReadOp>(
-      currentLocation, fromAddress, toAddress, mlir::TypeAttr::get(elementType),
-      hint, memoryOrder);
-}
-
-static void genOmpAtomicWriteStatement(
-    Fortran::lower::AbstractConverter &converter,
-    Fortran::lower::pft::Evaluation &eval, mlir::Value lhsAddr,
-    mlir::Value rhsExpr,
-    const Fortran::parser::OmpAtomicClauseList *leftHandClauseList,
-    const Fortran::parser::OmpAtomicClauseList *rightHandClauseList,
-    mlir::Value *evaluatedExprValue = nullptr) {
-  // Generate `omp.atomic.write` operation for atomic assignment statements
-  fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
-  mlir::Location currentLocation = converter.getCurrentLocation();
-  // If no hint clause is specified, the effect is as if
-  // hint(omp_sync_hint_none) had been specified.
-  mlir::IntegerAttr hint = nullptr;
-  mlir::omp::ClauseMemoryOrderKindAttr memoryOrder = nullptr;
-  if (leftHandClauseList)
-    genOmpAtomicHintAndMemoryOrderClauses(converter, *leftHandClauseList, hint,
-                                          memoryOrder);
-  if (rightHandClauseList)
-    genOmpAtomicHintAndMemoryOrderClauses(converter, *rightHandClauseList, hint,
-                                          memoryOrder);
-  firOpBuilder.create<mlir::omp::AtomicWriteOp>(currentLocation, lhsAddr,
-                                                rhsExpr, hint, memoryOrder);
-}
-
-static void genOmpAtomicUpdateStatement(
-    Fortran::lower::AbstractConverter &converter,
-    Fortran::lower::pft::Evaluation &eval, mlir::Value lhsAddr,
-    mlir::Type varType, const Fortran::parser::Variable &assignmentStmtVariable,
-    const Fortran::parser::Expr &assignmentStmtExpr,
-    const Fortran::parser::OmpAtomicClauseList *leftHandClauseList,
-    const Fortran::parser::OmpAtomicClauseList *rightHandClauseList) {
-  // Generate `omp.atomic.update` operation for atomic assignment statements
-  fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
-  mlir::Location currentLocation = converter.getCurrentLocation();
-
-  // If no hint clause is specified, the effect is as if
-  // hint(omp_sync_hint_none) had been specified.
-  mlir::IntegerAttr hint = nullptr;
-  mlir::omp::ClauseMemoryOrderKindAttr memoryOrder = nullptr;
-  if (leftHandClauseList)
-    genOmpAtomicHintAndMemoryOrderClauses(converter, *leftHandClauseList, hint,
-                                          memoryOrder);
-  if (rightHandClauseList)
-    genOmpAtomicHintAndMemoryOrderClauses(converter, *rightHandClauseList, hint,
-                                          memoryOrder);
-  auto atomicUpdateOp = firOpBuilder.create<mlir::omp::AtomicUpdateOp>(
-      currentLocation, lhsAddr, hint, memoryOrder);
-
-  //// Generate body of Atomic Update operation
-  // If an argument for the region is provided then create the block with that
-  // argument. Also update the symbol's address with the argument mlir value.
-  llvm::SmallVector<mlir::Type> varTys = {varType};
-  llvm::SmallVector<mlir::Location> locs = {currentLocation};
-  firOpBuilder.createBlock(&atomicUpdateOp.getRegion(), {}, varTys, locs);
-  mlir::Value val =
-      fir::getBase(atomicUpdateOp.getRegion().front().getArgument(0));
-  const auto *varDesignator =
-      std::get_if<Fortran::common::Indirection<Fortran::parser::Designator>>(
-          &assignmentStmtVariable.u);
-  assert(varDesignator && "Variable designator for atomic update assignment "
-                          "statement does not exist");
-  const Fortran::parser::Name *name =
-      Fortran::semantics::getDesignatorNameIfDataRef(varDesignator->value());
-  if (!name)
-    TODO(converter.getCurrentLocation(),
-         "Array references as atomic update variable");
-  assert(name && name->symbol &&
-         "No symbol attached to atomic update variable");
-  converter.bindSymbol(*name->symbol, val);
-  // Set the insert for the terminator operation to go at the end of the
-  // block.
-  mlir::Block &block = atomicUpdateOp.getRegion().back();
-  firOpBuilder.setInsertionPointToEnd(&block);
-
-  Fortran::lower::StatementContext stmtCtx;
-  mlir::Value rhsExpr = fir::getBase(converter.genExprValue(
-      *Fortran::semantics::GetExpr(assignmentStmtExpr), stmtCtx));
-  mlir::Value convertResult =
-      firOpBuilder.createConvert(currentLocation, varType, rhsExpr);
-  // Insert the terminator: YieldOp.
-  firOpBuilder.create<mlir::omp::YieldOp>(currentLocation, convertResult);
-  // Reset the insert point to before the terminator.
-  firOpBuilder.setInsertionPointToStart(&block);
-}
-
-static void
-genOmpAtomicWrite(Fortran::lower::AbstractConverter &converter,
-                  Fortran::lower::pft::Evaluation &eval,
-                  const Fortran::parser::OmpAtomicWrite &atomicWrite) {
-  // Get the value and address of atomic write operands.
-  const Fortran::parser::OmpAtomicClauseList &rightHandClauseList =
-      std::get<2>(atomicWrite.t);
-  const Fortran::parser::OmpAtomicClauseList &leftHandClauseList =
-      std::get<0>(atomicWrite.t);
-  const Fortran::parser::AssignmentStmt &stmt =
-      std::get<3>(atomicWrite.t).statement;
-  const Fortran::evaluate::Assignment &assign = *stmt.typedAssignment->v;
-  Fortran::lower::StatementContext stmtCtx;
-  // Get the value and address of atomic write operands.
-  mlir::Value rhsExpr =
-      fir::getBase(converter.genExprValue(assign.rhs, stmtCtx));
-  mlir::Value lhsAddr =
-      fir::getBase(converter.genExprAddr(assign.lhs, stmtCtx));
-  genOmpAtomicWriteStatement(converter, eval, lhsAddr, rhsExpr,
-                             &leftHandClauseList, &rightHandClauseList);
-}
-
-static void genOmpAtomicRead(Fortran::lower::AbstractConverter &converter,
-                             Fortran::lower::pft::Evaluation &eval,
-                             const Fortran::parser::OmpAtomicRead &atomicRead) {
-  // Get the address of atomic read operands.
-  const Fortran::parser::OmpAtomicClauseList &rightHandClauseList =
-      std::get<2>(atomicRead.t);
-  const Fortran::parser::OmpAtomicClauseList &leftHandClauseList =
-      std::get<0>(atomicRead.t);
-  const auto &assignmentStmtExpr =
-      std::get<Fortran::parser::Expr>(std::get<3>(atomicRead.t).statement.t);
-  const auto &assignmentStmtVariable = std::get<Fortran::parser::Variable>(
-      std::get<3>(atomicRead.t).statement.t);
-
-  Fortran::lower::StatementContext stmtCtx;
-  const Fortran::semantics::SomeExpr &fromExpr =
-      *Fortran::semantics::GetExpr(assignmentStmtExpr);
-  mlir::Type elementType = converter.genType(fromExpr);
-  mlir::Value fromAddress =
-      fir::getBase(converter.genExprAddr(fromExpr, stmtCtx));
-  mlir::Value toAddress = fir::getBase(converter.genExprAddr(
-      *Fortran::semantics::GetExpr(assignmentStmtVariable), stmtCtx));
-  genOmpAtomicCaptureStatement(converter, eval, fromAddress, toAddress,
-                               &leftHandClauseList, &rightHandClauseList,
-                               elementType);
-}
-
-static void
-genOmpAtomicUpdate(Fortran::lower::AbstractConverter &converter,
-                   Fortran::lower::pft::Evaluation &eval,
-                   const Fortran::parser::OmpAtomicUpdate &atomicUpdate) {
-  const Fortran::parser::OmpAtomicClauseList &rightHandClauseList =
-      std::get<2>(atomicUpdate.t);
-  const Fortran::parser::OmpAtomicClauseList &leftHandClauseList =
-      std::get<0>(atomicUpdate.t);
-  const auto &assignmentStmtExpr =
-      std::get<Fortran::parser::Expr>(std::get<3>(atomicUpdate.t).statement.t);
-  const auto &assignmentStmtVariable = std::get<Fortran::parser::Variable>(
-      std::get<3>(atomicUpdate.t).statement.t);
-
-  Fortran::lower::StatementContext stmtCtx;
-  mlir::Value lhsAddr = fir::getBase(converter.genExprAddr(
-      *Fortran::semantics::GetExpr(assignmentStmtVariable), stmtCtx));
-  mlir::Type varType =
-      fir::getBase(
-          converter.genExprValue(
-              *Fortran::semantics::GetExpr(assignmentStmtVariable), stmtCtx))
-          .getType();
-  genOmpAtomicUpdateStatement(converter, eval, lhsAddr, varType,
-                              assignmentStmtVariable, assignmentStmtExpr,
-                              &leftHandClauseList, &rightHandClauseList);
-}
-
-static void genOmpAtomic(Fortran::lower::AbstractConverter &converter,
-                         Fortran::lower::pft::Evaluation &eval,
-                         const Fortran::parser::OmpAtomic &atomicConstruct) {
-  const Fortran::parser::OmpAtomicClauseList &atomicClauseList =
-      std::get<Fortran::parser::OmpAtomicClauseList>(atomicConstruct.t);
-  const auto &assignmentStmtExpr = std::get<Fortran::parser::Expr>(
-      std::get<Fortran::parser::Statement<Fortran::parser::AssignmentStmt>>(
-          atomicConstruct.t)
-          .statement.t);
-  const auto &assignmentStmtVariable = std::get<Fortran::parser::Variable>(
-      std::get<Fortran::parser::Statement<Fortran::parser::AssignmentStmt>>(
-          atomicConstruct.t)
-          .statement.t);
-  Fortran::lower::StatementContext stmtCtx;
-  mlir::Value lhsAddr = fir::getBase(converter.genExprAddr(
-      *Fortran::semantics::GetExpr(assignmentStmtVariable), stmtCtx));
-  mlir::Type varType =
-      fir::getBase(
-          converter.genExprValue(
-              *Fortran::semantics::GetExpr(assignmentStmtVariable), stmtCtx))
-          .getType();
-  // If atomic-clause is not present on the construct, the behaviour is as if
-  // the update clause is specified
-  genOmpAtomicUpdateStatement(converter, eval, lhsAddr, varType,
-                              assignmentStmtVariable, assignmentStmtExpr,
-                              &atomicClauseList, nullptr);
-}
-
-static void
-genOmpAtomicCapture(Fortran::lower::AbstractConverter &converter,
-                    Fortran::lower::pft::Evaluation &eval,
-                    const Fortran::parser::OmpAtomicCapture &atomicCapture) {
-  fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
-  mlir::Location currentLocation = converter.getCurrentLocation();
-
-  mlir::IntegerAttr hint = nullptr;
-  mlir::omp::ClauseMemoryOrderKindAttr memoryOrder = nullptr;
-  const Fortran::parser::OmpAtomicClauseList &rightHandClauseList =
-      std::get<2>(atomicCapture.t);
-  const Fortran::parser::OmpAtomicClauseList &leftHandClauseList =
-      std::get<0>(atomicCapture.t);
-  genOmpAtomicHintAndMemoryOrderClauses(converter, leftHandClauseList, hint,
-                                        memoryOrder);
-  genOmpAtomicHintAndMemoryOrderClauses(converter, rightHandClauseList, hint,
-                                        memoryOrder);
-
-  const Fortran::parser::AssignmentStmt &stmt1 =
-      std::get<3>(atomicCapture.t).v.statement;
-  const auto &stmt1Var{std::get<Fortran::parser::Variable>(stmt1.t)};
-  const auto &stmt1Expr{std::get<Fortran::parser::Expr>(stmt1.t)};
-  const Fortran::parser::AssignmentStmt &stmt2 =
-      std::get<4>(atomicCapture.t).v.statement;
-  const auto &stmt2Var{std::get<Fortran::parser::Variable>(stmt2.t)};
-  const auto &stmt2Expr{std::get<Fortran::parser::Expr>(stmt2.t)};
-
-  // Pre-evaluate expressions to be used in the various operations inside
-  // `omp.atomic.capture` since it is not desirable to have anything other than
-  // a `omp.atomic.read`, `omp.atomic.write`, or `omp.atomic.update` operation
-  // inside `omp.atomic.capture`
-  Fortran::lower::StatementContext stmtCtx;
-  mlir::Value stmt1LHSArg, stmt1RHSArg, stmt2LHSArg, stmt2RHSArg;
-  mlir::Type elementType;
-  // LHS evaluations are common to all combinations of `omp.atomic.capture`
-  stmt1LHSArg = fir::getBase(
-      converter.genExprAddr(*Fortran::semantics::GetExpr(stmt1Var), stmtCtx));
-  stmt2LHSArg = fir::getBase(
-      converter.genExprAddr(*Fortran::semantics::GetExpr(stmt2Var), stmtCtx));
-
-  // Operation specific RHS evaluations
-  if (checkForSingleVariableOnRHS(stmt1)) {
-    // Atomic capture construct is of the form [capture-stmt, update-stmt] or
-    // of the form [capture-stmt, write-stmt]
-    stmt1RHSArg = fir::getBase(converter.genExprAddr(
-        *Fortran::semantics::GetExpr(stmt1Expr), stmtCtx));
-    stmt2RHSArg = fir::getBase(converter.genExprValue(
-        *Fortran::semantics::GetExpr(stmt2Expr), stmtCtx));
-
-  } else {
-    // Atomic capture construct is of the form [update-stmt, capture-stmt]
-    stmt1RHSArg = fir::getBase(converter.genExprValue(
-        *Fortran::semantics::GetExpr(stmt1Expr), stmtCtx));
-    stmt2RHSArg = fir::getBase(converter.genExprAddr(
-        *Fortran::semantics::GetExpr(stmt2Expr), stmtCtx));
-  }
-  // Type information used in generation of `omp.atomic.update` operation
-  mlir::Type stmt1VarType =
-      fir::getBase(converter.genExprValue(
-                       *Fortran::semantics::GetExpr(stmt1Var), stmtCtx))
-          .getType();
-  mlir::Type stmt2VarType =
-      fir::getBase(converter.genExprValue(
-                       *Fortran::semantics::GetExpr(stmt2Var), stmtCtx))
-          .getType();
-
-  auto atomicCaptureOp = firOpBuilder.create<mlir::omp::AtomicCaptureOp>(
-      currentLocation, hint, memoryOrder);
-  firOpBuilder.createBlock(&atomicCaptureOp.getRegion());
-  mlir::Block &block = atomicCaptureOp.getRegion().back();
-  firOpBuilder.setInsertionPointToStart(&block);
-  if (checkForSingleVariableOnRHS(stmt1)) {
-    if (checkForSymbolMatch(stmt2)) {
-      // Atomic capture construct is of the form [capture-stmt, update-stmt]
-      const Fortran::semantics::SomeExpr &fromExpr =
-          *Fortran::semantics::GetExpr(stmt1Expr);
-      elementType = converter.genType(fromExpr);
-      genOmpAtomicCaptureStatement(converter, eval, stmt1RHSArg, stmt1LHSArg,
-                                   /*leftHandClauseList=*/nullptr,
-                                   /*rightHandClauseList=*/nullptr,
-                                   elementType);
-      genOmpAtomicUpdateStatement(converter, eval, stmt1RHSArg, stmt2VarType,
-                                  stmt2Var, stmt2Expr,
-                                  /*leftHandClauseList=*/nullptr,
-                                  /*rightHandClauseList=*/nullptr);
-    } else {
-      // Atomic capture construct is of the form [capture-stmt, write-stmt]
-      const Fortran::semantics::SomeExpr &fromExpr =
-          *Fortran::semantics::GetExpr(stmt1Expr);
-      elementType = converter.genType(fromExpr);
-      genOmpAtomicCaptureStatement(converter, eval, stmt1RHSArg, stmt1LHSArg,
-                                   /*leftHandClauseList=*/nullptr,
-                                   /*rightHandClauseList=*/nullptr,
-                                   elementType);
-      genOmpAtomicWriteStatement(converter, eval, stmt1RHSArg, stmt2RHSArg,
-                                 /*leftHandClauseList=*/nullptr,
-                                 /*rightHandClauseList=*/nullptr);
-    }
-  } else {
-    // Atomic capture construct is of the form [update-stmt, capture-stmt]
-    firOpBuilder.setInsertionPointToEnd(&block);
-    const Fortran::semantics::SomeExpr &fromExpr =
-        *Fortran::semantics::GetExpr(stmt2Expr);
-    elementType = converter.genType(fromExpr);
-    genOmpAtomicCaptureStatement(converter, eval, stmt1LHSArg, stmt2LHSArg,
-                                 /*leftHandClauseList=*/nullptr,
-                                 /*rightHandClauseList=*/nullptr, elementType);
-    firOpBuilder.setInsertionPointToStart(&block);
-    genOmpAtomicUpdateStatement(converter, eval, stmt1LHSArg, stmt1VarType,
-                                stmt1Var, stmt1Expr,
-                                /*leftHandClauseList=*/nullptr,
-                                /*rightHandClauseList=*/nullptr);
-  }
-  firOpBuilder.setInsertionPointToEnd(&block);
-  firOpBuilder.create<mlir::omp::TerminatorOp>(currentLocation);
-  firOpBuilder.setInsertionPointToStart(&block);
-}
-
 static void
 genOMP(Fortran::lower::AbstractConverter &converter,
        Fortran::lower::pft::Evaluation &eval,
        const Fortran::parser::OpenMPAtomicConstruct &atomicConstruct) {
-  std::visit(Fortran::common::visitors{
-                 [&](const Fortran::parser::OmpAtomicRead &atomicRead) {
-                   genOmpAtomicRead(converter, eval, atomicRead);
-                 },
-                 [&](const Fortran::parser::OmpAtomicWrite &atomicWrite) {
-                   genOmpAtomicWrite(converter, eval, atomicWrite);
-                 },
-                 [&](const Fortran::parser::OmpAtomic &atomicConstruct) {
-                   genOmpAtomic(converter, eval, atomicConstruct);
-                 },
-                 [&](const Fortran::parser::OmpAtomicUpdate &atomicUpdate) {
-                   genOmpAtomicUpdate(converter, eval, atomicUpdate);
-                 },
-                 [&](const Fortran::parser::OmpAtomicCapture &atomicCapture) {
-                   genOmpAtomicCapture(converter, eval, atomicCapture);
-                 },
-             },
-             atomicConstruct.u);
+  std::visit(
+      Fortran::common::visitors{
+            Fortran::lower::genOmpAccAtomicRead<
+                Fortran::parser::OmpAtomicRead,
+                Fortran::parser::OmpAtomicClauseList>(converter, eval,
+                                                      atomicRead);
+            Fortran::lower::genOmpAccAtomicWrite<
+                Fortran::parser::OmpAtomicWrite,
+                Fortran::parser::OmpAtomicClauseList>(converter, eval,
+                                                      atomicWrite);
+            Fortran::lower::genOmpAtomic<Fortran::parser::OmpAtomic,
+                                         Fortran::parser::OmpAtomicClauseList>(
+                converter, eval, atomicConstruct);
+            Fortran::lower::genOmpAccAtomicUpdate<
+                Fortran::parser::OmpAtomicUpdate,
+                Fortran::parser::OmpAtomicClauseList>(converter, eval,
+                                                      atomicUpdate);
+            Fortran::lower::genOmpAccAtomicCapture<
+                Fortran::parser::OmpAtomicCapture,
+                Fortran::parser::OmpAtomicClauseList>(converter, eval,
+                                                      atomicCapture);
 }
 
 static void genOMP(Fortran::lower::AbstractConverter &converter,
                    Fortran::lower::pft::Evaluation &eval,
                    const Fortran::parser::OpenMPDeclareTargetConstruct
                        &declareTargetConstruct) {
   llvm::SmallVector<DeclareTargetCapturePair, 0> symbolAndClause;
   mlir::ModuleOp mod = converter.getFirOpBuilder().getModule();

flang/test/Lower/OpenACC/acc-atomic-capture.f90

+! RUN: %flang_fc1 -emit-fir -fopenacc %s -o - | FileCheck %s
+
+! This test checks the lowering of atomic capture
+
+program acc_atomic_capture_test                                                                                                     
+    integer :: x, y
+
+!CHECK: %[[X:.*]] = fir.alloca i32 {bindc_name = "x", uniq_name = "_QFEx"}
+!CHECK: %[[Y:.*]] = fir.alloca i32 {bindc_name = "y", uniq_name = "_QFEy"}
+!CHECK: acc.atomic.capture {
+!CHECK: acc.atomic.read %[[X]] = %[[Y]] : !fir.ref<i32>
+!CHECK: acc.atomic.update %[[Y]] : !fir.ref<i32> {
+!CHECK: ^bb0(%[[ARG:.*]]: i32):
+!CHECK: %[[temp:.*]] = fir.load %[[X]] : !fir.ref<i32>
+!CHECK: %[[result:.*]] = arith.addi %[[temp]], %[[ARG]] : i32
+!CHECK: acc.yield %[[result]] : i32
+!CHECK: }
+!CHECK: }
+
+    !$acc atomic capture
+        x = y
+        y = x + y
+    !$acc end atomic
+
+
+!CHECK: acc.atomic.capture {
+!CHECK: acc.atomic.update %[[Y]] : !fir.ref<i32> {
+!CHECK: ^bb0(%[[ARG:.*]]: i32):
+!CHECK: %[[temp:.*]] = fir.load %[[X]] : !fir.ref<i32>
+!CHECK: %[[result:.*]] = arith.muli %[[temp]], %[[ARG]] : i32
+!CHECK: acc.yield %[[result]] : i32
+!CHECK: }
+!CHECK: acc.atomic.read %[[X]] = %[[Y]] : !fir.ref<i32>
+!CHECK: }
+
+    !$acc atomic capture
+        y = x * y 
+        x = y
+    !$acc end atomic
+
+!CHECK: %[[constant_20:.*]] = arith.constant 20 : i32
+!CHECK: %[[constant_8:.*]] = arith.constant 8 : i32
+!CHECK: %[[temp:.*]] = fir.load %[[X]] : !fir.ref<i32>
+!CHECK: %[[result:.*]] = arith.subi %[[constant_8]], %[[temp]] : i32
+!CHECK: %[[result_noreassoc:.*]] = fir.no_reassoc %[[result]] : i32
+!CHECK: %[[result:.*]] = arith.addi %[[constant_20]], %[[result_noreassoc]] : i32
+!CHECK: acc.atomic.capture {
+!CHECK: acc.atomic.read %[[X]] = %[[Y]] : !fir.ref<i32>
+!CHECK: acc.atomic.write %[[Y]] = %[[result]] : !fir.ref<i32>, i32
+!CHECK: }
+
+    !$acc atomic capture
+        x = y
+        y = 2 * 10 + (8 - x) 
+    !$acc end atomic 
+end program
+
+
+
+subroutine pointers_in_atomic_capture()
+!CHECK: %[[A:.*]] = fir.alloca !fir.box<!fir.ptr<i32>> {bindc_name = "a", uniq_name = "_QFpointers_in_atomic_captureEa"}
+!CHECK: {{.*}} = fir.zero_bits !fir.ptr<i32>
+!CHECK: {{.*}} = fir.embox {{.*}} : (!fir.ptr<i32>) -> !fir.box<!fir.ptr<i32>>
+!CHECK: fir.store {{.*}} to %[[A]] : !fir.ref<!fir.box<!fir.ptr<i32>>>
+!CHECK: %[[B:.*]] = fir.alloca !fir.box<!fir.ptr<i32>> {bindc_name = "b", uniq_name = "_QFpointers_in_atomic_captureEb"}
+!CHECK: {{.*}} = fir.zero_bits !fir.ptr<i32>
+!CHECK: {{.*}} = fir.embox {{.*}} : (!fir.ptr<i32>) -> !fir.box<!fir.ptr<i32>>
+!CHECK: fir.store {{.*}} to %[[B]] : !fir.ref<!fir.box<!fir.ptr<i32>>>
+!CHECK: %[[C:.*]] = fir.alloca i32 {bindc_name = "c", fir.target, uniq_name = "_QFpointers_in_atomic_captureEc"}
+!CHECK: %[[D:.*]] = fir.alloca i32 {bindc_name = "d", fir.target, uniq_name = "_QFpointers_in_atomic_captureEd"}
+!CHECK: {{.*}} = fir.embox {{.*}} : (!fir.ref<i32>) -> !fir.box<!fir.ptr<i32>>
+!CHECK: fir.store {{.*}} to %[[A]] : !fir.ref<!fir.box<!fir.ptr<i32>>>
+!CHECK: {{.*}} = fir.embox {{.*}} : (!fir.ref<i32>) -> !fir.box<!fir.ptr<i32>>
+!CHECK: fir.store {{.*}} to %[[B]] : !fir.ref<!fir.box<!fir.ptr<i32>>>
+!CHECK: %[[loaded_A:.*]] = fir.load %[[A]] : !fir.ref<!fir.box<!fir.ptr<i32>>>
+!CHECK: %[[loaded_A_addr:.*]] = fir.box_addr %[[loaded_A]] : (!fir.box<!fir.ptr<i32>>) -> !fir.ptr<i32>
+!CHECK: %[[loaded_B:.*]] = fir.load %[[B]] : !fir.ref<!fir.box<!fir.ptr<i32>>>
+!CHECK: %[[loaded_B_addr:.*]] = fir.box_addr %[[loaded_B]] : (!fir.box<!fir.ptr<i32>>) -> !fir.ptr<i32>
+!CHECK: acc.atomic.capture   {
+!CHECK: acc.atomic.update %[[loaded_A_addr]] : !fir.ptr<i32> {
+!CHECK: ^bb0(%[[ARG:.*]]: i32):
+!CHECK: %[[PRIVATE_LOADED_B:.*]] = fir.load %[[B]] : !fir.ref<!fir.box<!fir.ptr<i32>>>
+!CHECK: %[[PRIVATE_LOADED_B_addr:.*]] = fir.box_addr %[[PRIVATE_LOADED_B]] : (!fir.box<!fir.ptr<i32>>) -> !fir.ptr<i32>
+!CHECK: %[[loaded_value:.*]] = fir.load %[[PRIVATE_LOADED_B_addr]] : !fir.ptr<i32>
+!CHECK: %[[result:.*]] = arith.addi %[[ARG]], %[[loaded_value]] : i32
+!CHECK: acc.yield %[[result]] : i32
+!CHECK: }
+!CHECK: acc.atomic.read %[[loaded_B_addr]] = %[[loaded_A_addr]] : !fir.ptr<i32>, i32
+!CHECK: }
+    integer, pointer :: a, b
+    integer, target :: c, d
+    a=>c
+    b=>d
+
+    !$acc atomic capture
+        a = a + b
+        b = a
+    !$acc end atomic
+end subroutine

flang/test/Lower/OpenACC/acc-atomic-read.f90

+! RUN: bbc --use-desc-for-alloc=false -fopenacc -emit-fir %s -o - | FileCheck %s
+
+! This test checks the lowering of atomic read
+
+!CHECK: func @_QQmain() attributes {fir.bindc_name = "acc_atomic_test"} {
+!CHECK: %[[VAR_G:.*]] = fir.alloca f32 {bindc_name = "g", uniq_name = "_QFEg"}
+!CHECK: %[[VAR_H:.*]] = fir.alloca f32 {bindc_name = "h", uniq_name = "_QFEh"}
+!CHECK: acc.atomic.read %[[VAR_G]] = %[[VAR_H]] : !fir.ref<f32>, f32
+!CHECK: return
+!CHECK: }
+
+program acc_atomic_test
+    real g, h
+    !$acc atomic read
+       g = h
+end program acc_atomic_test
+
+! Test lowering atomic read for pointer variables.
+! Please notice to use %[[VAL_4]] and %[[VAL_1]] for operands of atomic
+! operation, instead of %[[VAL_3]] and %[[VAL_0]].
+
+!CHECK-LABEL: func.func @_QPatomic_read_pointer() {
+!CHECK:         %[[VAL_0:.*]] = fir.alloca !fir.box<!fir.ptr<i32>> {bindc_name = "x", uniq_name = "_QFatomic_read_pointerEx"}
+!CHECK:         %[[VAL_1:.*]] = fir.alloca !fir.ptr<i32> {uniq_name = "_QFatomic_read_pointerEx.addr"}
+!CHECK:         %[[VAL_2:.*]] = fir.zero_bits !fir.ptr<i32>
+!CHECK:         fir.store %[[VAL_2]] to %[[VAL_1]] : !fir.ref<!fir.ptr<i32>>
+!CHECK:         %[[VAL_3:.*]] = fir.alloca !fir.box<!fir.ptr<i32>> {bindc_name = "y", uniq_name = "_QFatomic_read_pointerEy"}
+!CHECK:         %[[VAL_4:.*]] = fir.alloca !fir.ptr<i32> {uniq_name = "_QFatomic_read_pointerEy.addr"}
+!CHECK:         %[[VAL_5:.*]] = fir.zero_bits !fir.ptr<i32>
+!CHECK:         fir.store %[[VAL_5]] to %[[VAL_4]] : !fir.ref<!fir.ptr<i32>>
+!CHECK:         %[[VAL_6:.*]] = fir.load %[[VAL_1]] : !fir.ref<!fir.ptr<i32>>
+!CHECK:         %[[VAL_7:.*]] = fir.load %[[VAL_4]] : !fir.ref<!fir.ptr<i32>>
+!CHECK:         acc.atomic.read %[[VAL_7]] = %[[VAL_6]]   : !fir.ptr<i32>, i32
+!CHECK:         %[[VAL_8:.*]] = fir.load %[[VAL_4]] : !fir.ref<!fir.ptr<i32>>
+!CHECK:         %[[VAL_9:.*]] = fir.load %[[VAL_8]] : !fir.ptr<i32>
+!CHECK:         %[[VAL_10:.*]] = fir.load %[[VAL_1]] : !fir.ref<!fir.ptr<i32>>
+!CHECK:         fir.store %[[VAL_9]] to %[[VAL_10]] : !fir.ptr<i32>
+!CHECK:         return
+!CHECK:       }
+
+subroutine atomic_read_pointer()
+  integer, pointer :: x, y
+
+  !$acc atomic read
+    y = x
+
+  x = y
+end
+

flang/test/Lower/OpenACC/acc-atomic-update.f90

+! This test checks lowering of atomic and atomic update constructs
+! RUN: bbc --use-desc-for-alloc=false -fopenacc -emit-fir %s -o - | FileCheck %s
+! RUN: %flang_fc1 -mllvm --use-desc-for-alloc=false -emit-fir -fopenacc %s -o - | FileCheck %s
+
+program acc_atomic_update_test
+    integer :: x, y, z
+    integer, pointer :: a, b
+    integer, target :: c, d
+    integer(1) :: i1
+
+    a=>c
+    b=>d
+
+!CHECK: func.func @_QQmain() attributes {fir.bindc_name = "acc_atomic_update_test"} {
+!CHECK: %[[A:.*]] = fir.alloca !fir.box<!fir.ptr<i32>> {bindc_name = "a", uniq_name = "_QFEa"}
+!CHECK: %[[A_ADDR:.*]] = fir.alloca !fir.ptr<i32> {uniq_name = "_QFEa.addr"}
+!CHECK: %{{.*}} = fir.zero_bits !fir.ptr<i32>
+!CHECK: fir.store %{{.*}} to %[[A_ADDR]] : !fir.ref<!fir.ptr<i32>>
+!CHECK: %[[B:.*]] = fir.alloca !fir.box<!fir.ptr<i32>> {bindc_name = "b", uniq_name = "_QFEb"}
+!CHECK: %[[B_ADDR:.*]] = fir.alloca !fir.ptr<i32> {uniq_name = "_QFEb.addr"}
+!CHECK: %{{.*}} = fir.zero_bits !fir.ptr<i32>
+!CHECK: fir.store %{{.*}} to %[[B_ADDR]] : !fir.ref<!fir.ptr<i32>>
+!CHECK: %[[C_ADDR:.*]] = fir.address_of(@_QFEc) : !fir.ref<i32>
+!CHECK: %[[D_ADDR:.*]] = fir.address_of(@_QFEd) : !fir.ref<i32>
+!CHECK: %[[I1:.*]] = fir.alloca i8 {bindc_name = "i1", uniq_name = "_QFEi1"}
+!CHECK: %[[X:.*]] = fir.alloca i32 {bindc_name = "x", uniq_name = "_QFEx"}
+!CHECK: %[[Y:.*]] = fir.alloca i32 {bindc_name = "y", uniq_name = "_QFEy"}
+!CHECK: %[[Z:.*]] = fir.alloca i32 {bindc_name = "z", uniq_name = "_QFEz"}
+!CHECK: %{{.*}} = fir.convert %[[C_ADDR]] : (!fir.ref<i32>) -> !fir.ptr<i32>
+!CHECK: fir.store %{{.*}} to %[[A_ADDR]] : !fir.ref<!fir.ptr<i32>>
+!CHECK: %{{.*}} = fir.convert %[[D_ADDR]] : (!fir.ref<i32>) -> !fir.ptr<i32>
+!CHECK: fir.store {{.*}} to %[[B_ADDR]] : !fir.ref<!fir.ptr<i32>>
+!CHECK: %[[LOADED_A:.*]] = fir.load %[[A_ADDR]] : !fir.ref<!fir.ptr<i32>>
+!CHECK:  acc.atomic.update   %[[LOADED_A]] : !fir.ptr<i32> {
+!CHECK:  ^bb0(%[[ARG:.*]]: i32):
+!CHECK:    %[[LOADED_B:.*]] = fir.load %[[B_ADDR]] : !fir.ref<!fir.ptr<i32>>
+!CHECK:    %{{.*}} = fir.load %[[LOADED_B]] : !fir.ptr<i32>
+!CHECK:    %[[RESULT:.*]] = arith.addi %[[ARG]], %{{.*}} : i32
+!CHECK:    acc.yield %[[RESULT]] : i32
+!CHECK: }
+    !$acc atomic update
+        a = a + b 
+
+!CHECK: acc.atomic.update   %[[Y]] : !fir.ref<i32> {
+!CHECK:  ^bb0(%[[ARG:.*]]: i32):
+!CHECK:    {{.*}} = arith.constant 1 : i32
+!CHECK:    %[[RESULT:.*]] = arith.addi %[[ARG]], {{.*}} : i32
+!CHECK:    acc.yield %[[RESULT]] : i32
+!CHECK:  }
+!CHECK:  acc.atomic.update   %[[Z]] : !fir.ref<i32> {
+!CHECK:  ^bb0(%[[ARG:.*]]: i32):
+!CHECK:    %[[LOADED_X:.*]] = fir.load %[[X]] : !fir.ref<i32>
+!CHECK:    %[[RESULT:.*]] = arith.muli %[[LOADED_X]], %[[ARG]] : i32
+!CHECK:    acc.yield %[[RESULT]] : i32
+!CHECK:  }
+    !$acc atomic 
+        y = y + 1
+    !$acc atomic update
+        z = x * z 
+
+!CHECK:  acc.atomic.update   %[[I1]] : !fir.ref<i8> {
+!CHECK:  ^bb0(%[[VAL:.*]]: i8):
+!CHECK:    %[[CVT_VAL:.*]] = fir.convert %[[VAL]] : (i8) -> i32
+!CHECK:    %[[C1_VAL:.*]] = arith.constant 1 : i32
+!CHECK:    %[[ADD_VAL:.*]] = arith.addi %[[CVT_VAL]], %[[C1_VAL]] : i32
+!CHECK:    %[[UPDATED_VAL:.*]] = fir.convert %[[ADD_VAL]] : (i32) -> i8
+!CHECK:    acc.yield %[[UPDATED_VAL]] : i8
+!CHECK:  }
+    !$acc atomic
+      i1 = i1 + 1
+    !$acc end atomic
+!CHECK:  return
+!CHECK: }
+end program acc_atomic_update_test

flang/test/Lower/OpenACC/acc-atomic-write.f90

+! RUN: bbc --use-desc-for-alloc=false -fopenacc -emit-fir %s -o - | FileCheck %s
+
+! This test checks the lowering of atomic write
+
+!CHECK: func @_QQmain() attributes {fir.bindc_name = "acc_atomic_write_test"} {
+!CHECK: %[[VAR_X:.*]] = fir.alloca i32 {bindc_name = "x", uniq_name = "_QFEx"}
+!CHECK: %[[VAR_Y:.*]] = fir.alloca i32 {bindc_name = "y", uniq_name = "_QFEy"}
+!CHECK: %[[CONST_7:.*]] = arith.constant 7 : i32
+!CHECK: {{.*}} = fir.load %[[VAR_Y]] : !fir.ref<i32>
+!CHECK: %[[VAR_7y:.*]] = arith.muli %[[CONST_7]], {{.*}} : i32
+!CHECK: acc.atomic.write %[[VAR_X]] = %[[VAR_7y]] : !fir.ref<i32>, i32
+!CHECK: return
+!CHECK: }
+
+program acc_atomic_write_test
+    integer :: x, y
+
+    !$acc atomic write
+        x = 7 * y
+
+end program acc_atomic_write_test
+
+! Test lowering atomic read for pointer variables.
+
+!CHECK-LABEL: func.func @_QPatomic_write_pointer() {
+!CHECK:         %[[VAL_0:.*]] = fir.alloca !fir.box<!fir.ptr<i32>> {bindc_name = "x", uniq_name = "_QFatomic_write_pointerEx"}
+!CHECK:         %[[VAL_1:.*]] = fir.alloca !fir.ptr<i32> {uniq_name = "_QFatomic_write_pointerEx.addr"}
+!CHECK:         %[[VAL_2:.*]] = fir.zero_bits !fir.ptr<i32>
+!CHECK:         fir.store %[[VAL_2]] to %[[VAL_1]] : !fir.ref<!fir.ptr<i32>>
+!CHECK:         %[[VAL_3:.*]] = arith.constant 1 : i32
+!CHECK:         %[[VAL_4:.*]] = fir.load %[[VAL_1]] : !fir.ref<!fir.ptr<i32>>
+!CHECK:         acc.atomic.write %[[VAL_4]] = %[[VAL_3]]   : !fir.ptr<i32>, i32
+!CHECK:         %[[VAL_5:.*]] = arith.constant 2 : i32
+!CHECK:         %[[VAL_6:.*]] = fir.load %[[VAL_1]] : !fir.ref<!fir.ptr<i32>>
+!CHECK:         fir.store %[[VAL_5]] to %[[VAL_6]] : !fir.ptr<i32>
+!CHECK:         return
+!CHECK:       }
+
+subroutine atomic_write_pointer()
+  integer, pointer :: x
+
+  !$acc atomic write
+    x = 1
+
+  x = 2
+end subroutine
+
+!CHECK-LABEL: func.func @_QPatomic_write_typed_assign
+!CHECK: %[[VAR:.*]] = fir.alloca f32 {bindc_name = "r2", uniq_name = "{{.*}}r2"}
+!CHECK: %[[CST:.*]] = arith.constant 0.000000e+00 : f32
+!CHECK: acc.atomic.write %[[VAR]] = %[[CST]]   : !fir.ref<f32>, f32
+
+subroutine atomic_write_typed_assign
+  real :: r2
+  !$acc atomic write
+  r2 = 0
+end subroutine