[MLIR][OpenMP] Added omp.atomic.update

[shraiysh]

Summary

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

[peixin]

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

604

According to OpenMP 5.0 Spec 2.17.7 at page 239, operator is one of +, *, -, /, .AND., .OR., .EQV., or .NEQV.. I think you referred to C/C++ operator for XOR, SHIFTR and SHIFTL?

MAX, MIN is instrinsic procedure name, not operator here. What you plan to do is to use the format x = x max y in MLIR for x = max(x, y) in fortran code, right? If yes, did you miss IAND, IOR, IEOR?

I think BINOP is not the right term since MAX and MIN is intrinsic procedure name instead of binary operator. You should give a new name here.

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

1441

How do you plan to handle the case x = max(x, max(y), min(z)) in fortran code? To make it as x = x max expr (lowering max(max(y), min(z)) to expr) in MLIR?

1482

In the else branch, the LHS X is also in RHS. isXLHSInRHSPart -> isXInRHSLeft? Or you can give one better name.

[shraiysh]

Thanks for the comments @peixin! I have tried to address them all, and will include the comments as mentioned if it is okay with you.

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

604

According to OpenMP 5.0 Spec 2.17.7 at page 239, operator is one of +, *, -, /, .AND., .OR., .EQV., or .NEQV.. I think you referred to C/C++ operator for XOR, SHIFTR and SHIFTL?

Yes, that is correct.

MAX, MIN is instrinsic procedure name, not operator here. What you plan to do is to use the format x = x max y in MLIR for x = max(x, y) in fortran code, right? If yes, did you miss IAND, IOR, IEOR?

The AND, OR and XOR are bitwise operations (will mention this as a comment for clarification). As logical values are essentially one-bit integer values, "logical and" is equivalent to "bitwise and for 1-bit integers". So we do not need both the sets AND, OR, XOR and IAND, IOR, IEOR.

I think BINOP is not the right term since MAX and MIN is intrinsic procedure name instead of binary operator. You should give a new name here.

I have borrowed the naming from OpenMP IR Builder (link). Although they do not handle MAX, MIN in createAtomicUpdate (link) right now, they have modeled it as a binary operation in llvm::AtomicRMWInst::BinOp (link). This is in line with C/C++, and prevents us from handling Fortran-specific features (intrinsic) in the common dialect. I am open to suggestions on another more appropriate naming though.

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

1441

Yes, more or less. The evaluation of the expr_list (e1, e2, ...) in max(x, e1, e2, ... ) is not atomic. Hence, all such operations can be lowered to the following before OpenMPDialect:

expr = max(e1, e2... )
x = x max expr // atomic update

Let me know if the lowering doesn't seem right and breaks any concurrency related OpenMP guarantees.

1482

Again, this naming is borrowed from OpenMP IRBuilder. I was a little hesitant while using this variable too as the naming was not very clear. What the name is trying to say is "is X an LHS op in the RHS part?". I will put this as a comment here, but I have to agree I could not think of a better name. Any suggestions on this are welcome 😄

[peixin]

Except for the variable name isXLHSInRHSPart, LGTM.

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

604

According to OpenMP 5.0 Spec 2.17.7 at page 239, operator is one of +, *, -, /, .AND., .OR., .EQV., or .NEQV.. I think you referred to C/C++ operator for XOR, SHIFTR and SHIFTL?

Yes, that is correct.

MAX, MIN is instrinsic procedure name, not operator here. What you plan to do is to use the format x = x max y in MLIR for x = max(x, y) in fortran code, right? If yes, did you miss IAND, IOR, IEOR?

The AND, OR and XOR are bitwise operations (will mention this as a comment for clarification). As logical values are essentially one-bit integer values, "logical and" is equivalent to "bitwise and for 1-bit integers". So we do not need both the sets AND, OR, XOR and IAND, IOR, IEOR.

I think BINOP is not the right term since MAX and MIN is intrinsic procedure name instead of binary operator. You should give a new name here.

I have borrowed the naming from OpenMP IR Builder (link). Although they do not handle MAX, MIN in createAtomicUpdate (link) right now, they have modeled it as a binary operation in llvm::AtomicRMWInst::BinOp (link). This is in line with C/C++, and prevents us from handling Fortran-specific features (intrinsic) in the common dialect. I am open to suggestions on another more appropriate naming though.

604

The this is fine with me. Thanks for the explanations.

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

1441

I think it's viable to do this when lowering from parse-tree to MLIR.

1482

Well. I think the variable name should be changed together with the one in OpenMPIRBuilder since its name is not what it means. Maybe isXBinopExpr?
true if X is Left H.S. in Right H.S. part of the update expression, false otherwise. (e.g. true for X = X BinOp Expr) -> true for X = X BinOp Expr, and false for X = Expr BinOp X. (In OpenMPIRBuilder)
You may wait for better suggestions from others.

[shraiysh]

Address comments and rebase with main

[shraiysh]

Removed extra return in one test case (was added by mistake during rebase)

[shraiysh]

Changed argument name in comments and arguments. This is ready for review now.

[kiranchandramohan]

Thanks for this patch. Had a quick look. Looks OK. I have a couple of points.

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

1454

Move to verifier (also)?

1466–1469

Should this be moved to the verifier (also)?

mlir/test/Dialect/OpenMP/ops.mlir

538–544

These could be for integers as well right? Can you add a test for regular integer types also?

[shraiysh]

Rebase with main. Addressed comments.

[shraiysh]

Rebase with main, addressed comments

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

1454

I believe this cannot be moved to the verifier, as it's value is used in the next line (while making the attribute)

1466–1469

This cannot be moved to the verifier because expr is used in the resolveOperand call later and it cannot be null.

mlir/test/Dialect/OpenMP/ops.mlir

538–544

This was a mistake. It was supposed to be a test for regular integers and I don't think it makes sense to right shift on one-bit integers. I have modified it for i32 now.

[shraiysh]

Modified the other test case for shiftr too (missed in previous diff)

[shraiysh]

Rebase with main

[kiranchandramohan]

What will prevent the builder from building operations which has the errors that you are catching during parsing?

[shraiysh]

Thanks for the review @kiranchandramohan.

In D112982#3160888, @kiranchandramohan wrote:

What will prevent the builder from building operations which has the errors that you are catching during parsing?

The BinOp error will be taken care of by the AtomicBinOpKindAttr. An integer value from the allowed ones will be forced because of that constraint.

The other error - about X being in RHS - that is taken care of by design as the in memory structure has X, expr, binOp and isXBinOpExpr. It is not possible to run into that error while building the operation.

[kiranchandramohan]

In D112982#3161031, @shraiysh wrote:

Thanks for the review @kiranchandramohan.

In D112982#3160888, @kiranchandramohan wrote:

What will prevent the builder from building operations which has the errors that you are catching during parsing?

The BinOp error will be taken care of by the AtomicBinOpKindAttr. An integer value from the allowed ones will be forced because of that constraint.

The other error - about X being in RHS - that is taken care of by design as the in memory structure has X, expr, binOp and isXBinOpExpr. It is not possible to run into that error while building the operation.

Thanks. I have one more question regarding binary ops.

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

597–636

Nit: Please add a TODO here to autogenerate from OMP.td in future if possible.

627

Are any non-binary operations supported currently? If not, can this be removed? If yes, can a test be added?

[shraiysh]

Addressed Comments.

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

627

This is the naming issue from earlier. The name here means "is it of the form X binop expr". The value false implies that it is of the form expr binop X. The previous name for this variable was isXLhsinRhsPart but that was confusing too. This is not a flag for it being a binop expression or not.

[kiranchandramohan]

Thanks for the clarification. LGTM.

[shraiysh]

Thanks for the review. Rebase with main.

llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h

private:
   ///                   instructions.
   /// \param RMWOp      The binary operation used for update. If
   ///                   operation is not supported by atomicRMW,
   ///                   or belong to {FADD, FSUB, BAD_BINOP}.
   ///                   Then a `cmpExch` based  atomic will be generated.
   /// \param UpdateOp   Code generator for complex expressions that cannot be
   ///                   expressed through atomicrmw instruction.
   /// \param VolatileX       true if \a X volatile?
-  /// \param IsXLHSInRHSPart true if \a X is Left H.S. in Right H.S. part of
-  ///                        the update expression, false otherwise.
-  ///                        (e.g. true for X = X BinOp Expr)
+  /// \param IsXBinopExpr true if \a X is Left H.S. in Right H.S. part of the
+  ///                     update expression, false otherwise.
+  ///                     (e.g. true for X = X BinOp Expr)
   ///
   /// \returns A pair of the old value of X before the update, and the value
   ///          used for the update.
   std::pair<Value *, Value *> emitAtomicUpdate(Instruction *AllocIP, Value *X,
                                                Value *Expr, AtomicOrdering AO,
                                                AtomicRMWInst::BinOp RMWOp,
                                                AtomicUpdateCallbackTy &UpdateOp,
                                                bool VolatileX,
-                                               bool IsXLHSInRHSPart);
+                                               bool IsXBinopExpr);
 
   /// Emit the binary op. described by \p RMWOp, using \p Src1 and \p Src2 .
   ///
   /// \Return The instruction
   Value *emitRMWOpAsInstruction(Value *Src1, Value *Src2,
                                 AtomicRMWInst::BinOp RMWOp);
 
 public:
public:
   /// \param Expr     The value to update X with.
   /// \param AO       Atomic ordering of the generated atomic instructions.
   /// \param RMWOp    The binary operation used for update. If operation
   ///                 is  not supported by atomicRMW, or belong to
   ///                 {FADD, FSUB, BAD_BINOP}. Then a `cmpExch` based
   ///                 atomic will be generated.
   /// \param UpdateOp   Code generator for complex expressions that cannot be
   ///                   expressed through atomicrmw instruction.
-  /// \param IsXLHSInRHSPart true if \a X is Left H.S. in Right H.S. part of
-  ///                        the update expression, false otherwise.
-  ///                        (e.g. true for X = X BinOp Expr)
+  /// \param IsXBinopExpr true if \a X is Left H.S. in Right H.S. part of the
+  ///                     update expression, false otherwise.
+  ///                     (e.g. true for X = X BinOp Expr)
   ///
   /// \return Insertion point after generated atomic update IR.
   InsertPointTy createAtomicUpdate(const LocationDescription &Loc,
                                    Instruction *AllocIP, AtomicOpValue &X,
                                    Value *Expr, AtomicOrdering AO,
                                    AtomicRMWInst::BinOp RMWOp,
                                    AtomicUpdateCallbackTy &UpdateOp,
-                                   bool IsXLHSInRHSPart);
+                                   bool IsXBinopExpr);
 
   /// Emit atomic update for constructs: --- Only Scalar data types
   /// V = X; X = X BinOp Expr ,
   /// X = X BinOp Expr; V = X,
   /// V = X; X = Expr BinOp X,
   /// X = Expr BinOp X; V = X,
   /// V = X; X = UpdateOp(X),
   /// X = UpdateOp(X); V = X,
   ///
   /// \param Loc        The insert and source location description.
   /// \param AllocIP    Instruction to create AllocaInst before.
   /// \param X          The target atomic pointer to be updated
   /// \param V          Memory address where to store captured value
   /// \param Expr       The value to update X with.
   /// \param AO         Atomic ordering of the generated atomic instructions
   /// \param RMWOp      The binary operation used for update. If
   ///                   operation is not supported by atomicRMW, or belong to
   ///                   {FADD, FSUB, BAD_BINOP}. Then a cmpExch based
   ///                   atomic will be generated.
   /// \param UpdateOp   Code generator for complex expressions that cannot be
   ///                   expressed through atomicrmw instruction.
   /// \param UpdateExpr true if X is an in place update of the form
   ///                   X = X BinOp Expr or X = Expr BinOp X
-  /// \param IsXLHSInRHSPart true if X is Left H.S. in Right H.S. part of the
-  ///                        update expression, false otherwise.
-  ///                        (e.g. true for X = X BinOp Expr)
+  /// \param IsXBinopExpr true if X is Left H.S. in Right H.S. part of the
+  ///                     update expression, false otherwise.
+  ///                     (e.g. true for X = X BinOp Expr)
   /// \param IsPostfixUpdate true if original value of 'x' must be stored in
   ///                        'v', not an updated one.
   ///
   /// \return Insertion point after generated atomic capture IR.
   InsertPointTy
   createAtomicCapture(const LocationDescription &Loc, Instruction *AllocIP,
                       AtomicOpValue &X, AtomicOpValue &V, Value *Expr,
                       AtomicOrdering AO, AtomicRMWInst::BinOp RMWOp,
                       AtomicUpdateCallbackTy &UpdateOp, bool UpdateExpr,
-                      bool IsPostfixUpdate, bool IsXLHSInRHSPart);
+                      bool IsPostfixUpdate, bool IsXBinopExpr);
 
   /// Create the control flow structure of a canonical OpenMP loop.
   ///
   /// The emitted loop will be disconnected, i.e. no edge to the loop's
   /// preheader and no terminator in the AfterBB. The OpenMPIRBuilder's
   /// IRBuilder location is not preserved.
   ///
   /// \param DL        DebugLoc used for the instructions in the skeleton.

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

OpenMPIRBuilder::createAtomicWrite(const LocationDescription &Loc,
 
   checkAndEmitFlushAfterAtomic(Loc, AO, AtomicKind::Write);
   return Builder.saveIP();
 }
 
 OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createAtomicUpdate(
     const LocationDescription &Loc, Instruction *AllocIP, AtomicOpValue &X,
     Value *Expr, AtomicOrdering AO, AtomicRMWInst::BinOp RMWOp,
-    AtomicUpdateCallbackTy &UpdateOp, bool IsXLHSInRHSPart) {
+    AtomicUpdateCallbackTy &UpdateOp, bool IsXBinopExpr) {
   if (!updateToLocation(Loc))
     return Loc.IP;
 
   LLVM_DEBUG({
     Type *XTy = X.Var->getType();
     assert(XTy->isPointerTy() &&
            "OMP Atomic expects a pointer to target memory");
     Type *XElemTy = XTy->getPointerElementType();
     assert((XElemTy->isFloatingPointTy() || XElemTy->isIntegerTy() ||
             XElemTy->isPointerTy()) &&
            "OMP atomic update expected a scalar type");
     assert((RMWOp != AtomicRMWInst::Max) && (RMWOp != AtomicRMWInst::Min) &&
            (RMWOp != AtomicRMWInst::UMax) && (RMWOp != AtomicRMWInst::UMin) &&
            "OpenMP atomic does not support LT or GT operations");
   });
 
   emitAtomicUpdate(AllocIP, X.Var, Expr, AO, RMWOp, UpdateOp, X.IsVolatile,
-                   IsXLHSInRHSPart);
+                   IsXBinopExpr);
   checkAndEmitFlushAfterAtomic(Loc, AO, AtomicKind::Update);
   return Builder.saveIP();
 }
 
 Value *OpenMPIRBuilder::emitRMWOpAsInstruction(Value *Src1, Value *Src2,
                                                AtomicRMWInst::BinOp RMWOp) {
   switch (RMWOp) {
   case AtomicRMWInst::Add:
Value *OpenMPIRBuilder::emitRMWOpAsInstruction(Value *Src1, Value *Src2,
   }
   llvm_unreachable("Unsupported atomic update operation");
 }
 
 std::pair<Value *, Value *>
 OpenMPIRBuilder::emitAtomicUpdate(Instruction *AllocIP, Value *X, Value *Expr,
                                   AtomicOrdering AO, AtomicRMWInst::BinOp RMWOp,
                                   AtomicUpdateCallbackTy &UpdateOp,
-                                  bool VolatileX, bool IsXLHSInRHSPart) {
+                                  bool VolatileX, bool IsXBinopExpr) {
   Type *XElemTy = X->getType()->getPointerElementType();
 
   bool DoCmpExch =
       ((RMWOp == AtomicRMWInst::BAD_BINOP) || (RMWOp == AtomicRMWInst::FAdd)) ||
       (RMWOp == AtomicRMWInst::FSub) ||
-      (RMWOp == AtomicRMWInst::Sub && !IsXLHSInRHSPart);
+      (RMWOp == AtomicRMWInst::Sub && !IsXBinopExpr);
 
   std::pair<Value *, Value *> Res;
   if (XElemTy->isIntegerTy() && !DoCmpExch) {
     Res.first = Builder.CreateAtomicRMW(RMWOp, X, Expr, llvm::MaybeAlign(), AO);
     // not needed except in case of postfix captures. Generate anyway for
     // consistency with the else part. Will be removed with any DCE pass.
     Res.second = emitRMWOpAsInstruction(Res.first, Expr, RMWOp);
   } else {
OpenMPIRBuilder::emitAtomicUpdate(Instruction *AllocIP, Value *X, Value *Expr,
 
   return Res;
 }
 
 OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createAtomicCapture(
     const LocationDescription &Loc, Instruction *AllocIP, AtomicOpValue &X,
     AtomicOpValue &V, Value *Expr, AtomicOrdering AO,
     AtomicRMWInst::BinOp RMWOp, AtomicUpdateCallbackTy &UpdateOp,
-    bool UpdateExpr, bool IsPostfixUpdate, bool IsXLHSInRHSPart) {
+    bool UpdateExpr, bool IsPostfixUpdate, bool IsXBinopExpr) {
   if (!updateToLocation(Loc))
     return Loc.IP;
 
   LLVM_DEBUG({
     Type *XTy = X.Var->getType();
     assert(XTy->isPointerTy() &&
            "OMP Atomic expects a pointer to target memory");
     Type *XElemTy = XTy->getPointerElementType();
     assert((XElemTy->isFloatingPointTy() || XElemTy->isIntegerTy() ||
             XElemTy->isPointerTy()) &&
            "OMP atomic capture expected a scalar type");
     assert((RMWOp != AtomicRMWInst::Max) && (RMWOp != AtomicRMWInst::Min) &&
            "OpenMP atomic does not support LT or GT operations");
   });
 
   // If UpdateExpr is 'x' updated with some `expr` not based on 'x',
   // 'x' is simply atomically rewritten with 'expr'.
   AtomicRMWInst::BinOp AtomicOp = (UpdateExpr ? RMWOp : AtomicRMWInst::Xchg);
-  std::pair<Value *, Value *> Result =
-      emitAtomicUpdate(AllocIP, X.Var, Expr, AO, AtomicOp, UpdateOp,
-                       X.IsVolatile, IsXLHSInRHSPart);
+  std::pair<Value *, Value *> Result = emitAtomicUpdate(
+      AllocIP, X.Var, Expr, AO, AtomicOp, UpdateOp, X.IsVolatile, IsXBinopExpr);
 
   Value *CapturedVal = (IsPostfixUpdate ? Result.first : Result.second);
   Builder.CreateStore(CapturedVal, V.Var, V.IsVolatile);
 
   checkAndEmitFlushAfterAtomic(Loc, AO, AtomicKind::Capture);
   return Builder.saveIP();
 }
 

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

def AtomicWriteOp : OpenMP_Op<"atomic.write"> {
   let arguments = (ins OpenMP_PointerLikeType:$address,
                        AnyType:$value,
                        DefaultValuedAttr<I64Attr, "0">:$hint,
                        OptionalAttr<MemoryOrderKind>:$memory_order);
   let parser = [{ return parseAtomicWriteOp(parser, result); }];
   let printer = [{ return printAtomicWriteOp(p, *this); }];
   let verifier = [{ return verifyAtomicWriteOp(*this); }];
 }
 
+// TODO: autogenerate from OMP.td in future if possible.
+def ATOMIC_BINOP_KIND_ADD : I64EnumAttrCase<"ADD", 0>;
+def ATOMIC_BINOP_KIND_MUL : I64EnumAttrCase<"MUL", 1>;
+def ATOMIC_BINOP_KIND_SUB : I64EnumAttrCase<"SUB", 2>;
+def ATOMIC_BINOP_KIND_DIV : I64EnumAttrCase<"DIV", 3>;
+def ATOMIC_BINOP_KIND_AND : I64EnumAttrCase<"AND", 4>;
+def ATOMIC_BINOP_KIND_OR : I64EnumAttrCase<"OR", 5>;
+def ATOMIC_BINOP_KIND_XOR : I64EnumAttrCase<"XOR", 6>;
+def ATOMIC_BINOP_KIND_SHIFT_RIGHT : I64EnumAttrCase<"SHIFTR", 7>;
+def ATOMIC_BINOP_KIND_SHIFT_LEFT : I64EnumAttrCase<"SHIFTL", 8>;
+def ATOMIC_BINOP_KIND_MAX : I64EnumAttrCase<"MAX", 9>;
+def ATOMIC_BINOP_KIND_MIN : I64EnumAttrCase<"MIN", 10>;
+def ATOMIC_BINOP_KIND_EQV : I64EnumAttrCase<"EQV", 11>;
+def ATOMIC_BINOP_KIND_NEQV : I64EnumAttrCase<"NEQV", 12>;
+
+def AtomicBinOpKindAttr : I64EnumAttr<
+    "AtomicBinOpKind", "BinOp for Atomic Updates",
+    [ATOMIC_BINOP_KIND_ADD, ATOMIC_BINOP_KIND_MUL, ATOMIC_BINOP_KIND_SUB,
+     ATOMIC_BINOP_KIND_DIV, ATOMIC_BINOP_KIND_AND, ATOMIC_BINOP_KIND_OR,
+     ATOMIC_BINOP_KIND_XOR, ATOMIC_BINOP_KIND_SHIFT_RIGHT,
+     ATOMIC_BINOP_KIND_SHIFT_LEFT, ATOMIC_BINOP_KIND_MAX,
+     ATOMIC_BINOP_KIND_MIN, ATOMIC_BINOP_KIND_EQV, ATOMIC_BINOP_KIND_NEQV]> {
+  let cppNamespace = "::mlir::omp";
+  let stringToSymbolFnName = "AtomicBinOpKindToEnum";
+  let symbolToStringFnName = "AtomicBinOpKindToString";
+}
+
+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); }];
+}
+
 //===----------------------------------------------------------------------===//
 // 2.19.5.7 declare reduction Directive
 //===----------------------------------------------------------------------===//
 
 def ReductionDeclareOp : OpenMP_Op<"reduction.declare", [Symbol]> {
   let summary = "declares a reduction kind";
 
   let description = [{

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

if (op.memory_order()) {
     StringRef memoryOrder = op.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 writes");
   }
   return verifySynchronizationHint(op, op.hint());
 }
 
+//===----------------------------------------------------------------------===//
+// AtomicUpdateOp
+//===----------------------------------------------------------------------===//
+
+/// Parser for AtomicUpdateOp
+///
+/// operation ::= `omp.atomic.update` atomic-clause-list region
+static ParseResult parseAtomicUpdateOp(OpAsmParser &parser,
+                                       OperationState &result) {
+  SmallVector<ClauseType> clauses = {memoryOrderClause, hintClause};
+  SmallVector<int> segments;
+  OpAsmParser::OperandType x, y, z;
+  Type xType, exprType;
+  StringRef binOp;
+
+  // x = y `op` z : xtype, exprtype
+  if (parser.parseOperand(x) || parser.parseEqual() || parser.parseOperand(y) ||
+      parser.parseKeyword(&binOp) || parser.parseOperand(z) ||
+      parseClauses(parser, result, clauses, segments) || parser.parseColon() ||
+      parser.parseType(xType) || parser.parseComma() ||
+      parser.parseType(exprType) ||
+      parser.resolveOperand(x, xType, result.operands)) {
+    return failure();
+  }
+
+  auto binOpEnum = AtomicBinOpKindToEnum(binOp.upper());
+  if (!binOpEnum)
+    return parser.emitError(parser.getNameLoc())
+           << "invalid atomic bin op in atomic update\n";
+  auto attr =
+      parser.getBuilder().getI64IntegerAttr((int64_t)binOpEnum.getValue());
+  result.addAttribute("binop", attr);
+
+  OpAsmParser::OperandType expr;
+  if (x.name == y.name && x.number == y.number) {
+    expr = z;
+    result.addAttribute("isXBinopExpr", parser.getBuilder().getUnitAttr());
+  } else if (x.name == z.name && x.number == z.number) {
+    expr = y;
+  } else {
+    return parser.emitError(parser.getNameLoc())
+           << "atomic update variable " << x.name
+           << " not found in the RHS of the assignment statement in an"
+              " atomic.update operation";
+  }
+  return parser.resolveOperand(expr, exprType, result.operands);
+}
+
+/// Printer for AtomicUpdateOp
+static void printAtomicUpdateOp(OpAsmPrinter &p, AtomicUpdateOp op) {
+  p << " " << op.x() << " = ";
+  Value y, z;
+  if (op.isXBinopExpr()) {
+    y = op.x();
+    z = op.expr();
+  } else {
+    y = op.expr();
+    z = op.x();
+  }
+  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();
+    if (memoryOrder.equals("acq_rel") || memoryOrder.equals("acquire"))
+      return op.emitError(
+          "memory-order must not be acq_rel or acquire for atomic updates");
+  }
+  return success();
+}
+
 #define GET_OP_CLASSES
 #include "mlir/Dialect/OpenMP/OpenMPOps.cpp.inc"

mlir/test/Dialect/OpenMP/invalid.mlir

 func @omp_atomic_write6(%addr : memref<i32>, %val : i32) {
   // expected-error @below {{attribute 'memory_order' failed to satisfy constraint: MemoryOrderKind Clause}}
   omp.atomic.write  %addr, %val memory_order(xyz) : memref<i32>, i32
   return
 }
 
 // -----
 
+func @omp_atomic_update1(%x: memref<i32>, %expr: i32, %foo: memref<i32>) {
+  // expected-error @below {{atomic update variable %x not found in the RHS of the assignment statement in an atomic.update operation}}
+  omp.atomic.update %x = %foo add %expr : memref<i32>, i32
+  return
+}
+
+// -----
+
+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}}
+  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}}
+  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_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_write(%addr : memref<i32>, %val : i32) {
   omp.atomic.write %addr, %val memory_order(release) : memref<i32>, i32
   // CHECK: omp.atomic.write %[[ADDR]], %[[VAL]] memory_order(relaxed) : memref<i32>, i32
   omp.atomic.write %addr, %val memory_order(relaxed) : memref<i32>, i32
   // CHECK: omp.atomic.write %[[ADDR]], %[[VAL]] hint(uncontended, speculative) : memref<i32>, i32
   omp.atomic.write %addr, %val hint(speculative, uncontended) : memref<i32>, i32
   return
 }
 
+// CHECK-LABEL: omp_atomic_update
+// CHECK-SAME: (%[[X:.*]]: memref<i32>, %[[EXPR:.*]]: i32, %[[XBOOL:.*]]: memref<i1>, %[[EXPRBOOL:.*]]: i1)
+func @omp_atomic_update(%x : memref<i32>, %expr : i32, %xBool : memref<i1>, %exprBool : i1) {
+  // CHECK: omp.atomic.update %[[X]] = %[[X]] add %[[EXPR]] : memref<i32>, i32
+  omp.atomic.update %x = %x add %expr : memref<i32>, i32
+  // CHECK: omp.atomic.update %[[X]] = %[[X]] sub %[[EXPR]] : memref<i32>, i32
+  omp.atomic.update %x = %x sub %expr : memref<i32>, i32
+  // CHECK: omp.atomic.update %[[X]] = %[[X]] mul %[[EXPR]] : memref<i32>, i32
+  omp.atomic.update %x = %x mul %expr : memref<i32>, i32
+  // CHECK: omp.atomic.update %[[X]] = %[[X]] div %[[EXPR]] : memref<i32>, i32
+  omp.atomic.update %x = %x div %expr : memref<i32>, i32
+  // CHECK: omp.atomic.update %[[XBOOL]] = %[[XBOOL]] and %[[EXPRBOOL]] : memref<i1>, i1
+  omp.atomic.update %xBool = %xBool and %exprBool : memref<i1>, i1
+  // CHECK: omp.atomic.update %[[XBOOL]] = %[[XBOOL]] or %[[EXPRBOOL]] : memref<i1>, i1
+  omp.atomic.update %xBool = %xBool or %exprBool : memref<i1>, i1
+  // CHECK: omp.atomic.update %[[XBOOL]] = %[[XBOOL]] xor %[[EXPRBOOL]] : memref<i1>, i1
+  omp.atomic.update %xBool = %xBool xor %exprBool : memref<i1>, i1
+  // CHECK: omp.atomic.update %[[X]] = %[[X]] shiftr %[[EXPR]] : memref<i32>, i32
+  omp.atomic.update %x = %x shiftr %expr : memref<i32>, i32
+  // CHECK: omp.atomic.update %[[X]] = %[[X]] shiftl %[[EXPR]] : memref<i32>, i32
+  omp.atomic.update %x = %x shiftl %expr : memref<i32>, i32
+  // CHECK: omp.atomic.update %[[X]] = %[[X]] max %[[EXPR]] : memref<i32>, i32
+  omp.atomic.update %x = %x max %expr : memref<i32>, i32
+  // CHECK: omp.atomic.update %[[X]] = %[[X]] min %[[EXPR]] : memref<i32>, i32
+  omp.atomic.update %x = %x min %expr : memref<i32>, i32
+  // CHECK: omp.atomic.update %[[XBOOL]] = %[[XBOOL]] eqv %[[EXPRBOOL]] : memref<i1>, i1
+  omp.atomic.update %xBool = %xBool eqv %exprBool : memref<i1>, i1
+  // CHECK: omp.atomic.update %[[XBOOL]] = %[[XBOOL]] neqv %[[EXPRBOOL]] : memref<i1>, i1
+  omp.atomic.update %xBool = %xBool neqv %exprBool : memref<i1>, i1
+
+  // CHECK: omp.atomic.update %[[X]] = %[[EXPR]] add %[[X]] : memref<i32>, i32
+  omp.atomic.update %x = %expr add %x : memref<i32>, i32
+  // CHECK: omp.atomic.update %[[X]] = %[[EXPR]] sub %[[X]] : memref<i32>, i32
+  omp.atomic.update %x = %expr sub %x : memref<i32>, i32
+  // CHECK: omp.atomic.update %[[X]] = %[[EXPR]] mul %[[X]] : memref<i32>, i32
+  omp.atomic.update %x = %expr mul %x : memref<i32>, i32
+  // CHECK: omp.atomic.update %[[X]] = %[[EXPR]] div %[[X]] : memref<i32>, i32
+  omp.atomic.update %x = %expr div %x : memref<i32>, i32
+  // CHECK: omp.atomic.update %[[XBOOL]] = %[[EXPRBOOL]] and %[[XBOOL]] : memref<i1>, i1
+  omp.atomic.update %xBool =  %exprBool and %xBool : memref<i1>, i1
+  // CHECK: omp.atomic.update %[[XBOOL]] = %[[EXPRBOOL]] or %[[XBOOL]] : memref<i1>, i1
+  omp.atomic.update %xBool =  %exprBool or %xBool : memref<i1>, i1
+  // CHECK: omp.atomic.update %[[XBOOL]] = %[[EXPRBOOL]] xor %[[XBOOL]] : memref<i1>, i1
+  omp.atomic.update %xBool =  %exprBool xor %xBool : memref<i1>, i1
+  // CHECK: omp.atomic.update %[[X]] = %[[EXPR]] shiftr %[[X]] : memref<i32>, i32
+  omp.atomic.update %x =  %expr shiftr %x : memref<i32>, i32
+  // CHECK: omp.atomic.update %[[X]] = %[[EXPR]] shiftl %[[X]] : memref<i32>, i32
+  omp.atomic.update %x = %expr shiftl %x : memref<i32>, i32
+  // CHECK: omp.atomic.update %[[X]] = %[[EXPR]] max %[[X]] : memref<i32>, i32
+  omp.atomic.update %x = %expr max %x : memref<i32>, i32
+  // 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_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