[mlir][OpenMP] Change omp.atomic.update to have generic updates

[shraiysh]

Summary

This patch changes the syntax of omp.atomic.update to allow the other
dialects to modify the variable with appropriate operations in the
region.

[shraiysh]

llvm.intr.maxnum(f32) -> llvm.intr.smax(i32) in test case and rebase with main.

[ftynse]

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

347–348

So this just dropped the type-match requirement for WsLoop and ReductionDeclare, hopefully only in documentation. In this case, the very least you could do is to move this documentation into that of the other ops that care. Or add "except for AtomicUpdate where it must match <...>" to this, but don't just drop important information on the floor.

671–674

Nit: I see that it was already there, but I would avoid such an non-specific reference to the standard. It's unclear which version of the standard this refers to, and there is no indication to anything more concrete than "the standard" that is hundreds of pages long. It is also subject to change.

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

1569

Nit: the "pointer to region argument" is confusing. Consider something like "the type of the operand must be a pointer type whose element type is the same as that of the region argument", maybe you can find something more concise.

[ftynse]

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

1567

There doesn't seem to be a check for the region to have at least one argument, so getArgument(0) may assert.

[shraiysh]

Thanks for the review @ftynse. Addressed comments.

[shraiysh]

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

347–348

Sure, added to the ReductionDeclareOp. WsLoopOp yields without any operands.

671–674

Added standard and section number. A one-line description of x is also added.

[ftynse]

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

1568

Could we have a test for this? It can be written using the generic syntax.

[shraiysh]

Thanks for the comment @ftynse. Added testcase.

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

let extraClassDeclaration = [{
     unsigned getNumReductionVars() { return reduction_vars().size(); }
   }];
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 }
 
 def YieldOp : OpenMP_Op<"yield",
     [NoSideEffect, ReturnLike, Terminator,
-     ParentOneOf<["WsLoopOp", "ReductionDeclareOp"]>]> {
+     ParentOneOf<["WsLoopOp", "ReductionDeclareOp", "AtomicUpdateOp"]>]> {
   let summary = "loop yield and termination operation";
   let description = [{
     "omp.yield" yields SSA values from the OpenMP dialect op region and
     terminates the region. The semantics of how the values are yielded is
     defined by the parent operation.
-    If "omp.yield" has any operands, the operands must match the parent
-    operation's results.
   }];
 
   let arguments = (ins Variadic<AnyType>:$results);
 
   let assemblyFormat = [{ ( `(` $results^ `:` type($results) `)` )? attr-dict}];
 }
 
 //===----------------------------------------------------------------------===//

 
 def AtomicUpdateOp : OpenMP_Op<"atomic.update"> {
 
   let summary = "performs an atomic update";
 
   let description = [{
     This operation performs an atomic update.
 
-    The operands `x` and `expr` are 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 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 `isXBinopExpr` is
-      - true when the expression is of the form `x binop expr` on RHS
-      - false when the expression is of the form `expr binop x` on RHS
-
-    The attribute `binop` is the binary operation being performed atomically.
+    The operand `x` is exactly the same as the operand `x` in the OpenMP
+    Standard. It is the address of the variable that is being updated. `x` is
+    atomically read/written.
 
     `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`.
+
+    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.
   }];
 
   let arguments = (ins OpenMP_PointerLikeType:$x,
-                       AnyType:$expr,
-                       UnitAttr:$isXBinopExpr,
-                       AtomicBinOpKindAttr:$binop,
                        DefaultValuedAttr<I64Attr, "0">:$hint,
                        OptionalAttr<MemoryOrderKindAttr>:$memory_order);
+  let regions = (region SizedRegion<1>:$region);
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 }
 
 def AtomicCaptureOp : OpenMP_Op<"atomic.capture",
     [SingleBlockImplicitTerminator<"TerminatorOp">]> {
   let summary = "performs an atomic capture";
   let description = [{

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

 }
 
 //===----------------------------------------------------------------------===//
 // AtomicUpdateOp
 //===----------------------------------------------------------------------===//
 
 /// Parser for AtomicUpdateOp
 ///
-/// operation ::= `omp.atomic.update` atomic-clause-list region
+/// operation ::= `omp.atomic.update` atomic-clause-list ssa-id-and-type region
 ParseResult AtomicUpdateOp::parse(OpAsmParser &parser, OperationState &result) {
   SmallVector<ClauseType> clauses = {memoryOrderClause, hintClause};
   SmallVector<int> segments;
-  OpAsmParser::OperandType x, y, z;
+  OpAsmParser::OperandType x, expr;
   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)) {
+  if (parseClauses(parser, result, clauses, segments) ||
+      parser.parseOperand(x) || parser.parseColon() ||
+      parser.parseType(xType) ||
+      parser.resolveOperand(x, xType, result.operands) ||
+      parser.parseRegion(*result.addRegion()))
     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);
+  return success();
 }
 
 void AtomicUpdateOp::print(OpAsmPrinter &p) {
-  p << " " << x() << " = ";
-  Value y, z;
-  if (isXBinopExpr()) {
-    y = x();
-    z = expr();
-  } else {
-    y = expr();
-    z = x();
-  }
-  p << y << " " << AtomicBinOpKindToString(binop()).lower() << " " << z << " ";
+  p << " ";
   if (auto mo = memory_order())
     p << "memory_order(" << stringifyClauseMemoryOrderKind(*mo) << ") ";
   if (hintAttr())
     printSynchronizationHint(p, *this, hintAttr());
-  p << ": " << x().getType() << ", " << expr().getType();
+  p << x() << " : " << x().getType();
+  p.printRegion(region());
 }
 
 /// Verifier for AtomicUpdateOp
 LogicalResult AtomicUpdateOp::verify() {
   if (auto mo = memory_order()) {
     if (*mo == ClauseMemoryOrderKind::acq_rel ||
         *mo == ClauseMemoryOrderKind::acquire) {
       return emitError(
           "memory-order must not be acq_rel or acquire for atomic updates");
     }
   }
+  if (x().getType().cast<PointerLikeType>().getElementType() !=
+      region().getArgument(0).getType()) {
+    return emitError(
+        "the operand `x` must be a pointer to region argument (type mismatch)");
+  }
+  auto &lastOp = region().front().getOperations().back();
+  YieldOp yieldOp = dyn_cast<YieldOp>(lastOp);
+  if (!yieldOp)
+    return emitError("region must end with an omp.yield operation");
+  if (yieldOp.results().size() != 1)
+    return emitError("only updated value must be returned");
+  if (yieldOp.results().front().getType() != region().getArgument(0).getType())
+    return emitError("input and yielded value must have the same type");
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // AtomicCaptureOp
 //===----------------------------------------------------------------------===//
 
 ParseResult AtomicCaptureOp::parse(OpAsmParser &parser,

mlir/test/Dialect/OpenMP/invalid.mlir

 func @omp_atomic_write6(%addr : memref<i32>, %val : i32) {
   // expected-error @below {{invalid memory order kind}}
   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
+func @omp_atomic_update1(%x: memref<i32>, %expr: f32) {
+  // expected-error @below {{the operand `x` must be a pointer to region argument (type mismatch)}}
+  omp.atomic.update %x : memref<i32> {
+  ^bb0(%xval: f32):
+    %newval = llvm.add %xval, %expr : f32
+    omp.yield (%newval : f32)
+  }
   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
+  // expected-error @below {{region must end with an omp.yield operation}}
+  omp.atomic.update %x : memref<i32> {
+  ^bb0(%xval: i32):
+    %newval = llvm.add %xval, %expr : i32
+    omp.terminator
+  }
   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
+  omp.atomic.update memory_order(acq_rel) %x : memref<i32> {
+  ^bb0(%xval: i32):
+    %newval = llvm.add %xval, %expr : i32
+    omp.yield (%newval : 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
+  omp.atomic.update memory_order(acquire) %x : memref<i32> {
+  ^bb0(%xval: i32):
+    %newval = llvm.add %xval, %expr : i32
+    omp.yield (%newval : 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>
+  omp.atomic.update %x : i32 {
+  ^bb0(%xval: i32):
+    %newval = llvm.add %xval, %expr : i32
+    omp.yield (%newval : i32)
+  }
+  return
+}
+
+// -----
+
+func @omp_atomic_update6(%x: memref<i32>, %expr: i32) {
+  // expected-error @below {{only updated value must be returned}}
+  omp.atomic.update %x : memref<i32> {
+  ^bb0(%xval: i32):
+    %newval = llvm.add %xval, %expr : i32
+    omp.yield (%newval, %expr : i32, i32)
+  }
+  return
+}
+
+// -----
+
+func @omp_atomic_update7(%x: memref<i32>, %expr: i32, %y: f32) {
+  // expected-error @below {{input and yielded value must have the same type}}
+  omp.atomic.update %x : memref<i32> {
+  ^bb0(%xval: i32):
+    %newval = llvm.add %xval, %expr : i32
+    omp.yield (%y: f32)
+  }
   return
 }
 
 // -----
 
 func @omp_atomic_capture(%x: memref<i32>, %v: memref<i32>, %expr: i32) {
   // expected-error @below {{expected three operations in omp.atomic.capture region}}
   omp.atomic.capture {
func @omp_atomic_capture(%x: memref<i32>, %v: memref<i32>, %expr: i32) {
   return
 }
 
 // -----
 
 func @omp_atomic_capture(%x: memref<i32>, %v: memref<i32>, %expr: i32) {
   omp.atomic.capture {
     // expected-error @below {{invalid sequence of operations in the capture region}}
-    omp.atomic.update %x = %x add %expr : memref<i32>, i32
-    omp.atomic.update %x = %x sub %expr : memref<i32>, i32
+    omp.atomic.update %x : memref<i32> {
+    ^bb0(%xval: i32):
+      %newval = llvm.add %xval, %expr : i32
+      omp.yield (%newval : i32)
+    }
+    omp.atomic.update %x : memref<i32> {
+    ^bb0(%xval: i32):
+      %newval = llvm.add %xval, %expr : i32
+      omp.yield (%newval : i32)
+    }
     omp.terminator
   }
   return
 }
 
 // -----
 
 func @omp_atomic_capture(%x: memref<i32>, %v: memref<i32>, %expr: i32) {
   omp.atomic.capture {
     // expected-error @below {{invalid sequence of operations in the capture region}}
     omp.atomic.write %x = %expr : memref<i32>, i32
     omp.atomic.write %x = %expr : memref<i32>, i32
     omp.terminator
   }
   return
 }
 
 // -----
 
 func @omp_atomic_capture(%x: memref<i32>, %v: memref<i32>, %expr: i32) {
   omp.atomic.capture {
     // expected-error @below {{invalid sequence of operations in the capture region}}
     omp.atomic.write %x = %expr : memref<i32>, i32
-    omp.atomic.update %x = %x add %expr : memref<i32>, i32
+    omp.atomic.update %x : memref<i32> {
+    ^bb0(%xval: i32):
+      %newval = llvm.add %xval, %expr : i32
+      omp.yield (%newval : i32)
+    }
     omp.terminator
   }
   return
 }
 
 // -----
 
 func @omp_atomic_capture(%x: memref<i32>, %v: memref<i32>, %expr: i32) {
   omp.atomic.capture {
     // expected-error @below {{invalid sequence of operations in the capture region}}
-    omp.atomic.update %x = %x add %expr : memref<i32>, i32
+    omp.atomic.update %x : memref<i32> {
+    ^bb0(%xval: i32):
+      %newval = llvm.add %xval, %expr : i32
+      omp.yield (%newval : i32)
+    }
     omp.atomic.write %x = %expr : memref<i32>, i32
     omp.terminator
   }
   return
 }
 
 // -----
 
 func @omp_atomic_capture(%x: memref<i32>, %v: memref<i32>, %expr: i32) {
   omp.atomic.capture {
     // expected-error @below {{invalid sequence of operations in the capture region}}
     omp.atomic.write %x = %expr : memref<i32>, i32
     omp.atomic.read %v = %x : memref<i32>
     omp.terminator
   }
   return
 }
 
 // -----
 
 func @omp_atomic_capture(%x: memref<i32>, %y: memref<i32>, %v: memref<i32>, %expr: i32) {
   omp.atomic.capture {
     // expected-error @below {{updated variable in omp.atomic.update must be captured in second operation}}
-    omp.atomic.update %x = %x add %expr : memref<i32>, i32
+    omp.atomic.update %x : memref<i32> {
+    ^bb0(%xval: i32):
+      %newval = llvm.add %xval, %expr : i32
+      omp.yield (%newval : i32)
+    }
     omp.atomic.read %v = %y : memref<i32>
     omp.terminator
   }
 }
 
 // -----
 
 func @omp_atomic_capture(%x: memref<i32>, %y: memref<i32>, %v: memref<i32>, %expr: i32) {
   omp.atomic.capture {
     // expected-error @below {{captured variable in omp.atomic.read must be updated in second operation}}
     omp.atomic.read %v = %y : memref<i32>
-    omp.atomic.update %x = %x add %expr : memref<i32>, i32
+    omp.atomic.update %x : memref<i32> {
+    ^bb0(%xval: i32):
+      %newval = llvm.add %xval, %expr : i32
+      omp.yield (%newval : i32)
+    }
     omp.terminator
   }
 }
 
 // -----
 
 func @omp_atomic_capture(%x: memref<i32>, %y: memref<i32>, %v: memref<i32>, %expr: i32) {
   omp.atomic.capture {

mlir/test/Dialect/OpenMP/ops.mlir

func @omp_atomic_write(%addr : memref<i32>, %val : 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
+  // CHECK: omp.atomic.update %[[X]] : memref<i32>
+  // CHECK-NEXT: (%[[XVAL:.*]]: i32):
+  // CHECK-NEXT:   %[[NEWVAL:.*]] = llvm.add %[[XVAL]], %[[EXPR]] : i32
+  // CHECK-NEXT:   omp.yield(%[[NEWVAL]] : i32)
+  omp.atomic.update %x : memref<i32> {
+  ^bb0(%xval: i32):
+    %newval = llvm.add %xval, %expr : i32
+    omp.yield(%newval : i32)
+  }
+  // CHECK: omp.atomic.update %[[XBOOL]] : memref<i1>
+  // CHECK-NEXT: (%[[XVAL:.*]]: i1):
+  // CHECK-NEXT:   %[[NEWVAL:.*]] = llvm.and %[[XVAL]], %[[EXPRBOOL]] : i1
+  // CHECK-NEXT:   omp.yield(%[[NEWVAL]] : i1)
+  omp.atomic.update %xBool : memref<i1> {
+  ^bb0(%xval: i1):
+    %newval = llvm.and %xval, %exprBool : i1
+    omp.yield(%newval : i1)
+  }
+  // CHECK: omp.atomic.update %[[X]] : memref<i32>
+  // CHECK-NEXT: (%[[XVAL:.*]]: i32):
+  // CHECK-NEXT:   %[[NEWVAL:.*]] = llvm.shl %[[XVAL]], %[[EXPR]] : i32
+  // CHECK-NEXT:   omp.yield(%[[NEWVAL]] : i32)
+  // CHECK-NEXT: }
+  omp.atomic.update %x : memref<i32> {
+  ^bb0(%xval: i32):
+    %newval = llvm.shl %xval, %expr : i32
+    omp.yield(%newval : i32)
+  }
+  // CHECK: omp.atomic.update %[[X]] : memref<i32>
+  // CHECK-NEXT: (%[[XVAL:.*]]: i32):
+  // CHECK-NEXT:   %[[NEWVAL:.*]] = "llvm.intr.maxnum"(%[[XVAL]], %[[EXPR]]) : (i32, i32) -> i32
+  // CHECK-NEXT:   omp.yield(%[[NEWVAL]] : i32)
+  // CHECK-NEXT: }
+  omp.atomic.update %x : memref<i32> {
+  ^bb0(%xval: i32):
+    %newval = "llvm.intr.maxnum"(%xval, %expr) : (i32, i32) -> i32
+    omp.yield(%newval : i32)
+  }
+
+  // CHECK: omp.atomic.update %[[XBOOL]] : memref<i1>
+  // CHECK-NEXT: (%[[XVAL:.*]]: i1):
+  // CHECK-NEXT:   %[[NEWVAL:.*]] = llvm.icmp "eq" %[[XVAL]], %[[EXPRBOOL]] : i1
+  // CHECK-NEXT:   omp.yield(%[[NEWVAL]] : i1)
+  // }
+  omp.atomic.update %xBool : memref<i1> {
+  ^bb0(%xval: i1):
+    %newval = llvm.icmp "eq" %xval, %exprBool : i1
+    omp.yield(%newval : i1)
+  }
   return
 }
 
 // CHECK-LABEL: omp_atomic_capture
 // CHECK-SAME: (%[[v:.*]]: memref<i32>, %[[x:.*]]: memref<i32>, %[[expr:.*]]: i32)
 func @omp_atomic_capture(%v: memref<i32>, %x: memref<i32>, %expr: i32) {
   // CHECK: omp.atomic.capture{
-  // CHECK-NEXT: omp.atomic.update %[[x]] = %[[expr]] add %[[x]] : memref<i32>, i32
+  // CHECK-NEXT: omp.atomic.update %[[x]] : memref<i32>
+  // CHECK-NEXT: (%[[xval:.*]]: i32):
+  // CHECK-NEXT:   %[[newval:.*]] = llvm.add %[[xval]], %[[expr]] : i32
+  // CHECK-NEXT:   omp.yield(%[[newval]] : i32)
+  // CHECK-NEXT: }
   // CHECK-NEXT: omp.atomic.read %[[v]] = %[[x]] : memref<i32>
   // CHECK-NEXT: omp.terminator
   // CHECK-NEXT: }
   omp.atomic.capture{
-    omp.atomic.update %x = %expr add %x : memref<i32>, i32
+    omp.atomic.update %x : memref<i32> {
+    ^bb0(%xval: i32):
+      %newval = llvm.add %xval, %expr : i32
+      omp.yield(%newval : i32)
+    }
     omp.atomic.read %v = %x : memref<i32>
     omp.terminator
   }
   // CHECK: omp.atomic.capture{
   // CHECK-NEXT: omp.atomic.read %[[v]] = %[[x]] : memref<i32>
-  // CHECK-NEXT: omp.atomic.update %[[x]] = %[[expr]] add %[[x]] : memref<i32>, i32
+  // CHECK-NEXT: omp.atomic.update %[[x]] : memref<i32>
+  // CHECK-NEXT: (%[[xval:.*]]: i32):
+  // CHECK-NEXT:   %[[newval:.*]] = llvm.add %[[xval]], %[[expr]] : i32
+  // CHECK-NEXT:   omp.yield(%[[newval]] : i32)
+  // CHECK-NEXT: }
   // CHECK-NEXT: omp.terminator
   // CHECK-NEXT: }
   omp.atomic.capture{
     omp.atomic.read %v = %x : memref<i32>
-    omp.atomic.update %x = %expr add %x : memref<i32>, i32
+    omp.atomic.update %x : memref<i32> {
+    ^bb0(%xval: i32):
+      %newval = llvm.add %xval, %expr : i32
+      omp.yield(%newval : i32)
+    }
     omp.terminator
   }
   // CHECK: omp.atomic.capture{
   // CHECK-NEXT: omp.atomic.read %[[v]] = %[[x]] : memref<i32>
   // CHECK-NEXT: omp.atomic.write %[[x]] = %[[expr]] : memref<i32>, i32
   // CHECK-NEXT: omp.terminator
   // CHECK-NEXT: }
   omp.atomic.capture{