[Flang][OpenMP] Fix for atomic lowering with HLFIR

[kiranchandramohan]

Summary

Atomic update operation is modelled in OpenMP dialect as
an operation that takes a reference to the operation being
updated. It also contains a region that will perform the
update. The block argument represents the loaded value from
the update location and the Yield operation is the value
that should be stored for the update.

OpenMP FIR lowering binds the value loaded from the update
address to the SymbolAddress. HLFIR lowering does not permit
SymbolAddresses to be a value. To work around this, the
lowering is now performed in two steps. First the body of
the atomic update is lowered into an SCF execute_region
operation. Then this is copied into the omp.atomic_update
as a second step that performs the following:
-> Create an omp.atomic_update with the block argument of
the correct type.
-> Copy the operations from the SCF execute_region. Convert
the scf.yield to an omp.yield.
-> Remove the loads of the update location and replace all
uses with the block argument.

[kiranchandramohan]

TODO:
-> Add a detailed explanation
-> Add a test for HLFIR.

flang/lib/Lower/OpenMP.cpp

3035

Lowering is in three steps :

subroutine sb
  integer :: a, b
  !$omp atomic update
    a = a + b
end subroutine

Lower to scf.execute_region_op

func.func @_QPsb() {
  %0 = fir.alloca i32 {bindc_name = "a", uniq_name = "_QFsbEa"}
  %1 = fir.alloca i32 {bindc_name = "b", uniq_name = "_QFsbEb"}
  %2 = scf.execute_region -> i32 {
    %3 = fir.load %0 : !fir.ref<i32>
    %4 = fir.load %1 : !fir.ref<i32>
    %5 = arith.addi %3, %4 : i32
    scf.yield %5 : i32
  }
  return
}

Move out all the non-update loads.

func.func @_QPsb() {
  %0 = fir.alloca i32 {bindc_name = "a", uniq_name = "_QFsbEa"}
  %1 = fir.alloca i32 {bindc_name = "b", uniq_name = "_QFsbEb"}
  %2 = fir.load %1 : !fir.ref<i32>
  %3 = scf.execute_region -> i32 {
    %4 = fir.load %0 : !fir.ref<i32>
    %5 = arith.addi %4, %2 : i32
    scf.yield %5 : i32
  }
  return
}

Convert to atomic.update.

func.func @_QPsb() {
  %0 = fir.alloca i32 {bindc_name = "a", uniq_name = "_QFsbEa"}
  %1 = fir.alloca i32 {bindc_name = "b", uniq_name = "_QFsbEb"}
  %2 = fir.load %1 : !fir.ref<i32>
  omp.atomic.update   %0 : !fir.ref<i32> {
  ^bb0(%arg0: i32):
    %3 = arith.addi %arg0, %2 : i32
    omp.yield(%3 : i32)
  }
  return
}

[tblah]

flang/lib/Lower/OpenMP.cpp

3035

Why did you decide to lower to scf.execute_region first then convert to an atomic.update?

[kiranchandramohan]

Reduce the scope of changes to only make omp.atomic_update work with
HLFIR. (Hoisting can come in a separate patch)
Add an hlfir test.

[kiranchandramohan]

flang/lib/Lower/OpenMP.cpp

3035

HLFIR lowering only permits a reference to be bound to SymbolAddresses. omp.atomic_update has a block argument that models the value loaded from the Address and this is used inside the atomic_update region. This was previously achieved by mapping the SymbolAddress to the block argument. This does not work anymore, so it is first lowered temporarily to an scf.execute_region and then the atomic_update operation is carefully constructed from the scf.execute_region.
Note: Have added some additional explanation to the summary.

[tblah]

LG from a HLFIR perspective

[vzakhari]

@razvanlupusoru, just FYI

[razvanlupusoru]

Innovative approach with the intermediate step!

The change seems reasonable to me - but I think your comment showing the multiple step transformation belongs as a comment in the code instead of just summarized in the description. Also, you may want to further break up genOmpAtomicUpdateStatement into multiple functions - feels like too much is happening in a single spot now.

(I am currently looking into reusing this logic for OpenACC so I will build my change on top of this one) :)

Thanks! Great work!

[kiranchandramohan]

Add code comments inline as suggested by @razvanlupusoru.
Will split the code in a follow-up patch.

[razvanlupusoru]

Thank you!

flang/lib/Lower/OpenMP.cpp

 #include "flang/Optimizer/Builder/BoxValue.h"
 #include "flang/Optimizer/Builder/FIRBuilder.h"
 #include "flang/Optimizer/Builder/Todo.h"
 #include "flang/Optimizer/HLFIR/HLFIROps.h"
 #include "flang/Parser/parse-tree.h"
 #include "flang/Semantics/openmp-directive-sets.h"
 #include "flang/Semantics/tools.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "llvm/Frontend/OpenMP/OMPConstants.h"
 
 using DeclareTargetCapturePair =
     std::pair<mlir::omp::DeclareTargetCaptureClause,
               Fortran::semantics::Symbol>;
 
 //===----------------------------------------------------------------------===//
 // Common helper functions
static void genOmpAtomicUpdateStatement(
   mlir::IntegerAttr hint = nullptr;
   mlir::omp::ClauseMemoryOrderKindAttr memoryOrder = nullptr;
   if (leftHandClauseList)
     genOmpAtomicHintAndMemoryOrderClauses(converter, *leftHandClauseList, hint,
                                           memoryOrder);
   if (rightHandClauseList)
     genOmpAtomicHintAndMemoryOrderClauses(converter, *rightHandClauseList, hint,
                                           memoryOrder);
-  auto atomicUpdateOp = firOpBuilder.create<mlir::omp::AtomicUpdateOp>(
-      currentLocation, lhsAddr, hint, memoryOrder);
-
-  //// Generate body of Atomic Update operation
-  // If an argument for the region is provided then create the block with that
-  // argument. Also update the symbol's address with the argument mlir value.
-  llvm::SmallVector<mlir::Type> varTys = {varType};
-  llvm::SmallVector<mlir::Location> locs = {currentLocation};
-  firOpBuilder.createBlock(&atomicUpdateOp.getRegion(), {}, varTys, locs);
-  mlir::Value val =
-      fir::getBase(atomicUpdateOp.getRegion().front().getArgument(0));
   const auto *varDesignator =
       std::get_if<Fortran::common::Indirection<Fortran::parser::Designator>>(
           &assignmentStmtVariable.u);
   assert(varDesignator && "Variable designator for atomic update assignment "
                           "statement does not exist");
   const Fortran::parser::Name *name =
       Fortran::semantics::getDesignatorNameIfDataRef(varDesignator->value());
   if (!name)
     TODO(converter.getCurrentLocation(),
          "Array references as atomic update variable");
   assert(name && name->symbol &&
          "No symbol attached to atomic update variable");
-  converter.bindSymbol(*name->symbol, val);
-  // Set the insert for the terminator operation to go at the end of the
-  // block.
-  mlir::Block &block = atomicUpdateOp.getRegion().back();
-  firOpBuilder.setInsertionPointToEnd(&block);
+  if (Fortran::semantics::IsAllocatableOrPointer(name->symbol->GetUltimate()))
+    converter.bindSymbol(*name->symbol, lhsAddr);
 
+  // Temporarily lower the atomic update into an ExecuteRegion Op
+  auto tempOp =
+      firOpBuilder.create<mlir::scf::ExecuteRegionOp>(currentLocation, varType);
+  firOpBuilder.createBlock(&tempOp.getRegion());
+  mlir::Block &block = tempOp.getRegion().back();
+  firOpBuilder.setInsertionPointToEnd(&block);
   Fortran::lower::StatementContext stmtCtx;
   mlir::Value rhsExpr = fir::getBase(converter.genExprValue(
       *Fortran::semantics::GetExpr(assignmentStmtExpr), stmtCtx));
   mlir::Value convertResult =
       firOpBuilder.createConvert(currentLocation, varType, rhsExpr);
   // Insert the terminator: YieldOp.
-  firOpBuilder.create<mlir::omp::YieldOp>(currentLocation, convertResult);
-  // Reset the insert point to before the terminator.
+  firOpBuilder.create<mlir::scf::YieldOp>(currentLocation, convertResult);
   firOpBuilder.setInsertionPointToStart(&block);
+
+  // Hoist all the operations that are not transitively dependent on the
+  // variable that is going to be updated.
+  mlir::Value updateVar = converter.getSymbolAddress(*name->symbol);
+  if (auto decl = updateVar.getDefiningOp<hlfir::DeclareOp>())
+    updateVar = decl.getBase();
+  llvm::SmallVector<mlir::OpOperand *> dependentOperands;
+  for (mlir::OpOperand &use : updateVar.getUses())
+    dependentOperands.push_back(&use);
+
+  llvm::SmallVector<mlir::Operation *> dependentOps;
+  while (!dependentOperands.empty()) {
+    mlir::OpOperand *use = dependentOperands.pop_back_val();
+    mlir::Operation *op = use->getOwner();
+    dependentOps.push_back(op);
+    for (mlir::OpResult result : op->getResults()) {
+      for (mlir::OpOperand &u : result.getUses()) {
+        dependentOperands.push_back(&u);
+      }
+    }
+  }
+
+  llvm::SmallVector<mlir::Operation *> opsToMove;
+  for (mlir::Operation &op : tempOp.getRegion().getOps())
+    if (llvm::find(dependentOps, &op) == dependentOps.end())
+      opsToMove.push_back(&op);
+
+  mlir::Operation *hoistPointOp =
+      tempOp->getParentOfType<mlir::omp::AtomicCaptureOp>();
+  if (!hoistPointOp)
+    hoistPointOp = tempOp.getOperation();
+  for (mlir::Operation *op : opsToMove)
+    op->moveBefore(hoistPointOp);
+
+  // Now create the AtomicUpdateOp using the Operations in the temporary
+  // SCF Execute Region Op.
+  firOpBuilder.setInsertionPointAfter(tempOp);
+  auto atomicUpdateOp = firOpBuilder.create<mlir::omp::AtomicUpdateOp>(
+      currentLocation, updateVar, hint, memoryOrder);
+
+  llvm::SmallVector<mlir::Type> varTys = {varType};
+  llvm::SmallVector<mlir::Location> locs = {currentLocation};
+  firOpBuilder.createBlock(&atomicUpdateOp.getRegion(), {}, varTys, locs);
+  mlir::Value val =
+      fir::getBase(atomicUpdateOp.getRegion().front().getArgument(0));
+
+  llvm::SmallVector<mlir::Operation *> ops;
+  for (mlir::Operation &op : tempOp.getRegion().getOps())
+    ops.push_back(&op);
+
+  // SCF Yield is converted to OMP Yield. All other operations are copied
+  for (mlir::Operation *op : ops) {
+    if (auto y = mlir::dyn_cast<mlir::scf::YieldOp>(op)) {
+      firOpBuilder.setInsertionPointToEnd(&atomicUpdateOp.getRegion().front());
+      firOpBuilder.create<mlir::omp::YieldOp>(currentLocation, y.getResults());
+      op->erase();
+    } else {
+      op->remove();
+      atomicUpdateOp.getRegion().front().push_back(op);
+    }
+  }
+
+  // Remove the load and replace all uses of load with the block argument
+  for (mlir::Operation &op : atomicUpdateOp.getRegion().getOps()) {
+    if (auto y = mlir::dyn_cast<fir::LoadOp>(&op)) {
+      y.getRes().replaceAllUsesWith(val);
+      break;
+    }
+  }
+
+  tempOp.erase();
 }
 
 static void
 genOmpAtomicWrite(Fortran::lower::AbstractConverter &converter,
                   Fortran::lower::pft::Evaluation &eval,
                   const Fortran::parser::OmpAtomicWrite &atomicWrite) {
   // Get the value and address of atomic write operands.
   const Fortran::parser::OmpAtomicClauseList &rightHandClauseList =

flang/test/Lower/OpenMP/atomic-capture.f90

 ! RUN: %flang_fc1 -emit-fir -fopenmp %s -o - | FileCheck %s
 
 ! This test checks the lowering of atomic capture
 
 program OmpAtomicCapture
     use omp_lib                                                                                                       
     integer :: x, y
 
 !CHECK: %[[X:.*]] = fir.alloca i32 {bindc_name = "x", uniq_name = "_QFEx"}
 !CHECK: %[[Y:.*]] = fir.alloca i32 {bindc_name = "y", uniq_name = "_QFEy"}
+!CHECK: %[[temp:.*]] = fir.load %[[X]] : !fir.ref<i32>
 !CHECK: omp.atomic.capture   memory_order(release) {
 !CHECK: omp.atomic.read %[[X]] = %[[Y]] : !fir.ref<i32>
 !CHECK: omp.atomic.update %[[Y]] : !fir.ref<i32> {
 !CHECK: ^bb0(%[[ARG:.*]]: i32):
-!CHECK: %[[temp:.*]] = fir.load %[[X]] : !fir.ref<i32>
 !CHECK: %[[result:.*]] = arith.addi %[[temp]], %[[ARG]] : i32
 !CHECK: omp.yield(%[[result]] : i32)
 !CHECK: }
 !CHECK: }
 
     !$omp atomic capture release
         x = y
         y = x + y
     !$omp end atomic
 
 
+!CHECK: %[[temp:.*]] = fir.load %[[X]] : !fir.ref<i32>
 !CHECK: omp.atomic.capture   hint(uncontended) {
 !CHECK: omp.atomic.update %[[Y]] : !fir.ref<i32> {
 !CHECK: ^bb0(%[[ARG:.*]]: i32):
-!CHECK: %[[temp:.*]] = fir.load %[[X]] : !fir.ref<i32>
 !CHECK: %[[result:.*]] = arith.muli %[[temp]], %[[ARG]] : i32
 !CHECK: omp.yield(%[[result]] : i32)
 !CHECK: }
 !CHECK: omp.atomic.read %[[X]] = %[[Y]] : !fir.ref<i32>
 !CHECK: }
 
     !$omp atomic hint(omp_sync_hint_uncontended) capture
         y = x * y 

 !CHECK: {{.*}} = fir.embox {{.*}} : (!fir.ref<i32>) -> !fir.box<!fir.ptr<i32>>
 !CHECK: fir.store {{.*}} to %[[A]] : !fir.ref<!fir.box<!fir.ptr<i32>>>
 !CHECK: {{.*}} = fir.embox {{.*}} : (!fir.ref<i32>) -> !fir.box<!fir.ptr<i32>>
 !CHECK: fir.store {{.*}} to %[[B]] : !fir.ref<!fir.box<!fir.ptr<i32>>>
 !CHECK: %[[loaded_A:.*]] = fir.load %[[A]] : !fir.ref<!fir.box<!fir.ptr<i32>>>
 !CHECK: %[[loaded_A_addr:.*]] = fir.box_addr %[[loaded_A]] : (!fir.box<!fir.ptr<i32>>) -> !fir.ptr<i32>
 !CHECK: %[[loaded_B:.*]] = fir.load %[[B]] : !fir.ref<!fir.box<!fir.ptr<i32>>>
 !CHECK: %[[loaded_B_addr:.*]] = fir.box_addr %[[loaded_B]] : (!fir.box<!fir.ptr<i32>>) -> !fir.ptr<i32>
-!CHECK: omp.atomic.capture   {
-!CHECK: omp.atomic.update %[[loaded_A_addr]] : !fir.ptr<i32> {
-!CHECK: ^bb0(%[[ARG:.*]]: i32):
 !CHECK: %[[PRIVATE_LOADED_B:.*]] = fir.load %[[B]] : !fir.ref<!fir.box<!fir.ptr<i32>>>
 !CHECK: %[[PRIVATE_LOADED_B_addr:.*]] = fir.box_addr %[[PRIVATE_LOADED_B]] : (!fir.box<!fir.ptr<i32>>) -> !fir.ptr<i32>
 !CHECK: %[[loaded_value:.*]] = fir.load %[[PRIVATE_LOADED_B_addr]] : !fir.ptr<i32>
+!CHECK: omp.atomic.capture   {
+!CHECK: omp.atomic.update %[[loaded_A_addr]] : !fir.ptr<i32> {
+!CHECK: ^bb0(%[[ARG:.*]]: i32):
 !CHECK: %[[result:.*]] = arith.addi %[[ARG]], %[[loaded_value]] : i32
 !CHECK: omp.yield(%[[result]] : i32)
 !CHECK: }
 !CHECK: omp.atomic.read %[[loaded_B_addr]] = %[[loaded_A_addr]] : !fir.ptr<i32>, i32
 !CHECK: }
     integer, pointer :: a, b
     integer, target :: c, d
     a=>c
     b=>d
 
     !$omp atomic capture
         a = a + b
         b = a
     !$omp end atomic
 end subroutine

flang/test/Lower/OpenMP/atomic-update.f90

 !CHECK: %[[X:.*]] = fir.alloca i32 {bindc_name = "x", uniq_name = "_QFEx"}
 !CHECK: %[[Y:.*]] = fir.alloca i32 {bindc_name = "y", uniq_name = "_QFEy"}
 !CHECK: %[[Z:.*]] = fir.alloca i32 {bindc_name = "z", uniq_name = "_QFEz"}
 !CHECK: %{{.*}} = fir.convert %[[C_ADDR]] : (!fir.ref<i32>) -> !fir.ptr<i32>
 !CHECK: fir.store %{{.*}} to %[[A_ADDR]] : !fir.ref<!fir.ptr<i32>>
 !CHECK: %{{.*}} = fir.convert %[[D_ADDR]] : (!fir.ref<i32>) -> !fir.ptr<i32>
 !CHECK: fir.store {{.*}} to %[[B_ADDR]] : !fir.ref<!fir.ptr<i32>>
 !CHECK: %[[LOADED_A:.*]] = fir.load %[[A_ADDR]] : !fir.ref<!fir.ptr<i32>>
-!CHECK:  omp.atomic.update   %[[LOADED_A]] : !fir.ptr<i32> {
-!CHECK:  ^bb0(%[[ARG:.*]]: i32):
-!CHECK:    %[[LOADED_B:.*]] = fir.load %[[B_ADDR]] : !fir.ref<!fir.ptr<i32>>
-!CHECK:    %{{.*}} = fir.load %[[LOADED_B]] : !fir.ptr<i32>
-!CHECK:    %[[RESULT:.*]] = arith.addi %[[ARG]], %{{.*}} : i32
-!CHECK:    omp.yield(%[[RESULT]] : i32)
+!CHECK: %[[LOADED_B:.*]] = fir.load %[[B_ADDR]] : !fir.ref<!fir.ptr<i32>>
+!CHECK: %[[BVAL:.*]] = fir.load %[[LOADED_B]] : !fir.ptr<i32>
+!CHECK: omp.atomic.update   %[[LOADED_A]] : !fir.ptr<i32> {
+!CHECK: ^bb0(%[[ARG:.*]]: i32):
+!CHECK:   %[[RESULT:.*]] = arith.addi %[[ARG]], %[[BVAL]] : i32
+!CHECK:   omp.yield(%[[RESULT]] : i32)
 !CHECK: }
     !$omp atomic update
         a = a + b 
 
+!CHECK: %[[CST_1:.*]] = arith.constant 1 : i32
 !CHECK: omp.atomic.update   %[[Y]] : !fir.ref<i32> {
 !CHECK:  ^bb0(%[[ARG:.*]]: i32):
-!CHECK:    {{.*}} = arith.constant 1 : i32
-!CHECK:    %[[RESULT:.*]] = arith.addi %[[ARG]], {{.*}} : i32
+!CHECK:    %[[RESULT:.*]] = arith.addi %[[ARG]], %[[CST_1]] : i32
 !CHECK:    omp.yield(%[[RESULT]] : i32)
 !CHECK:  }
+    !$omp atomic
+        y = y + 1
+
+!CHECK:  %[[LOADED_X:.*]] = fir.load %[[X]] : !fir.ref<i32>
 !CHECK:  omp.atomic.update   %[[Z]] : !fir.ref<i32> {
 !CHECK:  ^bb0(%[[ARG:.*]]: i32):
-!CHECK:    %[[LOADED_X:.*]] = fir.load %[[X]] : !fir.ref<i32>
 !CHECK:    %[[RESULT:.*]] = arith.muli %[[LOADED_X]], %[[ARG]] : i32
 !CHECK:    omp.yield(%[[RESULT]] : i32)
 !CHECK:  }
-    !$omp atomic 
-        y = y + 1
     !$omp atomic update
         z = x * z 
 
+!CHECK:  %[[CST_1:.*]] = arith.constant 1 : i32
 !CHECK:  omp.atomic.update   memory_order(relaxed) hint(uncontended) %[[X]] : !fir.ref<i32> {
 !CHECK:  ^bb0(%[[ARG:.*]]: i32):
-!CHECK:    %{{.*}} = arith.constant 1 : i32
-!CHECK:    %[[RESULT:.*]] = arith.subi %[[ARG]], {{.*}} : i32
+!CHECK:    %[[RESULT:.*]] = arith.subi %[[ARG]], %[[CST_1]] : i32
 !CHECK:    omp.yield(%[[RESULT]] : i32)
 !CHECK:  }
-!CHECK:  omp.atomic.update   memory_order(relaxed) %[[Y]] : !fir.ref<i32> {
-!CHECK:  ^bb0(%[[ARG:.*]]: i32):
-!CHECK:    %[[LOADED_X:.*]] = fir.load %[[X]] : !fir.ref<i32>
-!CHECK:    %[[LOADED_Z:.*]] = fir.load %[[Z]] : !fir.ref<i32>
+    !$omp atomic relaxed update hint(omp_sync_hint_uncontended)
+        x = x - 1
+
+!CHECK:  %[[LOADED_X:.*]] = fir.load %[[X]] : !fir.ref<i32>
+!CHECK:  %[[LOADED_Z:.*]] = fir.load %[[Z]] : !fir.ref<i32>
+!CHECK:  omp.atomic.update   memory_order(relaxed) %[[Y]] : !fir.ref<i32> {
+!CHECK:  ^bb0(%[[ARG:.*]]: i32):
 !CHECK:    %{{.*}} = arith.cmpi sgt, %[[LOADED_X]], %[[ARG]] : i32
 !CHECK:    %{{.*}} = arith.select %{{.*}}, %[[LOADED_X]], %[[ARG]] : i32
 !CHECK:    %{{.*}} = arith.cmpi sgt, %{{.*}}, %[[LOADED_Z]] : i32
 !CHECK:    %[[RESULT:.*]] = arith.select %{{.*}}, %{{.*}}, %[[LOADED_Z]] : i32
 !CHECK:    omp.yield(%[[RESULT]] : i32)
 !CHECK:  }
+    !$omp atomic update relaxed
+        y = max(x, y, z)
+
+!CHECK:  %[[LOADED_X:.*]] = fir.load %[[X]] : !fir.ref<i32>
 !CHECK:  omp.atomic.update   memory_order(relaxed) hint(contended) %[[Z]] : !fir.ref<i32> {
 !CHECK:  ^bb0(%[[ARG:.*]]: i32):
-!CHECK:    %[[LOADED_X:.*]] = fir.load %[[X]] : !fir.ref<i32>
 !CHECK:    %[[RESULT:.*]] = arith.addi %[[ARG]], %[[LOADED_X]] : i32
 !CHECK:    omp.yield(%[[RESULT]] : i32)
 !CHECK:  }
-    !$omp atomic relaxed update hint(omp_sync_hint_uncontended)
-        x = x - 1
-    !$omp atomic update relaxed 
-        y = max(x, y, z)
     !$omp atomic relaxed hint(omp_sync_hint_contended)
         z = z + x
 
+!CHECK:  %[[CST_10:.*]] = arith.constant 10 : i32
 !CHECK:  omp.atomic.update   memory_order(release) hint(contended) %[[Z]] : !fir.ref<i32> {
 !CHECK:  ^bb0(%[[ARG:.*]]: i32):
-!CHECK:    %{{.*}} = arith.constant 10 : i32
-!CHECK:   %[[RESULT:.*]] = arith.muli {{.*}}, %[[ARG]] : i32
+!CHECK:   %[[RESULT:.*]] = arith.muli %[[CST_10]], %[[ARG]] : i32
 !CHECK:    omp.yield(%[[RESULT]] : i32)
 !CHECK:  }
+    !$omp atomic release update hint(omp_lock_hint_contended)
+        z = z * 10
+
+
+!CHECK:  %[[LOADED_Z:.*]] = fir.load %[[Z]] : !fir.ref<i32>
 !CHECK:  omp.atomic.update   memory_order(release) hint(speculative) %[[X]] : !fir.ref<i32> {
 !CHECK:  ^bb0(%[[ARG:.*]]: i32):
-!CHECK:    %[[LOADED_Z:.*]] = fir.load %[[Z]] : !fir.ref<i32>
 !CHECK:    %[[RESULT:.*]] = arith.divsi %[[ARG]], %[[LOADED_Z]] : i32
 !CHECK:    omp.yield(%[[RESULT]] : i32)
 !CHECK:  }
-
-    !$omp atomic release update hint(omp_lock_hint_contended)
-        z = z * 10
     !$omp atomic hint(omp_lock_hint_speculative) update release
         x = x / z
 
+!CHECK:  %[[CST_10:.*]] = arith.constant 10 : i32
 !CHECK:  omp.atomic.update   memory_order(seq_cst) hint(nonspeculative) %[[Y]] : !fir.ref<i32> {
 !CHECK:  ^bb0(%[[ARG:.*]]: i32):
-!CHECK:    %{{.*}} = arith.constant 10 : i32
-!CHECK:    %[[RESULT:.*]] = arith.addi %{{.*}}, %[[ARG]] : i32
+!CHECK:    %[[RESULT:.*]] = arith.addi %[[CST_10]], %[[ARG]] : i32
 !CHECK:   omp.yield(%[[RESULT]] : i32)
 !CHECK:  }
+    !$omp atomic hint(omp_sync_hint_nonspeculative) seq_cst
+        y = 10 + y
+
+!CHECK:  %[[LOADED_Y:.*]] = fir.load %[[Y]] : !fir.ref<i32>
 !CHECK:  omp.atomic.update   memory_order(seq_cst) %[[Z]] : !fir.ref<i32> {
 !CHECK:  ^bb0(%[[ARG:.*]]: i32):
-!CHECK:    %[[LOADED_Y:.*]] = fir.load %[[Y]] : !fir.ref<i32>
 !CHECK:    %[[RESULT:.*]] = arith.addi %[[LOADED_Y]], %[[ARG]] : i32
 !CHECK:    omp.yield(%[[RESULT]] : i32)
 !CHECK:  }
-    !$omp atomic hint(omp_sync_hint_nonspeculative) seq_cst
-        y = 10 + y
     !$omp atomic seq_cst update
         z = y + z
 
+!CHECK:  %[[C1_VAL:.*]] = arith.constant 1 : i32
 !CHECK:  omp.atomic.update   %[[I1]] : !fir.ref<i8> {
 !CHECK:  ^bb0(%[[VAL:.*]]: i8):
 !CHECK:    %[[CVT_VAL:.*]] = fir.convert %[[VAL]] : (i8) -> i32
-!CHECK:    %[[C1_VAL:.*]] = arith.constant 1 : i32
 !CHECK:    %[[ADD_VAL:.*]] = arith.addi %[[CVT_VAL]], %[[C1_VAL]] : i32
 !CHECK:    %[[UPDATED_VAL:.*]] = fir.convert %[[ADD_VAL]] : (i32) -> i8
 !CHECK:    omp.yield(%[[UPDATED_VAL]] : i8)
 !CHECK:  }
     !$omp atomic
       i1 = i1 + 1
     !$omp end atomic
 !CHECK:  return
 !CHECK: }
 end program OmpAtomicUpdate