[openacc] Add dialect representation for atomic operations

[razvanlupusoru]

Summary

The OpenACC standard specifies an atomic construct in section 2.12
(of 3.3 spec), used to ensure that a specific location is accessed or
updated atomically. Four different clauses are allowed: read, write,
update, or capture. If no clause appears, it is as if update
is used.

The OpenMP specification defines the same clauses for omp atomic. The
types of expression and the clauses in the OpenACC spec match the OpenMP
spec exactly. The main difference is that the OpenMP specification is a
superset - it includes clauses for hint and memory order. It also
allows conditional expression statements. But otherwise, the expression
definition matches.

Thus, for OpenACC, we reuse the OpenMP implementation as follows:

• The dialect operations are duplicated and the not needed parts removed

(hint and memory order). The goal is no dependence or coupling between
the two dialects so they can continue to evolve without hindrance from
the other. - The verification operations are also duplicated with the
not needed parts removed. This part could possibly be shared by having a
common verifier pieces that are templated. Having it shared could be
beneficial when fixing errors or adding further constraints - the
disadvantage is that it would introduce some coupling on design.

The frontend lowering necessary to generate the dialect can be reused.
This will be done in a follow up change.

[clementval]

mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td

1273–1283

What is the use of these? I don't see any use in the lowering as well. Same question for the other ops.

[razvanlupusoru]

mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td

1273–1283

My intent was to ensure that the operations provide same-ish interface as the OpenMP ones. But I agree this is not needed for now.

[clementval]

mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td

1273–1283

If needed it would be better to achieve this through a shared interface attached to the op itself.

[razvanlupusoru]

mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td

1273–1283

Great idea. I actually like the idea of this approach since it doesn't couple the two dialects! Nice. I will investigate what needs shared early next week.

[razvanlupusoru]

Rebase. Remove several unused methods from new operations.

[clementval]

Looks mostly good for me. Just one comment about the empty verifier.

mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp

1101

It's better to disable the verifier in the op declaration so we don't need an empty verifier.

[razvanlupusoru]

Merged via github PR instead: https://github.com/llvm/llvm-project/pull/65493

mlir/include/mlir/Dialect/OpenACC/OpenACC.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_OPENACC_OPENACC_H_
 #define MLIR_DIALECT_OPENACC_OPENACC_H_
 
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/SymbolTable.h"
 
 #include "mlir/Bytecode/BytecodeOpInterface.h"
 #include "mlir/Dialect/OpenACC/OpenACCOpsDialect.h.inc"
 #include "mlir/Dialect/OpenACC/OpenACCOpsEnums.h.inc"
 #include "mlir/Dialect/OpenACC/OpenACCTypeInterfaces.h.inc"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"

mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td

let assemblyFormat = [{
     attr-dict-with-keyword
   }];
 
   let hasVerifier = 1;
 }
 
 // Yield operation for the acc.loop and acc.parallel operations.
 def OpenACC_YieldOp : OpenACC_Op<"yield", [ReturnLike, Terminator,
-    ParentOneOf<["FirstprivateRecipeOp, LoopOp, ParallelOp, PrivateRecipeOp, ReductionRecipeOp, SerialOp"]>]> {
+    ParentOneOf<["FirstprivateRecipeOp, LoopOp, ParallelOp, PrivateRecipeOp,"
+                 "ReductionRecipeOp, SerialOp, AtomicUpdateOp"]>]> {
   let summary = "Acc yield and termination operation";
 
   let description = [{
     `acc.yield` is a special terminator operation for block inside regions in
-    acc ops (parallel and loop). It returns values to the immediately enclosing
-    acc op.
+    various acc ops (including parallel, loop, atomic.update). It returns values
+    to the immediately enclosing acc op.
   }];
 
   let arguments = (ins Variadic<AnyType>:$operands);
 
   let builders = [OpBuilder<(ins), [{ /* nothing to do */ }]>];
 
   let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
 }
 
 //===----------------------------------------------------------------------===//
+// 2.12 atomic construct
+//===----------------------------------------------------------------------===//
+
+def AtomicReadOp : OpenACC_Op<"atomic.read", [AllTypesMatch<["x", "v"]>]> {
+
+  let summary = "performs an atomic read";
+
+  let description = [{
+    This operation performs an atomic read.
+
+    The operand `x` is the address from where the value is atomically read.
+    The operand `v` is the address where the value is stored after reading.
+  }];
+
+  let arguments = (ins OpenACC_PointerLikeType:$x,
+                       OpenACC_PointerLikeType:$v,
+                       TypeAttr:$element_type);
+  let assemblyFormat = [{
+    $v `=` $x
+    `:` type($x) `,` $element_type attr-dict
+  }];
+  let hasVerifier = 1;
+}
+
+def AtomicWriteOp : OpenACC_Op<"atomic.write"> {
+
+  let summary = "performs an atomic write";
+
+  let description = [{
+    This operation performs an atomic write.
+
+    The operand `address` is the address to where the `value` is atomically
+    written w.r.t. multiple threads. The evaluation of `value` need not be
+    atomic w.r.t. the write to address. In general, the type(address) must
+    dereference to type(value).
+  }];
+
+  let arguments = (ins OpenACC_PointerLikeType:$address,
+                       AnyType:$value);
+  let assemblyFormat = [{
+    $address `=` $value
+    `:` type($address) `,` type($value)
+    attr-dict
+  }];
+  let hasVerifier = 1;
+}
+
+def AtomicUpdateOp : OpenACC_Op<"atomic.update",
+                               [SingleBlockImplicitTerminator<"YieldOp">,
+                                RecursiveMemoryEffects]> {
+
+  let summary = "performs an atomic update";
+
+  let description = [{
+    This operation performs an atomic update.
+
+    The operand `x` is exactly the same as the operand `x` in the OpenACC
+    Standard (OpenACC 3.3, section 2.12). It is the address of the variable
+    that is being updated. `x` is atomically read/written.
+
+    The region describes how to update the value of `x`. It takes the value at
+    `x` as an input and must yield the updated value. Only the update to `x` is
+    atomic. Generally the region must have only one instruction, but can
+    potentially have more than one instructions too. The update is sematically
+    similar to a compare-exchange loop based atomic update.
+
+    The syntax of atomic update operation is different from atomic read and
+    atomic write operations. This is because only the host dialect knows how to
+    appropriately update a value. For example, while generating LLVM IR, if
+    there are no special `atomicrmw` instructions for the operation-type
+    combination in atomic update, a compare-exchange loop is generated, where
+    the core update operation is directly translated like regular operations by
+    the host dialect. The front-end must handle semantic checks for allowed
+    operations.
+  }];
+
+  let arguments = (ins Arg<OpenACC_PointerLikeType,
+                           "Address of variable to be updated",
+                           [MemRead, MemWrite]>:$x);
+  let regions = (region SizedRegion<1>:$region);
+  let assemblyFormat = [{
+    $x `:` type($x) $region attr-dict
+  }];
+  let hasVerifier = 1;
+  let hasRegionVerifier = 1;
+  let hasCanonicalizeMethod = 1;
+  let extraClassDeclaration = [{
+    Operation* getFirstOp() {
+      return &getRegion().front().getOperations().front();
+    }
+
+    /// Returns true if the new value is same as old value and the operation is
+    /// a no-op, false otherwise.
+    bool isNoOp();
+
+    /// Returns the new value if the operation is equivalent to just a write
+    /// operation. Otherwise, returns nullptr.
+    Value getWriteOpVal();
+  }];
+}
+
+def AtomicCaptureOp : OpenACC_Op<"atomic.capture",
+    [SingleBlockImplicitTerminator<"TerminatorOp">]> {
+  let summary = "performs an atomic capture";
+  let description = [{
+    This operation performs an atomic capture.
+
+    The region has the following allowed forms:
+
+    ```
+      acc.atomic.capture {
+        acc.atomic.update ...
+        acc.atomic.read ...
+        acc.terminator
+      }
+
+      acc.atomic.capture {
+        acc.atomic.read ...
+        acc.atomic.update ...
+        acc.terminator
+      }
+
+      acc.atomic.capture {
+        acc.atomic.read ...
+        acc.atomic.write ...
+        acc.terminator
+      }
+    ```
+
+  }];
+
+  let regions = (region SizedRegion<1>:$region);
+  let assemblyFormat = [{
+    $region attr-dict
+  }];
+  let hasRegionVerifier = 1;
+  let hasVerifier = 1;
+  let extraClassDeclaration = [{
+    /// Returns the first operation in atomic capture region
+    Operation* getFirstOp();
+
+    /// Returns the second operation in atomic capture region
+    Operation* getSecondOp();
+
+    /// Returns the `atomic.read` operation inside the region, if any.
+    /// Otherwise, it returns nullptr.
+    AtomicReadOp getAtomicReadOp();
+
+    /// Returns the `atomic.write` operation inside the region, if any.
+    /// Otherwise, it returns nullptr.
+    AtomicWriteOp getAtomicWriteOp();
+
+    /// Returns the `atomic.update` operation inside the region, if any.
+    /// Otherwise, it returns nullptr.
+    AtomicUpdateOp getAtomicUpdateOp();
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // 2.13 Declare Directive
 //===----------------------------------------------------------------------===//
 
 def OpenACC_DeclareEnterOp : OpenACC_Op<"declare_enter", []> {
   let summary = "declare directive - entry to implicit data region";
 
   let description = [{
     The "acc.declare_enter" operation represents the OpenACC declare directive

mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp

 
 using namespace mlir;
 using namespace acc;
 
 #include "mlir/Dialect/OpenACC/OpenACCOpsDialect.cpp.inc"
 #include "mlir/Dialect/OpenACC/OpenACCOpsEnums.cpp.inc"
 #include "mlir/Dialect/OpenACC/OpenACCTypeInterfaces.cpp.inc"
 
+namespace {
+/// Model for pointer-like types that already provide a `getElementType` method.
+template <typename T>
+struct PointerLikeModel
+    : public PointerLikeType::ExternalModel<PointerLikeModel<T>, T> {
+  Type getElementType(Type pointer) const {
+    return llvm::cast<T>(pointer).getElementType();
+  }
+};
+} // namespace
+
 //===----------------------------------------------------------------------===//
 // OpenACC operations
 //===----------------------------------------------------------------------===//
 
 void OpenACCDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/OpenACC/OpenACCOps.cpp.inc"
       >();
   addAttributes<
 #define GET_ATTRDEF_LIST
 #include "mlir/Dialect/OpenACC/OpenACCOpsAttributes.cpp.inc"
       >();
   addTypes<
 #define GET_TYPEDEF_LIST
 #include "mlir/Dialect/OpenACC/OpenACCOpsTypes.cpp.inc"
       >();
 
   // By attaching interfaces here, we make the OpenACC dialect dependent on
   // the other dialects. This is probably better than having dialects like LLVM
   // and memref be dependent on OpenACC.
-  LLVM::LLVMPointerType::attachInterface<PointerLikeType>(*getContext());
-  MemRefType::attachInterface<PointerLikeType>(*getContext());
+  LLVM::LLVMPointerType::attachInterface<
+      PointerLikeModel<LLVM::LLVMPointerType>>(*getContext());
+  MemRefType::attachInterface<PointerLikeModel<MemRefType>>(*getContext());
 }
 
 //===----------------------------------------------------------------------===//
 // DataBoundsOp
 //===----------------------------------------------------------------------===//
 LogicalResult acc::DataBoundsOp::verify() {
   auto extent = getExtent();
   auto upperbound = getUpperbound();

 }
 
 void EnterDataOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                               MLIRContext *context) {
   results.add<RemoveConstantIfCondition<EnterDataOp>>(context);
 }
 
 //===----------------------------------------------------------------------===//
+// AtomicReadOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult AtomicReadOp::verify() {
+  if (getX() == getV())
+    return emitError(
+        "read and write must not be to the same location for atomic reads");
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// AtomicWriteOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult AtomicWriteOp::verify() {
+  Type elementType = getAddress().getType().getElementType();
+  if (elementType && elementType != getValue().getType())
+    return emitError("address must dereference to value type");
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// AtomicUpdateOp
+//===----------------------------------------------------------------------===//
+
+bool AtomicUpdateOp::isNoOp() {
+  YieldOp yieldOp = dyn_cast<YieldOp>(getFirstOp());
+  return (yieldOp &&
+          yieldOp.getOperands().front() == getRegion().front().getArgument(0));
+}
+
+Value AtomicUpdateOp::getWriteOpVal() {
+  YieldOp yieldOp = dyn_cast<YieldOp>(getFirstOp());
+  if (yieldOp &&
+      yieldOp.getOperands().front() != getRegion().front().getArgument(0))
+    return yieldOp.getOperands().front();
+  return nullptr;
+}
+
+LogicalResult AtomicUpdateOp::canonicalize(AtomicUpdateOp op,
+                                           PatternRewriter &rewriter) {
+  if (op.isNoOp()) {
+    rewriter.eraseOp(op);
+    return success();
+  }
+
+  if (Value writeVal = op.getWriteOpVal()) {
+    rewriter.replaceOpWithNewOp<AtomicWriteOp>(op, op.getX(), writeVal);
+    return success();
+  }
+
+  return failure();
+}
+
+LogicalResult AtomicUpdateOp::verify() {
+  if (getRegion().getNumArguments() != 1)
+    return emitError("the region must accept exactly one argument");
+
+  Type elementType = getX().getType().getElementType();
+  if (elementType && elementType != getRegion().getArgument(0).getType()) {
+    return emitError("the type of the operand must be a pointer type whose "
+                     "element type is the same as that of the region argument");
+  }
+
+  return success();
+}
+
+LogicalResult AtomicUpdateOp::verifyRegions() {
+  YieldOp yieldOp = *getRegion().getOps<YieldOp>().begin();
+
+  if (yieldOp.getOperands().size() != 1)
+    return emitError("only updated value must be returned");
+  if (yieldOp.getOperands().front().getType() !=
+      getRegion().getArgument(0).getType())
+    return emitError("input and yielded value must have the same type");
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// AtomicCaptureOp
+//===----------------------------------------------------------------------===//
+
+Operation *AtomicCaptureOp::getFirstOp() {
+  return &getRegion().front().getOperations().front();
+}
+
+Operation *AtomicCaptureOp::getSecondOp() {
+  auto &ops = getRegion().front().getOperations();
+  return ops.getNextNode(ops.front());
+}
+
+AtomicReadOp AtomicCaptureOp::getAtomicReadOp() {
+  if (auto op = dyn_cast<AtomicReadOp>(getFirstOp()))
+    return op;
+  return dyn_cast<AtomicReadOp>(getSecondOp());
+}
+
+AtomicWriteOp AtomicCaptureOp::getAtomicWriteOp() {
+  if (auto op = dyn_cast<AtomicWriteOp>(getFirstOp()))
+    return op;
+  return dyn_cast<AtomicWriteOp>(getSecondOp());
+}
+
+AtomicUpdateOp AtomicCaptureOp::getAtomicUpdateOp() {
+  if (auto op = dyn_cast<AtomicUpdateOp>(getFirstOp()))
+    return op;
+  return dyn_cast<AtomicUpdateOp>(getSecondOp());
+}
+
+LogicalResult AtomicCaptureOp::verify() { return success(); }
+
+LogicalResult AtomicCaptureOp::verifyRegions() {
+  Block::OpListType &ops = getRegion().front().getOperations();
+  if (ops.size() != 3)
+    return emitError()
+           << "expected three operations in acc.atomic.capture region (one "
+              "terminator, and two atomic ops)";
+  auto &firstOp = ops.front();
+  auto &secondOp = *ops.getNextNode(firstOp);
+  auto firstReadStmt = dyn_cast<AtomicReadOp>(firstOp);
+  auto firstUpdateStmt = dyn_cast<AtomicUpdateOp>(firstOp);
+  auto secondReadStmt = dyn_cast<AtomicReadOp>(secondOp);
+  auto secondUpdateStmt = dyn_cast<AtomicUpdateOp>(secondOp);
+  auto secondWriteStmt = dyn_cast<AtomicWriteOp>(secondOp);
+
+  if (!((firstUpdateStmt && secondReadStmt) ||
+        (firstReadStmt && secondUpdateStmt) ||
+        (firstReadStmt && secondWriteStmt)))
+    return ops.front().emitError()
+           << "invalid sequence of operations in the capture region";
+  if (firstUpdateStmt && secondReadStmt &&
+      firstUpdateStmt.getX() != secondReadStmt.getX())
+    return firstUpdateStmt.emitError()
+           << "updated variable in acc.atomic.update must be captured in "
+              "second operation";
+  if (firstReadStmt && secondUpdateStmt &&
+      firstReadStmt.getX() != secondUpdateStmt.getX())
+    return firstReadStmt.emitError()
+           << "captured variable in acc.atomic.read must be updated in second "
+              "operation";
+  if (firstReadStmt && secondWriteStmt &&
+      firstReadStmt.getX() != secondWriteStmt.getAddress())
+    return firstReadStmt.emitError()
+           << "captured variable in acc.atomic.read must be updated in "
+              "second operation";
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // DeclareEnterOp
 //===----------------------------------------------------------------------===//
 
 template <typename Op>
 static LogicalResult checkDeclareOperands(Op &op,
                                           const mlir::ValueRange &operands) {
   if (operands.empty())
     return emitError(

mlir/test/Dialect/OpenACC/ops.mlir

 // CHECK: [[IDXVALUE:%.*]] = arith.constant 1 : index
 // CHECK: [[IFCOND:%.*]] = arith.constant true
 // CHECK: acc.set device_type([[I32VALUE]] : i32)
 // CHECK: acc.set device_num([[I64VALUE]] : i64)
 // CHECK: acc.set device_num([[I32VALUE]] : i32)
 // CHECK: acc.set device_num([[IDXVALUE]] : index)
 // CHECK: acc.set device_num([[IDXVALUE]] : index) if([[IFCOND]])
 // CHECK: acc.set default_async([[I32VALUE]] : i32)
+
+// -----
+
+// CHECK-LABEL: func.func @acc_atomic_read
+// CHECK-SAME: (%[[v:.*]]: memref<i32>, %[[x:.*]]: memref<i32>)
+func.func @acc_atomic_read(%v: memref<i32>, %x: memref<i32>) {
+  // CHECK: acc.atomic.read %[[v]] = %[[x]] : memref<i32>, i32
+  acc.atomic.read %v = %x : memref<i32>, i32
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func.func @acc_atomic_write
+// CHECK-SAME: (%[[ADDR:.*]]: memref<i32>, %[[VAL:.*]]: i32)
+func.func @acc_atomic_write(%addr : memref<i32>, %val : i32) {
+  // CHECK: acc.atomic.write %[[ADDR]] = %[[VAL]] : memref<i32>, i32
+  acc.atomic.write %addr = %val : memref<i32>, i32
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func.func @acc_atomic_update
+// CHECK-SAME: (%[[X:.*]]: memref<i32>, %[[EXPR:.*]]: i32, %[[XBOOL:.*]]: memref<i1>, %[[EXPRBOOL:.*]]: i1)
+func.func @acc_atomic_update(%x : memref<i32>, %expr : i32, %xBool : memref<i1>, %exprBool : i1) {
+  // CHECK: acc.atomic.update %[[X]] : memref<i32>
+  // CHECK-NEXT: (%[[XVAL:.*]]: i32):
+  // CHECK-NEXT:   %[[NEWVAL:.*]] = llvm.add %[[XVAL]], %[[EXPR]] : i32
+  // CHECK-NEXT:   acc.yield %[[NEWVAL]] : i32
+  acc.atomic.update %x : memref<i32> {
+  ^bb0(%xval: i32):
+    %newval = llvm.add %xval, %expr : i32
+    acc.yield %newval : i32
+  }
+  // CHECK: acc.atomic.update %[[XBOOL]] : memref<i1>
+  // CHECK-NEXT: (%[[XVAL:.*]]: i1):
+  // CHECK-NEXT:   %[[NEWVAL:.*]] = llvm.and %[[XVAL]], %[[EXPRBOOL]] : i1
+  // CHECK-NEXT:   acc.yield %[[NEWVAL]] : i1
+  acc.atomic.update %xBool : memref<i1> {
+  ^bb0(%xval: i1):
+    %newval = llvm.and %xval, %exprBool : i1
+    acc.yield %newval : i1
+  }
+  // CHECK: acc.atomic.update %[[X]] : memref<i32>
+  // CHECK-NEXT: (%[[XVAL:.*]]: i32):
+  // CHECK-NEXT:   %[[NEWVAL:.*]] = llvm.shl %[[XVAL]], %[[EXPR]] : i32
+  // CHECK-NEXT:   acc.yield %[[NEWVAL]] : i32
+  // CHECK-NEXT: }
+  acc.atomic.update %x : memref<i32> {
+  ^bb0(%xval: i32):
+    %newval = llvm.shl %xval, %expr : i32
+    acc.yield %newval : i32
+  }
+  // CHECK: acc.atomic.update %[[X]] : memref<i32>
+  // CHECK-NEXT: (%[[XVAL:.*]]: i32):
+  // CHECK-NEXT:   %[[NEWVAL:.*]] = llvm.intr.smax(%[[XVAL]], %[[EXPR]]) : (i32, i32) -> i32
+  // CHECK-NEXT:   acc.yield %[[NEWVAL]] : i32
+  // CHECK-NEXT: }
+  acc.atomic.update %x : memref<i32> {
+  ^bb0(%xval: i32):
+    %newval = llvm.intr.smax(%xval, %expr) : (i32, i32) -> i32
+    acc.yield %newval : i32
+  }
+
+  // CHECK: acc.atomic.update %[[XBOOL]] : memref<i1>
+  // CHECK-NEXT: (%[[XVAL:.*]]: i1):
+  // CHECK-NEXT:   %[[NEWVAL:.*]] = llvm.icmp "eq" %[[XVAL]], %[[EXPRBOOL]] : i1
+  // CHECK-NEXT:   acc.yield %[[NEWVAL]] : i1
+  // }
+  acc.atomic.update %xBool : memref<i1> {
+  ^bb0(%xval: i1):
+    %newval = llvm.icmp "eq" %xval, %exprBool : i1
+    acc.yield %newval : i1
+  }
+
+  // CHECK: acc.atomic.update %[[X]] : memref<i32> {
+  // CHECK-NEXT: (%[[XVAL:.*]]: i32):
+  // CHECK-NEXT:   acc.yield %[[XVAL]] : i32
+  // CHECK-NEXT: }
+  acc.atomic.update %x : memref<i32> {
+  ^bb0(%xval:i32):
+    acc.yield %xval : i32
+  }
+
+  // CHECK: acc.atomic.update %[[X]] : memref<i32> {
+  // CHECK-NEXT: (%[[XVAL:.*]]: i32):
+  // CHECK-NEXT:   acc.yield %{{.+}} : i32
+  // CHECK-NEXT: }
+  %const = arith.constant 42 : i32
+  acc.atomic.update %x : memref<i32> {
+  ^bb0(%xval:i32):
+    acc.yield %const : i32
+  }
+
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func.func @acc_atomic_capture
+// CHECK-SAME: (%[[v:.*]]: memref<i32>, %[[x:.*]]: memref<i32>, %[[expr:.*]]: i32)
+func.func @acc_atomic_capture(%v: memref<i32>, %x: memref<i32>, %expr: i32) {
+  // CHECK: acc.atomic.capture {
+  // CHECK-NEXT: acc.atomic.update %[[x]] : memref<i32>
+  // CHECK-NEXT: (%[[xval:.*]]: i32):
+  // CHECK-NEXT:   %[[newval:.*]] = llvm.add %[[xval]], %[[expr]] : i32
+  // CHECK-NEXT:   acc.yield %[[newval]] : i32
+  // CHECK-NEXT: }
+  // CHECK-NEXT: acc.atomic.read %[[v]] = %[[x]] : memref<i32>, i32
+  // CHECK-NEXT: }
+  acc.atomic.capture {
+    acc.atomic.update %x : memref<i32> {
+    ^bb0(%xval: i32):
+      %newval = llvm.add %xval, %expr : i32
+      acc.yield %newval : i32
+    }
+    acc.atomic.read %v = %x : memref<i32>, i32
+  }
+  // CHECK: acc.atomic.capture {
+  // CHECK-NEXT: acc.atomic.read %[[v]] = %[[x]] : memref<i32>, i32
+  // CHECK-NEXT: acc.atomic.update %[[x]] : memref<i32>
+  // CHECK-NEXT: (%[[xval:.*]]: i32):
+  // CHECK-NEXT:   %[[newval:.*]] = llvm.add %[[xval]], %[[expr]] : i32
+  // CHECK-NEXT:   acc.yield %[[newval]] : i32
+  // CHECK-NEXT: }
+  // CHECK-NEXT: }
+  acc.atomic.capture {
+    acc.atomic.read %v = %x : memref<i32>, i32
+    acc.atomic.update %x : memref<i32> {
+    ^bb0(%xval: i32):
+      %newval = llvm.add %xval, %expr : i32
+      acc.yield %newval : i32
+    }
+  }
+  // CHECK: acc.atomic.capture {
+  // CHECK-NEXT: acc.atomic.read %[[v]] = %[[x]] : memref<i32>, i32
+  // CHECK-NEXT: acc.atomic.write %[[x]] = %[[expr]] : memref<i32>, i32
+  // CHECK-NEXT: }
+  acc.atomic.capture {
+    acc.atomic.read %v = %x : memref<i32>, i32
+    acc.atomic.write %x = %expr : memref<i32>, i32
+  }
+
+  return
+}