[Mlir-commits] [mlir] [OpenACC] update location of recipes when materializing (PR #205915)
Scott Manley
llvmlistbot at llvm.org
Fri Jun 26 05:08:25 PDT 2026
https://github.com/rscottmanley updated https://github.com/llvm/llvm-project/pull/205915
>From 9e6e6bd79068ca9bafbddf851929c49c418f6bd1 Mon Sep 17 00:00:00 2001
From: Scott Manley <scmanley at nvidia.com>
Date: Fri, 26 Jun 2026 05:05:24 -0700
Subject: [PATCH] [OpenACC] update location of recipes when materializing
As recipes are generated per type and not per variable, we can end up
with the same location for multiple private/firstprivate/reduction
variables. When materializing the recipes, set the Location of all
Operations within the recipe region to be that of the op that is being
materialized.
---
.../Transforms/ACCRecipeMaterialization.cpp | 32 +++-
.../acc-recipe-materialization-loc.mlir | 142 ++++++++++++++++++
2 files changed, 168 insertions(+), 6 deletions(-)
create mode 100644 mlir/test/Dialect/OpenACC/acc-recipe-materialization-loc.mlir
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp
index 04f2a996dadc9..05a4257737e78 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp
@@ -73,6 +73,15 @@ namespace {
using namespace mlir;
+static void setLocation(Region ®ion, Location loc) {
+ // Since recipes are generated per type and not per variable, the location
+ // of the recipe operations which get inlined will not necessarily be the
+ // same as the location of the op that is being materialized. Force an update
+ // of the location of the recipe operations to the location of the op that is
+ // being materialized.
+ region.walk([&](Operation *op) { op->setLoc(loc); });
+}
+
static void saveVarName(StringRef name, Value dst) {
if (name.empty())
return;
@@ -130,7 +139,7 @@ static void resolveVarNamePlaceholders(Block *block, Block::iterator ip,
// block. Values must be provided for the destroy region block arguments
// according to the recipe specifications.
template <typename RecipeOpTy>
-static void cloneDestroy(RecipeOpTy recipe, mlir::Block *block,
+static void cloneDestroy(Location loc, RecipeOpTy recipe, mlir::Block *block,
Block::iterator ip,
const llvm::SmallVector<mlir::Value> &arguments) {
IRMapping mapping{};
@@ -138,6 +147,9 @@ static void cloneDestroy(RecipeOpTy recipe, mlir::Block *block,
assert(destroyRegion.getBlocks().front().getNumArguments() ==
arguments.size() &&
"unexpected acc recipe destroy block arguments");
+
+ setLocation(destroyRegion, loc);
+
mapping.map(destroyRegion.getBlocks().front().getArguments(), arguments);
acc::cloneACCRegionInto(&destroyRegion, block, ip, mapping,
/*resultsToReplace=*/{});
@@ -267,6 +279,9 @@ ACCRecipeMaterialization::materialize(OpTy op, RecipeOpTy recipe, AccOpTy accOp,
initArgs.append(triples);
mapping.map(initRegion.getBlocks().front().getArguments(), initArgs);
+ Location loc = op.getLoc();
+ setLocation(initRegion, loc);
+
if constexpr (std::is_same_v<OpTy, acc::PrivateOp>) {
// Clone the init region for a private.
Block *block = ®ion.front();
@@ -279,7 +294,7 @@ ACCRecipeMaterialization::materialize(OpTy op, RecipeOpTy recipe, AccOpTy accOp,
if (!recipe.getDestroyRegion().empty()) {
results.insert(results.begin(), origPtr);
results.append(triples);
- cloneDestroy(recipe, block, std::prev(block->end()), results);
+ cloneDestroy(loc, recipe, block, std::prev(block->end()), results);
}
} else if constexpr (std::is_same_v<OpTy, acc::FirstprivateOp>) {
// Clone the init region for a firstprivate.
@@ -297,16 +312,19 @@ ACCRecipeMaterialization::materialize(OpTy op, RecipeOpTy recipe, AccOpTy accOp,
mapping.clear();
mapping.map(recipe.getCopyRegion().front().getArguments(), results);
// Clone the copy region for a firstprivate.
- acc::cloneACCRegionInto(&recipe.getCopyRegion(), block, std::next(ip),
- mapping, {});
+ Region ©Region = recipe.getCopyRegion();
+ setLocation(copyRegion, loc);
+ acc::cloneACCRegionInto(©Region, block, std::next(ip), mapping, {});
if (!recipe.getDestroyRegion().empty()) {
// origPtr was already pushed.
- cloneDestroy(recipe, block, std::prev(block->end()), results);
+ cloneDestroy(loc, recipe, block, std::prev(block->end()), results);
}
} else if constexpr (std::is_same_v<OpTy, acc::ReductionOp>) {
auto cloneRegionIntoAccRegion = [&](Region *src, Region *dest,
bool hasResult) {
src->cloneInto(dest, mapping);
+ // TODO: update location of the cloned operations to the location of the
+ // op
Block *block = &dest->front();
Operation *terminator = block->getTerminator();
b.setInsertionPoint(terminator);
@@ -339,6 +357,8 @@ ACCRecipeMaterialization::materialize(OpTy op, RecipeOpTy recipe, AccOpTy accOp,
// Clone the combiner region into acc.reduction_combine_region.
Region &combinerRegion = recipe.getCombinerRegion();
+ setLocation(combinerRegion, loc);
+
Block *entryBlock = &combinerRegion.front();
if constexpr (std::is_same_v<AccOpTy, acc::ParallelOp>)
@@ -375,7 +395,7 @@ ACCRecipeMaterialization::materialize(OpTy op, RecipeOpTy recipe, AccOpTy accOp,
if (!recipe.getDestroyRegion().empty()) {
SmallVector<Value> results{origPtr, reductionOp.getResult()};
Block::iterator ip = std::next(Block::iterator(combineRegionOp));
- cloneDestroy(recipe, combineRegionOp->getBlock(), ip, results);
+ cloneDestroy(loc, recipe, combineRegionOp->getBlock(), ip, results);
}
} else {
llvm_unreachable("unexpected op type");
diff --git a/mlir/test/Dialect/OpenACC/acc-recipe-materialization-loc.mlir b/mlir/test/Dialect/OpenACC/acc-recipe-materialization-loc.mlir
new file mode 100644
index 0000000000000..87f07da1483fe
--- /dev/null
+++ b/mlir/test/Dialect/OpenACC/acc-recipe-materialization-loc.mlir
@@ -0,0 +1,142 @@
+// RUN: mlir-opt %s -acc-recipe-materialization \
+// RUN: -mlir-print-debuginfo -mlir-pretty-debuginfo | FileCheck %s
+
+acc.private.recipe @privatization_memref_index : memref<index> init {
+^bb0(%arg0: memref<index>):
+ %alloca = memref.alloca() : memref<index>
+ acc.yield %alloca : memref<index>
+} destroy {
+^bb0(%arg0: memref<index>, %arg1: memref<index>):
+ memref.dealloc %arg1 : memref<index>
+ acc.terminator
+}
+
+acc.firstprivate.recipe @firstprivatization_memref_index : memref<index> init {
+^bb0(%arg0: memref<index>):
+ %alloca = memref.alloca() : memref<index>
+ acc.yield %alloca : memref<index>
+} copy {
+^bb0(%arg0: memref<index>, %arg1: memref<index>):
+ %0 = memref.load %arg0[] : memref<index>
+ memref.store %0, %arg1[] : memref<index>
+ acc.terminator
+} destroy {
+^bb0(%arg0: memref<index>, %arg1: memref<index>):
+ memref.dealloc %arg1 : memref<index>
+ acc.terminator
+}
+
+acc.reduction.recipe @reduction_add_memref_index : memref<index> reduction_operator <add> init {
+^bb0(%arg0: memref<index>):
+ %c0 = arith.constant 0 : index
+ %alloca = memref.alloca() : memref<index>
+ memref.store %c0, %alloca[] : memref<index>
+ acc.yield %alloca : memref<index>
+} combiner {
+^bb0(%arg0: memref<index>, %arg1: memref<index>):
+ %0 = memref.load %arg0[] : memref<index>
+ %1 = memref.load %arg1[] : memref<index>
+ %2 = arith.addi %0, %1 : index
+ memref.store %2, %arg0[] : memref<index>
+ acc.yield %arg0 : memref<index>
+} destroy {
+^bb0(%arg0: memref<index>, %arg1: memref<index>):
+ memref.dealloc %arg1 : memref<index>
+ acc.terminator
+}
+
+func.func @private_loc(%arg0: memref<index>) {
+ %c1 = arith.constant 1 : index
+ %0 = acc.private varPtr(%arg0 : memref<index>) recipe(@privatization_memref_index) -> memref<index> {implicit = true, name = "priv0"}
+ acc.parallel private(%0 : memref<index>) {
+ memref.store %c1, %0[] : memref<index>
+ acc.yield
+ } attributes {independent = [#acc.device_type<none>]}
+ %1 = acc.private varPtr(%arg0 : memref<index>) recipe(@privatization_memref_index) -> memref<index> {implicit = true, name = "priv1"}
+ acc.parallel private(%1 : memref<index>) {
+ memref.store %c1, %1[] : memref<index>
+ acc.yield
+ } attributes {independent = [#acc.device_type<none>]}
+
+ // CHECK-LABEL: func.func @private_loc
+ // CHECK: acc.parallel
+ // CHECK: [[ALLOCA0:%.*]] = memref.alloca() {{.*}}"priv0"{{.*}}:50
+ // CHECK: memref.dealloc [[ALLOCA0]] {{.*}}:50
+ // CHECK: acc.parallel
+ // CHECK: [[ALLOCA1:%.*]] = memref.alloca() {{.*}}"priv1"{{.*}}:55
+ // CHECK: memref.dealloc [[ALLOCA1]] {{.*}}:55
+ return
+}
+
+func.func @firstprivate_loc() {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %alloca = memref.alloca() : memref<index>
+ memref.store %c0, %alloca[] : memref<index>
+ %0 = acc.firstprivate varPtr(%alloca : memref<index>) recipe(@firstprivatization_memref_index) -> memref<index> {implicit = true, name = "firstpriv0"}
+ acc.parallel firstprivate(%0 : memref<index>) {
+ memref.store %c1, %0[] : memref<index>
+ acc.yield
+ }
+ %1 = acc.firstprivate varPtr(%alloca : memref<index>) recipe(@firstprivatization_memref_index) -> memref<index> {implicit = true, name = "firstpriv1"}
+ acc.parallel firstprivate(%1 : memref<index>) {
+ memref.store %c1, %1[] : memref<index>
+ acc.yield
+ }
+ // CHECK-LABEL: func.func @firstprivate_loc
+ // CHECK: acc.parallel
+ // CHECK: [[ALLOCA2:%.*]] = memref.alloca() {{.*}}"firstpriv0"{{.*}}:76
+ // CHECK-NEXT: [[LOAD2:%.*]] = memref.load{{.*}}:76
+ // CHECK-NEXT: memref.store [[LOAD2]], [[ALLOCA2]][]{{.*}}:76
+ // CHECK: memref.dealloc [[ALLOCA2]] {{.*}}:76
+ // CHECK: acc.parallel
+ // CHECK: [[ALLOCA3:%.*]] = memref.alloca() {{.*}}"firstpriv1"{{.*}}:81
+ // CHECK-NEXT: [[LOAD3:%.*]] = memref.load{{.*}}:81
+ // CHECK-NEXT: memref.store [[LOAD3]], [[ALLOCA3]][]{{.*}}:81
+ // CHECK: memref.dealloc [[ALLOCA3]] {{.*}}:81
+ return
+}
+
+func.func @reduction_loc(%arg0: memref<index>) {
+ %c1 = arith.constant 1 : index
+ %0 = acc.reduction varPtr(%arg0 : memref<index>) recipe(@reduction_add_memref_index) -> memref<index> {name = "r0"}
+ acc.parallel reduction(%0 : memref<index>) {
+ memref.store %c1, %0[] : memref<index>
+ acc.yield
+ }
+ %1 = acc.reduction varPtr(%arg0 : memref<index>) recipe(@reduction_add_memref_index) -> memref<index> {name = "r1"}
+ acc.parallel reduction(%1 : memref<index>) {
+ memref.store %c1, %1[] : memref<index>
+ acc.yield
+ }
+ // CHECK-LABEL: func.func @reduction_loc
+ // CHECK: acc.parallel
+ // CHECK: [[INIT0:%.*]] = acc.reduction_init {{.*}} {
+ // CHECK: [[ZERO0:%.*]] = arith.constant 0 : index {{.*}}:102
+ // CHECK-NEXT: [[ALLOCA4:%.*]] = memref.alloca() {{.*}}:102
+ // CHECK-NEXT: memref.store [[ZERO0]], [[ALLOCA4]][]{{.*}}:102
+ // CHECK-NEXT: acc.yield [[ALLOCA4]] {{.*}}:102
+ // CHECK-NEXT: } {{.*}}"r0"{{.*}}:102
+ // CHECK: acc.reduction_combine_region {{.*}} {
+ // CHECK-NEXT: [[LOADX:%.*]] = memref.load{{.*}}:102
+ // CHECK-NEXT: [[LOADY:%.*]] = memref.load{{.*}}:102
+ // CHECK-NEXT: [[ADD0:%.*]] = arith.addi [[LOADX]], [[LOADY]]{{.*}}:102
+ // CHECK-NEXT: memref.store [[ADD0]]{{.*}}:102
+ // CHECK-NEXT: } {{.*}}:102
+ // CHECK: memref.dealloc [[INIT0]]{{.*}}:102
+ // CHECK: acc.parallel
+ // CHECK: [[INIT1:%.*]] = acc.reduction_init {{.*}} {
+ // CHECK: [[ZERO1:%.*]] = arith.constant 0 : index {{.*}}:107
+ // CHECK-NEXT: [[ALLOCA5:%.*]] = memref.alloca() {{.*}}:107
+ // CHECK-NEXT: memref.store [[ZERO1]], [[ALLOCA5]][]{{.*}}:107
+ // CHECK-NEXT: acc.yield [[ALLOCA5]] {{.*}}:107
+ // CHECK-NEXT: } {{.*}}"r1"{{.*}}:107
+ // CHECK: acc.reduction_combine_region {{.*}} {
+ // CHECK-NEXT: [[LOADA:%.*]] = memref.load{{.*}}:107
+ // CHECK-NEXT: [[LOADB:%.*]] = memref.load{{.*}}:107
+ // CHECK-NEXT: [[ADD1:%.*]] = arith.addi [[LOADA]], [[LOADB]]{{.*}}:107
+ // CHECK-NEXT: memref.store [[ADD1]]{{.*}}:107
+ // CHECK-NEXT: } {{.*}}:107
+ // CHECK: memref.dealloc [[INIT1]]{{.*}}:107
+ return
+}
More information about the Mlir-commits
mailing list