[OpenACC][CIR] Implement atomic-write lowering

clang/include/clang/AST/StmtOpenACC.h

@@ -821,6 +821,7 @@ class OpenACCAtomicConstruct final
   struct StmtInfo {
     const Expr *V;
     const Expr *X;
+    const Expr *Expr;
     // TODO: OpenACC: We should expand this as we're implementing the other
     // atomic construct kinds.
   };

clang/lib/AST/StmtOpenACC.cpp

@@ -324,6 +324,18 @@ OpenACCAtomicConstruct *OpenACCAtomicConstruct::Create(
   return Inst;
 }
 
+static std::pair<const Expr *, const Expr *> getBinaryOpArgs(const Expr *Op) {
+  if (const auto *BO = dyn_cast<BinaryOperator>(Op)) {
+    assert(BO->getOpcode() == BO_Assign);
+    return {BO->getLHS(), BO->getRHS()};
+  }
+
+  const auto *OO = cast<CXXOperatorCallExpr>(Op);
+  assert(OO->getOperator() == OO_Equal);
+
+  return {OO->getArg(0), OO->getArg(1)};
+}
+
 const OpenACCAtomicConstruct::StmtInfo
 OpenACCAtomicConstruct::getAssociatedStmtInfo() const {
   // This ends up being a vastly simplified version of SemaOpenACCAtomic, since
@@ -333,27 +345,35 @@ OpenACCAtomicConstruct::getAssociatedStmtInfo() const {
 
   switch (AtomicKind) {
   case OpenACCAtomicKind::None:
-  case OpenACCAtomicKind::Write:
   case OpenACCAtomicKind::Update:
   case OpenACCAtomicKind::Capture:
-    assert(false && "Only 'read' has been implemented here");
+    assert(false && "Only 'read'/'write' has been implemented here");
     return {};
   case OpenACCAtomicKind::Read: {
     // Read only supports the format 'v = x'; where both sides are a scalar
     // expression. This can come in 2 forms; BinaryOperator or
     // CXXOperatorCallExpr (rarely).
-    const Expr *AssignExpr = cast<const Expr>(getAssociatedStmt());
-    if (const auto *BO = dyn_cast<BinaryOperator>(AssignExpr)) {
-      assert(BO->getOpcode() == BO_Assign);
-      return {BO->getLHS()->IgnoreImpCasts(), BO->getRHS()->IgnoreImpCasts()};
-    }
-
-    const auto *OO = cast<CXXOperatorCallExpr>(AssignExpr);
-    assert(OO->getOperator() == OO_Equal);
-
-    return {OO->getArg(0)->IgnoreImpCasts(), OO->getArg(1)->IgnoreImpCasts()};
+    std::pair<const Expr *, const Expr *> BinaryArgs =
+        getBinaryOpArgs(cast<const Expr>(getAssociatedStmt()));
+    // We want the L-value for each side, so we ignore implicit casts.
+    return {BinaryArgs.first->IgnoreImpCasts(),
+            BinaryArgs.second->IgnoreImpCasts(), /*expr=*/nullptr};
   }
+  case OpenACCAtomicKind::Write: {
+    // Write supports only the format 'x = expr', where the expression is scalar
+    // type, and 'x' is a scalar l value. As above, this can come in 2 forms;
+    // Binary Operator or CXXOperatorCallExpr.
+    std::pair<const Expr *, const Expr *> BinaryArgs =
+        getBinaryOpArgs(cast<const Expr>(getAssociatedStmt()));
+    // We want the L-value for ONLY the X side, so we ignore implicit casts. For
+    // the right side (the expr), we emit it as an r-value so we need to
+    // maintain implicit casts.
+    return {/*v=*/nullptr, BinaryArgs.first->IgnoreImpCasts(),
+            BinaryArgs.second};
   }
+  }
+
+  llvm_unreachable("unknown OpenACC atomic kind");
 }
 
 OpenACCCacheConstruct *OpenACCCacheConstruct::CreateEmpty(const ASTContext &C,

clang/lib/CIR/CodeGen/CIRGenStmtOpenACC.cpp

@@ -306,9 +306,10 @@ CIRGenFunction::emitOpenACCCacheConstruct(const OpenACCCacheConstruct &s) {
 
 mlir::LogicalResult
 CIRGenFunction::emitOpenACCAtomicConstruct(const OpenACCAtomicConstruct &s) {
-  // For now, we are only support 'read', so diagnose. We can switch on the kind
-  // later once we start implementing the other 3 forms.
-  if (s.getAtomicKind() != OpenACCAtomicKind::Read) {
+  // For now, we are only support 'read'/'write', so diagnose. We can switch on
+  // the kind later once we start implementing the other 3 forms. While we
+  if (s.getAtomicKind() != OpenACCAtomicKind::Read &&
+      s.getAtomicKind() != OpenACCAtomicKind::Write) {
     cgm.errorNYI(s.getSourceRange(), "OpenACC Atomic Construct");
     return mlir::failure();
   }
@@ -318,17 +319,41 @@ CIRGenFunction::emitOpenACCAtomicConstruct(const OpenACCAtomicConstruct &s) {
   // it has custom emit logic.
   mlir::Location start = getLoc(s.getSourceRange().getBegin());
   OpenACCAtomicConstruct::StmtInfo inf = s.getAssociatedStmtInfo();
-  // Atomic 'read' only permits 'v = x', where v and x are both scalar L values.
-  // The getAssociatedStmtInfo strips off implicit casts, which includes
-  // implicit conversions and L-to-R-Value conversions, so we can just emit it
-  // as an L value.  The Flang implementation has no problem with different
-  // types, so it appears that the dialect can handle the conversions.
-  mlir::Value v = emitLValue(inf.V).getPointer();
-  mlir::Value x = emitLValue(inf.X).getPointer();
-  mlir::Type resTy = convertType(inf.V->getType());
-  auto op = mlir::acc::AtomicReadOp::create(builder, start, x, v, resTy,
-                                            /*ifCond=*/{});
-  emitOpenACCClauses(op, s.getDirectiveKind(), s.getDirectiveLoc(),
-                     s.clauses());
-  return mlir::success();
+
+  switch (s.getAtomicKind()) {
+  case OpenACCAtomicKind::None:
+  case OpenACCAtomicKind::Update:
+  case OpenACCAtomicKind::Capture:
+    llvm_unreachable("Unimplemented atomic construct type, should have "
+                     "diagnosed/returned above");
+    return mlir::failure();
+  case OpenACCAtomicKind::Read: {
+
+    // Atomic 'read' only permits 'v = x', where v and x are both scalar L
+    // values. The getAssociatedStmtInfo strips off implicit casts, which
+    // includes implicit conversions and L-to-R-Value conversions, so we can
+    // just emit it as an L value.  The Flang implementation has no problem with
+    // different types, so it appears that the dialect can handle the
+    // conversions.
+    mlir::Value v = emitLValue(inf.V).getPointer();
+    mlir::Value x = emitLValue(inf.X).getPointer();
+    mlir::Type resTy = convertType(inf.V->getType());
+    auto op = mlir::acc::AtomicReadOp::create(builder, start, x, v, resTy,
+                                              /*ifCond=*/{});
+    emitOpenACCClauses(op, s.getDirectiveKind(), s.getDirectiveLoc(),
+                       s.clauses());
+    return mlir::success();
+  }
+  case OpenACCAtomicKind::Write: {
+    mlir::Value x = emitLValue(inf.X).getPointer();
+    mlir::Value expr = emitAnyExpr(inf.Expr).getValue();
+    auto op = mlir::acc::AtomicWriteOp::create(builder, start, x, expr,
+                                               /*ifCond=*/{});
+    emitOpenACCClauses(op, s.getDirectiveKind(), s.getDirectiveLoc(),
+                       s.clauses());
+    return mlir::success();
+  }
+  }
+
+  llvm_unreachable("unknown OpenACC atomic kind");
 }

clang/test/CIR/CodeGenOpenACC/atomic-write.cpp

@@ -0,0 +1,55 @@
+// RUN: %clang_cc1 -fopenacc -triple x86_64-linux-gnu -Wno-openacc-self-if-potential-conflict -emit-cir -fclangir -triple x86_64-linux-pc %s -o - | FileCheck %s
+
+extern "C" bool condition(int x, unsigned int y, float f);
+extern "C" double do_thing(float f);
+
+struct ConvertsToScalar {
+  operator float();
+};
+
+void use(int x, unsigned int y, float f, ConvertsToScalar cts) {
+  // CHECK: cir.func{{.*}}(%[[X_ARG:.*]]: !s32i{{.*}}, %[[Y_ARG:.*]]: !u32i{{.*}}, %[[F_ARG:.*]]: !cir.float{{.*}}){{.*}}, %[[CTS_ARG:.*]]: !rec_ConvertsToScalar{{.*}}) {
+  // CHECK-NEXT: %[[X_ALLOC:.*]] = cir.alloca !s32i, !cir.ptr<!s32i>, ["x", init]
+  // CHECK-NEXT: %[[Y_ALLOC:.*]] = cir.alloca !u32i, !cir.ptr<!u32i>, ["y", init]
+  // CHECK-NEXT: %[[F_ALLOC:.*]] = cir.alloca !cir.float, !cir.ptr<!cir.float>, ["f", init]
+  // CHECK-NEXT: %[[CTS_ALLOC:.*]] = cir.alloca !rec_ConvertsToScalar, !cir.ptr<!rec_ConvertsToScalar>, ["cts", init]
+  //
+  // CHECK-NEXT: cir.store %[[X_ARG]], %[[X_ALLOC]] : !s32i, !cir.ptr<!s32i>
+  // CHECK-NEXT: cir.store %[[Y_ARG]], %[[Y_ALLOC]] : !u32i, !cir.ptr<!u32i>
+  // CHECK-NEXT: cir.store %[[F_ARG]], %[[F_ALLOC]] : !cir.float, !cir.ptr<!cir.float>
+  // CHECK-NEXT: cir.store %[[CTS_ARG]], %[[CTS_ALLOC]] : !rec_ConvertsToScalar, !cir.ptr<!rec_ConvertsToScalar>
+
+  // CHECK-NEXT: %[[Y_LOAD:.*]] = cir.load {{.*}}%[[Y_ALLOC]] : !cir.ptr<!u32i>, !u32i
+  // CHECK-NEXT: %[[Y_TO_FLOAT:.*]] = cir.cast int_to_float %[[Y_LOAD]] : !u32i -> !cir.float
+  // CHECK-NEXT: %[[F_LOAD:.*]] = cir.load {{.*}}%[[F_ALLOC]] : !cir.ptr<!cir.float>, !cir.float
+  // CHECK-NEXT: %[[MUL:.*]] = cir.binop(mul, %[[Y_TO_FLOAT]], %[[F_LOAD]]) : !cir.float
+  // CHECK-NEXT: %[[RHS_CAST:.*]] = cir.cast float_to_int %[[MUL]] : !cir.float -> !s32i
+  // CHECK-NEXT: acc.atomic.write %[[X_ALLOC]] = %[[RHS_CAST]] : !cir.ptr<!s32i>, !s32i
+#pragma acc atomic write
+  x = y * f;
+
+  // CHECK-NEXT: %[[F_LOAD:.*]] = cir.load {{.*}}%[[F_ALLOC]] : !cir.ptr<!cir.float>, !cir.float
+  // CHECK-NEXT: %[[CALL:.*]] = cir.call @do_thing(%[[F_LOAD]]) : (!cir.float) -> !cir.double
+  // CHECK-NEXT: %[[CALL_CAST:.*]] = cir.cast float_to_int %[[CALL]] : !cir.double -> !u32i
+  // CHECK-NEXT: acc.atomic.write %[[Y_ALLOC]] = %[[CALL_CAST]] : !cir.ptr<!u32i>, !u32i
+#pragma acc atomic write
+  y = do_thing(f);
+
+  // CHECK-NEXT: %[[X_LOAD:.*]] = cir.load {{.*}}%[[X_ALLOC]] : !cir.ptr<!s32i>, !s32i
+  // CHECK-NEXT: %[[X_CAST:.*]] = cir.cast int_to_float %[[X_LOAD]] : !s32i -> !cir.float
+  // CHECK-NEXT: %[[THING_CALL:.*]] = cir.call @do_thing(%[[X_CAST]]) : (!cir.float) -> !cir.double
+  // CHECK-NEXT: %[[THING_CAST:.*]] = cir.cast floating %[[THING_CALL]] : !cir.double -> !cir.float
+  // CHECK-NEXT: %[[X_LOAD:.*]] = cir.load {{.*}}%[[X_ALLOC]] : !cir.ptr<!s32i>, !s32i
+  // CHECK-NEXT: %[[Y_LOAD:.*]] = cir.load {{.*}}%[[Y_ALLOC]] : !cir.ptr<!u32i>, !u32i
+  // CHECK-NEXT: %[[F_LOAD:.*]] = cir.load {{.*}}%[[F_ALLOC]] : !cir.ptr<!cir.float>, !cir.float
+  // CHECK-NEXT: %[[COND_CALL:.*]] = cir.call @condition(%[[X_LOAD]], %[[Y_LOAD]], %[[F_LOAD]]) : (!s32i, !u32i, !cir.float) -> !cir.bool
+  // CHECK-NEXT: %[[COND_CAST:.*]] = builtin.unrealized_conversion_cast %[[COND_CALL]] : !cir.bool to i1
+  // CHECK-NEXT: acc.atomic.write if(%[[COND_CAST]]) %[[F_ALLOC]] = %[[THING_CAST]] : !cir.ptr<!cir.float>, !cir.float
+#pragma acc atomic write if (condition(x, y, f))
+  f = do_thing(x);
+
+  // CHECK-NEXT: %[[CTS_CONV_CALL:.*]] = cir.call @{{.*}}(%[[CTS_ALLOC]]) : (!cir.ptr<!rec_ConvertsToScalar>) -> !cir.float
+  // CHECK-NEXT: acc.atomic.write %[[F_ALLOC]] = %[[CTS_CONV_CALL]] : !cir.ptr<!cir.float>, !cir.float
+#pragma acc atomic write
+  f = cts;
+}