[flang-commits] [flang] [flang][OpenACC] Handle declare actions on CUF allocation ops (PR #214568)

[via flang-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-openacc

Author: Andre Kuhlenschmidt (akuhlens)

<details>
<summary>Changes</summary>

Recover the typed descriptor from cuf.allocate and cuf.deallocate when inserting calls to OpenACC declare-action recipes. Add conversion and lowering coverage for pinned allocations.

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


3 Files Affected:

- (modified) flang/lib/Optimizer/OpenACC/Transforms/ACCDeclareActionConversion.cpp (+5) 
- (modified) flang/test/Fir/OpenACC/declare-action-conversion.fir (+33) 
- (added) flang/test/Lower/OpenACC/acc-declare-cuda-pinned.f90 (+29) 


``````````diff
diff --git a/flang/lib/Optimizer/OpenACC/Transforms/ACCDeclareActionConversion.cpp b/flang/lib/Optimizer/OpenACC/Transforms/ACCDeclareActionConversion.cpp
index 0973e18f7241f..cd01e34ae503e 100644
--- a/flang/lib/Optimizer/OpenACC/Transforms/ACCDeclareActionConversion.cpp
+++ b/flang/lib/Optimizer/OpenACC/Transforms/ACCDeclareActionConversion.cpp
@@ -52,6 +52,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "flang/Optimizer/Dialect/CUF/CUFOps.h"
 #include "flang/Optimizer/Dialect/FIROps.h"
 #include "flang/Optimizer/Dialect/FIRType.h"
 #include "flang/Optimizer/OpenACC/Passes.h"
@@ -175,6 +176,10 @@ class ACCDeclareActionConversion
                         [&](auto store) { return store.getMemref(); })
                     .Case<fir::BoxAddrOp>(
                         [&](auto boxAddr) { return boxAddr.getVal(); })
+                    .Case<cuf::AllocateOp>(
+                        [&](auto allocate) { return allocate.getBox(); })
+                    .Case<cuf::DeallocateOp>(
+                        [&](auto deallocate) { return deallocate.getBox(); })
                     .Case<fir::CallOp>([&](fir::CallOp call) -> Value {
                       if (auto callee = call.getCalleeAttr()) {
                         StringRef funcName =
diff --git a/flang/test/Fir/OpenACC/declare-action-conversion.fir b/flang/test/Fir/OpenACC/declare-action-conversion.fir
index fe54dae0f5fb3..d4f82cf9045b6 100644
--- a/flang/test/Fir/OpenACC/declare-action-conversion.fir
+++ b/flang/test/Fir/OpenACC/declare-action-conversion.fir
@@ -171,3 +171,36 @@ module {
   func.func private @_FortranAAllocatableDeallocate(!fir.ref<!fir.box<none>>, i1, !fir.box<none>, !fir.ref<i8>, i32) -> i32
   func.func private @_FortranAAllocatableDeallocatePolymorphic(!fir.ref<!fir.box<none>>, !fir.ref<none>, i1, !fir.box<none>, !fir.ref<i8>, i32) -> i32
 }
+
+// -----
+
+// Test declare_action on CUF allocation and deallocation operations.
+module {
+  func.func private @_QFpre_alloc(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+    return
+  }
+  func.func private @_QFpost_alloc(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+    return
+  }
+  func.func private @_QFpre_dealloc(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+    return
+  }
+  func.func private @_QFpost_dealloc(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+    return
+  }
+
+  func.func @_QPcuf_declare_actions(%box: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+    %0 = cuf.allocate %box : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {acc.declare_action = #acc.declare_action<preAlloc = @_QFpre_alloc, postAlloc = @_QFpost_alloc>, data_attr = #cuf.cuda<pinned>} -> i32
+    %1 = cuf.deallocate %box : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {acc.declare_action = #acc.declare_action<preDealloc = @_QFpre_dealloc, postDealloc = @_QFpost_dealloc>, data_attr = #cuf.cuda<pinned>} -> i32
+    return
+  }
+
+// CHECK-LABEL: func.func @_QPcuf_declare_actions(
+// CHECK-SAME: %[[BOX:.*]]: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>)
+// CHECK: fir.call @_QFpre_alloc(%[[BOX]])
+// CHECK-NEXT: %{{.*}} = cuf.allocate %[[BOX]]
+// CHECK-NEXT: fir.call @_QFpost_alloc(%[[BOX]])
+// CHECK: fir.call @_QFpre_dealloc(%[[BOX]])
+// CHECK-NEXT: %{{.*}} = cuf.deallocate %[[BOX]]
+// CHECK-NEXT: fir.call @_QFpost_dealloc(%[[BOX]])
+}
diff --git a/flang/test/Lower/OpenACC/acc-declare-cuda-pinned.f90 b/flang/test/Lower/OpenACC/acc-declare-cuda-pinned.f90
new file mode 100644
index 0000000000000..c6993431c9036
--- /dev/null
+++ b/flang/test/Lower/OpenACC/acc-declare-cuda-pinned.f90
@@ -0,0 +1,29 @@
+! RUN: bbc -emit-hlfir -fopenacc -fcuda -gpu=pinned %s -o - | FileCheck %s
+
+program declare_pinned
+  real, allocatable :: a
+  !$acc declare device_resident(a)
+  allocate(a)
+end program
+
+! CHECK-LABEL: func.func @_QQmain()
+! CHECK: %[[A_ALLOC:.*]] = cuf.alloc !fir.box<!fir.heap<f32>> {{.*}}data_attr = #cuf.cuda<pinned>
+! CHECK: %[[A:.*]]:2 = hlfir.declare %[[A_ALLOC]]
+! CHECK: cuf.allocate %[[A]]#0 : !fir.ref<!fir.box<!fir.heap<f32>>> {acc.declare_action = #acc.declare_action<postAlloc = @{{.*}}_acc_declare_post_alloc>, data_attr = #cuf.cuda<pinned>} -> i32
+! CHECK: cuf.deallocate %[[A]]#0 : !fir.ref<!fir.box<!fir.heap<f32>>> {acc.declare_action = #acc.declare_action<postDealloc = @{{.*}}_acc_declare_post_dealloc>, data_attr = #cuf.cuda<pinned>} -> i32
+
+! CHECK-LABEL: func.func private @{{.*}}_acc_declare_post_alloc(
+! CHECK-SAME: %[[POST_ALLOC_ARG:.*]]: !fir.ref<!fir.box<!fir.heap<f32>>>)
+! CHECK: acc.declare_device_resident varPtr(%[[POST_ALLOC_ARG]] : !fir.ref<!fir.box<!fir.heap<f32>>>)
+! CHECK: acc.declare_enter
+
+! CHECK-LABEL: func.func private @{{.*}}_acc_declare_pre_dealloc(
+! CHECK-SAME: %[[PRE_DEALLOC_ARG:.*]]: !fir.ref<!fir.box<!fir.heap<f32>>>)
+! CHECK: acc.getdeviceptr varPtr(%[[PRE_DEALLOC_ARG]] : !fir.ref<!fir.box<!fir.heap<f32>>>)
+! CHECK: acc.declare_exit
+! CHECK: acc.delete
+
+! CHECK-LABEL: func.func private @{{.*}}_acc_declare_post_dealloc(
+! CHECK-SAME: %[[POST_DEALLOC_ARG:.*]]: !fir.ref<!fir.box<!fir.heap<f32>>>)
+! CHECK: acc.getdeviceptr varPtr(%[[POST_DEALLOC_ARG]] : !fir.ref<!fir.box<!fir.heap<f32>>>){{.*}}implicit = true
+! CHECK: acc.declare_exit

``````````

</details>


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