[Mlir-commits] [mlir] [mlir][acc] Use atomicrmw for simple atomic captures (PR #219552)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Aug 28 13:07:51 PDT 2026
https://github.com/khaki3 updated https://github.com/llvm/llvm-project/pull/219552
>From 93abd32e53ff5efffa35c016b1a1cecdc3b4a43a Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Fri, 28 Aug 2026 01:41:50 -0700
Subject: [PATCH 1/2] [mlir][acc] Use atomicrmw for simple atomic captures
```fortran
!$acc atomic capture
nSmall = nSmall + 1
indx = nSmall
!$acc end atomic
```
acc.atomic.capture always generated a cmpxchg loop, while acc.atomic.update
already mapped a simple binop to atomicrmw. On a partition loop with 204800
threads contending on one scalar, the CAS retries dominate: the loop runs
1.64s versus 8.4e-05s with the classic compiler.
Fix: give the capture conversion the same atomicrmw path. atomicrmw returns
the old value, so `{read, update}` stores it directly and `{update, read}`
reapplies the binop to it.
An expression that reads memory keeps the cmpxchg loop: it may load the
captured value itself, which only the loop can re-point at the atomically
loaded value.
The loop now runs 9.5e-05s.
---
.../OpenACCToLLVM/ACCAtomicPatterns.cpp | 174 ++++++++++++++----
.../test/Conversion/OpenACCToLLVM/atomic.mlir | 53 +++++-
2 files changed, 189 insertions(+), 38 deletions(-)
diff --git a/mlir/lib/Conversion/OpenACCToLLVM/ACCAtomicPatterns.cpp b/mlir/lib/Conversion/OpenACCToLLVM/ACCAtomicPatterns.cpp
index 3629bd9be32eb..24ea912950d58 100644
--- a/mlir/lib/Conversion/OpenACCToLLVM/ACCAtomicPatterns.cpp
+++ b/mlir/lib/Conversion/OpenACCToLLVM/ACCAtomicPatterns.cpp
@@ -22,6 +22,7 @@
#include "mlir/Dialect/OpenACC/OpenACC.h"
#include "mlir/Dialect/OpenACC/OpenACCUtilsType.h"
#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/Debug.h"
@@ -497,6 +498,41 @@ Value ACCAtomicOpConversion<AtomicOpTy>::genUpdateCmpxchgLoop(
llvm_unreachable("invalid cmpxchg loop");
}
+static std::optional<LLVM::AtomicBinOp> getAtomicBinOp(Operation *op,
+ bool updateIsLhs) {
+ return TypeSwitch<Operation *, std::optional<LLVM::AtomicBinOp>>(op)
+ .Case<arith::AddFOp>([](auto) { return LLVM::AtomicBinOp::fadd; })
+ .Case<arith::AddIOp>([](auto) { return LLVM::AtomicBinOp::add; })
+ .Case<arith::SubFOp>(
+ [updateIsLhs](auto) -> std::optional<LLVM::AtomicBinOp> {
+ // atomicrmw fsub is always `*ptr = *ptr - val`.
+ if (!updateIsLhs)
+ return std::nullopt;
+ return LLVM::AtomicBinOp::fsub;
+ })
+ .Case<arith::SubIOp>(
+ [updateIsLhs](auto) -> std::optional<LLVM::AtomicBinOp> {
+ // atomicrmw sub is always `*ptr = *ptr - val`.
+ if (!updateIsLhs)
+ return std::nullopt;
+ return LLVM::AtomicBinOp::sub;
+ })
+ .Case<arith::AndIOp>([](auto) { return LLVM::AtomicBinOp::_and; })
+ .Case<arith::OrIOp>([](auto) { return LLVM::AtomicBinOp::_or; })
+ .Case<arith::XOrIOp>([](auto) { return LLVM::AtomicBinOp::_xor; })
+ .Case<arith::MaxSIOp>([](auto) { return LLVM::AtomicBinOp::max; })
+ .Case<arith::MinSIOp>([](auto) { return LLVM::AtomicBinOp::min; })
+ .Case<arith::MaxUIOp>([](auto) { return LLVM::AtomicBinOp::umax; })
+ .Case<arith::MinUIOp>([](auto) { return LLVM::AtomicBinOp::umin; })
+ .Case<arith::MaximumFOp>([](auto) { return LLVM::AtomicBinOp::fmaximum; })
+ .Case<arith::MinimumFOp>([](auto) { return LLVM::AtomicBinOp::fminimum; })
+ .Case<arith::MaxNumFOp>(
+ [](auto) { return LLVM::AtomicBinOp::fmaximumnum; })
+ .Case<arith::MinNumFOp>(
+ [](auto) { return LLVM::AtomicBinOp::fminimumnum; })
+ .Default([](Operation *) { return std::nullopt; });
+}
+
/// Generate llvm.atomicrmw or an llvm.cmpxchg loop.
template <>
LogicalResult ACCAtomicOpConversion<AtomicUpdateOp>::matchAndRewrite(
@@ -549,42 +585,6 @@ LogicalResult ACCAtomicOpConversion<AtomicUpdateOp>::matchAndRewrite(
// https://llvm.org/docs/LangRef.html#floating-point-min-max-intrinsics-comparison
// https://mlir.llvm.org/docs/Dialects/ArithOps/#arithmaximumf-arithmaximumfop
// https://mlir.llvm.org/docs/Dialects/ArithOps/#arithmaxnumf-arithmaxnumfop
- auto getAtomicBinOp =
- [](Operation *op, bool updateIsLhs) -> std::optional<LLVM::AtomicBinOp> {
- return TypeSwitch<Operation *, std::optional<LLVM::AtomicBinOp>>(op)
- .Case<arith::AddFOp>([](auto) { return LLVM::AtomicBinOp::fadd; })
- .Case<arith::AddIOp>([](auto) { return LLVM::AtomicBinOp::add; })
- .Case<arith::SubFOp>(
- [updateIsLhs](auto) -> std::optional<LLVM::AtomicBinOp> {
- // atomicrmw fsub is always `*ptr = *ptr - val`.
- if (!updateIsLhs)
- return std::nullopt;
- return LLVM::AtomicBinOp::fsub;
- })
- .Case<arith::SubIOp>(
- [updateIsLhs](auto) -> std::optional<LLVM::AtomicBinOp> {
- // atomicrmw sub is always `*ptr = *ptr - val`.
- if (!updateIsLhs)
- return std::nullopt;
- return LLVM::AtomicBinOp::sub;
- })
- .Case<arith::AndIOp>([](auto) { return LLVM::AtomicBinOp::_and; })
- .Case<arith::OrIOp>([](auto) { return LLVM::AtomicBinOp::_or; })
- .Case<arith::XOrIOp>([](auto) { return LLVM::AtomicBinOp::_xor; })
- .Case<arith::MaxSIOp>([](auto) { return LLVM::AtomicBinOp::max; })
- .Case<arith::MinSIOp>([](auto) { return LLVM::AtomicBinOp::min; })
- .Case<arith::MaxUIOp>([](auto) { return LLVM::AtomicBinOp::umax; })
- .Case<arith::MinUIOp>([](auto) { return LLVM::AtomicBinOp::umin; })
- .Case<arith::MaximumFOp>(
- [](auto) { return LLVM::AtomicBinOp::fmaximum; })
- .Case<arith::MinimumFOp>(
- [](auto) { return LLVM::AtomicBinOp::fminimum; })
- .Case<arith::MaxNumFOp>(
- [](auto) { return LLVM::AtomicBinOp::fmaximumnum; })
- .Case<arith::MinNumFOp>(
- [](auto) { return LLVM::AtomicBinOp::fminimumnum; })
- .Default([](Operation *) { return std::nullopt; });
- };
// Select the kind and the val of atomicrmw.
std::optional<Value> val = std::nullopt;
@@ -691,6 +691,63 @@ LogicalResult ACCAtomicOpConversion<AtomicUpdateOp>::matchAndRewrite(
return success();
}
+/// The cmpxchg loop re-points reads of `v`/`x` inside `expr` at the atomically
+/// loaded value (see moveDependency). An atomicrmw cannot do that, so any
+/// `expr` that loads from memory keeps the loop: the load may be the captured
+/// value itself, and the address is not reliably comparable here.
+static bool exprReadsMemory(Value expr) {
+ SmallVector<Value> worklist{expr};
+ llvm::DenseSet<Operation *> seen;
+ while (!worklist.empty()) {
+ Operation *def = worklist.pop_back_val().getDefiningOp();
+ if (!def || !seen.insert(def).second)
+ continue;
+ if (!isMemoryEffectFree(def))
+ return true;
+ worklist.append(def->getOperands().begin(), def->getOperands().end());
+ }
+ return false;
+}
+
+/// Match an update region holding a single `x = x <binop> expr`. Returns the
+/// atomicrmw kind and the binary operation, or nullopt when the region needs a
+/// cmpxchg loop.
+static std::optional<std::pair<LLVM::AtomicBinOp, Operation *>>
+matchSimpleAtomicUpdate(AtomicUpdateOp update, const TypeConverter *converter) {
+ Block &block = update.getRegion().front();
+ if (!llvm::hasNItems(block.getOperations(), 2))
+ return std::nullopt;
+ Operation &binOp = block.front();
+ Operation &yield = block.back();
+ if (binOp.getNumOperands() != 2 || binOp.getNumResults() != 1)
+ return std::nullopt;
+ if (yield.getNumOperands() != 1 || yield.getOperand(0) != binOp.getResult(0))
+ return std::nullopt;
+
+ // The updated value must be used exactly once, as an operand of the binop.
+ Value arg = block.getArgument(0);
+ bool updateIsLhs = binOp.getOperand(0) == arg;
+ if (updateIsLhs == (binOp.getOperand(1) == arg))
+ return std::nullopt;
+ if (!arg.hasOneUse())
+ return std::nullopt;
+
+ // Keep serialized and aggregate types on the cmpxchg path.
+ Type argTy = arg.getType();
+ if (!argTy.isIntOrFloat() || converter->convertType(argTy) != argTy)
+ return std::nullopt;
+
+ // The other operand has to be available outside the region.
+ Operation *exprDef = binOp.getOperand(updateIsLhs ? 1 : 0).getDefiningOp();
+ if (exprDef && exprDef->getBlock() == &block)
+ return std::nullopt;
+
+ std::optional<LLVM::AtomicBinOp> kind = getAtomicBinOp(&binOp, updateIsLhs);
+ if (!kind)
+ return std::nullopt;
+ return std::make_pair(*kind, &binOp);
+}
+
/// Generate an llvm.cmpxchg loop.
template <>
LogicalResult ACCAtomicOpConversion<AtomicCaptureOp>::matchAndRewrite(
@@ -700,6 +757,51 @@ LogicalResult ACCAtomicOpConversion<AtomicCaptureOp>::matchAndRewrite(
Operation *secondOp = capture.getSecondOp();
Value vPtr = nullptr;
Value storeVal = nullptr;
+
+ // A single `x = x <binop> expr` capture becomes one atomicrmw. The cmpxchg
+ // loop below serializes retries and collapses under contention.
+ auto update = dyn_cast<AtomicUpdateOp>(firstOp)
+ ? cast<AtomicUpdateOp>(firstOp)
+ : dyn_cast<AtomicUpdateOp>(secondOp);
+ auto read = dyn_cast<AtomicReadOp>(firstOp)
+ ? cast<AtomicReadOp>(firstOp)
+ : dyn_cast<AtomicReadOp>(secondOp);
+ if (update && read && read.getX() == update.getX()) {
+ std::optional<std::pair<LLVM::AtomicBinOp, Operation *>> matched =
+ matchSimpleAtomicUpdate(update, this->getTypeConverter());
+ Value arg = update.getRegion().front().getArgument(0);
+ Value xRef = update.getX();
+ Value vRef = read.getV();
+ if (matched) {
+ auto [kind, binOp] = *matched;
+ bool updateIsLhs = binOp->getOperand(0) == arg;
+ Value expr = binOp->getOperand(updateIsLhs ? 1 : 0);
+ if (exprReadsMemory(expr))
+ matched = std::nullopt;
+ else {
+ Location loc = capture.getLoc();
+ Value xPtr = getAtomicPointer(xRef, rewriter.getRemappedValue(xRef),
+ loc, rewriter);
+ Value vFastPtr = getAtomicPointer(vRef, rewriter.getRemappedValue(vRef),
+ loc, rewriter);
+ rewriter.setInsertionPoint(capture);
+ auto rmw = LLVM::AtomicRMWOp::create(rewriter, loc, kind, xPtr,
+ rewriter.getRemappedValue(expr),
+ LLVM::AtomicOrdering::monotonic);
+ // atomicrmw yields the old value; `{update, read}` captures the new
+ // one.
+ Value captured = rmw.getRes();
+ if (isa<AtomicUpdateOp>(firstOp)) {
+ rewriter.moveOpAfter(binOp, rmw);
+ binOp->replaceUsesOfWith(arg, captured);
+ captured = binOp->getResult(0);
+ }
+ rewriter.replaceOpWithNewOp<LLVM::StoreOp>(capture, captured, vFastPtr);
+ return success();
+ }
+ }
+ }
+
if (auto firstReadStmt = dyn_cast<AtomicReadOp>(firstOp)) {
Location loc = capture.getLoc();
Value xRef = firstReadStmt.getX();
diff --git a/mlir/test/Conversion/OpenACCToLLVM/atomic.mlir b/mlir/test/Conversion/OpenACCToLLVM/atomic.mlir
index 21bdbc5aaa8e3..dbb5ea438b443 100644
--- a/mlir/test/Conversion/OpenACCToLLVM/atomic.mlir
+++ b/mlir/test/Conversion/OpenACCToLLVM/atomic.mlir
@@ -140,9 +140,14 @@ module {
// -----
+// A simple binop capture uses atomicrmw; the read follows the update, so the
+// captured value is recomputed from the old one.
+
// CHECK-LABEL: llvm.func @convert_capture_ur
-// CHECK: llvm.cmpxchg %{{.*}}, %{{.*}}, %{{.*}} acq_rel monotonic : !llvm.ptr, i32
-// CHECK: llvm.store %{{.*}}, %{{.*}} : i32, !llvm.ptr
+// CHECK-NOT: llvm.cmpxchg
+// CHECK: %[[OLD:.*]] = llvm.atomicrmw add %{{.*}}, %[[VAL:.*]] monotonic : !llvm.ptr, i32
+// CHECK: %[[NEW:.*]] = llvm.add %[[OLD]], %[[VAL]] : i32
+// CHECK: llvm.store %[[NEW]], %{{.*}} : i32, !llvm.ptr
module {
func.func @convert_capture_ur(%v: memref<i32>, %x: memref<i32>, %val: i32) {
@@ -160,6 +165,50 @@ module {
// -----
+// The read precedes the update, so the old value is captured directly.
+
+// CHECK-LABEL: llvm.func @convert_capture_ru_atomicrmw
+// CHECK-NOT: llvm.cmpxchg
+// CHECK: %[[OLD:.*]] = llvm.atomicrmw sub %{{.*}}, %{{.*}} monotonic : !llvm.ptr, i32
+// CHECK: llvm.store %[[OLD]], %{{.*}} : i32, !llvm.ptr
+
+module {
+ func.func @convert_capture_ru_atomicrmw(%v: memref<i32>, %x: memref<i32>, %val: i32) {
+ acc.atomic.capture {
+ acc.atomic.read %v = %x : memref<i32>, memref<i32>, i32
+ acc.atomic.update %x : memref<i32> {
+ ^bb0(%arg: i32):
+ %0 = arith.subi %arg, %val : i32
+ acc.yield %0 : i32
+ }
+ }
+ return
+ }
+}
+
+// -----
+
+// `expr - x` is not an atomicrmw sub and must keep the cmpxchg loop.
+
+// CHECK-LABEL: llvm.func @convert_capture_sub_rhs
+// CHECK: llvm.cmpxchg
+
+module {
+ func.func @convert_capture_sub_rhs(%v: memref<i32>, %x: memref<i32>, %val: i32) {
+ acc.atomic.capture {
+ acc.atomic.update %x : memref<i32> {
+ ^bb0(%arg: i32):
+ %0 = arith.subi %val, %arg : i32
+ acc.yield %0 : i32
+ }
+ acc.atomic.read %v = %x : memref<i32>, memref<i32>, i32
+ }
+ return
+ }
+}
+
+// -----
+
// Test per-component atomicrmw for double complex (complex<f64>) atomic update.
// CHECK-LABEL: llvm.func @double_complex_atomic_add
>From c366cdb7f0aefcf25f4e07e9ca0de23ffd97f181 Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Fri, 28 Aug 2026 13:07:34 -0700
Subject: [PATCH 2/2] [mlir][acc] Share the update shape check with the capture
path
Use getAtomicUpdateOp()/getAtomicReadOp(), and share the
`x = x <binop> expr` match between the update and capture conversions so
the same region cannot take atomicrmw in one and the cmpxchg loop in the
other. Move the capture fast path into a helper that fails early rather
than falling through.
---
.../OpenACCToLLVM/ACCAtomicPatterns.cpp | 161 +++++++++---------
1 file changed, 84 insertions(+), 77 deletions(-)
diff --git a/mlir/lib/Conversion/OpenACCToLLVM/ACCAtomicPatterns.cpp b/mlir/lib/Conversion/OpenACCToLLVM/ACCAtomicPatterns.cpp
index 24ea912950d58..760154649df42 100644
--- a/mlir/lib/Conversion/OpenACCToLLVM/ACCAtomicPatterns.cpp
+++ b/mlir/lib/Conversion/OpenACCToLLVM/ACCAtomicPatterns.cpp
@@ -178,6 +178,13 @@ class ACCAtomicOpConversion : public ConvertOpToLLVMPattern<AtomicOpTy> {
Block *constructCmpxchgLoop(Value ptr, Type type, Value expr,
ConversionPatternRewriter &rewriter) const;
+ /// Emit a single atomicrmw for a `x = x <binop> expr` capture, or fail so
+ /// the caller falls back to the cmpxchg loop.
+ LogicalResult
+ tryEmitCaptureAtomicRMW(AtomicCaptureOp capture, AtomicUpdateOp update,
+ AtomicReadOp read,
+ ConversionPatternRewriter &rewriter) const;
+
Value genUpdateCmpxchgLoop(AtomicUpdateOp update,
ConversionPatternRewriter &rewriter) const;
};
@@ -533,6 +540,31 @@ static std::optional<LLVM::AtomicBinOp> getAtomicBinOp(Operation *op,
.Default([](Operation *) { return std::nullopt; });
}
+/// Match an update region computing `x = x <binop> expr`, and return the
+/// atomicrmw kind with the binary operation. Shared by the update and capture
+/// conversions so the same region cannot take atomicrmw in one and the cmpxchg
+/// loop in the other.
+static std::optional<std::pair<LLVM::AtomicBinOp, Operation *>>
+matchAtomicBinOpUpdate(AtomicUpdateOp update) {
+ Block &block = update.getRegion().front();
+ Value arg = block.getArgument(0);
+ Operation *yield = block.getTerminator();
+ if (!yield || yield->getNumOperands() != 1)
+ return std::nullopt;
+ Operation *binOp = yield->getOperand(0).getDefiningOp();
+ if (!binOp || binOp->getBlock() != &block || binOp->getNumOperands() != 2 ||
+ binOp->getNumResults() != 1)
+ return std::nullopt;
+ // The updated value has to feed the binop and nothing else.
+ if (!arg.hasOneUse() || arg.use_begin()->getOwner() != binOp)
+ return std::nullopt;
+ bool updateIsLhs = binOp->getOperand(0) == arg;
+ std::optional<LLVM::AtomicBinOp> kind = getAtomicBinOp(binOp, updateIsLhs);
+ if (!kind)
+ return std::nullopt;
+ return std::make_pair(*kind, binOp);
+}
+
/// Generate llvm.atomicrmw or an llvm.cmpxchg loop.
template <>
LogicalResult ACCAtomicOpConversion<AtomicUpdateOp>::matchAndRewrite(
@@ -594,11 +626,11 @@ LogicalResult ACCAtomicOpConversion<AtomicUpdateOp>::matchAndRewrite(
Operation &firstOp = ops.front();
Operation &yield = ops.back();
- if (dependents.size() == 2 && firstOp.getResult(0) == yield.getOperand(0)) {
- bool updateIsLhs = firstOp.getOperand(0) == updateArgument;
- kind = getAtomicBinOp(&firstOp, updateIsLhs);
- if (kind)
- val = firstOp.getOperand(updateIsLhs ? 1 : 0);
+ if (auto matched = matchAtomicBinOpUpdate(update)) {
+ Operation *binOp = matched->second;
+ bool updateIsLhs = binOp->getOperand(0) == updateArgument;
+ kind = matched->first;
+ val = binOp->getOperand(updateIsLhs ? 1 : 0);
}
// Per-component atomicrmw info for complex type decomposition.
@@ -709,43 +741,55 @@ static bool exprReadsMemory(Value expr) {
return false;
}
-/// Match an update region holding a single `x = x <binop> expr`. Returns the
-/// atomicrmw kind and the binary operation, or nullopt when the region needs a
-/// cmpxchg loop.
-static std::optional<std::pair<LLVM::AtomicBinOp, Operation *>>
-matchSimpleAtomicUpdate(AtomicUpdateOp update, const TypeConverter *converter) {
- Block &block = update.getRegion().front();
- if (!llvm::hasNItems(block.getOperations(), 2))
- return std::nullopt;
- Operation &binOp = block.front();
- Operation &yield = block.back();
- if (binOp.getNumOperands() != 2 || binOp.getNumResults() != 1)
- return std::nullopt;
- if (yield.getNumOperands() != 1 || yield.getOperand(0) != binOp.getResult(0))
- return std::nullopt;
+/// Emit a single atomicrmw for a `x = x <binop> expr` capture.
+template <typename AtomicOpTy>
+LogicalResult ACCAtomicOpConversion<AtomicOpTy>::tryEmitCaptureAtomicRMW(
+ AtomicCaptureOp capture, AtomicUpdateOp update, AtomicReadOp read,
+ ConversionPatternRewriter &rewriter) const {
+ if (read.getX() != update.getX())
+ return failure();
+ auto matched = matchAtomicBinOpUpdate(update);
+ if (!matched)
+ return failure();
+ auto [kind, binOp] = *matched;
- // The updated value must be used exactly once, as an operand of the binop.
- Value arg = block.getArgument(0);
- bool updateIsLhs = binOp.getOperand(0) == arg;
- if (updateIsLhs == (binOp.getOperand(1) == arg))
- return std::nullopt;
- if (!arg.hasOneUse())
- return std::nullopt;
+ Value arg = update.getRegion().front().getArgument(0);
+ bool updateIsLhs = binOp->getOperand(0) == arg;
+ Value expr = binOp->getOperand(updateIsLhs ? 1 : 0);
// Keep serialized and aggregate types on the cmpxchg path.
Type argTy = arg.getType();
- if (!argTy.isIntOrFloat() || converter->convertType(argTy) != argTy)
- return std::nullopt;
-
- // The other operand has to be available outside the region.
- Operation *exprDef = binOp.getOperand(updateIsLhs ? 1 : 0).getDefiningOp();
- if (exprDef && exprDef->getBlock() == &block)
- return std::nullopt;
+ if (!argTy.isIntOrFloat() ||
+ this->getTypeConverter()->convertType(argTy) != argTy)
+ return failure();
+ // The operand must already be available, and must not be the captured value.
+ Operation *exprDef = expr.getDefiningOp();
+ if (exprDef && exprDef->getBlock() == &update.getRegion().front())
+ return failure();
+ if (exprReadsMemory(expr))
+ return failure();
+
+ Location loc = capture.getLoc();
+ Value xRef = update.getX();
+ Value vRef = read.getV();
+ Value xPtr =
+ getAtomicPointer(xRef, rewriter.getRemappedValue(xRef), loc, rewriter);
+ Value vPtr =
+ getAtomicPointer(vRef, rewriter.getRemappedValue(vRef), loc, rewriter);
- std::optional<LLVM::AtomicBinOp> kind = getAtomicBinOp(&binOp, updateIsLhs);
- if (!kind)
- return std::nullopt;
- return std::make_pair(*kind, &binOp);
+ rewriter.setInsertionPoint(capture);
+ auto rmw = LLVM::AtomicRMWOp::create(rewriter, loc, kind, xPtr,
+ rewriter.getRemappedValue(expr),
+ LLVM::AtomicOrdering::monotonic);
+ // atomicrmw yields the old value; `{update, read}` captures the new one.
+ Value captured = rmw.getRes();
+ if (capture.getFirstOp() == update.getOperation()) {
+ rewriter.moveOpAfter(binOp, rmw);
+ binOp->replaceUsesOfWith(arg, captured);
+ captured = binOp->getResult(0);
+ }
+ rewriter.replaceOpWithNewOp<LLVM::StoreOp>(capture, captured, vPtr);
+ return success();
}
/// Generate an llvm.cmpxchg loop.
@@ -760,47 +804,10 @@ LogicalResult ACCAtomicOpConversion<AtomicCaptureOp>::matchAndRewrite(
// A single `x = x <binop> expr` capture becomes one atomicrmw. The cmpxchg
// loop below serializes retries and collapses under contention.
- auto update = dyn_cast<AtomicUpdateOp>(firstOp)
- ? cast<AtomicUpdateOp>(firstOp)
- : dyn_cast<AtomicUpdateOp>(secondOp);
- auto read = dyn_cast<AtomicReadOp>(firstOp)
- ? cast<AtomicReadOp>(firstOp)
- : dyn_cast<AtomicReadOp>(secondOp);
- if (update && read && read.getX() == update.getX()) {
- std::optional<std::pair<LLVM::AtomicBinOp, Operation *>> matched =
- matchSimpleAtomicUpdate(update, this->getTypeConverter());
- Value arg = update.getRegion().front().getArgument(0);
- Value xRef = update.getX();
- Value vRef = read.getV();
- if (matched) {
- auto [kind, binOp] = *matched;
- bool updateIsLhs = binOp->getOperand(0) == arg;
- Value expr = binOp->getOperand(updateIsLhs ? 1 : 0);
- if (exprReadsMemory(expr))
- matched = std::nullopt;
- else {
- Location loc = capture.getLoc();
- Value xPtr = getAtomicPointer(xRef, rewriter.getRemappedValue(xRef),
- loc, rewriter);
- Value vFastPtr = getAtomicPointer(vRef, rewriter.getRemappedValue(vRef),
- loc, rewriter);
- rewriter.setInsertionPoint(capture);
- auto rmw = LLVM::AtomicRMWOp::create(rewriter, loc, kind, xPtr,
- rewriter.getRemappedValue(expr),
- LLVM::AtomicOrdering::monotonic);
- // atomicrmw yields the old value; `{update, read}` captures the new
- // one.
- Value captured = rmw.getRes();
- if (isa<AtomicUpdateOp>(firstOp)) {
- rewriter.moveOpAfter(binOp, rmw);
- binOp->replaceUsesOfWith(arg, captured);
- captured = binOp->getResult(0);
- }
- rewriter.replaceOpWithNewOp<LLVM::StoreOp>(capture, captured, vFastPtr);
+ if (AtomicUpdateOp update = capture.getAtomicUpdateOp())
+ if (AtomicReadOp read = capture.getAtomicReadOp())
+ if (succeeded(tryEmitCaptureAtomicRMW(capture, update, read, rewriter)))
return success();
- }
- }
- }
if (auto firstReadStmt = dyn_cast<AtomicReadOp>(firstOp)) {
Location loc = capture.getLoc();
More information about the Mlir-commits
mailing list