[MLIR][OpenACC] Defer lowering atomic capture children

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();

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