[mlir][OpenMP] Added omp.atomic.capture operation

[shraiysh]

Summary

This patch supports the atomic construct (capture) following section 2.17.7 of OpenMP 5.0 standard. Also added tests for the same.

[peixin]

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

714

For fortran code, expr here refers to both expr and expr_list.

718

The evaluation of `expr` need not be atomic w.r.t the read or write of the location designated by `x`.

This is for update clause. For capture clause, the standard states that Neither the evaluation of expr or expr_list, nor the write to the location designated by v, need be atomic with respect to the read or write of the location designated by x.

718

In general, type(x) must dereference to type(expr).

Not clear what this means.

720

Not only binary operation. Since this is used for both c/c++ and fortran, binop is the binary operator or intrinsic (MAX, MIN, IAND, IOR, or IEOR).

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

1504

The restriction states that If atomic-clause is update or not present then memory-order-clause must not be acq_rel or acquire, which is not for capture clause. One interesting thing is that the description in the standard states that If the write, update, or capture clause is specified and the release, acq_rel, or seq_cst clause is specified then the strong flush on entry to the atomic operation is also a release flush. Checked with OpenMP 5.1, and there is no this kind of restriction. Don't understand why acq_rel and acquire are not allowed in update or capture clause.

1517

Nit: How about x_type , v_type , expr_type?

1535

%x = %expr binop %x | %x = %x binop %expr would be clear than %x = ...

1569

The same as above.

1600

Nit: remove the brace.

1641

Remove the brace. https://llvm.org/docs/CodingStandards.html#don-t-use-braces-on-simple-single-statement-bodies-of-if-else-loop-statements

mlir/test/Dialect/OpenMP/ops.mlir

602

lines 593-602 ->

v = x = x binop expr;
x = x binop expr; v = x;

608

v = x = expr binop x;
x = expr binop x; v = x;

620

lines 613-620 ->

v = x; x = x binop expr;
v = x; x binop= expr;

634

Redundant test cases. Move the comment fortran expressions in front of the above test cases.

[shraiysh]

Thanks for the detailed review @peixin. I will work on this patch and address comments based on the discussion from patch D116396.

[shraiysh]

Updated patch to use omp.atomic.read and omp.atomic.update

[shraiysh]

Depends on D116396

[shraiysh]

Fix missing semi-colon error

[peixin]

This design is much easier to review. Nice job. LGTM.

[kiranchandramohan]

For discussion,
-> Generally, in a single block region, if the terminator has no special semantics then it is better to use SingleBlockImplicitTerminator. In this case more so since it is just a couple of instructions.
-> Would it be better to model to capture x, expr, v, y etc as operands of the capture operation?
-> Would it be better to model this is a complex operation without a region?

[peixin]

In D115851#3224776, @kiranchandramohan wrote:

-> Would it be better to model to capture x, expr, v, y etc as operands of the capture operation?
-> Would it be better to model this is a complex operation without a region?

The first version of this patch model the capture construct using one standalone operation (https://reviews.llvm.org/D115851?id=394750), but it is kind of hard to read.

[kiranchandramohan]

In D115851#3224817, @peixin wrote:

In D115851#3224776, @kiranchandramohan wrote:

-> Would it be better to model to capture x, expr, v, y etc as operands of the capture operation?
-> Would it be better to model this is a complex operation without a region?

The first version of this patch model the capture construct using one standalone operation (https://reviews.llvm.org/D115851?id=394750), but it is kind of hard to read.

OK. The only concern I have is that the operands of the operation are not captured in the operation definition. I don't know whether that will cause any issues. I think we can go ahead with this for now. Can you address the terminator comment?

[shraiysh]

Thank you for the reviews @kiranchandramohan and @peixin. I will address the terminator comment soon. I have been busy with some other work recently. Apologies for the delay.

[shraiysh]

Addressed SingleBlockImplicitTerminator comment.

[shraiysh]

Gentle ping!

[kiranchandramohan]

1. If the terminator is implicit will it create difficulties in lowering?
2. Should we update the examples in the description and the tests to be without the terminator?

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

1537

If the terminator is implicit, does this hold?

[shraiysh]

Thanks for the comments @kiranchandramohan. IIUC SingleBlockImplicitTerminator just adds an error message if the region does not end with omp.terminator. I thought it would accept the following code but it gives a missing omp.terminator error instead

omp.atomic.capture {
  omp.atomic.read ...
  omp.atomic.update ...
}

I think lowering should not be a problem because the in-memory structure still has an omp.terminator at the end of the region.

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

1537

Yes, it still holds (there is a test for this error message). I expected that with an implicit terminator, it will only have two, but this is not true. It looks like having SingleBlockImplicitTerminator just makes sure that we end a region with omp.terminator.

[kiranchandramohan]

Ahh, sorry about that. I kind of thought the terminator will not be visible.

LGTM.

[shraiysh]

Thanks! Rebase with main.

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

def AtomicUpdateOp : OpenMP_Op<"atomic.update"> {
   let arguments = (ins OpenMP_PointerLikeType:$x,
                        AnyType:$expr,
                        UnitAttr:$isXBinopExpr,
                        AtomicBinOpKindAttr:$binop,
                        DefaultValuedAttr<I64Attr, "0">:$hint,
                        OptionalAttr<MemoryOrderKind>:$memory_order);
   let parser = [{ return parseAtomicUpdateOp(parser, result); }];
   let printer = [{ return printAtomicUpdateOp(p, *this); }];
-  let verifier = [{ return verifyAtomicUpdateOp(*this); }];
+  let verifier = [{ return verifyAtomicUpdateOrCaptureOp(*this); }];
+}
+
+def AtomicCaptureOp : OpenMP_Op<"atomic.capture"> {
+
+  let description = "performs an atomic capture";
+
+  let summary = [{
+    This operation performs an atomic capture.
+
+    An atomic capture is one of the following forms:
+      %[[X]] = %[[X]] binop %[[expr]], %[[V]] = %[[X]]
+      %[[X]] = %[[expr]] binop %[[X]], %[[V]] = %[[X]]
+      %[[V]] = %[[X]], %[[X]] = %[[X]] binop %[[expr]]
+      %[[V]] = %[[X]], %[[X]] = %[[expr]] binop %[[X]]
+      %[[V]] = %[[X]], %[[X]] = %[[expr]]
+
+    The operands `x`, `v` and `expr` exactly the same as the operands `x` and
+    `expr` in the OpenMP Standard. The operand `x` is the address of the
+    variable that is being captured and updated. `x` is atomically
+    read/written. The evaluation of `expr` need not be atomic w.r.t the read or
+    write of the location designated by `x`. In general, type(x) must
+    dereference to type(expr).
+
+    The attribute `binop` is the binary operation being performed atomically.
+    It is an optional attribute - it is absent in the fifth form of this
+    operation.
+
+    The attribute `is_x_binop_expr` is
+      - true when the update expression is of the form `x binop expr` on RHS
+      - false when the update expression is of the form `expr binop x` on RHS
+
+    The attribute `is_postfix_update` is
+      - true when the capture is before update
+      - false when the capture is after update
+
+    `hint` is the value of hint (as used in the hint clause). It is a compile
+    time constant. As the name suggests, this is just a hint for optimization.
+
+    `memory_order` indicates the memory ordering behavior of the construct. It
+    can be one of `seq_cst`, `acq_rel`, `release`, `acquire` or `relaxed`.
+  }];
+  let arguments = (ins OpenMP_PointerLikeType:$x,
+                       OpenMP_PointerLikeType:$v,
+                       AnyType:$expr,
+                       OptionalAttr<AtomicBinOpKindAttr>:$binop,
+                       UnitAttr:$is_x_binop_expr,
+                       UnitAttr:$is_postfix_update,
+                       DefaultValuedAttr<I64Attr, "0">:$hint,
+                       OptionalAttr<MemoryOrderKind>:$memory_order);
+  let parser = [{ return parseAtomicCaptureOp(parser, result); }];
+  let printer = [{ return printAtomicCaptureOp(p, *this); }];
+  let verifier = [{ return verifyAtomicUpdateOrCaptureOp(*this); }];
 }
 
 //===----------------------------------------------------------------------===//
 // 2.19.5.7 declare reduction Directive
 //===----------------------------------------------------------------------===//
 
 def ReductionDeclareOp : OpenMP_Op<"reduction.declare", [Symbol]> {
   let summary = "declares a reduction kind";

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

 #include "mlir/Dialect/OpenMP/OpenMPOps.cpp.inc"
       >();
 
   LLVM::LLVMPointerType::attachInterface<
       PointerLikeModel<LLVM::LLVMPointerType>>(*getContext());
   MemRefType::attachInterface<PointerLikeModel<MemRefType>>(*getContext());
 }
 
+bool equal(OpAsmParser::OperandType &a, OpAsmParser::OperandType &b) {
+  return a.name == b.name && a.number == b.number;
+}
+
 //===----------------------------------------------------------------------===//
 // ParallelOp
 //===----------------------------------------------------------------------===//
 
 void ParallelOp::build(OpBuilder &builder, OperationState &state,
                        ArrayRef<NamedAttribute> attributes) {
   ParallelOp::build(
       builder, state, /*if_expr_var=*/nullptr, /*num_threads_var=*/nullptr,
p << y << " " << AtomicBinOpKindToString(op.binop()).lower() << " " << z
     << " ";
   if (op.memory_order())
     p << "memory_order(" << op.memory_order() << ") ";
   if (op.hintAttr())
     printSynchronizationHint(p, op, op.hintAttr());
   p << ": " << op.x().getType() << ", " << op.expr().getType();
 }
 
-/// Verifier for AtomicUpdateOp
-static LogicalResult verifyAtomicUpdateOp(AtomicUpdateOp op) {
-  if (op.memory_order()) {
-    StringRef memoryOrder = op.memory_order().getValue();
+/// Verifier for AtomicUpdateOp and AtomicCaptureOp
+static LogicalResult verifyAtomicUpdateOrCaptureOp(Operation *op) {
+  if (op->hasAttrOfType<StringAttr>("memory_order")) {
+    StringRef memoryOrder =
+        op->getAttrOfType<StringAttr>("memory_order").getValue();
     if (memoryOrder.equals("acq_rel") || memoryOrder.equals("acquire"))
-      return op.emitError(
-          "memory-order must not be acq_rel or acquire for atomic updates");
+      return op->emitError("memory-order must not be acq_rel or acquire for "
+                           "atomic update/capture");
+  }
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// AtomicCaptureOp
+//===----------------------------------------------------------------------===//
+
+/// Parse an atomic capture op
+/// operation ::= `omp.atomic.capture` expr, expr : types
+/// expr ::= operand `=` operand | operand `=` operand binop operand
+/// types ::= captured_variable_type `,` capturing_variable_type `,` expr_type
+static ParseResult parseAtomicCaptureOp(OpAsmParser &parser,
+                                        OperationState &result) {
+  SmallVector<ClauseType> clauses = {hintClause, memoryOrderClause};
+  SmallVector<int> segments;
+  // Parse expressions %a = %b [binop|,]
+  OpAsmParser::OperandType x, v, expr;
+  StringRef binOp;
+  bool isPostfixUpdate = false;
+  bool isXLhsInRhsPart = false;
+  bool hasBinop = false;
+
+  OpAsmParser::OperandType a, b;
+  if (parser.parseOperand(a) || parser.parseEqual() || parser.parseOperand(b))
+    return failure();
+
+  // The next operand can either be a comma or a binary operation.
+  if (succeeded(parser.parseOptionalComma())) {
+    // %v = %x, %x = ...
+    v = a;
+    isPostfixUpdate = true;
+    OpAsmParser::OperandType lhs, rhs;
+    if (parser.parseOperand(x) || parser.parseEqual() ||
+        parser.parseOperand(lhs))
+      return failure();
+    if (!equal(x, b))
+      return parser.emitError(parser.getCurrentLocation())
+             << "captured variable not found in LHS of second operation";
+    if (succeeded(parser.parseOptionalKeyword(&binOp))) {
+      hasBinop = true;
+      if (parser.parseOperand(rhs))
+        return failure();
+      if (equal(lhs, x)) {
+        isXLhsInRhsPart = true;
+        expr = rhs;
+      } else if (equal(rhs, x)) {
+        isXLhsInRhsPart = false;
+        expr = lhs;
+      } else {
+        return parser.emitError(parser.getCurrentLocation())
+               << "captured variable not found in RHS of second operation";
+      }
+    } else {
+      // %v = %x, %x = %expr
+      hasBinop = false;
+      expr = lhs;
+      if (equal(x, lhs))
+        return parser.emitError(parser.getCurrentLocation())
+               << "self assignment without update to captured variable in "
+                  "second operation";
+    }
+  } else if (succeeded(parser.parseOptionalKeyword(&binOp))) {
+    // %x = ..., %v = %x
+    isPostfixUpdate = false;
+    hasBinop = true;
+    x = a;
+    OpAsmParser::OperandType lhs, rhs;
+    lhs = b;
+    if (parser.parseOperand(rhs) || parser.parseComma() ||
+        parser.parseOperand(v) || parser.parseEqual() || parser.parseOperand(b))
+      return failure();
+    if (equal(lhs, x)) {
+      isXLhsInRhsPart = true;
+      expr = rhs;
+    } else if (equal(rhs, x)) {
+      isXLhsInRhsPart = false;
+      expr = lhs;
+    } else {
+      return parser.emitError(parser.getCurrentLocation())
+             << "captured variable not found in RHS of first operation";
     }
+    if (!equal(x, b)) {
+      return parser.emitError(parser.getCurrentLocation())
+             << "updated variable not captured in second operation";
+    }
+  } else {
+    return parser.emitError(parser.getCurrentLocation())
+           << "expected comma or valid binary operation";
+  }
+
+  if (isPostfixUpdate)
+    result.addAttribute("is_postfix_update",
+                        UnitAttr::get(parser.getContext()));
+  if (hasBinop) {
+    result.addAttribute(
+        "binop", parser.getBuilder().getI64IntegerAttr(
+                     (int64_t)AtomicBinOpKindToEnum(binOp.upper()).getValue()));
+  }
+  if (isXLhsInRhsPart)
+    result.addAttribute("is_x_binop_expr", UnitAttr::get(parser.getContext()));
+
+  SmallVector<Type> types;
+  if (parseClauses(parser, result, clauses, segments) ||
+      parser.parseColonTypeList(types))
+    return failure();
+  if (types.size() != 3)
+    return parser.emitError(parser.getNameLoc()) << "expected three types";
+  if (parser.resolveOperand(x, types[0], result.operands) ||
+      parser.resolveOperand(v, types[1], result.operands) ||
+      parser.resolveOperand(expr, types[2], result.operands))
+    return failure();
   return success();
 }
 
+/// Print an atomic capture op
+static void printAtomicCaptureOp(OpAsmPrinter &p, AtomicCaptureOp op) {
+  p << " ";
+  if (op.is_postfix_update()) {
+    p << op.v() << " = " << op.x() << ", " << op.x() << " = ";
+    if (op.binop()) {
+      StringRef binop = AtomicBinOpKindToString(op.binop().getValue()).lower();
+      if (op.is_x_binop_expr())
+        p << op.x() << " " << binop << " " << op.expr() << " ";
+      else
+        p << op.expr() << " " << binop << " " << op.x() << " ";
+    } else {
+      p << op.expr() << " ";
+    }
+  } else {
+    p << op.x() << " = ";
+    if (op.binop()) {
+      StringRef binop = AtomicBinOpKindToString(op.binop().getValue()).lower();
+      if (op.is_x_binop_expr())
+        p << op.x() << " " << binop << " " << op.expr();
+      else {
+        p << op.expr() << " " << binop << " " << op.x();
+      }
+    }
+    p << ", " << op.v() << " = " << op.x() << " ";
+  }
+  if (op.memory_order())
+    p << "memory_order(" << op.memory_order() << ") ";
+  if (op.hintAttr())
+    printSynchronizationHint(p, op, op.hintAttr());
+  p << ": " << op.x().getType() << ", " << op.v().getType() << ", "
+    << op.expr().getType();
+}
+
 #define GET_OP_CLASSES
 #include "mlir/Dialect/OpenMP/OpenMPOps.cpp.inc"

mlir/test/Dialect/OpenMP/invalid.mlir

func @omp_atomic_update2(%x: memref<i32>, %expr: i32) {
   // expected-error @below {{invalid atomic bin op in atomic update}}
   omp.atomic.update %x = %x invalid %expr : memref<i32>, i32
   return
 }
 
 // -----
 
 func @omp_atomic_update3(%x: memref<i32>, %expr: i32) {
-  // expected-error @below {{memory-order must not be acq_rel or acquire for atomic updates}}
+  // expected-error @below {{memory-order must not be acq_rel or acquire for atomic update/capture}}
   omp.atomic.update %x = %x add %expr memory_order(acq_rel) : memref<i32>, i32
   return
 }
 
 // -----
 
 func @omp_atomic_update4(%x: memref<i32>, %expr: i32) {
-  // expected-error @below {{memory-order must not be acq_rel or acquire for atomic updates}}
+  // expected-error @below {{memory-order must not be acq_rel or acquire for atomic update/capture}}
   omp.atomic.update %x = %x add %expr memory_order(acquire) : memref<i32>, i32
   return
 }
 
 // -----
 
 // expected-note @below {{prior use here}}
 func @omp_atomic_update5(%x: memref<i32>, %expr: i32) {
   // expected-error @below {{use of value '%x' expects different type than prior uses: 'i32' vs 'memref<i32>'}}
   omp.atomic.update %x = %x add %expr : i32, memref<i32>
   return
 }
 
 // -----
 
+func @omp_atomic_capture1(%x: memref<i32>, %expr: i32, %v: memref<i32>) {
+  // expected-error @below {{op operand #0 must be OpenMP-compatible variable type, but got 'i32'}}
+  omp.atomic.capture %expr = %expr add %x, %v = %expr : i32, memref<i32>, memref<i32>
+  return
+}
+
+// -----
+
+func @omp_atomic_capture2(%x: memref<i32>, %y: memref<i32>, %expr: i32, %v: memref<i32>) {
+  // expected-error @below {{captured variable not found in LHS of second operation}}
+  omp.atomic.capture %v = %x, %y = %y add %expr : memref<i32>, memref<i32>, i32
+  return
+}
+
+// -----
+
+func @omp_atomic_capture3(%x: memref<i32>, %y: memref<i32>, %expr: i32, %v: memref<i32>) {
+  // expected-error @below {{captured variable not found in RHS of second operation}}
+  omp.atomic.capture %v = %x, %x = %y add %expr : memref<i32>, memref<i32>, i32
+  return
+}
+
+// -----
+
+func @omp_atomic_capture4(%x: memref<i32>, %y: memref<i32>, %expr: i32, %v: memref<i32>) {
+  // expected-error @below {{self assignment without update to captured variable in second operation}}
+  omp.atomic.capture %v = %x, %x = %x : memref<i32>, memref<i32>, i32
+  return
+}
+
+// -----
+
+func @omp_atomic_capture5(%x: memref<i32>, %y: memref<i32>, %expr: i32, %v: memref<i32>) {
+  // expected-error @below {{captured variable not found in RHS of first operation}}
+  omp.atomic.capture %x = %y add %expr, %v = %x : memref<i32>, memref<i32>, i32
+  return
+}
+
+// -----
+
+func @omp_atomic_capture6(%x: memref<i32>, %expr: i32, %v: memref<i32>, %y: memref<i32>) {
+  // expected-error @below {{updated variable not captured in second operation}}
+  omp.atomic.capture %x = %x add %expr, %v = %y : memref<i32>, memref<i32>, i32
+  return
+}
+
+// -----
+
+func @omp_atomic_capture7(%x: memref<i32>, %v: memref<i32>) {
+  // expected-error @below {{expected comma or valid binary operation}}
+  omp.atomic.capture %v = %x : memref<i32>, memref<i32>, i32
+  return
+}
+
+// -----
+
+func @omp_atomic_capture8(%x: memref<i32>, %expr: i32, %v: memref<i32>) {
+  // expected-error @below {{expected three types}}
+  omp.atomic.capture %x = %x add %expr, %v = %x memory_order(acquire) : memref<i32>, i32
+  return
+}
+
+// -----
+
+func @omp_atomic_capture9(%x: memref<i32>, %expr: i32, %v: memref<i32>) {
+  // expected-error @below {{memory-order must not be acq_rel or acquire for atomic update/capture}}
+  omp.atomic.capture %x = %x add %expr, %v = %x memory_order(acq_rel) : memref<i32>, memref<i32>, i32
+  return
+}
+
+// -----
+
+func @omp_atomic_capture10(%x: memref<i32>, %expr: i32, %v: memref<i32>) {
+  // expected-error @below {{memory-order must not be acq_rel or acquire for atomic update/capture}}
+  omp.atomic.capture %x = %x add %expr, %v = %x memory_order(acquire) : memref<i32>, memref<i32>, i32
+  return
+}
+
+// -----
+
 func @omp_sections(%data_var1 : memref<i32>, %data_var2 : memref<i32>, %data_var3 : memref<i32>) -> () {
   // expected-error @below {{operand used in both private and firstprivate clauses}}
   omp.sections private(%data_var1 : memref<i32>) firstprivate(%data_var1 : memref<i32>) {
     omp.terminator
   }
   return
 }
 

mlir/test/Dialect/OpenMP/ops.mlir

func @omp_atomic_update(%x : memref<i32>, %expr : i32, %xBool : memref<i1>, %exprBool : i1) {
   // CHECK: omp.atomic.update %[[X]] = %[[EXPR]] min %[[X]] : memref<i32>, i32
   omp.atomic.update %x = %expr min %x : memref<i32>, i32
   // CHECK: omp.atomic.update %[[XBOOL]] = %[[EXPRBOOL]] eqv %[[XBOOL]] : memref<i1>, i1
   omp.atomic.update %xBool =  %exprBool eqv %xBool : memref<i1>, i1
   // CHECK: omp.atomic.update %[[XBOOL]] = %[[EXPRBOOL]] neqv %[[XBOOL]] : memref<i1>, i1
   omp.atomic.update %xBool =  %exprBool neqv %xBool : memref<i1>, i1
   // CHECK: omp.atomic.update %[[X]] = %[[EXPR]] add %[[X]] memory_order(seq_cst) hint(speculative) : memref<i32>, i32
   omp.atomic.update %x = %expr add %x hint(speculative) memory_order(seq_cst) : memref<i32>, i32
     return
 }
 
+// CHECK-LABEL: @omp_atomic_capture
+// CHECK-SAME: (%[[X:.*]]: memref<i32>, %[[V:.*]]: memref<i32>, %[[E:.*]]: i32)
+func @omp_atomic_capture(%x : memref<i32>, %v : memref<i32>, %e : i32) {
+
+    // C/C++ expressions:
+    //   v = ++x;
+    //   v = --x;
+    //   v = x binop= expr;
+    //   v = x = x binop expr;
+    //   x binop= expr; v = x;
+    //   x = x binop expr; v = x;
+    //   ++x; v = x;
+    //   x++; v = x;
+    //   --x; v = x;
+    //   x--; v = x;
+    // CHECK: omp.atomic.capture %[[X]] = %[[X]] add %[[E]], %[[V]] = %[[X]] : memref<i32>, memref<i32>, i32
+    omp.atomic.capture %x = %x add %e, %v = %x : memref<i32>, memref<i32>, i32
+
+    // C/C++ expressions:
+    //   v = x = expr binop x;
+    //   x = expr binop x; v = x;
+    // CHECK: omp.atomic.capture %[[X]] = %[[E]] add %[[X]], %[[V]] = %[[X]] : memref<i32>, memref<i32>, i32
+    omp.atomic.capture %x = %e add %x, %v = %x : memref<i32>, memref<i32>, i32
+
+    // C/C++ expressions:
+    //   v = x++;
+    //   v = x--;
+    //   v = x; x binop= expr;
+    //   v = x; x = x binop expr;
+    //   v = x; x++;
+    //   v = x; ++x;
+    //   v = x; x--;
+    //   v = x; --x;
+    // CHECK: omp.atomic.capture %[[V]] = %[[X]], %[[X]] = %[[X]] sub %[[E]] : memref<i32>, memref<i32>, i32
+    omp.atomic.capture %v = %x, %x = %x sub %e : memref<i32>, memref<i32>, i32
+
+    // C/C++ expressions:
+    //   v = x; x = expr binop x;
+    // CHECK: omp.atomic.capture %[[V]] = %[[X]], %[[X]] = %[[E]] sub %[[X]] : memref<i32>, memref<i32>, i32
+    omp.atomic.capture %v = %x, %x = %e sub %x : memref<i32>, memref<i32>, i32
+
+    // C/C++ expressions:
+    //   v = x; x = expr;
+    // CHECK: omp.atomic.capture %[[V]] = %[[X]], %[[X]] = %[[E]] : memref<i32>, memref<i32>, i32
+    omp.atomic.capture %v = %x, %x = %e : memref<i32>, memref<i32>, i32
+
+    // Fortran expressions:
+    //   update-statement; capture-statement
+    // CHECK: omp.atomic.capture %[[X]] = %[[X]] add %[[E]], %[[V]] = %[[X]] : memref<i32>, memref<i32>, i32
+    // CHECK: omp.atomic.capture %[[X]] = %[[E]] add %[[X]], %[[V]] = %[[X]] : memref<i32>, memref<i32>, i32
+    omp.atomic.capture %x = %x add %e, %v = %x : memref<i32>, memref<i32>, i32
+    omp.atomic.capture %x = %e add %x, %v = %x : memref<i32>, memref<i32>, i32
+
+    // Fortran expressions:
+    //   capture-statement; update-statement
+    // CHECK: omp.atomic.capture %[[V]] = %[[X]], %[[X]] = %[[X]] sub %[[E]] : memref<i32>, memref<i32>, i32
+    // CHECK: omp.atomic.capture %[[V]] = %[[X]], %[[X]] = %[[E]] sub %[[X]] : memref<i32>, memref<i32>, i32
+    omp.atomic.capture %v = %x, %x = %x sub %e : memref<i32>, memref<i32>, i32
+    omp.atomic.capture %v = %x, %x = %e sub %x : memref<i32>, memref<i32>, i32
+
+    // Fortran expressions:
+    //   capture-statement; write-statement
+    // CHECK: omp.atomic.capture %[[V]] = %[[X]], %[[X]] = %[[E]] : memref<i32>, memref<i32>, i32
+    omp.atomic.capture %v = %x, %x = %e : memref<i32>, memref<i32>, i32
+
+  return
+}
+
 // CHECK-LABEL: omp_sectionsop
 func @omp_sectionsop(%data_var1 : memref<i32>, %data_var2 : memref<i32>,
                      %data_var3 : memref<i32>, %redn_var : !llvm.ptr<f32>) {
 
   // CHECK: omp.sections private(%{{.*}} : memref<i32>) {
   "omp.sections" (%data_var1) ({
     // CHECK: omp.terminator
     omp.terminator