[Flang][OpenMP] Use typed assignment in Atomic Write lowering

Use typed assignment in Atomic Write lowering to better handle
type conversions of allowed types.

Note: We should make similar changes for other constructs in
later patches.

Reviewed By: NimishMishra

Differential Revision: https://reviews.llvm.org/D154163

Author: kiranchandramohan

flang/lib/Lower/OpenMP.cpp

@@ -2358,16 +2358,16 @@ genOmpAtomicWrite(Fortran::lower::AbstractConverter &converter,
       std::get<2>(atomicWrite.t);
   const Fortran::parser::OmpAtomicClauseList &leftHandClauseList =
       std::get<0>(atomicWrite.t);
-  const auto &assignmentStmtExpr =
-      std::get<Fortran::parser::Expr>(std::get<3>(atomicWrite.t).statement.t);
-  const auto &assignmentStmtVariable = std::get<Fortran::parser::Variable>(
-      std::get<3>(atomicWrite.t).statement.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 rhs_expr = fir::getBase(converter.genExprValue(
-      *Fortran::semantics::GetExpr(assignmentStmtExpr), stmtCtx));
-  mlir::Value lhs_addr = fir::getBase(converter.genExprAddr(
-      *Fortran::semantics::GetExpr(assignmentStmtVariable), stmtCtx));
+  mlir::Value rhs_expr =
+      fir::getBase(converter.genExprValue(assign.rhs, stmtCtx));
+
+  mlir::Value lhs_addr =
+      fir::getBase(converter.genExprAddr(assign.lhs, stmtCtx));
   genOmpAtomicWriteStatement(converter, eval, lhs_addr, rhs_expr,
                              &leftHandClauseList, &rightHandClauseList);
 }

flang/test/Lower/OpenMP/atomic-write.f90

@@ -63,3 +63,13 @@ subroutine atomic_write_pointer()
   x = 2
 end
 
+!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: omp.atomic.write %[[VAR]] = %[[CST]]   : !fir.ref<f32>, f32
+
+subroutine atomic_write_typed_assign
+  real :: r2
+  !$omp atomic write
+  r2 = 0
+end subroutine