[flang-commits] [flang] [flang][OpenMP] Common lowering flow for atomic update (PR #69866)

[Kiran Chandramohan via flang-commits]

================
@@ -301,38 +287,36 @@ static inline void genOmpAccAtomicUpdateStatement(
   mlir::Value val =
       fir::getBase(atomicUpdateOp->getRegion(0).front().getArgument(0));
 
-  llvm::SmallVector<mlir::Operation *> ops;
-  for (mlir::Operation &op : tempOp.getRegion().getOps())
-    ops.push_back(&op);
-
-  // SCF Yield is converted to OMP Yield. All other operations are copied
-  for (mlir::Operation *op : ops) {
-    if (auto y = mlir::dyn_cast<mlir::scf::YieldOp>(op)) {
-      firOpBuilder.setInsertionPointToEnd(
-          &atomicUpdateOp->getRegion(0).front());
-      if constexpr (std::is_same<AtomicListT,
-                                 Fortran::parser::OmpAtomicClauseList>()) {
-        firOpBuilder.create<mlir::omp::YieldOp>(currentLocation,
-                                                y.getResults());
-      } else {
-        firOpBuilder.create<mlir::acc::YieldOp>(currentLocation,
-                                                y.getResults());
-      }
-      op->erase();
-    } else {
-      op->remove();
-      atomicUpdateOp->getRegion(0).front().push_back(op);
-    }
-  }
-
-  // Remove the load and replace all uses of load with the block argument
-  for (mlir::Operation &op : atomicUpdateOp->getRegion(0).getOps()) {
-    fir::LoadOp y = mlir::dyn_cast<fir::LoadOp>(&op);
-    if (y && y.getMemref() == updateVar)
-      y.getRes().replaceAllUsesWith(val);
+  mlir::Value op = nullptr;
+  if (std::get_if<Fortran::parser::Expr::Add>(&assignmentStmtExpr.u)) {
+    op = firOpBuilder.create<mlir::arith::AddIOp>(currentLocation, val,
+                                                  convertRhs);
+  } else if (std::get_if<Fortran::parser::Expr::Subtract>(
+                 &assignmentStmtExpr.u)) {
+    op = firOpBuilder.create<mlir::arith::SubIOp>(currentLocation, val,
+                                                  convertRhs);
+  } else if (std::get_if<Fortran::parser::Expr::Multiply>(
+                 &assignmentStmtExpr.u)) {
+    op = firOpBuilder.create<mlir::arith::MulIOp>(currentLocation, val,
+                                                  convertRhs);
+  } else if (std::get_if<Fortran::parser::Expr::Divide>(
+                 &assignmentStmtExpr.u)) {
+    op = firOpBuilder.create<mlir::arith::DivUIOp>(currentLocation, val,
+                                                   convertRhs);
+  } else if (std::get_if<Fortran::parser::Expr::AND>(&assignmentStmtExpr.u)) {
+    op = firOpBuilder.create<mlir::arith::AndIOp>(currentLocation, val,
+                                                  convertRhs);
+  } else if (std::get_if<Fortran::parser::Expr::OR>(&assignmentStmtExpr.u)) {
+    op = firOpBuilder.create<mlir::arith::OrIOp>(currentLocation, val,
+                                                 convertRhs);
+  } else if (std::get_if<Fortran::parser::Expr::EQV>(&assignmentStmtExpr.u)) {
+    op = firOpBuilder.create<mlir::arith::CmpIOp>(
+        currentLocation, mlir::arith::CmpIPredicate::eq, val, convertRhs);
+  } else if (std::get_if<Fortran::parser::Expr::NEQV>(&assignmentStmtExpr.u)) {
+    op = firOpBuilder.create<mlir::arith::CmpIOp>(
+        currentLocation, mlir::arith::CmpIPredicate::ne, val, convertRhs);
   }
-
-  tempOp.erase();
+  firOpBuilder.create<mlir::omp::YieldOp>(currentLocation, op);
----------------
kiranchandramohan wrote:

Should there be an `if` to decide between OpenMP and OpenACC Yields?

https://github.com/llvm/llvm-project/pull/69866