[Mlir-commits] [mlir] [mlir][acc] Preserve shared data operations during if-clause lowering (PR #208771)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Jul 10 09:30:50 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-mlir-openacc
Author: Matsu (khaki3)
<details>
<summary>Changes</summary>
Example:
```fortran
!$acc kernels present(grid, c) if(offload_on)
do nc = 1, grid%rc%ncol
c(1,nc) = 0
end do
c(1,1) = 0
!$acc end kernels
```
In this code, kernels restructuring can create compute constructs that share data-entry operations with an enclosing data region. If-clause lowering could rewrite or erase these shared operations, producing invalid IR or duplicate erasure.
Fix: Preserve externally owned data operations, deduplicate cleanup, and retain duplicate operand ordering while cloning local operations once.
---
Full diff: https://github.com/llvm/llvm-project/pull/208771.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/OpenACC/Transforms/ACCIfClauseLowering.cpp (+74-27)
- (modified) mlir/test/Dialect/OpenACC/acc-if-clause-lowering.mlir (+99)
``````````diff
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCIfClauseLowering.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCIfClauseLowering.cpp
index 71df75958a134..8adf441f7bd44 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCIfClauseLowering.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCIfClauseLowering.cpp
@@ -66,6 +66,7 @@
#include "mlir/IR/PatternMatch.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SetVector.h"
#include "llvm/Support/Debug.h"
namespace mlir {
@@ -92,8 +93,10 @@ class ACCIfClauseLowering
void convertHostRegion(Operation *computeOp, Region ®ion);
template <typename OpTy>
- void lowerIfClauseForComputeConstruct(OpTy computeConstructOp,
- SmallVector<Operation *> &eraseOps);
+ void
+ lowerIfClauseForComputeConstruct(OpTy computeConstructOp,
+ llvm::SetVector<Operation *> &eraseOps,
+ llvm::SetVector<Operation *> &condEraseOps);
public:
void runOnOperation() override;
@@ -122,7 +125,8 @@ void ACCIfClauseLowering::convertHostRegion(Operation *computeOp,
// constructs
template <typename OpTy>
void ACCIfClauseLowering::lowerIfClauseForComputeConstruct(
- OpTy computeConstructOp, SmallVector<Operation *> &eraseOps) {
+ OpTy computeConstructOp, llvm::SetVector<Operation *> &eraseOps,
+ llvm::SetVector<Operation *> &condEraseOps) {
Value ifCond = computeConstructOp.getIfCond();
if (!ifCond)
return;
@@ -135,24 +139,36 @@ void ACCIfClauseLowering::lowerIfClauseForComputeConstruct(
// Collect data clause operations that need to be recreated in the if
// condition
- SmallVector<Operation *> dataEntryOps;
- SmallVector<Operation *> dataExitOps;
+ llvm::SetVector<Operation *> dataEntryOps;
+ llvm::SetVector<Operation *> dataExitOps;
SmallVector<Operation *> firstprivateOps;
SmallVector<Operation *> privateOps;
SmallVector<Operation *> reductionOps;
+ // Entries used outside this construct belong to the surrounding scope.
+ // Their entry and exit operations must remain outside the conditional.
+ auto isExternallyOwned = [&](Operation *dataOp) {
+ for (Operation *user : dataOp->getUsers())
+ if (!isa<ACC_DATA_EXIT_OPS>(user) && user != computeConstructOp &&
+ !computeConstructOp->isAncestor(user))
+ return true;
+ return false;
+ };
+
// Collect data entry operations
for (Value operand : computeConstructOp.getDataClauseOperands())
if (Operation *defOp = operand.getDefiningOp())
if (isa<ACC_DATA_ENTRY_OPS>(defOp))
- dataEntryOps.push_back(defOp);
+ dataEntryOps.insert(defOp);
- // Find corresponding exit operations for each entry operation.
+ // Find corresponding exit operations for each local entry operation. Exit ops
+ // of externally owned entry ops belong to the enclosing construct.
// Iterate backwards through entry ops since exit ops appear in reverse order.
for (Operation *dataEntryOp : llvm::reverse(dataEntryOps))
- for (Operation *user : dataEntryOp->getUsers())
- if (isa<ACC_DATA_EXIT_OPS>(user))
- dataExitOps.push_back(user);
+ if (!isExternallyOwned(dataEntryOp))
+ for (Operation *user : dataEntryOp->getUsers())
+ if (isa<ACC_DATA_EXIT_OPS>(user))
+ dataExitOps.insert(user);
// Collect firstprivate, private, and reduction operations
auto collectOps = [&](SmallVector<Operation *> &ops, OperandRange operands) {
@@ -191,7 +207,18 @@ void ACCIfClauseLowering::lowerIfClauseForComputeConstruct(
deviceMapping.map(op->getResult(0), clonedOp->getResult(0));
}
};
- cloneAndMapOps(dataEntryOps, deviceDataOperands);
+ // Clone each local data entry op once. Externally owned ops (mapped by a
+ // surrounding scope) are referenced directly.
+ for (Operation *op : dataEntryOps) {
+ if (isExternallyOwned(op))
+ continue;
+ Operation *clonedOp = rewriter.clone(*op, deviceMapping);
+ deviceMapping.map(op->getResult(0), clonedOp->getResult(0));
+ }
+ // Preserve the original operand order and multiplicity while using the
+ // cloned value for local entries and the original value for external ones.
+ for (Value operand : computeConstructOp.getDataClauseOperands())
+ deviceDataOperands.push_back(deviceMapping.lookupOrDefault(operand));
cloneAndMapOps(firstprivateOps, firstprivateOperands);
cloneAndMapOps(privateOps, privateOperands);
cloneAndMapOps(reductionOps, reductionOperands);
@@ -240,43 +267,63 @@ void ACCIfClauseLowering::lowerIfClauseForComputeConstruct(
}
// The original op is now empty and can be erased
- eraseOps.push_back(computeConstructOp);
+ eraseOps.insert(computeConstructOp);
// TODO: Can probably 'move' the data ops instead of cloning them
// which would eliminate need to explicitly erase
for (Operation *dataOp : dataExitOps)
- eraseOps.push_back(dataOp);
+ eraseOps.insert(dataOp);
- // The new host code may contain uses of the acc variables. Replace them by
- // the host values.
- auto replaceAndEraseOps = [&](SmallVector<Operation *> &ops) {
+ // Redirect host-side uses while preserving uses outside this construct.
+ // Shared operations are erased only after their users are gone.
+ Region &elseRegion = ifOp.getElseRegion();
+ auto replaceHostUsesAndScheduleErase = [&](auto &ops) {
for (Operation *op : ops) {
- getAccVar(op).replaceAllUsesWith(getVar(op));
- eraseOps.push_back(op);
+ getAccVar(op).replaceUsesWithIf(getVar(op), [&](OpOperand &use) {
+ return elseRegion.isAncestor(use.getOwner()->getParentRegion());
+ });
+ condEraseOps.insert(op);
}
};
- replaceAndEraseOps(dataEntryOps);
- replaceAndEraseOps(firstprivateOps);
- replaceAndEraseOps(privateOps);
- replaceAndEraseOps(reductionOps);
+ replaceHostUsesAndScheduleErase(dataEntryOps);
+ replaceHostUsesAndScheduleErase(firstprivateOps);
+ replaceHostUsesAndScheduleErase(privateOps);
+ replaceHostUsesAndScheduleErase(reductionOps);
}
void ACCIfClauseLowering::runOnOperation() {
func::FuncOp funcOp = getOperation();
accSupport = &getAnalysis<OpenACCSupport>();
- SmallVector<Operation *> eraseOps;
+ llvm::SetVector<Operation *> eraseOps;
+ llvm::SetVector<Operation *> condEraseOps;
funcOp.walk([&](Operation *op) {
if (auto parallelOp = dyn_cast<acc::ParallelOp>(op))
- lowerIfClauseForComputeConstruct(parallelOp, eraseOps);
+ lowerIfClauseForComputeConstruct(parallelOp, eraseOps, condEraseOps);
else if (auto kernelsOp = dyn_cast<acc::KernelsOp>(op))
- lowerIfClauseForComputeConstruct(kernelsOp, eraseOps);
+ lowerIfClauseForComputeConstruct(kernelsOp, eraseOps, condEraseOps);
else if (auto serialOp = dyn_cast<acc::SerialOp>(op))
- lowerIfClauseForComputeConstruct(serialOp, eraseOps);
+ lowerIfClauseForComputeConstruct(serialOp, eraseOps, condEraseOps);
});
- for (Operation *op : eraseOps)
+ for (Operation *op : llvm::reverse(eraseOps))
op->erase();
+ // Shared entry/private/reduction ops can become dead in stages.
+ // Revisit deferred producers after their consumers are erased.
+ SmallVector<Operation *> pendingEraseOps(condEraseOps.begin(),
+ condEraseOps.end());
+ bool erased;
+ do {
+ erased = false;
+ for (size_t i = pendingEraseOps.size(); i > 0; --i) {
+ Operation *op = pendingEraseOps[i - 1];
+ if (!op->use_empty())
+ continue;
+ pendingEraseOps.erase(pendingEraseOps.begin() + i - 1);
+ op->erase();
+ erased = true;
+ }
+ } while (erased);
}
} // namespace
diff --git a/mlir/test/Dialect/OpenACC/acc-if-clause-lowering.mlir b/mlir/test/Dialect/OpenACC/acc-if-clause-lowering.mlir
index af4cc72a42645..9cc0df1dfb107 100644
--- a/mlir/test/Dialect/OpenACC/acc-if-clause-lowering.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-if-clause-lowering.mlir
@@ -406,3 +406,102 @@ func.func @test_parallel_if_atomic_capture(%x: memref<i32>, %v: memref<i32>, %co
}
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
+// ends up with a non data-entry op as its data operand and fails to verify).
+// CHECK-LABEL: func.func @test_kernels_if_shared_present
+func.func @test_kernels_if_shared_present(%arg0: memref<10xi32>, %cond: i1) {
+ %c0_i32 = arith.constant 0 : i32
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ // CHECK: %[[PRESENT:.*]] = acc.present varPtr(%arg0 : memref<10xi32>) -> memref<10xi32>
+ %present = acc.present varPtr(%arg0 : memref<10xi32>) -> memref<10xi32>
+ // The enclosing acc.data must keep referencing the original present op.
+ // CHECK: acc.data dataOperands(%[[PRESENT]] : memref<10xi32>) {
+ acc.data dataOperands(%present : memref<10xi32>) {
+ // CHECK: scf.if %{{.*}} {
+ // The externally owned present op is referenced directly, not cloned/erased.
+ // CHECK-NOT: acc.present
+ // CHECK: acc.kernels dataOperands(%[[PRESENT]] : memref<10xi32>)
+ // CHECK: } else {
+ // Host path uses the original variable, not the present result.
+ // CHECK: memref.store %{{.*}}, %arg0[%{{.*}}] : memref<10xi32>
+ // CHECK: }
+ acc.kernels dataOperands(%present : memref<10xi32>) if(%cond) {
+ scf.for %i = %c0 to %c10 step %c1 {
+ memref.store %c0_i32, %present[%i] : memref<10xi32>
+ }
+ acc.terminator
+ }
+ acc.terminator
+ }
+ return
+}
+
+// -----
+
+// A data entry and exit shared by sibling compute constructs belong to their
+// surrounding scope. Neither conditional lowering may clone or erase them.
+// CHECK-LABEL: func.func @test_sibling_compute_shared_data
+func.func @test_sibling_compute_shared_data(%arg0: memref<10xi32>, %cond0: i1, %cond1: i1) {
+ %c0_i32 = arith.constant 0 : i32
+ %c0 = arith.constant 0 : index
+ // CHECK: %[[COPYIN:.*]] = acc.copyin varPtr(%arg0 : memref<10xi32>) -> memref<10xi32>
+ %copyin = acc.copyin varPtr(%arg0 : memref<10xi32>) -> memref<10xi32>
+ // CHECK: scf.if %{{.*}} {
+ // CHECK: acc.parallel dataOperands(%[[COPYIN]] : memref<10xi32>)
+ // CHECK-NOT: acc.copyout
+ // CHECK: } else {
+ // CHECK: memref.store %{{.*}}, %arg0[%{{.*}}] : memref<10xi32>
+ // CHECK: }
+ acc.parallel dataOperands(%copyin : memref<10xi32>) if(%cond0) {
+ memref.store %c0_i32, %copyin[%c0] : memref<10xi32>
+ acc.yield
+ }
+ // CHECK: scf.if %{{.*}} {
+ // CHECK: acc.serial dataOperands(%[[COPYIN]] : memref<10xi32>)
+ // CHECK-NOT: acc.copyout
+ // CHECK: } else {
+ // CHECK: memref.store %{{.*}}, %arg0[%{{.*}}] : memref<10xi32>
+ // CHECK: }
+ acc.serial dataOperands(%copyin : memref<10xi32>) if(%cond1) {
+ memref.store %c0_i32, %copyin[%c0] : memref<10xi32>
+ acc.yield
+ }
+ // CHECK: acc.copyout accPtr(%[[COPYIN]] : memref<10xi32>) to varPtr(%arg0 : memref<10xi32>)
+ acc.copyout accPtr(%copyin : memref<10xi32>) to varPtr(%arg0 : memref<10xi32>)
+ return
+}
+
+// -----
+
+// Duplicate data operands must preserve their multiplicity without cloning or
+// erasing the same entry/exit operation more than once.
+// CHECK-LABEL: func.func @test_duplicate_data_operand
+func.func @test_duplicate_data_operand(%arg0: memref<10xi32>, %cond: i1) {
+ %c0_i32 = arith.constant 0 : i32
+ %c0 = arith.constant 0 : index
+ %copyin = acc.copyin varPtr(%arg0 : memref<10xi32>) -> memref<10xi32>
+ // CHECK-NOT: acc.copyin
+ // CHECK: scf.if %{{.*}} {
+ // CHECK: %[[COPYIN:.*]] = acc.copyin varPtr(%arg0 : memref<10xi32>) -> memref<10xi32>
+ // CHECK: acc.parallel dataOperands(%[[COPYIN]], %[[COPYIN]] : memref<10xi32>, memref<10xi32>)
+ // CHECK: acc.copyout accPtr(%[[COPYIN]] : memref<10xi32>) to varPtr(%arg0 : memref<10xi32>)
+ // CHECK-NOT: acc.copyout
+ // CHECK: } else {
+ // CHECK: memref.store %{{.*}}, %arg0[%{{.*}}] : memref<10xi32>
+ // CHECK: }
+ // CHECK-NOT: acc.copyout
+ // CHECK: return
+ acc.parallel dataOperands(%copyin, %copyin : memref<10xi32>, memref<10xi32>) if(%cond) {
+ memref.store %c0_i32, %copyin[%c0] : memref<10xi32>
+ acc.yield
+ }
+ acc.copyout accPtr(%copyin : memref<10xi32>) to varPtr(%arg0 : memref<10xi32>)
+ return
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/208771
More information about the Mlir-commits
mailing list