[OpenMP][mlir] Lowering for omp.atomic.update

[shraiysh]

Summary

This patch adds lowering from omp.atomic.update to LLVM IR. Whenever a
special LLVM IR instruction is available for the operation, atomicrmw
instruction is emitted, otherwise a compare-exchange loop based update
is emitted.

Depends on D119522

[ftynse]

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

954–955

Please do not commit commented-out code.

968

This comment is not helpful, we can see it is boilerplate. I would rather have most of it factored out into a function like getOpenMPLoc(IRBuilderBase &).

978

There doesn't seem to be a check for the region to only have two operations (binop and terminator), and this translation does not account for it but just ignores the remaining operations.

980

There doesn't seem to be any check that updateOp does have operands at all anywhere.

989

Please use better names than x and X in the same function, these are extremely hard to keep in one's head for more than a couple of lines.

993

Please consider a better name than ao.

998

MLIR uses camelBack names, IRB isn't something that it would use.

mlir/test/Target/LLVMIR/openmp-llvm.mlir

936

Note that MLIR tends to use "lowering" for _conversions_ between dialects rather than for _translations_ between MLIR and other IRs.

[shraiysh]

Thanks for the review @ftynse. Addressed comments.

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

968

For the function getOpenMPLoc, we would need opInst and moduleTranslation too. Plus, llvm::OpenMPIRBuilder::LocationDescription doesn't have a copy constructor, so that would also have to be passed. Also, the ompBuilder is used later in this function, so that cannot be pushed into the function. At this point, a function like the following doesn't feel like it is worth it. I have removed the comment as suggested. Let me know if you have any suggestions for an alternative.

getOpenMPLoc(Operation &opInst, llvm::IRBuilderBase &builder, LLVM::ModuleTranslation moduleTranslation, llvm::OpenMPIRBuilder::LocationDescription &ompLoc)

[ftynse]

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

968

I don't have a problem with multi-argument signatures. ModuleTranslation needs to be passed by reference as well. Can't we just make LocationDescription copy-constructible (or movable)?

978

Please add a test for user-visible error messages.

[shraiysh]

Handle multiple operations in update region.
Added tests for same.
Addressed comments.
Rebase with main.

[shraiysh]

Requesting review again, because added translation for multiple operations in update region.

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

968

Now the constructor directly accepts builder, so this is not needed. Thanks for the discussion.

[shraiysh]

Rebase with main. Ping!

[shraiysh]

Rebase with main. Ping!

[peixin]

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

961

OpenMP update -> OpenMP atomic update

mlir/test/Target/LLVMIR/openmp-llvm.mlir

1028

I don't think you should process the lowering of the "expr" in MLIR stage. For different programing languages, the expr may have different structure. Can you try to analyze the "expr" when lowering parse-tree to MLIR? That is, the input for the "expr" in atomic update is always one mlir value and the LLVM IR generated for analyzing the "expr" is outside the atomic update region. Another reason for doing this is that "expr" cannot access the storage location designated by x, and the operations inside the atomic update region will be very clean. That is, you will always only have two instructions in atomic update region, one for update "xval", the other one for yield the "newval".

[peixin]

mlir/test/Target/LLVMIR/openmp-llvm.mlir

1028

Similar to https://github.com/flang-compiler/f18-llvm-project/blob/9e96e8e4b1f84db0c30eaf4798f094f571b1c565/flang/lib/Lower/OpenMP.cpp#L320-L322. Using genExprValue should provide you the result value of the "expr".

[shraiysh]

Thanks for the review @peixin. Addressed comments.

mlir/test/Target/LLVMIR/openmp-llvm.mlir

1028

The lowering of %expr is not handled in the MLIR region. It is assumed to be a single lowered value, with lowering handled in Flang as you suggested. The idea behind allowing more than one operations was to have the following in MLIR

omp.atomic.update %x {
^bb0(%xval):
  %xnew = f(%xval, %expr)
  omp.yield %xnew
}

when we cannot express x = f(x, expr) in MLIR with a single instruction (for example f(x,expr) = (x^expr)^expr)) as we might not have the appropriate single op in the dialect (which can be anything including but not limited to LLVM IR Dialect and FIR Dialect) and that should be okay. In this case, we cannot have f(x, expr) such that the operation can be done in one op. By allowing multiple operations in the region, we provide the liberty of handling this in multiple ops. The generated LLVM IR will work appropriately. Another thing is that the MLIR Operation cannot depend on Flang, because this operation will be used by multiple dialects, so it has to be generic enough to support simple and complicated Dialects. Should I put a comment about this in the code itself?

[shraiysh]

Rebase with main.

[peixin]

mlir/test/Target/LLVMIR/openmp-llvm.mlir

1028

However, this is for OpenMP dialect. Currently, OpenMP standard only defines C, C++, and fortran, and all the update statements for them can be transformed as one single instruction and one omp.yield operation. f(x,expr) = (x^expr)^expr is obviously not allowed in OpenMP standard. If someone want to use this generic atomic update, they should implement one in other dialects in somewhere else, instead of OpenMP dialect here. What you currently implemented is one generic atomic update, not restricted to OpenMP atomic update.

[shraiysh]

mlir/test/Target/LLVMIR/openmp-llvm.mlir

1028

can be transformed as one single instruction and one omp.yield operation

Yes, in the LLVM IR Dialect. We cannot say the same for all the other dialects with certainty. (or can we? Feel free to correct me if this sounds absurd.) Maybe we can put a check during lowering to LLVM IR for having exactly two operations, but enforcing it in the operation verifier itself is not entirely correct as the internal op can be in any dialect (In some example dialect, a few atomic update statements for some types might need more than a single op in that dialect).

What you currently implemented is one generic atomic update, not restricted to OpenMP atomic update.

That is correct, because only the dialect handling the operation knows how to update it appropriately (including how many ops to take for that update).

[peixin]

mlir/test/Target/LLVMIR/openmp-llvm.mlir

1028

I think we can say that. Actually, OpenMP standard says that. Please notice that OpenMPToLLVMIRTranslation.cpp is for OpenMP. In semantics, having more than one single op is not correct. If someone needs it, they should not use OpenMP stomic update since it is not allowed in OpenMP atomic update.

[peixin]

mlir/test/Target/LLVMIR/openmp-llvm.mlir

1028

OK. I noticed that you changed the omp.atomic.update dialect. https://github.com/llvm/llvm-project/blob/0e38b295435bf49c21e9ff97354a855ac8c78c7e/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td#L637-L638 Here, you referred to OpenMP standard, but you make it too generic so that this opertion is beyond the OpenMP atomic update. I think we need to discuss about the change in D119522. @kiranchandramohan What do you think?

[shraiysh]

Addressed reviewer comments. Rebase with main.

[peixin]

The lowering for memory-order-clause and hint clause will be in future work, right?

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

961

Can you add some comments to describe your design method to tell others why you don't do it similar to atomic read/write? You have expalined this in slack. Can you move those info here? Also, comment that the input is checked in semantic analysis for each programming language.

980

Should this be one if condition or one assert? If I understand correctly, the input not satisfying the if condition is not one standard OpenMP atomic update statement. Right?

983

Will this return BAD_BINOP for int128, fp128, x86_fp80? They seems not to be supported in AtomicRMWInst. Can you add one test case for one of them to represent the corner case?

Also, what about the complex data type?

[shraiysh]

Thanks for the review @peixin. Addressed comments.

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

961

Sure. Added in the description for omp.atomic.update operation in OpenMPOps.td.

980

Yes, thank you. Added.

983

I have added a TODO in OMPIRBuilder.cpp for this. It goes into error at the moment. This is a bug in OmpIRBuilder.

// TODO: handle the case where XElemTy is not byte-sized or not a power of 2.

[shraiysh]

Rebase with main

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

Value *OpenMPIRBuilder::emitRMWOpAsInstruction(Value *Src1, Value *Src2,
   }
   llvm_unreachable("Unsupported atomic update operation");
 }
 
 std::pair<Value *, Value *> OpenMPIRBuilder::emitAtomicUpdate(
     InsertPointTy AllocaIP, Value *X, Type *XElemTy, Value *Expr,
     AtomicOrdering AO, AtomicRMWInst::BinOp RMWOp,
     AtomicUpdateCallbackTy &UpdateOp, bool VolatileX, bool IsXBinopExpr) {
+  // TODO: handle the case where XElemTy is not byte-sized or not a power of 2
+  // or a complex datatype.
   bool DoCmpExch = (RMWOp == AtomicRMWInst::BAD_BINOP) ||
                    (RMWOp == AtomicRMWInst::FAdd) ||
                    (RMWOp == AtomicRMWInst::FSub) ||
                    (RMWOp == AtomicRMWInst::Sub && !IsXBinopExpr) || !XElemTy;
 
   std::pair<Value *, Value *> Res;
   if (XElemTy->isIntegerTy() && !DoCmpExch) {
     Res.first = Builder.CreateAtomicRMW(RMWOp, X, Expr, llvm::MaybeAlign(), AO);

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

let description = [{
     `memory_order` indicates the memory ordering behavior of the construct. It
     can be one of `seq_cst`, `release` 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.
+
+    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 OpenMP_PointerLikeType:$x,
                        DefaultValuedAttr<I64Attr, "0">:$hint_val,
                        OptionalAttr<MemoryOrderKindAttr>:$memory_order_val);
   let regions = (region SizedRegion<1>:$region);
   let assemblyFormat = [{
     oilist( `memory_order` `(` custom<ClauseAttr>($memory_order_val) `)`

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

if (region().getNumArguments() != 1)
     return emitError("the region must accept exactly one argument");
 
   if (x().getType().cast<PointerLikeType>().getElementType() !=
       region().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");
   }
 
+  if (region().front().getOperations().size() < 2)
+    return emitError() << "the update region must have at least two operations "
+                          "(binop and terminator)";
+
   YieldOp yieldOp = *region().getOps<YieldOp>().begin();
 
   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();
 }

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

convertOmpAtomicWrite(Operation &opInst, llvm::IRBuilderBase &builder,
   llvm::Value *dest = moduleTranslation.lookupValue(writeOp.address());
   llvm::Type *ty = moduleTranslation.convertType(writeOp.value().getType());
   llvm::OpenMPIRBuilder::AtomicOpValue x = {dest, ty, /*isSigned=*/false,
                                             /*isVolatile=*/false};
   builder.restoreIP(ompBuilder->createAtomicWrite(ompLoc, x, expr, ao));
   return success();
 }
 
+/// Converts an LLVM dialect binary operation to the corresponding enum value
+/// for `atomicrmw` supported binary operation.
+llvm::AtomicRMWInst::BinOp convertBinOpToAtomic(Operation &op) {
+  return llvm::TypeSwitch<Operation *, llvm::AtomicRMWInst::BinOp>(&op)
+      .Case([&](LLVM::AddOp) { return llvm::AtomicRMWInst::BinOp::Add; })
+      .Case([&](LLVM::SubOp) { return llvm::AtomicRMWInst::BinOp::Sub; })
+      .Case([&](LLVM::AndOp) { return llvm::AtomicRMWInst::BinOp::And; })
+      .Case([&](LLVM::OrOp) { return llvm::AtomicRMWInst::BinOp::Or; })
+      .Case([&](LLVM::XOrOp) { return llvm::AtomicRMWInst::BinOp::Xor; })
+      .Case([&](LLVM::UMaxOp) { return llvm::AtomicRMWInst::BinOp::UMax; })
+      .Case([&](LLVM::UMinOp) { return llvm::AtomicRMWInst::BinOp::UMin; })
+      .Case([&](LLVM::FAddOp) { return llvm::AtomicRMWInst::BinOp::FAdd; })
+      .Case([&](LLVM::FSubOp) { return llvm::AtomicRMWInst::BinOp::FSub; })
+      .Default(llvm::AtomicRMWInst::BinOp::BAD_BINOP);
+}
+
+/// Converts an OpenMP atomic update operation using OpenMPIRBuilder.
+static LogicalResult
+convertOmpAtomicUpdate(omp::AtomicUpdateOp &opInst,
+                       llvm::IRBuilderBase &builder,
+                       LLVM::ModuleTranslation &moduleTranslation) {
+  llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
+  llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder);
+
+  // Convert values and types.
+  auto &innerOpList = opInst.region().front().getOperations();
+  if (innerOpList.size() != 2)
+    return opInst.emitError("exactly two operations are allowed inside an "
+                            "atomic update region while lowering to LLVM IR");
+
+  Operation &innerUpdateOp = innerOpList.front();
+
+  if (innerUpdateOp.getNumOperands() != 2 ||
+      !llvm::is_contained(innerUpdateOp.getOperands(),
+                          opInst.getRegion().getArgument(0)))
+    return opInst.emitError(
+        "the update operation inside the region must be a binary operation and "
+        "that update operation must have the region argument as an operand");
+
+  llvm::AtomicRMWInst::BinOp binop = convertBinOpToAtomic(innerUpdateOp);
+
+  bool isXBinopExpr =
+      innerUpdateOp.getNumOperands() > 0 &&
+      innerUpdateOp.getOperand(0) == opInst.getRegion().getArgument(0);
+
+  mlir::Value mlirExpr = (isXBinopExpr ? innerUpdateOp.getOperand(1)
+                                       : innerUpdateOp.getOperand(0));
+  llvm::Value *llvmExpr = moduleTranslation.lookupValue(mlirExpr);
+  llvm::Value *llvmX = moduleTranslation.lookupValue(opInst.x());
+  LLVM::LLVMPointerType mlirXType =
+      opInst.x().getType().cast<LLVM::LLVMPointerType>();
+  llvm::Type *llvmXElementType =
+      moduleTranslation.convertType(mlirXType.getElementType());
+  llvm::OpenMPIRBuilder::AtomicOpValue llvmAtomicX = {llvmX, llvmXElementType,
+                                                      /*isSigned=*/false,
+                                                      /*isVolatile=*/false};
+
+  llvm::AtomicOrdering atomicOrdering =
+      convertAtomicOrdering(opInst.memory_order_val());
+
+  // Generate update code.
+  LogicalResult updateGenStatus = success();
+  auto updateFn = [&opInst, &moduleTranslation, &updateGenStatus](
+                      llvm::Value *atomicx,
+                      llvm::IRBuilder<> &builder) -> llvm::Value * {
+    Block &bb = *opInst.region().begin();
+    moduleTranslation.mapValue(*opInst.region().args_begin(), atomicx);
+    moduleTranslation.mapBlock(&bb, builder.GetInsertBlock());
+    if (failed(moduleTranslation.convertBlock(bb, true, builder))) {
+      updateGenStatus = (opInst.emitError()
+                         << "unable to convert update operation to llvm IR");
+      return nullptr;
+    }
+    omp::YieldOp yieldop = dyn_cast<omp::YieldOp>(bb.getTerminator());
+    assert(yieldop && yieldop.results().size() == 1 &&
+           "terminator must be omp.yield op and it must have exactly one "
+           "argument");
+    return moduleTranslation.lookupValue(yieldop.results()[0]);
+  };
+
+  // Handle ambiguous alloca, if any.
+  auto allocaIP = findAllocaInsertPoint(builder, moduleTranslation);
+  llvm::UnreachableInst *unreachableInst;
+  if (allocaIP.getPoint() == ompLoc.IP.getPoint()) {
+    // Same point => split basic block and make them unambigous.
+    unreachableInst = builder.CreateUnreachable();
+    builder.SetInsertPoint(builder.GetInsertBlock()->splitBasicBlock(
+        unreachableInst, "alloca_split"));
+    ompLoc.IP = builder.saveIP();
+    unreachableInst->removeFromParent();
+  }
+  builder.restoreIP(ompBuilder->createAtomicUpdate(
+      ompLoc, findAllocaInsertPoint(builder, moduleTranslation), llvmAtomicX,
+      llvmExpr, atomicOrdering, binop, updateFn, isXBinopExpr));
+  return updateGenStatus;
+}
+
 /// Converts an OpenMP reduction operation using OpenMPIRBuilder. Expects the
 /// mapping between reduction variables and their private equivalents to have
 /// been stored on the ModuleTranslation stack. Currently only supports
 /// reduction within WsLoopOp, but can be easily extended.
 static LogicalResult
 convertOmpReductionOp(omp::ReductionOp reductionOp,
                       llvm::IRBuilderBase &builder,
                       LLVM::ModuleTranslation &moduleTranslation) {
return llvm::TypeSwitch<Operation *, LogicalResult>(op)
         return convertOmpWsLoop(*op, builder, moduleTranslation);
       })
       .Case([&](omp::AtomicReadOp) {
         return convertOmpAtomicRead(*op, builder, moduleTranslation);
       })
       .Case([&](omp::AtomicWriteOp) {
         return convertOmpAtomicWrite(*op, builder, moduleTranslation);
       })
+      .Case([&](omp::AtomicUpdateOp op) {
+        return convertOmpAtomicUpdate(op, builder, moduleTranslation);
+      })
       .Case([&](omp::SectionsOp) {
         return convertOmpSections(*op, builder, moduleTranslation);
       })
       .Case<omp::YieldOp, omp::TerminatorOp, omp::ReductionDeclareOp,
             omp::CriticalDeclareOp>([](auto op) {
         // `yield` and `terminator` can be just omitted. The block structure
         // was created in the region that handles their parent operation.
         // `reduction.declare` will be used by reductions and is not

mlir/test/Dialect/OpenMP/invalid.mlir

^bb0(%xval: i32, %tmp: i32):
     %newval = llvm.add %xval, %expr : i32
     omp.yield (%newval : i32)
   }
   return
 }
 
 // -----
 
+func @omp_atomic_update9(%x: memref<i32>, %expr: i32) {
+  // expected-error @below {{the update region must have at least two operations (binop and terminator)}}
+  omp.atomic.update %x : memref<i32> {
+  ^bb0(%xval: i32):
+    omp.yield (%xval : i32)
+  }
+  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 {
     omp.atomic.read %v = %x : memref<i32>
     omp.terminator
   }
   return
 }

mlir/test/Target/LLVMIR/openmp-llvm-invalid.mlir

+// RUN: mlir-translate -mlir-to-llvmir %s -split-input-file -verify-diagnostics
+
+// Checking translation when the update is carried out by using more than one op
+// in the region.
+llvm.func @omp_atomic_update_multiple_step_update(%x: !llvm.ptr<i32>, %expr: i32) {
+  // expected-error @+2 {{exactly two operations are allowed inside an atomic update region while lowering to LLVM IR}}
+  // expected-error @+1 {{LLVM Translation failed for operation: omp.atomic.update}}
+  omp.atomic.update %x : !llvm.ptr<i32> {
+  ^bb0(%xval: i32):
+    %t1 = llvm.mul %xval, %expr : i32
+    %t2 = llvm.sdiv %t1, %expr : i32
+    %newval = llvm.add %xval, %t2 : i32
+    omp.yield(%newval : i32)
+  }
+  llvm.return
+}
+
+// -----
+
+// Checking translation when the captured variable is not used in the inner
+// update operation
+llvm.func @omp_atomic_update_multiple_step_update(%x: !llvm.ptr<i32>, %expr: i32) {
+  // expected-error @+2 {{the update operation inside the region must be a binary operation and that update operation must have the region argument as an operand}}
+  // expected-error @+1 {{LLVM Translation failed for operation: omp.atomic.update}}
+  omp.atomic.update %x : !llvm.ptr<i32> {
+  ^bb0(%xval: i32):
+    %newval = llvm.mul %expr, %expr : i32
+    omp.yield(%newval : i32)
+  }
+  llvm.return
+}

mlir/test/Target/LLVMIR/openmp-llvm.mlir

llvm.func @omp_atomic_write(%x: !llvm.ptr<i32>, %expr: i32) -> () {
   omp.atomic.write %x = %expr memory_order(release) : !llvm.ptr<i32>, i32
   // CHECK: store atomic i32 %[[expr]], i32* %[[x]] monotonic, align 4
   omp.atomic.write %x = %expr memory_order(relaxed) : !llvm.ptr<i32>, i32
   llvm.return
 }
 
 // -----
 
+// Checking simple atomicrmw and cmpxchg based translation. This also checks for
+// ambigous alloca insert point by putting llvm.mul as the first update operation.
+// CHECK-LABEL: @omp_atomic_update
+// CHECK-SAME: (i32* %[[x:.*]], i32 %[[expr:.*]], i1* %[[xbool:.*]], i1 %[[exprbool:.*]])
+llvm.func @omp_atomic_update(%x:!llvm.ptr<i32>, %expr: i32, %xbool: !llvm.ptr<i1>, %exprbool: i1) {
+  // CHECK: %[[t1:.*]] = mul i32 %[[x_old:.*]], %[[expr]]
+  // CHECK: store i32 %[[t1]], i32* %[[x_new:.*]]
+  // CHECK: %[[t2:.*]] = load i32, i32* %[[x_new]]
+  // CHECK: cmpxchg i32* %[[x]], i32 %[[x_old]], i32 %[[t2]]
+  omp.atomic.update %x : !llvm.ptr<i32> {
+  ^bb0(%xval: i32):
+    %newval = llvm.mul %xval, %expr : i32
+    omp.yield(%newval : i32)
+  }
+  // CHECK: atomicrmw add i32* %[[x]], i32 %[[expr]] monotonic
+  omp.atomic.update %x : !llvm.ptr<i32> {
+  ^bb0(%xval: i32):
+    %newval = llvm.add %xval, %expr : i32
+    omp.yield(%newval : i32)
+  }
+  llvm.return
+}
+
+// -----
+
+// Checking an order-dependent operation when the order is `expr binop x`
+// CHECK-LABEL: @omp_atomic_update_ordering
+// CHECK-SAME: (i32* %[[x:.*]], i32 %[[expr:.*]])
+llvm.func @omp_atomic_update_ordering(%x:!llvm.ptr<i32>, %expr: i32) {
+  // CHECK: %[[t1:.*]] = shl i32 %[[expr]], %[[x_old:[^ ,]*]]
+  // CHECK: store i32 %[[t1]], i32* %[[x_new:.*]]
+  // CHECK: %[[t2:.*]] = load i32, i32* %[[x_new]]
+  // CHECK: cmpxchg i32* %[[x]], i32 %[[x_old]], i32 %[[t2]]
+  omp.atomic.update %x : !llvm.ptr<i32> {
+  ^bb0(%xval: i32):
+    %newval = llvm.shl %expr, %xval : i32
+    omp.yield(%newval : i32)
+  }
+  llvm.return
+}
+
+// -----
+
+// Checking an order-dependent operation when the order is `x binop expr`
+// CHECK-LABEL: @omp_atomic_update_ordering
+// CHECK-SAME: (i32* %[[x:.*]], i32 %[[expr:.*]])
+llvm.func @omp_atomic_update_ordering(%x:!llvm.ptr<i32>, %expr: i32) {
+  // CHECK: %[[t1:.*]] = shl i32 %[[x_old:.*]], %[[expr]]
+  // CHECK: store i32 %[[t1]], i32* %[[x_new:.*]]
+  // CHECK: %[[t2:.*]] = load i32, i32* %[[x_new]]
+  // CHECK: cmpxchg i32* %[[x]], i32 %[[x_old]], i32 %[[t2]] monotonic
+  omp.atomic.update %x : !llvm.ptr<i32> {
+  ^bb0(%xval: i32):
+    %newval = llvm.shl %xval, %expr : i32
+    omp.yield(%newval : i32)
+  }
+  llvm.return
+}
+
+// -----
+
+// Checking intrinsic translation.
+// CHECK-LABEL: @omp_atomic_update_intrinsic
+// CHECK-SAME: (i32* %[[x:.*]], i32 %[[expr:.*]])
+llvm.func @omp_atomic_update_intrinsic(%x:!llvm.ptr<i32>, %expr: i32) {
+  // CHECK: %[[t1:.*]] = call i32 @llvm.smax.i32(i32 %[[x_old:.*]], i32 %[[expr]])
+  // CHECK: store i32 %[[t1]], i32* %[[x_new:.*]]
+  // CHECK: %[[t2:.*]] = load i32, i32* %[[x_new]]
+  // CHECK: cmpxchg i32* %[[x]], i32 %[[x_old]], i32 %[[t2]]
+  omp.atomic.update %x : !llvm.ptr<i32> {
+  ^bb0(%xval: i32):
+    %newval = "llvm.intr.smax"(%xval, %expr) : (i32, i32) -> i32
+    omp.yield(%newval : i32)
+  }
+  llvm.return
+}
+
+// -----
+
 // CHECK-LABEL: @omp_sections_empty
 llvm.func @omp_sections_empty() -> () {
   omp.sections {
     omp.terminator
   }
   // CHECK-NEXT: ret void
   llvm.return
 }
 
 // -----
 
 // Check IR generation for simple empty sections. This only checks the overall
 // shape of the IR, detailed checking is done by the OpenMPIRBuilder.
 
 // CHECK-LABEL: @omp_sections_trivial
 llvm.func @omp_sections_trivial() -> () {
   // CHECK:   br label %[[PREHEADER:.*]]
 
   // CHECK: [[PREHEADER]]:
   // CHECK:   %{{.*}} = call i32 @__kmpc_global_thread_num({{.*}})
   // CHECK:   call void @__kmpc_for_static_init_4u({{.*}})
   // CHECK:   br label %[[HEADER:.*]]