[Mlir-commits] [mlir] [mlir][OpenACC] Map the initial value when a private recipe reads its… (PR #219215)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-openacc

@llvm/pr-subscribers-openacc

Author: Matsu (khaki3)

<details>
<summary>Changes</summary>

Example:
```fortran
real*8, allocatable :: dep(:)
allocate(dep(n))
!$acc parallel loop private(dep) async(1)
do i = 1, n
   dep(i) = 1.0d0
end do
!$acc wait(1)
deallocate(dep)
```

In this code, the private recipe's init region reads `dep`'s descriptor for its bounds, so the descriptor stays live-in to the compute region with no data clause. `acc-implicit-data` then maps it with copy semantics, and the runtime copies the array payload in and out — data the kernel never reads, on storage freed right after the wait.

Fix: `acc-recipe-materialization` maps such a variable through `acc.firstprivate_map`, the same way firstprivate already does, so its initial value is captured rather than left to implicit mapping.

---
Full diff: https://github.com/llvm/llvm-project/pull/219215.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp (+24-14) 
- (added) mlir/test/Dialect/OpenACC/acc-recipe-materialization-private-init-reads-var.mlir (+53) 


``````````diff
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp
index e1775ce3706e0..e2df498178df2 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp
@@ -167,9 +167,9 @@ class ACCRecipeMaterialization
   void runOnOperation() override;
 
 private:
-  // When handling firstprivate, the initial value needs to be available on
-  // the GPU. One way to get that value there is to map the variable through
-  // global memory.
+  // When the recipe reads the original variable, its initial value needs to be
+  // available on the GPU. One way to get that value there is to map the
+  // variable through global memory.
   // Thus, when we materialize a firstprivate, we materialize it into
   // a mapping action first. This function ends up with doing the following:
   // %dev = acc.firstprivate var(%var)
@@ -180,7 +180,8 @@ class ACCRecipeMaterialization
   // being removed. But because of the way we chain it to the
   // `acc.firstprivate_map`, then its result becomes live-in to the
   // compute region and used as the variable the initial value is loaded from.
-  void handleFirstprivateMapping(acc::FirstprivateOp firstprivateOp) const;
+  template <typename OpTy>
+  void handleInitialValueMapping(OpTy op) const;
   template <typename OpTy>
   void removeRecipe(OpTy op, ModuleOp moduleOp) const;
   template <typename OpTy, typename RecipeOpTy, typename AccOpTy>
@@ -192,15 +193,22 @@ class ACCRecipeMaterialization
                                     acc::ACCToGPUMappingPolicy &policy) const;
 };
 
-void ACCRecipeMaterialization::handleFirstprivateMapping(
-    acc::FirstprivateOp firstprivateOp) const {
-  OpBuilder builder(firstprivateOp);
-  auto mapFirstprivateOp = acc::FirstprivateMapInitialOp::create(
-      builder, firstprivateOp.getLoc(), firstprivateOp.getVar(),
-      firstprivateOp.getStructured(), firstprivateOp.getImplicit(),
-      firstprivateOp.getBounds());
-  mapFirstprivateOp.setName(firstprivateOp.getName());
-  firstprivateOp.getVarMutable().assign(mapFirstprivateOp.getAccVar());
+template <typename OpTy>
+void ACCRecipeMaterialization::handleInitialValueMapping(OpTy op) const {
+  OpBuilder builder(op);
+  auto mapInitialOp = acc::FirstprivateMapInitialOp::create(
+      builder, op.getLoc(), op.getVar(), op.getStructured(), op.getImplicit(),
+      op.getBounds());
+  mapInitialOp.setName(op.getName());
+  op.getVarMutable().assign(mapInitialOp.getAccVar());
+}
+
+// Whether the init region reads the variable it privatizes - a descriptor
+// recipe loads it for the bounds, while a scalar one ignores it.
+static bool initReadsVar(Region &initRegion) {
+  if (initRegion.empty() || initRegion.getNumArguments() == 0)
+    return false;
+  return !initRegion.getArgument(0).use_empty();
 }
 
 template <typename OpTy>
@@ -447,7 +455,7 @@ LogicalResult ACCRecipeMaterialization::materializeForACCOp(
       auto recipeOp = cast<acc::FirstprivateRecipeOp>(decl);
       LLVM_DEBUG(llvm::dbgs() << "materializing: " << firstprivateOp << "\n"
                               << symbolRef << "\n");
-      handleFirstprivateMapping(firstprivateOp);
+      handleInitialValueMapping(firstprivateOp);
       if (failed(
               materialize(firstprivateOp, recipeOp, accOp, accSupport, policy)))
         return failure();
@@ -466,6 +474,8 @@ LogicalResult ACCRecipeMaterialization::materializeForACCOp(
       auto recipeOp = cast<acc::PrivateRecipeOp>(decl);
       LLVM_DEBUG(llvm::dbgs() << "materializing: " << privateOp << "\n"
                               << symbolRef << "\n");
+      if (initReadsVar(recipeOp.getInitRegion()))
+        handleInitialValueMapping(privateOp);
       if (failed(materialize(privateOp, recipeOp, accOp, accSupport, policy)))
         return failure();
     }
diff --git a/mlir/test/Dialect/OpenACC/acc-recipe-materialization-private-init-reads-var.mlir b/mlir/test/Dialect/OpenACC/acc-recipe-materialization-private-init-reads-var.mlir
new file mode 100644
index 0000000000000..1f677e7b3ce67
--- /dev/null
+++ b/mlir/test/Dialect/OpenACC/acc-recipe-materialization-private-init-reads-var.mlir
@@ -0,0 +1,53 @@
+// RUN: mlir-opt %s -acc-recipe-materialization | FileCheck %s
+
+// A private recipe whose init region reads the original variable needs that
+// variable on the device, so the materialization has to map its initial value.
+acc.private.recipe @privatization_memref_dyn : memref<?xi32> init {
+^bb0(%arg0: memref<?xi32>):
+  %c0 = arith.constant 0 : index
+  %dim = memref.dim %arg0, %c0 : memref<?xi32>
+  %0 = memref.alloca(%dim) : memref<?xi32>
+  acc.yield %0 : memref<?xi32>
+} destroy {
+^bb0(%arg0: memref<?xi32>, %arg1: memref<?xi32>):
+  acc.terminator
+}
+
+// A private recipe that ignores the original variable must not map anything.
+acc.private.recipe @privatization_memref_i32 : memref<i32> init {
+^bb0(%arg0: memref<i32>):
+  %0 = memref.alloca() : memref<i32>
+  acc.yield %0 : memref<i32>
+}
+
+// CHECK-LABEL: func.func @private_dyn
+// CHECK: %[[MAP:.*]] = acc.firstprivate_map varPtr(%{{.*}} : memref<?xi32>)
+// CHECK: acc.parallel
+// CHECK: %[[DIM:.*]] = memref.dim %[[MAP]]
+// CHECK: %[[ALLOCA:.*]] = memref.alloca(%[[DIM]])
+// CHECK: memref.store %{{.*}}, %[[ALLOCA]]
+
+func.func @private_dyn(%arg0 : memref<?xi32>) {
+  %c0 = arith.constant 0 : index
+  %c1336 = arith.constant 1336 : i32
+  %priv = acc.private varPtr(%arg0 : memref<?xi32>) recipe(@privatization_memref_dyn) implicit(true) name("t") -> memref<?xi32>
+  acc.parallel private(%priv : memref<?xi32>) {
+    memref.store %c1336, %priv[%c0] : memref<?xi32>
+    acc.yield
+  }
+  return
+}
+
+// CHECK-LABEL: func.func @private_scalar
+// CHECK-NOT: acc.firstprivate_map
+// CHECK: acc.parallel
+
+func.func @private_scalar(%arg0 : memref<i32>) {
+  %c1336 = arith.constant 1336 : i32
+  %priv = acc.private varPtr(%arg0 : memref<i32>) recipe(@privatization_memref_i32) implicit(true) name("t") -> memref<i32>
+  acc.parallel private(%priv : memref<i32>) {
+    memref.store %c1336, %priv[] : memref<i32>
+    acc.yield
+  }
+  return
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/219215