[Mlir-commits] [flang] [mlir] [Flang] [OpenMP] atomic compare (PR #184761)

[Krzysztof Parzyszek]

================
@@ -521,8 +542,121 @@ void Fortran::lower::omp::lowerAtomic(
     memOrder = makeValidForAction(memOrder, action0, action1, version);
 
   if (auto *cond = get(analysis.cond)) {
-    (void)cond;
-    TODO(loc, "OpenMP ATOMIC COMPARE");
+    // atomic compare: if (x == e) x = d
+    // e : expecteVal
+    // d : desiredVal
+
+    // Check for compound clauses (fail, capture, weak) that are not yet
+    // supported with atomic compare.
+    if (llvm::any_of(clauses, [](const omp::Clause &clause) {
+          return clause.id == llvm::omp::Clause::OMPC_fail ||
+                 clause.id == llvm::omp::Clause::OMPC_capture ||
+                 clause.id == llvm::omp::Clause::OMPC_weak;
+        })) {
+      TODO(loc, "Compound clauses of OpenMP ATOMIC COMPARE");
+    }
+
+    Fortran::common::RelationalOperator relOpr =
+        Fortran::common::RelationalOperator::EQ;
+    std::optional<semantics::SomeExpr> expectedExprStorage;
+
+    if (const auto *rel = Fortran::evaluate::UnwrapExpr<
+            Fortran::evaluate::Relational<Fortran::evaluate::SomeType>>(
+            *cond)) {
+      std::visit(
+          [&](const auto &relImpl) {
+            relOpr = relImpl.opr;
+            using Operand = typename std::decay_t<decltype(relImpl)>::Operand;
+            auto leftExpr = Fortran::evaluate::AsGenericExpr(
+                Fortran::evaluate::Expr<Operand>{relImpl.left()});
+            auto rightExpr = Fortran::evaluate::AsGenericExpr(
+                Fortran::evaluate::Expr<Operand>{relImpl.right()});
+            if (Fortran::evaluate::IsSameOrConvertOf(rightExpr, atom)) {
+              // e.g. e == x  (atom is on the right)
+              // left operand is expected value (e)
+              // reverse the operator so that the comparison becomes
+              // x <reversed-op> e.
+              expectedExprStorage = std::move(leftExpr);
+              relOpr = reverseRelOp(relOpr);
+            } else {
+              // Form: x == e  (atom is on the left, or default)
+              expectedExprStorage = std::move(rightExpr);
+            }
+          },
+          rel->u);
+    }
+    if (!expectedExprStorage) {
+      mlir::emitError(loc, "internal error: atomic compare condition is not a "
+                           "recognized relational expression");
+      return;
+    }
+
+    mlir::UnitAttr weakAttr = nullptr;
+    mlir::Operation *atomicOp = mlir::omp::AtomicCompareOp::create(
+        builder, loc, atomAddr, weakAttr, hint,
+        makeMemOrderAttr(converter, memOrder));
+    mlir::Type elemTypeOfX = fir::unwrapRefType(atomAddr.getType());
+    mlir::Block *block = builder.createBlock(&atomicOp->getRegion(0));
+    mlir::Value blockArg = block->addArgument(elemTypeOfX, loc);
+    builder.setInsertionPointToEnd(block);
+
+    mlir::Value expectedVal = fir::getBase(
+        converter.genExprValue(*expectedExprStorage, stmtCtx, &loc));
+    if (expectedVal.getType() != elemTypeOfX) {
+      expectedVal = builder.createConvert(loc, elemTypeOfX, expectedVal);
+    }
----------------
kparzysz wrote:

I think this should be done before the atomic op.

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