lib/Lower/OpenMP/Atomic.cpp

//===-- Atomic.cpp -- Lowering of atomic constructs -----------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#include "Atomic.h"
#include "flang/Evaluate/expression.h"
#include "flang/Evaluate/fold.h"
#include "flang/Evaluate/tools.h"
#include "flang/Evaluate/traverse.h"
#include "flang/Evaluate/type.h"
#include "flang/Lower/AbstractConverter.h"
#include "flang/Lower/ConvertType.h"
#include "flang/Lower/OpenMP/Clauses.h"
#include "flang/Lower/PFTBuilder.h"
#include "flang/Lower/StatementContext.h"
#include "flang/Lower/SymbolMap.h"
#include "flang/Optimizer/Builder/FIRBuilder.h"
#include "flang/Optimizer/Builder/Todo.h"
#include "flang/Parser/parse-tree.h"
#include "flang/Semantics/openmp-utils.h"
#include "flang/Semantics/semantics.h"
#include "flang/Semantics/type.h"
#include "flang/Support/Fortran.h"
#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/raw_ostream.h"

#include <optional>
#include <string>
#include <type_traits>
#include <variant>

static llvm::cl::opt<bool> DumpAtomicAnalysis("fdebug-dump-atomic-analysis");

using namespace Fortran;

// Don't import the entire Fortran::lower.
namespace omp {
using namespace Fortran::lower::omp;
}

[[maybe_unused]] static void
dumpAtomicAnalysis(const parser::OpenMPAtomicConstruct::Analysis &analysis) {
  auto whatStr = [](int k) {
    std::string txt = "?";
    switch (k & parser::OpenMPAtomicConstruct::Analysis::Action) {
    case parser::OpenMPAtomicConstruct::Analysis::None:
      txt = "None";
      break;
    case parser::OpenMPAtomicConstruct::Analysis::Read:
      txt = "Read";
      break;
    case parser::OpenMPAtomicConstruct::Analysis::Write:
      txt = "Write";
      break;
    case parser::OpenMPAtomicConstruct::Analysis::Update:
      txt = "Update";
      break;
    }
    switch (k & parser::OpenMPAtomicConstruct::Analysis::Condition) {
    case parser::OpenMPAtomicConstruct::Analysis::IfTrue:
      txt += " | IfTrue";
      break;
    case parser::OpenMPAtomicConstruct::Analysis::IfFalse:
      txt += " | IfFalse";
      break;
    }
    return txt;
  };

  auto exprStr = [&](const parser::TypedExpr &expr) {
    if (auto *maybe = expr.get()) {
      if (maybe->v)
        return maybe->v->AsFortran();
    }
    return "<null>"s;
  };
  auto assignStr = [&](const parser::TypedAssignment &assign) {
    if (auto *maybe = assign.get(); maybe && maybe->v) {
      std::string str;
      llvm::raw_string_ostream os(str);
      maybe->v->AsFortran(os);
      return str;
    }
    return "<null>"s;
  };

  const semantics::SomeExpr &atom = *analysis.atom.get()->v;

  llvm::errs() << "Analysis {\n";
  llvm::errs() << "  atom: " << atom.AsFortran() << "\n";
  llvm::errs() << "  cond: " << exprStr(analysis.cond) << "\n";
  llvm::errs() << "  op0 {\n";
  llvm::errs() << "    what: " << whatStr(analysis.op0.what) << "\n";
  llvm::errs() << "    assign: " << assignStr(analysis.op0.assign) << "\n";
  llvm::errs() << "  }\n";
  llvm::errs() << "  op1 {\n";
  llvm::errs() << "    what: " << whatStr(analysis.op1.what) << "\n";
  llvm::errs() << "    assign: " << assignStr(analysis.op1.assign) << "\n";
  llvm::errs() << "  }\n";
  llvm::errs() << "}\n";
}

static bool isPointerAssignment(const evaluate::Assignment &assign) {
  return common::visit(
      common::visitors{
          [](const evaluate::Assignment::BoundsSpec &) { return true; },
          [](const evaluate::Assignment::BoundsRemapping &) { return true; },
          [](const auto &) { return false; },
      },
      assign.u);
}

static fir::FirOpBuilder::InsertPoint
getInsertionPointBefore(mlir::Operation *op) {
  return fir::FirOpBuilder::InsertPoint(op->getBlock(),
                                        mlir::Block::iterator(op));
}

static fir::FirOpBuilder::InsertPoint
getInsertionPointAfter(mlir::Operation *op) {
  return fir::FirOpBuilder::InsertPoint(op->getBlock(),
                                        ++mlir::Block::iterator(op));
}

static mlir::IntegerAttr getAtomicHint(lower::AbstractConverter &converter,
                                       const omp::List<omp::Clause> &clauses) {
  fir::FirOpBuilder &builder = converter.getFirOpBuilder();
  for (const omp::Clause &clause : clauses) {
    if (clause.id != llvm::omp::Clause::OMPC_hint)
      continue;
    auto &hint = std::get<omp::clause::Hint>(clause.u);
    auto maybeVal = evaluate::ToInt64(hint.v);
    CHECK(maybeVal);
    return builder.getI64IntegerAttr(*maybeVal);
  }
  return nullptr;
}

static mlir::omp::ClauseMemoryOrderKind
getMemoryOrderKind(common::OmpMemoryOrderType kind) {
  switch (kind) {
  case common::OmpMemoryOrderType::Acq_Rel:
    return mlir::omp::ClauseMemoryOrderKind::Acq_rel;
  case common::OmpMemoryOrderType::Acquire:
    return mlir::omp::ClauseMemoryOrderKind::Acquire;
  case common::OmpMemoryOrderType::Relaxed:
    return mlir::omp::ClauseMemoryOrderKind::Relaxed;
  case common::OmpMemoryOrderType::Release:
    return mlir::omp::ClauseMemoryOrderKind::Release;
  case common::OmpMemoryOrderType::Seq_Cst:
    return mlir::omp::ClauseMemoryOrderKind::Seq_cst;
  }
  llvm_unreachable("Unexpected kind");
}

static mlir::omp::ClauseMemoryOrderKind
getMemoryOrderKind(omp::clause::Fail::MemoryOrder kind) {
  using MemoryOrder = omp::clause::Fail::MemoryOrder;
  switch (kind) {
  case MemoryOrder::AcqRel:
    return mlir::omp::ClauseMemoryOrderKind::Acq_rel;
  case MemoryOrder::Acquire:
    return mlir::omp::ClauseMemoryOrderKind::Acquire;
  case MemoryOrder::Relaxed:
    return mlir::omp::ClauseMemoryOrderKind::Relaxed;
  case MemoryOrder::Release:
    return mlir::omp::ClauseMemoryOrderKind::Release;
  case MemoryOrder::SeqCst:
    return mlir::omp::ClauseMemoryOrderKind::Seq_cst;
  }
  llvm_unreachable("Unexpected memory order");
}

static std::optional<mlir::omp::ClauseMemoryOrderKind>
getMemoryOrderKind(llvm::omp::Clause clauseId) {
  switch (clauseId) {
  case llvm::omp::Clause::OMPC_acq_rel:
    return mlir::omp::ClauseMemoryOrderKind::Acq_rel;
  case llvm::omp::Clause::OMPC_acquire:
    return mlir::omp::ClauseMemoryOrderKind::Acquire;
  case llvm::omp::Clause::OMPC_relaxed:
    return mlir::omp::ClauseMemoryOrderKind::Relaxed;
  case llvm::omp::Clause::OMPC_release:
    return mlir::omp::ClauseMemoryOrderKind::Release;
  case llvm::omp::Clause::OMPC_seq_cst:
    return mlir::omp::ClauseMemoryOrderKind::Seq_cst;
  default:
    return std::nullopt;
  }
}

static std::optional<mlir::omp::ClauseMemoryOrderKind>
getMemoryOrderFromRequires(const semantics::Scope &scope) {
  // The REQUIRES construct is only allowed in the main program scope
  // and module scope, but seems like we also accept it in a subprogram
  // scope.
  // For safety, traverse all enclosing scopes and check if their symbol
  // contains REQUIRES.
  const semantics::Scope &unitScope = semantics::omp::GetProgramUnit(scope);
  if (auto *symbol = unitScope.symbol()) {
    const common::OmpMemoryOrderType *admo = common::visit(
        [](auto &&s) {
          using WithOmpDeclarative = semantics::WithOmpDeclarative;
          if constexpr (std::is_convertible_v<decltype(s),
                                              const WithOmpDeclarative &>) {
            if (auto &admo{s.ompAtomicDefaultMemOrder()}) {
              return &*admo;
            }
          }
          return static_cast<const common::OmpMemoryOrderType *>(nullptr);
        },
        symbol->details());

    if (admo)
      return getMemoryOrderKind(*admo);
  }

  return std::nullopt;
}

static std::optional<mlir::omp::ClauseMemoryOrderKind>
getDefaultAtomicMemOrder(semantics::SemanticsContext &semaCtx) {
  llvm::omp::Version version = semaCtx.langOptions().getOpenMPVersion();
  if (version > 50)
    return mlir::omp::ClauseMemoryOrderKind::Relaxed;
  return std::nullopt;
}

static std::pair<std::optional<mlir::omp::ClauseMemoryOrderKind>, bool>
getAtomicMemoryOrder(semantics::SemanticsContext &semaCtx,
                     const omp::List<omp::Clause> &clauses,
                     const semantics::Scope &scope) {
  for (const omp::Clause &clause : clauses) {
    if (auto maybeKind = getMemoryOrderKind(clause.id))
      return std::make_pair(*maybeKind, /*canOverride=*/false);
  }

  if (auto maybeKind = getMemoryOrderFromRequires(scope))
    return std::make_pair(*maybeKind, /*canOverride=*/true);

  return std::make_pair(getDefaultAtomicMemOrder(semaCtx),
                        /*canOverride=*/false);
}

static std::optional<mlir::omp::ClauseMemoryOrderKind>
makeValidForAction(std::optional<mlir::omp::ClauseMemoryOrderKind> memOrder,
                   int action0, int action1, llvm::omp::Version version) {
  // When the atomic default memory order specified on a REQUIRES directive is
  // disallowed on a given ATOMIC operation, and it's not ACQ_REL, the order
  // reverts to RELAXED. ACQ_REL decays to either ACQUIRE or RELEASE, depending
  // on the operation.

  if (!memOrder) {
    return memOrder;
  }

  using Analysis = parser::OpenMPAtomicConstruct::Analysis;
  // Figure out the main action (i.e. disregard a potential capture operation)
  int action = action0;
  bool isCapture = action1 != Analysis::None;
  if (isCapture)
    action = action0 == Analysis::Read ? action1 : action0;

  // All orderings are valid for capture operations per the OpenMP spec.
  // The individual sub-operations (read/write/update) inside the capture
  // will have their orderings handled separately.
  if (isCapture)
    return memOrder;

  // Avaliable orderings: acquire, acq_rel, relaxed, release, seq_cst

  if (version == 50) {
    if (action == Analysis::Read) {
      // "acq_rel" decays to "acquire" for read
      if (*memOrder == mlir::omp::ClauseMemoryOrderKind::Acq_rel)
        return mlir::omp::ClauseMemoryOrderKind::Acquire;
    } else if (action == Analysis::Write) {
      // "acq_rel" decays to "release" for write
      if (*memOrder == mlir::omp::ClauseMemoryOrderKind::Acq_rel)
        return mlir::omp::ClauseMemoryOrderKind::Release;
    } else if (action == Analysis::Update) {
      // "acquire" decays to "relaxed", "acq_rel" decays to "release"
      if (*memOrder == mlir::omp::ClauseMemoryOrderKind::Acquire)
        return mlir::omp::ClauseMemoryOrderKind::Relaxed;
      if (*memOrder == mlir::omp::ClauseMemoryOrderKind::Acq_rel)
        return mlir::omp::ClauseMemoryOrderKind::Release;
    }
  }

  if (version >= 50) {
    if (action == Analysis::Read) {
      // "release" prohibited
      if (*memOrder == mlir::omp::ClauseMemoryOrderKind::Release)
        return mlir::omp::ClauseMemoryOrderKind::Relaxed;
    }
    if (action == Analysis::Write) {
      // "acquire" prohibited
      if (*memOrder == mlir::omp::ClauseMemoryOrderKind::Acquire)
        return mlir::omp::ClauseMemoryOrderKind::Relaxed;
    }
  } else {
    if (action == Analysis::Read) {
      // "release" prohibited
      if (*memOrder == mlir::omp::ClauseMemoryOrderKind::Release)
        return mlir::omp::ClauseMemoryOrderKind::Relaxed;
    } else {
      if (action & Analysis::Write) { // include "update"
        // "acquire" prohibited
        if (*memOrder == mlir::omp::ClauseMemoryOrderKind::Acquire)
          return mlir::omp::ClauseMemoryOrderKind::Relaxed;
        if (action == Analysis::Update) {
          // "acq_rel" prohibited
          if (*memOrder == mlir::omp::ClauseMemoryOrderKind::Acq_rel)
            return mlir::omp::ClauseMemoryOrderKind::Relaxed;
        }
      }
    }
  }

  return memOrder;
}

static mlir::omp::ClauseMemoryOrderKindAttr
makeMemOrderAttr(lower::AbstractConverter &converter,
                 std::optional<mlir::omp::ClauseMemoryOrderKind> maybeKind) {
  if (maybeKind) {
    return mlir::omp::ClauseMemoryOrderKindAttr::get(
        converter.getFirOpBuilder().getContext(), *maybeKind);
  }
  return nullptr;
}

static mlir::Operation * //
genAtomicRead(lower::AbstractConverter &converter,
              semantics::SemanticsContext &semaCtx, mlir::Location loc,
              lower::StatementContext &stmtCtx, mlir::Value atomAddr,
              const semantics::SomeExpr &atom,
              const evaluate::Assignment &assign, mlir::IntegerAttr hint,
              std::optional<mlir::omp::ClauseMemoryOrderKind> memOrder,
              fir::FirOpBuilder::InsertPoint preAt,
              fir::FirOpBuilder::InsertPoint atomicAt,
              fir::FirOpBuilder::InsertPoint postAt) {
  fir::FirOpBuilder &builder = converter.getFirOpBuilder();
  builder.restoreInsertionPoint(preAt);

  // If the atomic clause is read then the memory-order clause must
  // not be release.
  if (memOrder) {
    if (*memOrder == mlir::omp::ClauseMemoryOrderKind::Release) {
      // Reset it back to the default.
      memOrder = getDefaultAtomicMemOrder(semaCtx);
    } else if (semaCtx.langOptions().getOpenMPVersion() <= 50 &&
               *memOrder == mlir::omp::ClauseMemoryOrderKind::Acq_rel) {
      // In OpenMP 5.0, acq_rel is not allowed on read; decay to acquire.
      // In OpenMP 5.1+, acq_rel is permitted on read.
      memOrder = mlir::omp::ClauseMemoryOrderKind::Acquire;
    }
  }

  mlir::Value storeAddr =
      fir::getBase(converter.genExprAddr(assign.lhs, stmtCtx, &loc));
  mlir::Type atomType = fir::unwrapRefType(atomAddr.getType());
  mlir::Type storeType = fir::unwrapRefType(storeAddr.getType());

  mlir::Value toAddr = [&]() {
    if (atomType == storeType)
      return storeAddr;
    return builder.createTemporary(loc, atomType, ".tmp.atomval");
  }();

  builder.restoreInsertionPoint(atomicAt);
  mlir::Operation *op = mlir::omp::AtomicReadOp::create(
      builder, loc, atomAddr, toAddr, mlir::TypeAttr::get(atomType), hint,
      makeMemOrderAttr(converter, memOrder));

  if (atomType != storeType) {
    lower::ExprToValueMap overrides;
    // The READ operation could be a part of UPDATE CAPTURE, so make sure
    // we don't emit extra code into the body of the atomic op.
    builder.restoreInsertionPoint(postAt);
    mlir::Value load = fir::LoadOp::create(builder, loc, toAddr);
    overrides.try_emplace(&atom, load);

    converter.overrideExprValues(&overrides);
    mlir::Value value =
        fir::getBase(converter.genExprValue(assign.rhs, stmtCtx, &loc));
    converter.resetExprOverrides();

    if (value.getType() != storeType)
      value = builder.createConvert(loc, storeType, value);
    fir::StoreOp::create(builder, loc, value, storeAddr);
  }
  return op;
}

static mlir::Operation * //
genAtomicWrite(lower::AbstractConverter &converter,
               semantics::SemanticsContext &semaCtx, mlir::Location loc,
               lower::StatementContext &stmtCtx, mlir::Value atomAddr,
               const semantics::SomeExpr &atom,
               const evaluate::Assignment &assign, mlir::IntegerAttr hint,
               std::optional<mlir::omp::ClauseMemoryOrderKind> memOrder,
               fir::FirOpBuilder::InsertPoint preAt,
               fir::FirOpBuilder::InsertPoint atomicAt,
               fir::FirOpBuilder::InsertPoint postAt) {
  fir::FirOpBuilder &builder = converter.getFirOpBuilder();
  builder.restoreInsertionPoint(preAt);

  // If the atomic clause is write then the memory-order clause must
  // not be acquire.
  if (memOrder) {
    if (*memOrder == mlir::omp::ClauseMemoryOrderKind::Acquire) {
      // Reset it back to the default.
      memOrder = getDefaultAtomicMemOrder(semaCtx);
    } else if (semaCtx.langOptions().getOpenMPVersion() <= 50 &&
               *memOrder == mlir::omp::ClauseMemoryOrderKind::Acq_rel) {
      // In OpenMP 5.0, acq_rel is not allowed on write; decay to release.
      // In OpenMP 5.1+, acq_rel is permitted on write.
      memOrder = mlir::omp::ClauseMemoryOrderKind::Release;
    }
  }

  mlir::Value value =
      fir::getBase(converter.genExprValue(assign.rhs, stmtCtx, &loc));
  mlir::Type atomType = fir::unwrapRefType(atomAddr.getType());
  mlir::Value converted = builder.createConvert(loc, atomType, value);

  builder.restoreInsertionPoint(atomicAt);
  mlir::Operation *op =
      mlir::omp::AtomicWriteOp::create(builder, loc, atomAddr, converted, hint,
                                       makeMemOrderAttr(converter, memOrder));
  return op;
}

static mlir::Operation *
genAtomicUpdate(lower::AbstractConverter &converter,
                semantics::SemanticsContext &semaCtx, mlir::Location loc,
                lower::StatementContext &stmtCtx, mlir::Value atomAddr,
                const semantics::SomeExpr &atom,
                const evaluate::Assignment &assign, mlir::IntegerAttr hint,
                std::optional<mlir::omp::ClauseMemoryOrderKind> memOrder,
                fir::FirOpBuilder::InsertPoint preAt,
                fir::FirOpBuilder::InsertPoint atomicAt,
                fir::FirOpBuilder::InsertPoint postAt) {
  lower::ExprToValueMap overrides;
  lower::StatementContext naCtx;
  fir::FirOpBuilder &builder = converter.getFirOpBuilder();
  builder.restoreInsertionPoint(preAt);

  mlir::Type atomType = fir::unwrapRefType(atomAddr.getType());

  // This must exist by now.
  semantics::SomeExpr rhs = assign.rhs;
  semantics::SomeExpr input = *evaluate::GetConvertInput(rhs);
  auto [opcode, args] = evaluate::GetTopLevelOperationIgnoreResizing(input);
  assert(!args.empty() && "Update operation without arguments");

  for (auto &arg : args) {
    if (!evaluate::IsSameOrConvertOf(arg, atom)) {
      mlir::Value val = fir::getBase(converter.genExprValue(arg, naCtx, &loc));
      overrides.try_emplace(&arg, val);
    }
  }

  mlir::ModuleOp module = builder.getModule();
  mlir::omp::AtomicControlAttr atomicControlAttr =
      mlir::omp::AtomicControlAttr::get(
          builder.getContext(), fir::getAtomicIgnoreDenormalMode(module),
          fir::getAtomicFineGrainedMemory(module),
          fir::getAtomicRemoteMemory(module));
  builder.restoreInsertionPoint(atomicAt);
  auto updateOp = mlir::omp::AtomicUpdateOp::create(
      builder, loc, atomAddr, atomicControlAttr, hint,
      makeMemOrderAttr(converter, memOrder));

  mlir::Region &region = updateOp->getRegion(0);
  mlir::Block *block = builder.createBlock(&region, {}, {atomType}, {loc});
  mlir::Value localAtom = fir::getBase(block->getArgument(0));
  overrides.try_emplace(&atom, localAtom);

  converter.overrideExprValues(&overrides);
  mlir::Value updated =
      fir::getBase(converter.genExprValue(rhs, stmtCtx, &loc));
  mlir::Value converted = builder.createConvert(loc, atomType, updated);
  mlir::omp::YieldOp::create(builder, loc, converted);
  converter.resetExprOverrides();

  builder.restoreInsertionPoint(postAt); // For naCtx cleanups
  return updateOp;
}

static mlir::Operation *
genAtomicOperation(lower::AbstractConverter &converter,
                   semantics::SemanticsContext &semaCtx, mlir::Location loc,
                   lower::StatementContext &stmtCtx, int action,
                   mlir::Value atomAddr, const semantics::SomeExpr &atom,
                   const evaluate::Assignment &assign, mlir::IntegerAttr hint,
                   std::optional<mlir::omp::ClauseMemoryOrderKind> memOrder,
                   fir::FirOpBuilder::InsertPoint preAt,
                   fir::FirOpBuilder::InsertPoint atomicAt,
                   fir::FirOpBuilder::InsertPoint postAt) {
  if (isPointerAssignment(assign)) {
    TODO(loc, "Code generation for pointer assignment is not implemented yet");
  }

  // This function and the functions called here do not preserve the
  // builder's insertion point, or set it to anything specific.
  switch (action) {
  case parser::OpenMPAtomicConstruct::Analysis::Read:
    return genAtomicRead(converter, semaCtx, loc, stmtCtx, atomAddr, atom,
                         assign, hint, memOrder, preAt, atomicAt, postAt);
  case parser::OpenMPAtomicConstruct::Analysis::Write:
    return genAtomicWrite(converter, semaCtx, loc, stmtCtx, atomAddr, atom,
                          assign, hint, memOrder, preAt, atomicAt, postAt);
  case parser::OpenMPAtomicConstruct::Analysis::Update:
    return genAtomicUpdate(converter, semaCtx, loc, stmtCtx, atomAddr, atom,
                           assign, hint, memOrder, preAt, atomicAt, postAt);
  default:
    return nullptr;
  }
}

/// Reverse a relational operator as if the operands were swapped.
/// e.g. LT becomes GT, LE becomes GE. Symmetric operators (EQ, NE)
/// are returned unchanged.
static common::RelationalOperator reverseRelOp(common::RelationalOperator op) {
  using RO = common::RelationalOperator;
  switch (op) {
  case RO::LT:
    return RO::GT;
  case RO::LE:
    return RO::GE;
  case RO::GT:
    return RO::LT;
  case RO::GE:
    return RO::LE;
  default:
    return op;
  }
}

void Fortran::lower::omp::lowerAtomic(
    AbstractConverter &converter, SymMap &symTable,
    semantics::SemanticsContext &semaCtx, pft::Evaluation &eval,
    const parser::OpenMPAtomicConstruct &construct) {
  auto get = [](auto &&typedWrapper) -> decltype(&*typedWrapper.get()->v) {
    if (auto *maybe = typedWrapper.get(); maybe && maybe->v) {
      return &*maybe->v;
    } else {
      return nullptr;
    }
  };

  fir::FirOpBuilder &builder = converter.getFirOpBuilder();
  const parser::OmpDirectiveSpecification &dirSpec = construct.BeginDir();
  omp::List<omp::Clause> clauses = makeClauses(dirSpec.Clauses(), semaCtx);
  lower::StatementContext stmtCtx;

  const parser::OpenMPAtomicConstruct::Analysis &analysis = construct.analysis;
  if (DumpAtomicAnalysis)
    dumpAtomicAnalysis(analysis);

  const semantics::SomeExpr &atom = *get(analysis.atom);
  mlir::Location loc = converter.genLocation(construct.source);
  mlir::Value atomAddr =
      fir::getBase(converter.genExprAddr(atom, stmtCtx, &loc));
  mlir::IntegerAttr hint = getAtomicHint(converter, clauses);
  auto [memOrder, canOverride] = getAtomicMemoryOrder(
      semaCtx, clauses, semaCtx.FindScope(construct.source));

  llvm::omp::Version version = semaCtx.langOptions().getOpenMPVersion();
  int action0 = analysis.op0.what & analysis.Action;
  int action1 = analysis.op1.what & analysis.Action;
  memOrder = makeValidForAction(memOrder, action0, action1, version);

  // --- Shared capture scaffolding ---
  mlir::Operation *captureOp = nullptr;
  fir::FirOpBuilder::InsertPoint preAt = builder.saveInsertionPoint();
  fir::FirOpBuilder::InsertPoint atomicAt, postAt;

  if (construct.IsCapture()) {
    assert(action0 != analysis.None && action1 != analysis.None &&
           "Expecting two actions");
    (void)action0;
    (void)action1;
    captureOp = mlir::omp::AtomicCaptureOp::create(
        builder, loc, hint, makeMemOrderAttr(converter, memOrder),
        /*fail_only=*/nullptr);
    // Set the non-atomic insertion point to before the atomic.capture.
    preAt = getInsertionPointBefore(captureOp);

    mlir::Block *block = builder.createBlock(&captureOp->getRegion(0));
    builder.setInsertionPointToEnd(block);
    // Set the atomic insertion point to before the terminator inside
    // atomic.capture.
    mlir::Operation *term = mlir::omp::TerminatorOp::create(builder, loc);
    atomicAt = getInsertionPointBefore(term);
    postAt = getInsertionPointAfter(captureOp);
    hint = nullptr;
    memOrder = std::nullopt;
  }

  if (auto *cond = get(analysis.cond)) {
    // atomic compare: if (x == e) x = d
    // e : expecteVal
    // d : desiredVal

    // Restore insertion point so pre-processing code (e.g. computing
    // expectedVal) is emitted before the capture op, not after the terminator.
    builder.restoreInsertionPoint(preAt);

    // The comparison-result forms of atomic compare
    //     (e.g. `r = x == e;
    //           if (r) x = d`)
    // are accepted by semantics, but the assignment to `r` is
    // not captured by the atomic analysis and would be silently dropped.
    // Detect the extra assignment here and emit an explicit TODO instead of
    // generating incorrect code.
    if (!construct.IsCapture()) {
      const parser::Block &body = std::get<parser::Block>(construct.t);
      if (body.size() > 1)
        TODO(loc, "atomic compare capturing the comparison result "
                  "(e.g. 'r = x == e')");
    }

    // The `fail` clause sets the memory ordering for a failed compare;
    // extract its argument to attach to the omp.atomic.compare op below.
    std::optional<mlir::omp::ClauseMemoryOrderKind> failMemOrder;
    for (const omp::Clause &clause : clauses) {
      if (const auto *fail = std::get_if<omp::clause::Fail>(&clause.u)) {
        failMemOrder = getMemoryOrderKind(fail->v);
        break;
      }
    }

    common::RelationalOperator relOpr = common::RelationalOperator::EQ;
    std::optional<semantics::SomeExpr> expectedExprStorage;
    bool isUnsigned = false;

    if (const auto *rel =
            evaluate::UnwrapExpr<evaluate::Relational<evaluate::SomeType>>(
                *cond)) {
      std::visit(
          [&](const auto &relImpl) {
            relOpr = relImpl.opr;
            using Operand = typename std::decay_t<decltype(relImpl)>::Operand;
            isUnsigned = Operand::category == common::TypeCategory::Unsigned;
            auto leftExpr = evaluate::AsGenericExpr(
                evaluate::Expr<Operand>{relImpl.left()});
            auto rightExpr = evaluate::AsGenericExpr(
                evaluate::Expr<Operand>{relImpl.right()});
            if (evaluate::IsSameOrConvertOf(rightExpr, atom)) {
              // e.g. e == x  (atom is on the right)
              // left operand is expected value (e)
              // reverse the operator so that the comparison becomes
              // x <reversed-op> e.
              expectedExprStorage = std::move(leftExpr);
              relOpr = reverseRelOp(relOpr);
            } else {
              // Form: x == e  (atom is on the left, or default)
              expectedExprStorage = std::move(rightExpr);
            }
          },
          rel->u);
    }

    // Fortran uses .eqv./.neqv. for logical equality/inequality, which are
    // LogicalOperation expressions rather than Relational expressions.
    if (!expectedExprStorage) {
      if (const auto *someLogical =
              evaluate::UnwrapExpr<evaluate::Expr<evaluate::SomeLogical>>(
                  *cond)) {
        common::visit(
            [&](const auto &kindLogical) {
              using LogicalExpr = std::decay_t<decltype(kindLogical)>;
              constexpr int K = LogicalExpr::Result::kind;
              if (const auto *logOp =
                      std::get_if<evaluate::LogicalOperation<K>>(
                          &kindLogical.u)) {
                if (logOp->logicalOperator == common::LogicalOperator::Eqv ||
                    logOp->logicalOperator == common::LogicalOperator::Neqv) {
                  relOpr =
                      (logOp->logicalOperator == common::LogicalOperator::Eqv)
                          ? common::RelationalOperator::EQ
                          : common::RelationalOperator::NE;
                  using Operand =
                      typename evaluate::LogicalOperation<K>::Operand;
                  auto leftExpr = evaluate::AsGenericExpr(
                      evaluate::Expr<Operand>{logOp->left()});
                  auto rightExpr = evaluate::AsGenericExpr(
                      evaluate::Expr<Operand>{logOp->right()});
                  if (evaluate::IsSameOrConvertOf(rightExpr, atom)) {
                    expectedExprStorage = std::move(leftExpr);
                    relOpr = reverseRelOp(relOpr);
                  } else {
                    expectedExprStorage = std::move(rightExpr);
                  }
                }
              }
            },
            someLogical->u);
      }
    }

    if (!expectedExprStorage) {
      mlir::emitError(loc, "internal error: atomic compare condition is not a "
                           "recognized relational expression");
      return;
    }

    mlir::Type elemTypeOfX = fir::unwrapRefType(atomAddr.getType());
    mlir::Value expectedVal = fir::getBase(
        converter.genExprValue(*expectedExprStorage, stmtCtx, &loc));
    if (expectedVal.getType() != elemTypeOfX) {
      expectedVal = builder.createConvert(loc, elemTypeOfX, expectedVal);
    }

    // For logical types, convert address and expected value to integer
    // type here (above the atomic compare region) so that the region
    // only contains arith.cmpi eq on integers. If done inside the region,
    // fir.convert logical<4> -> i32 would lower to `icmp ne %val, 0`
    // which violates the atomic compare verifier's expectations.
    if (mlir::isa<fir::LogicalType>(elemTypeOfX)) {
      unsigned kind = mlir::cast<fir::LogicalType>(elemTypeOfX).getFKind();
      mlir::Type intTy = builder.getIntegerType(kind * 8);
      mlir::Type intRefTy = builder.getRefType(intTy);
      atomAddr = builder.createConvert(loc, intRefTy, atomAddr);
      expectedVal = builder.createConvert(loc, intTy, expectedVal);
      elemTypeOfX = intTy;
    }

    // If this is a compare+capture, determine the ordering of ops.
    // Pattern 1 (prefix): v = x; if (x == e) x = d  → read first
    // Pattern 2 (postfix): if (x == e) x = d; v = x → compare first
    // Pattern 3 (fail-only): if (x == e) x = d; else v = x → read on failure
    bool isPostfixCapture = false;
    bool isFailOnly = false;
    const evaluate::Assignment *readAssign = nullptr;
    if (construct.IsCapture()) {
      // Determine which op is the read and check for fail-only (IfFalse).
      int readWhat = 0;
      if (analysis.op0.what & analysis.Read) {
        readAssign = get(analysis.op0.assign);
        readWhat = analysis.op0.what;
      } else if (analysis.op1.what & analysis.Read) {
        readAssign = get(analysis.op1.assign);
        readWhat = analysis.op1.what;
        isPostfixCapture = true;
      }
      assert(readAssign && "Expected a read assignment for compare capture");

      // Check if the read is conditioned on comparison failure (else branch).
      if (readWhat & analysis.IfFalse)
        isFailOnly = true;

      if (!isPostfixCapture && !isFailOnly) {
        // Pattern 1 (prefix): read is first, generate it before compare.
        mlir::Operation *readOp =
            genAtomicRead(converter, semaCtx, loc, stmtCtx, atomAddr, atom,
                          *readAssign, hint, memOrder, preAt, atomicAt, postAt);
        assert(readOp && "Should have created an atomic read operation");
        builder.setInsertionPointAfter(readOp);
      } else {
        // Pattern 2 (postfix) or 3 (fail-only): compare first, read after.
        builder.restoreInsertionPoint(atomicAt);
      }

      // Set the fail_only attribute on the capture op.
      if (isFailOnly && captureOp)
        mlir::cast<mlir::omp::AtomicCaptureOp>(captureOp).setFailOnly(true);
    }

    mlir::UnitAttr weakAttr = nullptr;
    if (llvm::any_of(clauses, [](const omp::Clause &clause) {
          return clause.id == llvm::omp::Clause::OMPC_weak;
        })) {
      weakAttr = builder.getUnitAttr();
    }

    // Extract write assignment (x = d) and generate desired value (d)
    // before creating the compare region, so that d is defined outside
    // the region and any intermediate conversions (e.g., logical-to-integer
    // truthiness normalization) don't appear inside the atomic compare block.
    [[maybe_unused]] int writeActionCond = 0;
    const evaluate::Assignment *writeAssign = nullptr;
    if (analysis.op0.what & analysis.Write) {
      writeAssign = get(analysis.op0.assign);
      writeActionCond = analysis.op0.what;
    }
    if (!writeAssign && (analysis.op1.what & analysis.Write)) {
      writeAssign = get(analysis.op1.assign);
      writeActionCond = analysis.op1.what;
    }
    if (!writeAssign) {
      mlir::emitError(loc,
                      "internal error: atomic compare has no write assignment");
      return;
    }
    assert((writeActionCond & analysis.IfTrue) &&
           "atomic compare write should be conditioned on IfTrue");

    // Generate desiredVal before the capture/compare region so any
    // intermediate ops (loads, conversions) don't pollute the atomic blocks.
    fir::FirOpBuilder::InsertPoint savedIP = builder.saveInsertionPoint();
    builder.restoreInsertionPoint(preAt);
    mlir::Value desiredVal =
        fir::getBase(converter.genExprValue(writeAssign->rhs, stmtCtx, &loc));
    if (desiredVal.getType() != elemTypeOfX)
      desiredVal = builder.createConvert(loc, elemTypeOfX, desiredVal);
    builder.restoreInsertionPoint(savedIP);

    mlir::Operation *atomicOp = mlir::omp::AtomicCompareOp::create(
        builder, loc, atomAddr, weakAttr, hint,
        makeMemOrderAttr(converter, memOrder),
        makeMemOrderAttr(converter, failMemOrder));
    mlir::Block *block = builder.createBlock(&atomicOp->getRegion(0));
    mlir::Value blockArg = block->addArgument(elemTypeOfX, loc);
    builder.setInsertionPointToEnd(block);

    // Generate comparison: e.g. x == e
    mlir::Value cmpResult;
    if (mlir::isa<mlir::IntegerType>(elemTypeOfX)) {
      auto pred = isUnsigned ? lower::translateUnsignedRelational(relOpr)
                             : lower::translateSignedRelational(relOpr);
      cmpResult = mlir::arith::CmpIOp::create(builder, loc, pred, blockArg,
                                              expectedVal);
    } else if (mlir::isa<mlir::FloatType>(elemTypeOfX)) {
      auto pred = lower::translateFloatRelational(relOpr);
      cmpResult = mlir::arith::CmpFOp::create(builder, loc, pred, blockArg,
                                              expectedVal);
    } else if (fir::isa_complex(elemTypeOfX)) {
      auto pred = lower::translateFloatRelational(relOpr);
      cmpResult =
          fir::CmpcOp::create(builder, loc, pred, blockArg, expectedVal);
    } else {
      mlir::emitError(loc, "unsupported type for atomic compare");
      return;
    }

    mlir::Value newVal = mlir::arith::SelectOp::create(builder, loc, cmpResult,
                                                       desiredVal, blockArg);

    // Generate omp.yield
    mlir::omp::YieldOp::create(builder, loc, newVal);
    builder.setInsertionPointAfter(atomicOp);

    // Pattern 2 (postfix) or 3 (fail-only): compare first, read second.
    // Generate read after compare for postfix or fail-only patterns.
    if (construct.IsCapture() && (isPostfixCapture || isFailOnly)) {
      fir::FirOpBuilder::InsertPoint afterCompareAt =
          builder.saveInsertionPoint();
      mlir::Operation *readOp = genAtomicRead(
          converter, semaCtx, loc, stmtCtx, atomAddr, atom, *readAssign, hint,
          memOrder, preAt, afterCompareAt, postAt);
      assert(readOp && "Should have created an atomic read operation");
      builder.setInsertionPointAfter(readOp);
    }
    // END omp atomic compare
  } else {
    if (!construct.IsCapture()) {
      // Non-capturing operation.
      assert(action0 != analysis.None && action1 == analysis.None &&
             "Expecting single action");
      assert(!(analysis.op0.what & analysis.Condition));
      postAt = atomicAt = preAt;
    }

    // The builder's insertion point needs to be specifically set before
    // each call to `genAtomicOperation`.
    mlir::Operation *firstOp = genAtomicOperation(
        converter, semaCtx, loc, stmtCtx, analysis.op0.what, atomAddr, atom,
        *get(analysis.op0.assign), hint, memOrder, preAt, atomicAt, postAt);
    assert(firstOp && "Should have created an atomic operation");
    atomicAt = getInsertionPointAfter(firstOp);

    mlir::Operation *secondOp = nullptr;
    if (analysis.op1.what != analysis.None) {
      secondOp = genAtomicOperation(
          converter, semaCtx, loc, stmtCtx, analysis.op1.what, atomAddr, atom,
          *get(analysis.op1.assign), hint, memOrder, preAt, atomicAt, postAt);
    }

    if (!construct.IsCapture()) {
      builder.setInsertionPointAfter(secondOp ? secondOp : firstOp);
    }
  }

  // Shared capture cleanup.
  if (construct.IsCapture()) {
    builder.restoreInsertionPoint(postAt);
  }
}