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(); +} + 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 { 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 { 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 { 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, %v: memref, %co // ----- +// CHECK-LABEL: func.func @test_parallel_if_atomic_capture_write( +// CHECK-SAME: %[[X:.*]]: memref, %[[V:.*]]: memref, %[[EXPR:.*]]: i32 +func.func @test_parallel_if_atomic_capture_write(%x: memref, %v: memref, %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, memref, i32 + acc.atomic.write %x = %expr : memref, 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