[Mlir-commits] [mlir] [OpenACC] add AutomaticAllocationScope to recipe ops (PR #124337)

[Scott Manley]

https://github.com/rscottmanley updated https://github.com/llvm/llvm-project/pull/124337

>From 87058144a81dda78a90d6ad73c657750dfb6cc34 Mon Sep 17 00:00:00 2001
From: Scott Manley <scmanley at nvidia.com>
Date: Fri, 24 Jan 2025 12:04:49 -0800
Subject: [PATCH] [OpenACC] add AutomaticAllocationScope to recipe ops

The recipe operators should have AutomaticAllocationScope so recipes
can be converted using operators that require parent ops to have
AutomaticAllocationScope
---
 .../mlir/Dialect/OpenACC/OpenACCOps.td        | 17 +++++---
 mlir/test/Dialect/OpenACC/ops.mlir            | 41 +++++++++++++++++++
 2 files changed, 52 insertions(+), 6 deletions(-)

diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td b/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
index c60eb5cc620a7d..7e9ed2c741cf76 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
@@ -1008,8 +1008,9 @@ def OpenACC_UpdateHostOp : OpenACC_DataExitOpWithVarPtr<"update_host",
 // 2.5.13 private clause
 //===----------------------------------------------------------------------===//
 
-def OpenACC_PrivateRecipeOp : OpenACC_Op<"private.recipe",
-    [IsolatedFromAbove, Symbol, RecipeInterface]> {
+def OpenACC_PrivateRecipeOp
+    : OpenACC_Op<"private.recipe", [IsolatedFromAbove, Symbol, RecipeInterface,
+                                    AutomaticAllocationScope]> {
   let summary = "privatization recipe";
 
   let description = [{
@@ -1065,8 +1066,10 @@ def OpenACC_PrivateRecipeOp : OpenACC_Op<"private.recipe",
 // 2.5.14 firstprivate clause
 //===----------------------------------------------------------------------===//
 
-def OpenACC_FirstprivateRecipeOp : OpenACC_Op<"firstprivate.recipe",
-    [IsolatedFromAbove, Symbol, RecipeInterface]> {
+def OpenACC_FirstprivateRecipeOp
+    : OpenACC_Op<"firstprivate.recipe", [IsolatedFromAbove, Symbol,
+                                         RecipeInterface,
+                                         AutomaticAllocationScope]> {
   let summary = "privatization recipe";
 
   let description = [{
@@ -1131,8 +1134,10 @@ def OpenACC_FirstprivateRecipeOp : OpenACC_Op<"firstprivate.recipe",
 // 2.5.15 reduction clause
 //===----------------------------------------------------------------------===//
 
-def OpenACC_ReductionRecipeOp : OpenACC_Op<"reduction.recipe",
-    [IsolatedFromAbove, Symbol, RecipeInterface]> {
+def OpenACC_ReductionRecipeOp
+    : OpenACC_Op<"reduction.recipe", [IsolatedFromAbove, Symbol,
+                                      RecipeInterface,
+                                      AutomaticAllocationScope]> {
   let summary = "reduction recipe";
 
   let description = [{
diff --git a/mlir/test/Dialect/OpenACC/ops.mlir b/mlir/test/Dialect/OpenACC/ops.mlir
index 2706792a263a82..28ab6f9fcfb4ce 100644
--- a/mlir/test/Dialect/OpenACC/ops.mlir
+++ b/mlir/test/Dialect/OpenACC/ops.mlir
@@ -1892,3 +1892,44 @@ func.func @acc_combined() {
 // CHECK: acc.loop combined(kernels)
 // CHECK: acc.serial combined(loop)
 // CHECK: acc.loop combined(serial)
+
+acc.firstprivate.recipe @firstprivatization_memref_i32 : memref<i32> init {
+^bb0(%arg0: memref<i32>):
+  %alloca = memref.alloca() : memref<i32>
+  acc.yield %alloca : memref<i32>
+} copy {
+^bb0(%arg0: memref<i32>, %arg1: memref<i32>):
+  %0 = memref.load %arg1[] : memref<i32>
+  memref.store %0, %arg0[] : memref<i32>
+  acc.terminator
+}
+
+// CHECK-LABEL: acc.firstprivate.recipe @firstprivatization_memref_i32
+// CHECK:       memref.alloca
+
+acc.reduction.recipe @reduction_add_memref_i32 : memref<i32> reduction_operator <add> init {
+^bb0(%arg0: memref<i32>):
+  %c0_i32 = arith.constant 0 : i32
+  %alloca = memref.alloca() : memref<i32>
+  memref.store %c0_i32, %alloca[] : memref<i32>
+  acc.yield %alloca : memref<i32>
+} combiner {
+^bb0(%arg0: memref<i32>, %arg1: memref<i32>):
+  %0 = memref.load %arg0[] : memref<i32>
+  %1 = memref.load %arg1[] : memref<i32>
+  %2 = arith.addi %0, %1 : i32
+  memref.store %2, %arg0[] : memref<i32>
+  acc.yield %arg0 : memref<i32>
+}
+
+// CHECK-LABEL: acc.reduction.recipe @reduction_add_memref_i32
+// CHECK:       memref.alloca
+
+acc.private.recipe @privatization_memref_i32 : memref<i32> init {
+^bb0(%arg0: memref<i32>):
+  %alloca = memref.alloca() : memref<i32>
+  acc.yield %alloca : memref<i32>
+}
+
+// CHECK-LABEL: acc.private.recipe @privatization_memref_i32
+// CHECK:       memref.alloca