[flang][openacc] Attach post allocate action on the correct operation

[clementval]

In some cases (when using stat), the action was attached to the invisible fir.result op. Apply same fix as in #89662.

[llvmbot]

@llvm/pr-subscribers-flang-fir-hlfir

@llvm/pr-subscribers-openacc

Author: Valentin Clement (バレンタイン クレメン) (clementval)

Changes

In some cases (when using stat), the action was attach to the invisible fir.result op. Apply same fix as in #89662.

---

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

2 Files Affected:

• (modified) flang/lib/Lower/OpenACC.cpp (+21-14)
• (modified) flang/test/Lower/OpenACC/acc-declare.f90 (+17)

diff --git a/flang/lib/Lower/OpenACC.cpp b/flang/lib/Lower/OpenACC.cpp
index 431fab52872d33..c97398fc43f923 100644
--- a/flang/lib/Lower/OpenACC.cpp
+++ b/flang/lib/Lower/OpenACC.cpp
@@ -4293,23 +4293,30 @@ void Fortran::lower::attachDeclarePostAllocAction(
     const Fortran::semantics::Symbol &sym) {
   std::stringstream fctName;
   fctName << converter.mangleName(sym) << declarePostAllocSuffix.str();
-  mlir::Operation &op = builder.getInsertionBlock()->back();
+  mlir::Operation *op = &builder.getInsertionBlock()->back();
+
+  if (auto resOp = mlir::dyn_cast<fir::ResultOp>(*op)) {
+    assert(resOp.getOperands().size() == 0 &&
+           "expect only fir.result op with no operand");
+    op = op->getPrevNode();
+  }
+  assert(op && "expect operation to attach the post allocation action");
 
-  if (op.hasAttr(mlir::acc::getDeclareActionAttrName())) {
-    auto attr = op.getAttrOfType<mlir::acc::DeclareActionAttr>(
+  if (op->hasAttr(mlir::acc::getDeclareActionAttrName())) {
+    auto attr = op->getAttrOfType<mlir::acc::DeclareActionAttr>(
         mlir::acc::getDeclareActionAttrName());
-    op.setAttr(mlir::acc::getDeclareActionAttrName(),
-               mlir::acc::DeclareActionAttr::get(
-                   builder.getContext(), attr.getPreAlloc(),
-                   /*postAlloc=*/builder.getSymbolRefAttr(fctName.str()),
-                   attr.getPreDealloc(), attr.getPostDealloc()));
+    op->setAttr(mlir::acc::getDeclareActionAttrName(),
+                mlir::acc::DeclareActionAttr::get(
+                    builder.getContext(), attr.getPreAlloc(),
+                    /*postAlloc=*/builder.getSymbolRefAttr(fctName.str()),
+                    attr.getPreDealloc(), attr.getPostDealloc()));
   } else {
-    op.setAttr(mlir::acc::getDeclareActionAttrName(),
-               mlir::acc::DeclareActionAttr::get(
-                   builder.getContext(),
-                   /*preAlloc=*/{},
-                   /*postAlloc=*/builder.getSymbolRefAttr(fctName.str()),
-                   /*preDealloc=*/{}, /*postDealloc=*/{}));
+    op->setAttr(mlir::acc::getDeclareActionAttrName(),
+                mlir::acc::DeclareActionAttr::get(
+                    builder.getContext(),
+                    /*preAlloc=*/{},
+                    /*postAlloc=*/builder.getSymbolRefAttr(fctName.str()),
+                    /*preDealloc=*/{}, /*postDealloc=*/{}));
   }
 }
 
diff --git a/flang/test/Lower/OpenACC/acc-declare.f90 b/flang/test/Lower/OpenACC/acc-declare.f90
index ff1e756c20e125..0066e712fbdcce 100644
--- a/flang/test/Lower/OpenACC/acc-declare.f90
+++ b/flang/test/Lower/OpenACC/acc-declare.f90
@@ -455,3 +455,20 @@ module acc_declare_allocatable_test3
 
 ! CHECK-LABEL: acc.global_ctor @_QMacc_declare_allocatable_test3Edata1_acc_ctor
 ! CHECK-LABEL: acc.global_ctor @_QMacc_declare_allocatable_test3Edata2_acc_ctor
+
+module acc_declare_post_action_stat
+  real, dimension(:), allocatable :: x, y
+  !$acc declare create(x,y)
+
+contains
+
+  subroutine init()
+    integer :: stat
+    allocate(x(10), y(10), stat=stat)
+  end subroutine
+end module
+
+! CHECK-LABEL: func.func @_QMacc_declare_post_action_statPinit()
+! CHECK: fir.call @_FortranAAllocatableAllocate({{.*}}) fastmath<contract> {acc.declare_action = #acc.declare_action<postAlloc = @_QMacc_declare_post_action_statEx_acc_declare_update_desc_post_alloc>} : (!fir.ref<!fir.box<none>>, i1, !fir.box<none>, !fir.ref<i8>, i32) -> i32
+! CHECK: fir.if
+! CHECK: fir.call @_FortranAAllocatableAllocate({{.*}}) fastmath<contract> {acc.declare_action = #acc.declare_action<postAlloc = @_QMacc_declare_post_action_statEy_acc_declare_update_desc_post_alloc>} : (!fir.ref<!fir.box<none>>, i1, !fir.box<none>, !fir.ref<i8>, i32) -> i32

[vzakhari]

Thank you, Valentin!