[OpenACC][CIR] Implement atomic-write lowering (llvm#164627)

This is a slightly more complicated variant of this, which supports 'x =
expr', so the right hand side is an r-value. This patch implements that,
adds some tests, and does some minor refactoring to the infrastructure
added for the 'atomic read' to make it more flexible for 'write'.

This is the second of four 'atomic' kinds.

Author: erichkeane

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' have 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 2 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;
+}