[MLIR][OpenACC] Defer lowering atomic capture children - #217475
Merged
Conversation
An atomic capture must retain its valid child sequence until its wrapper is inlined. Make child conversions decline while nested in a capture. Moving the wrapper block requeues the existing children, which can then lower without worklist-order assumptions. Assisted-by: Codex
|
@llvm/pr-subscribers-mlir Author: Mehdi Amini (joker-eph) ChangesAn atomic capture must retain its valid child sequence until its wrapper is inlined. Make child conversions decline while nested in a capture. Moving the wrapper block requeues the existing children, which can then lower without worklist-order assumptions. Assisted-by: Codex Full diff: https://github.com/llvm/llvm-project/pull/217475.diff 2 Files Affected:
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCSpecializeForHost.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCSpecializeForHost.cpp
index aaa364bf86684..9019827ac3964 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCSpecializeForHost.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCSpecializeForHost.cpp
@@ -110,11 +110,18 @@ static bool isInsideACCComputeConstruct(Operation *op) {
return false;
}
+/// Return true if an enclosing compute construct or capture must be removed
+/// before converting an atomic operation.
+static bool isAtomicConversionDeferred(Operation *op) {
+ return isInsideACCComputeConstruct(op) ||
+ op->getParentOfType<acc::AtomicCaptureOp>();
+}
+
namespace {
// Lower orphan acc.atomic.update by: load from addr, clone region expr with
// the loaded value, then store the computed result back to addr.
-// Only matches if NOT inside a compute region.
+// Only matches outside compute regions and atomic captures.
class ACCOrphanAtomicUpdateOpConversion
: public OpRewritePattern<acc::AtomicUpdateOp> {
public:
@@ -123,8 +130,7 @@ class ACCOrphanAtomicUpdateOpConversion
LogicalResult matchAndRewrite(acc::AtomicUpdateOp atomicUpdateOp,
PatternRewriter &rewriter) const override {
- // Only convert if this op is not inside an ACC compute construct
- if (isInsideACCComputeConstruct(atomicUpdateOp))
+ if (isAtomicConversionDeferred(atomicUpdateOp))
return failure();
Value x = atomicUpdateOp.getX();
@@ -167,7 +173,7 @@ class ACCOrphanAtomicUpdateOpConversion
};
// Lower orphan acc.atomic.read by: load from src, then store into dst.
-// Only matches if NOT inside an ACC compute construct.
+// Only matches outside compute regions and atomic captures.
class ACCOrphanAtomicReadOpConversion
: public OpRewritePattern<acc::AtomicReadOp> {
public:
@@ -176,8 +182,7 @@ class ACCOrphanAtomicReadOpConversion
LogicalResult matchAndRewrite(acc::AtomicReadOp readOp,
PatternRewriter &rewriter) const override {
- // Only convert if this op is not inside an ACC compute construct
- if (isInsideACCComputeConstruct(readOp))
+ if (isAtomicConversionDeferred(readOp))
return failure();
Value x = readOp.getX();
@@ -208,7 +213,7 @@ class ACCOrphanAtomicReadOpConversion
};
// Lower orphan acc.atomic.write by: store value into addr.
-// Only matches if NOT inside an ACC compute construct.
+// Only matches outside compute regions and atomic captures.
class ACCOrphanAtomicWriteOpConversion
: public OpRewritePattern<acc::AtomicWriteOp> {
public:
@@ -217,8 +222,7 @@ class ACCOrphanAtomicWriteOpConversion
LogicalResult matchAndRewrite(acc::AtomicWriteOp writeOp,
PatternRewriter &rewriter) const override {
- // Only convert if this op is not inside an ACC compute construct
- if (isInsideACCComputeConstruct(writeOp))
+ if (isAtomicConversionDeferred(writeOp))
return failure();
Value x = writeOp.getX();
diff --git a/mlir/test/Dialect/OpenACC/acc-if-clause-lowering.mlir b/mlir/test/Dialect/OpenACC/acc-if-clause-lowering.mlir
index cbdf9a1018804..e1c748f78b174 100644
--- a/mlir/test/Dialect/OpenACC/acc-if-clause-lowering.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-if-clause-lowering.mlir
@@ -1,7 +1,5 @@
// RUN: mlir-opt %s -acc-if-clause-lowering -split-input-file | FileCheck %s
-// XFAIL: mlir-expensive-checks
-
// Test acc.parallel with if condition
// CHECK-LABEL: func.func @test_parallel_if
func.func @test_parallel_if(%arg0: memref<10xi32>, %cond: i1) {
@@ -411,6 +409,30 @@ func.func @test_parallel_if_atomic_capture(%x: memref<i32>, %v: memref<i32>, %co
// -----
+// CHECK-LABEL: func.func @test_parallel_if_atomic_capture_write(
+// CHECK-SAME: %[[X:.*]]: memref<i32>, %[[V:.*]]: memref<i32>, %[[EXPR:.*]]: i32
+func.func @test_parallel_if_atomic_capture_write(%x: memref<i32>, %v: memref<i32>, %expr: i32, %cond: i1) {
+ // CHECK: scf.if %{{.*}} {
+ // CHECK: acc.parallel {
+ // CHECK: acc.atomic.capture {
+ // CHECK: acc.atomic.read %[[V]] = %[[X]]
+ // CHECK: acc.atomic.write %[[X]] = %[[EXPR]]
+ // CHECK: } else {
+ // CHECK-NOT: acc.atomic
+ // CHECK: memref.copy %[[X]], %[[V]]
+ // CHECK: memref.store %[[EXPR]], %[[X]][]
+ acc.parallel if(%cond) {
+ acc.atomic.capture {
+ acc.atomic.read %v = %x : memref<i32>, memref<i32>, i32
+ acc.atomic.write %x = %expr : memref<i32>, i32
+ }
+ acc.yield
+ }
+ return
+}
+
+// -----
+
// A data entry op (acc.present) shared by an enclosing acc.data and a nested
// acc.kernels that has an if clause. Lowering the kernels' if clause must not
// erase or rewrite the present op that acc.data still uses (otherwise acc.data
|
|
@llvm/pr-subscribers-openacc Author: Mehdi Amini (joker-eph) ChangesAn atomic capture must retain its valid child sequence until its wrapper is inlined. Make child conversions decline while nested in a capture. Moving the wrapper block requeues the existing children, which can then lower without worklist-order assumptions. Assisted-by: Codex Full diff: https://github.com/llvm/llvm-project/pull/217475.diff 2 Files Affected:
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCSpecializeForHost.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCSpecializeForHost.cpp
index aaa364bf86684..9019827ac3964 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCSpecializeForHost.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCSpecializeForHost.cpp
@@ -110,11 +110,18 @@ static bool isInsideACCComputeConstruct(Operation *op) {
return false;
}
+/// Return true if an enclosing compute construct or capture must be removed
+/// before converting an atomic operation.
+static bool isAtomicConversionDeferred(Operation *op) {
+ return isInsideACCComputeConstruct(op) ||
+ op->getParentOfType<acc::AtomicCaptureOp>();
+}
+
namespace {
// Lower orphan acc.atomic.update by: load from addr, clone region expr with
// the loaded value, then store the computed result back to addr.
-// Only matches if NOT inside a compute region.
+// Only matches outside compute regions and atomic captures.
class ACCOrphanAtomicUpdateOpConversion
: public OpRewritePattern<acc::AtomicUpdateOp> {
public:
@@ -123,8 +130,7 @@ class ACCOrphanAtomicUpdateOpConversion
LogicalResult matchAndRewrite(acc::AtomicUpdateOp atomicUpdateOp,
PatternRewriter &rewriter) const override {
- // Only convert if this op is not inside an ACC compute construct
- if (isInsideACCComputeConstruct(atomicUpdateOp))
+ if (isAtomicConversionDeferred(atomicUpdateOp))
return failure();
Value x = atomicUpdateOp.getX();
@@ -167,7 +173,7 @@ class ACCOrphanAtomicUpdateOpConversion
};
// Lower orphan acc.atomic.read by: load from src, then store into dst.
-// Only matches if NOT inside an ACC compute construct.
+// Only matches outside compute regions and atomic captures.
class ACCOrphanAtomicReadOpConversion
: public OpRewritePattern<acc::AtomicReadOp> {
public:
@@ -176,8 +182,7 @@ class ACCOrphanAtomicReadOpConversion
LogicalResult matchAndRewrite(acc::AtomicReadOp readOp,
PatternRewriter &rewriter) const override {
- // Only convert if this op is not inside an ACC compute construct
- if (isInsideACCComputeConstruct(readOp))
+ if (isAtomicConversionDeferred(readOp))
return failure();
Value x = readOp.getX();
@@ -208,7 +213,7 @@ class ACCOrphanAtomicReadOpConversion
};
// Lower orphan acc.atomic.write by: store value into addr.
-// Only matches if NOT inside an ACC compute construct.
+// Only matches outside compute regions and atomic captures.
class ACCOrphanAtomicWriteOpConversion
: public OpRewritePattern<acc::AtomicWriteOp> {
public:
@@ -217,8 +222,7 @@ class ACCOrphanAtomicWriteOpConversion
LogicalResult matchAndRewrite(acc::AtomicWriteOp writeOp,
PatternRewriter &rewriter) const override {
- // Only convert if this op is not inside an ACC compute construct
- if (isInsideACCComputeConstruct(writeOp))
+ if (isAtomicConversionDeferred(writeOp))
return failure();
Value x = writeOp.getX();
diff --git a/mlir/test/Dialect/OpenACC/acc-if-clause-lowering.mlir b/mlir/test/Dialect/OpenACC/acc-if-clause-lowering.mlir
index cbdf9a1018804..e1c748f78b174 100644
--- a/mlir/test/Dialect/OpenACC/acc-if-clause-lowering.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-if-clause-lowering.mlir
@@ -1,7 +1,5 @@
// RUN: mlir-opt %s -acc-if-clause-lowering -split-input-file | FileCheck %s
-// XFAIL: mlir-expensive-checks
-
// Test acc.parallel with if condition
// CHECK-LABEL: func.func @test_parallel_if
func.func @test_parallel_if(%arg0: memref<10xi32>, %cond: i1) {
@@ -411,6 +409,30 @@ func.func @test_parallel_if_atomic_capture(%x: memref<i32>, %v: memref<i32>, %co
// -----
+// CHECK-LABEL: func.func @test_parallel_if_atomic_capture_write(
+// CHECK-SAME: %[[X:.*]]: memref<i32>, %[[V:.*]]: memref<i32>, %[[EXPR:.*]]: i32
+func.func @test_parallel_if_atomic_capture_write(%x: memref<i32>, %v: memref<i32>, %expr: i32, %cond: i1) {
+ // CHECK: scf.if %{{.*}} {
+ // CHECK: acc.parallel {
+ // CHECK: acc.atomic.capture {
+ // CHECK: acc.atomic.read %[[V]] = %[[X]]
+ // CHECK: acc.atomic.write %[[X]] = %[[EXPR]]
+ // CHECK: } else {
+ // CHECK-NOT: acc.atomic
+ // CHECK: memref.copy %[[X]], %[[V]]
+ // CHECK: memref.store %[[EXPR]], %[[X]][]
+ acc.parallel if(%cond) {
+ acc.atomic.capture {
+ acc.atomic.read %v = %x : memref<i32>, memref<i32>, i32
+ acc.atomic.write %x = %expr : memref<i32>, i32
+ }
+ acc.yield
+ }
+ return
+}
+
+// -----
+
// A data entry op (acc.present) shared by an enclosing acc.data and a nested
// acc.kernels that has an if clause. Lowering the kernels' if clause must not
// erase or rewrite the present op that acc.data still uses (otherwise acc.data
|
|
@llvm/pr-subscribers-mlir-openacc Author: Mehdi Amini (joker-eph) ChangesAn atomic capture must retain its valid child sequence until its wrapper is inlined. Make child conversions decline while nested in a capture. Moving the wrapper block requeues the existing children, which can then lower without worklist-order assumptions. Assisted-by: Codex Full diff: https://github.com/llvm/llvm-project/pull/217475.diff 2 Files Affected:
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCSpecializeForHost.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCSpecializeForHost.cpp
index aaa364bf86684..9019827ac3964 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCSpecializeForHost.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCSpecializeForHost.cpp
@@ -110,11 +110,18 @@ static bool isInsideACCComputeConstruct(Operation *op) {
return false;
}
+/// Return true if an enclosing compute construct or capture must be removed
+/// before converting an atomic operation.
+static bool isAtomicConversionDeferred(Operation *op) {
+ return isInsideACCComputeConstruct(op) ||
+ op->getParentOfType<acc::AtomicCaptureOp>();
+}
+
namespace {
// Lower orphan acc.atomic.update by: load from addr, clone region expr with
// the loaded value, then store the computed result back to addr.
-// Only matches if NOT inside a compute region.
+// Only matches outside compute regions and atomic captures.
class ACCOrphanAtomicUpdateOpConversion
: public OpRewritePattern<acc::AtomicUpdateOp> {
public:
@@ -123,8 +130,7 @@ class ACCOrphanAtomicUpdateOpConversion
LogicalResult matchAndRewrite(acc::AtomicUpdateOp atomicUpdateOp,
PatternRewriter &rewriter) const override {
- // Only convert if this op is not inside an ACC compute construct
- if (isInsideACCComputeConstruct(atomicUpdateOp))
+ if (isAtomicConversionDeferred(atomicUpdateOp))
return failure();
Value x = atomicUpdateOp.getX();
@@ -167,7 +173,7 @@ class ACCOrphanAtomicUpdateOpConversion
};
// Lower orphan acc.atomic.read by: load from src, then store into dst.
-// Only matches if NOT inside an ACC compute construct.
+// Only matches outside compute regions and atomic captures.
class ACCOrphanAtomicReadOpConversion
: public OpRewritePattern<acc::AtomicReadOp> {
public:
@@ -176,8 +182,7 @@ class ACCOrphanAtomicReadOpConversion
LogicalResult matchAndRewrite(acc::AtomicReadOp readOp,
PatternRewriter &rewriter) const override {
- // Only convert if this op is not inside an ACC compute construct
- if (isInsideACCComputeConstruct(readOp))
+ if (isAtomicConversionDeferred(readOp))
return failure();
Value x = readOp.getX();
@@ -208,7 +213,7 @@ class ACCOrphanAtomicReadOpConversion
};
// Lower orphan acc.atomic.write by: store value into addr.
-// Only matches if NOT inside an ACC compute construct.
+// Only matches outside compute regions and atomic captures.
class ACCOrphanAtomicWriteOpConversion
: public OpRewritePattern<acc::AtomicWriteOp> {
public:
@@ -217,8 +222,7 @@ class ACCOrphanAtomicWriteOpConversion
LogicalResult matchAndRewrite(acc::AtomicWriteOp writeOp,
PatternRewriter &rewriter) const override {
- // Only convert if this op is not inside an ACC compute construct
- if (isInsideACCComputeConstruct(writeOp))
+ if (isAtomicConversionDeferred(writeOp))
return failure();
Value x = writeOp.getX();
diff --git a/mlir/test/Dialect/OpenACC/acc-if-clause-lowering.mlir b/mlir/test/Dialect/OpenACC/acc-if-clause-lowering.mlir
index cbdf9a1018804..e1c748f78b174 100644
--- a/mlir/test/Dialect/OpenACC/acc-if-clause-lowering.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-if-clause-lowering.mlir
@@ -1,7 +1,5 @@
// RUN: mlir-opt %s -acc-if-clause-lowering -split-input-file | FileCheck %s
-// XFAIL: mlir-expensive-checks
-
// Test acc.parallel with if condition
// CHECK-LABEL: func.func @test_parallel_if
func.func @test_parallel_if(%arg0: memref<10xi32>, %cond: i1) {
@@ -411,6 +409,30 @@ func.func @test_parallel_if_atomic_capture(%x: memref<i32>, %v: memref<i32>, %co
// -----
+// CHECK-LABEL: func.func @test_parallel_if_atomic_capture_write(
+// CHECK-SAME: %[[X:.*]]: memref<i32>, %[[V:.*]]: memref<i32>, %[[EXPR:.*]]: i32
+func.func @test_parallel_if_atomic_capture_write(%x: memref<i32>, %v: memref<i32>, %expr: i32, %cond: i1) {
+ // CHECK: scf.if %{{.*}} {
+ // CHECK: acc.parallel {
+ // CHECK: acc.atomic.capture {
+ // CHECK: acc.atomic.read %[[V]] = %[[X]]
+ // CHECK: acc.atomic.write %[[X]] = %[[EXPR]]
+ // CHECK: } else {
+ // CHECK-NOT: acc.atomic
+ // CHECK: memref.copy %[[X]], %[[V]]
+ // CHECK: memref.store %[[EXPR]], %[[X]][]
+ acc.parallel if(%cond) {
+ acc.atomic.capture {
+ acc.atomic.read %v = %x : memref<i32>, memref<i32>, i32
+ acc.atomic.write %x = %expr : memref<i32>, i32
+ }
+ acc.yield
+ }
+ return
+}
+
+// -----
+
// A data entry op (acc.present) shared by an enclosing acc.data and a nested
// acc.kernels that has an if clause. Lowering the kernels' if clause must not
// erase or rewrite the present op that acc.data still uses (otherwise acc.data
|
khaki3
approved these changes
Aug 19, 2026
Contributor
|
Thanks Mehdi! LGTM. |
kieroxide
pushed a commit
to kieroxide/llvm-project
that referenced
this pull request
Aug 21, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
An atomic capture must retain its valid child sequence until its wrapper is inlined. Make child conversions decline while nested in a capture. Moving the wrapper block requeues the existing children, which can then lower without worklist-order assumptions.
Assisted-by: Codex