[flang][openacc] Add acc.declare_action attributes on operation

[clementval]

Summary

This patches adds the acc.declare_action attrbites on
post allocate operation and pre/post deallocate operations.

[razvanlupusoru]

Looks great! Nice!

flang/include/flang/Lower/OpenACC.h

 namespace Fortran {
 namespace parser {
 struct OpenACCConstruct;
 struct OpenACCDeclarativeConstruct;
 } // namespace parser
 
 namespace semantics {
 class SemanticsContext;
+class Symbol;
 }
 
 namespace lower {
 
 class AbstractConverter;
 class StatementContext;
 
 namespace pft {

 
 /// Get a acc.firstprivate.recipe op for the given type or create it if it does
 /// not exist yet.
 mlir::acc::FirstprivateRecipeOp createOrGetFirstprivateRecipe(mlir::OpBuilder &,
                                                               llvm::StringRef,
                                                               mlir::Location,
                                                               mlir::Type);
 
+void attachDeclarePostAllocAction(AbstractConverter &, fir::FirOpBuilder &,
+                                  const Fortran::semantics::Symbol &);
+void attachDeclarePreDeallocAction(AbstractConverter &, fir::FirOpBuilder &,
+                                   mlir::Value beginOpValue,
+                                   const Fortran::semantics::Symbol &);
+void attachDeclarePostDeallocAction(AbstractConverter &, fir::FirOpBuilder &,
+                                    const Fortran::semantics::Symbol &);
+
 } // namespace lower
 } // namespace Fortran
 
 #endif // FORTRAN_LOWER_OPENACC_H

flang/include/flang/Optimizer/Builder/MutableBox.h

void genFinalization(fir::FirOpBuilder &builder, mlir::Location loc,
                      const fir::MutableBoxValue &box);
 
 void genInlinedAllocation(fir::FirOpBuilder &builder, mlir::Location loc,
                           const fir::MutableBoxValue &box,
                           mlir::ValueRange lbounds, mlir::ValueRange extents,
                           mlir::ValueRange lenParams, llvm::StringRef allocName,
                           bool mustBeHeap = false);
 
-void genInlinedDeallocate(fir::FirOpBuilder &builder, mlir::Location loc,
+mlir::Value genInlinedDeallocate(fir::FirOpBuilder &builder, mlir::Location loc,
                                  const fir::MutableBoxValue &box);
 
 /// When the MutableBoxValue was passed as a fir.ref<fir.box> to a call that may
 /// have modified it, update the MutableBoxValue according to the
 /// fir.ref<fir.box> value.
 void syncMutableBoxFromIRBox(fir::FirOpBuilder &builder, mlir::Location loc,
                              const fir::MutableBoxValue &box);
 
 /// Read all mutable properties into a normal symbol box.

flang/lib/Lower/Allocatable.cpp

 //===----------------------------------------------------------------------===//
 
 #include "flang/Lower/Allocatable.h"
 #include "flang/Evaluate/tools.h"
 #include "flang/Lower/AbstractConverter.h"
 #include "flang/Lower/ConvertType.h"
 #include "flang/Lower/IterationSpace.h"
 #include "flang/Lower/Mangler.h"
+#include "flang/Lower/OpenACC.h"
 #include "flang/Lower/PFTBuilder.h"
 #include "flang/Lower/Runtime.h"
 #include "flang/Lower/StatementContext.h"
 #include "flang/Optimizer/Builder/FIRBuilder.h"
 #include "flang/Optimizer/Builder/Runtime/RTBuilder.h"
 #include "flang/Optimizer/Builder/Todo.h"
 #include "flang/Optimizer/Dialect/FIROps.h"
 #include "flang/Optimizer/Dialect/FIROpsSupport.h"
fir::MutableBoxValue boxAddr =
         genMutableBoxValue(converter, loc, alloc.getAllocObj());
 
     if (sourceExpr)
       genSourceMoldAllocation(alloc, boxAddr, /*isSource=*/true);
     else if (moldExpr)
       genSourceMoldAllocation(alloc, boxAddr, /*isSource=*/false);
     else
       genSimpleAllocation(alloc, boxAddr);
+    postAllocationAction(alloc);
   }
 
   static bool lowerBoundsAreOnes(const Allocation &alloc) {
     for (const Fortran::parser::AllocateShapeSpec &shapeSpec :
          alloc.getShapeSpecs())
       if (std::get<0>(shapeSpec.t))
         return false;
     return true;
for (const Fortran::parser::AllocateShapeSpec &shapeSpec :
         extents.emplace_back(ub);
       }
     }
     fir::factory::genInlinedAllocation(builder, loc, box, lbounds, extents,
                                        lenParams, mangleAlloc(alloc),
                                        /*mustBeHeap=*/true);
   }
 
+  void postAllocationAction(const Allocation &alloc) {
+    if (alloc.getSymbol().test(Fortran::semantics::Symbol::Flag::AccDeclare))
+      Fortran::lower::attachDeclarePostAllocAction(converter, builder,
+                                                   alloc.getSymbol());
+  }
+
   void genSimpleAllocation(const Allocation &alloc,
                            const fir::MutableBoxValue &box) {
     if (!box.isDerived() && !errorManager.hasStatSpec() &&
         !alloc.type.IsPolymorphic() && !alloc.hasCoarraySpec() &&
         !useAllocateRuntime) {
       genInlinedAllocation(alloc, box);
       return;
     }
void Fortran::lower::genAllocateStmt(
   AllocateStmtHelper{converter, stmt, loc}.lower();
 }
 
 //===----------------------------------------------------------------------===//
 // Deallocate statement implementation
 //===----------------------------------------------------------------------===//
 
 // Generate deallocation of a pointer/allocatable.
-static void genDeallocate(fir::FirOpBuilder &builder, mlir::Location loc,
+static mlir::Value genDeallocate(fir::FirOpBuilder &builder, mlir::Location loc,
                                  const fir::MutableBoxValue &box,
                                  ErrorManager &errorManager,
                                  mlir::Value declaredTypeDesc = {}) {
   // Deallocate intrinsic types inline.
   if (!box.isDerived() && !box.isPolymorphic() &&
       !box.isUnlimitedPolymorphic() && !errorManager.hasStatSpec() &&
       !useAllocateRuntime) {
-    fir::factory::genInlinedDeallocate(builder, loc, box);
-    return;
+    return fir::factory::genInlinedDeallocate(builder, loc, box);
   }
   // Use runtime calls to deallocate descriptor cases. Sync MutableBoxValue
   // with its descriptor before and after calls if needed.
   errorManager.genStatCheck(builder, loc);
   mlir::Value stat =
       genRuntimeDeallocate(builder, loc, box, errorManager, declaredTypeDesc);
   fir::factory::syncMutableBoxFromIRBox(builder, loc, box);
   errorManager.assignStat(builder, loc, stat);
+  return stat;
 }
 
 void Fortran::lower::genDeallocateBox(
     Fortran::lower::AbstractConverter &converter,
     const fir::MutableBoxValue &box, mlir::Location loc,
     mlir::Value declaredTypeDesc) {
   const Fortran::lower::SomeExpr *statExpr = nullptr;
   const Fortran::lower::SomeExpr *errMsgExpr = nullptr;
   ErrorManager errorManager;
   errorManager.init(converter, loc, statExpr, errMsgExpr);
   fir::FirOpBuilder &builder = converter.getFirOpBuilder();
   genDeallocate(builder, loc, box, errorManager, declaredTypeDesc);
 }
 
+static void preDeallocationAction(Fortran::lower::AbstractConverter &converter,
+                                  fir::FirOpBuilder &builder,
+                                  mlir::Value beginOpValue,
+                                  const Fortran::semantics::Symbol &sym) {
+  if (sym.test(Fortran::semantics::Symbol::Flag::AccDeclare))
+    Fortran::lower::attachDeclarePreDeallocAction(converter, builder,
+                                                  beginOpValue, sym);
+}
+
+static void postDeallocationAction(Fortran::lower::AbstractConverter &converter,
+                                   fir::FirOpBuilder &builder,
+                                   const Fortran::semantics::Symbol &sym) {
+  if (sym.test(Fortran::semantics::Symbol::Flag::AccDeclare))
+    Fortran::lower::attachDeclarePostDeallocAction(converter, builder, sym);
+}
+
 void Fortran::lower::genDeallocateStmt(
     Fortran::lower::AbstractConverter &converter,
     const Fortran::parser::DeallocateStmt &stmt, mlir::Location loc) {
   const Fortran::lower::SomeExpr *statExpr = nullptr;
   const Fortran::lower::SomeExpr *errMsgExpr = nullptr;
   for (const Fortran::parser::StatOrErrmsg &statOrErr :
        std::get<std::list<Fortran::parser::StatOrErrmsg>>(stmt.t))
     std::visit(Fortran::common::visitors{
                    [&](const Fortran::parser::StatVariable &statVar) {
                      statExpr = Fortran::semantics::GetExpr(statVar);
                    },
                    [&](const Fortran::parser::MsgVariable &errMsgVar) {
                      errMsgExpr = Fortran::semantics::GetExpr(errMsgVar);
                    },
                },
                statOrErr.u);
   ErrorManager errorManager;
   errorManager.init(converter, loc, statExpr, errMsgExpr);
   fir::FirOpBuilder &builder = converter.getFirOpBuilder();
   mlir::OpBuilder::InsertPoint insertPt = builder.saveInsertionPoint();
   for (const Fortran::parser::AllocateObject &allocateObject :
        std::get<std::list<Fortran::parser::AllocateObject>>(stmt.t)) {
+    const Fortran::semantics::Symbol &symbol = unwrapSymbol(allocateObject);
     fir::MutableBoxValue box =
         genMutableBoxValue(converter, loc, allocateObject);
-
     mlir::Value declaredTypeDesc = {};
     if (box.isPolymorphic()) {
-      const Fortran::semantics::Symbol &symbol = unwrapSymbol(allocateObject);
       assert(symbol.GetType());
       if (const Fortran::semantics::DerivedTypeSpec *derivedTypeSpec =
               symbol.GetType()->AsDerived()) {
         declaredTypeDesc =
             Fortran::lower::getTypeDescAddr(converter, loc, *derivedTypeSpec);
       }
     }
+    mlir::Value beginOpValue =
         genDeallocate(builder, loc, box, errorManager, declaredTypeDesc);
+    preDeallocationAction(converter, builder, beginOpValue, symbol);
+    postDeallocationAction(converter, builder, symbol);
   }
   builder.restoreInsertionPoint(insertPt);
 }
 
 //===----------------------------------------------------------------------===//
 // MutableBoxValue creation implementation
 //===----------------------------------------------------------------------===//
 

flang/lib/Lower/OpenACC.cpp

std::visit(
           [&](const Fortran::parser::OpenACCRoutineConstruct
                   &routineConstruct) {
             TODO(converter.genLocation(routineConstruct.source),
                  "OpenACC Routine construct not lowered yet!");
           },
       },
       accDeclConstruct.u);
 }
+
+void Fortran::lower::attachDeclarePostAllocAction(
+    AbstractConverter &converter, fir::FirOpBuilder &builder,
+    const Fortran::semantics::Symbol &sym) {
+  std::stringstream fctName;
+  fctName << converter.mangleName(sym) << declarePostAllocSuffix.str();
+  mlir::Operation &op = builder.getInsertionBlock()->back();
+  op.setAttr(mlir::acc::getDeclareActionAttrName(),
+             mlir::acc::DeclareActionAttr::get(
+                 builder.getContext(),
+                 /*preAlloc=*/{},
+                 /*postAlloc=*/builder.getSymbolRefAttr(fctName.str()),
+                 /*preDealloc=*/{}, /*postDealloc=*/{}));
+}
+
+void Fortran::lower::attachDeclarePreDeallocAction(
+    AbstractConverter &converter, fir::FirOpBuilder &builder,
+    mlir::Value beginOpValue, const Fortran::semantics::Symbol &sym) {
+  std::stringstream fctName;
+  fctName << converter.mangleName(sym) << declarePreDeallocSuffix.str();
+  beginOpValue.getDefiningOp()->setAttr(
+      mlir::acc::getDeclareActionAttrName(),
+      mlir::acc::DeclareActionAttr::get(
+          builder.getContext(),
+          /*preAlloc=*/{}, /*postAlloc=*/{},
+          /*preDealloc=*/builder.getSymbolRefAttr(fctName.str()),
+          /*postDealloc=*/{}));
+}
+
+void Fortran::lower::attachDeclarePostDeallocAction(
+    AbstractConverter &converter, fir::FirOpBuilder &builder,
+    const Fortran::semantics::Symbol &sym) {
+  std::stringstream fctName;
+  fctName << converter.mangleName(sym) << declarePostAllocSuffix.str();
+  mlir::Operation &op = builder.getInsertionBlock()->back();
+  op.setAttr(mlir::acc::getDeclareActionAttrName(),
+             mlir::acc::DeclareActionAttr::get(
+                 builder.getContext(),
+                 /*preAlloc=*/{}, /*postAlloc=*/{}, /*preDealloc=*/{},
+                 /*postDealloc=*/builder.getSymbolRefAttr(fctName.str())));
+}

flang/lib/Optimizer/Builder/MutableBox.cpp

if (box.getEleTy().isa<fir::RecordType>()) {
     mlir::Value irBox = fir::factory::getMutableIRBox(builder, loc, box);
     fir::runtime::genDerivedTypeInitialize(builder, loc, irBox);
   }
 
   heap->setAttr(fir::MustBeHeapAttr::getAttrName(),
                 fir::MustBeHeapAttr::get(builder.getContext(), mustBeHeap));
 }
 
-void fir::factory::genInlinedDeallocate(fir::FirOpBuilder &builder,
-                                        mlir::Location loc,
-                                        const fir::MutableBoxValue &box) {
+mlir::Value
+fir::factory::genInlinedDeallocate(fir::FirOpBuilder &builder,
   auto addr = MutablePropertyReader(builder, loc, box).readBaseAddress();
   genFinalizeAndFree(builder, loc, addr);
   MutablePropertyWriter{builder, loc, box}.setUnallocatedStatus();
+  return addr;
 }
 
 fir::factory::MutableBoxReallocation fir::factory::genReallocIfNeeded(
     fir::FirOpBuilder &builder, mlir::Location loc,
     const fir::MutableBoxValue &box, mlir::ValueRange shape,
     mlir::ValueRange lengthParams,
     fir::factory::ReallocStorageHandlerFunc storageHandler) {
   // Implement 10.2.1.3 point 3 logic when lhs is an array.

flang/test/Lower/OpenACC/acc-declare.f90

subroutine acc_declare_deviceptr2()
     !$acc declare deviceptr(dataparam)
   end subroutine
 
 ! CHECK-LABEL: func.func @_QMacc_declarePacc_declare_deviceptr2()
 ! CHECK: %[[ALLOCA:.*]] = fir.alloca !fir.array<100xf32> {acc.declare = #acc.declare<dataClause =  acc_deviceptr>, bindc_name = "dataparam", uniq_name = "_QMacc_declareFacc_declare_deviceptr2Edataparam"}
 ! CHECK: %[[DEVICEPTR:.*]] = acc.deviceptr varPtr(%[[ALLOCA]] : !fir.ref<!fir.array<100xf32>>) bounds(%{{.*}}) -> !fir.ref<!fir.array<100xf32>> {name = "dataparam"}
 ! CHECK: acc.declare_enter dataOperands(%[[DEVICEPTR]] : !fir.ref<!fir.array<100xf32>>)
 
+  subroutine acc_declare_allocate()
+    integer, allocatable :: a(:)
+    !$acc declare create(a)
+
+    allocate(a(100))
+
+! CHECK: %{{.*}} = fir.allocmem !fir.array<?xi32>, %{{.*}} {fir.must_be_heap = true, uniq_name = "_QMacc_declareFacc_declare_allocateEa.alloc"}
+! CHECK: fir.store %{{.*}} to %{{.*}} {acc.declare_action = #acc.declare_action<postAlloc = @_QMacc_declareFacc_declare_allocateEa_acc_declare_update_desc_post_alloc>} : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
+
+    deallocate(a)
+
+! CHECK: %{{.*}} = fir.box_addr %{{.*}} {acc.declare_action = #acc.declare_action<preDealloc = @_QMacc_declareFacc_declare_allocateEa_acc_declare_update_desc_pre_dealloc>} : (!fir.box<!fir.heap<!fir.array<?xi32>>>) -> !fir.heap<!fir.array<?xi32>>
+
+! CHECK: fir.freemem %{{.*}} : !fir.heap<!fir.array<?xi32>>
+! CHECK: fir.store %{{.*}} to %{{.*}} {acc.declare_action = #acc.declare_action<postDealloc = @_QMacc_declareFacc_declare_allocateEa_acc_declare_update_desc_post_alloc>} : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
+
+  end subroutine
+
 end module
 
 module acc_declare_allocatable_test
  integer, allocatable :: data1(:)
  !$acc declare create(data1)
 end module
 
 ! CHECK-LABEL: acc.global_ctor @_QMacc_declare_allocatable_testEdata1_acc_ctor {

 
 ! CHECK-LABEL: acc.global_dtor @_QMacc_declare_allocatable_testEdata1_acc_dtor {
 ! CHECK:         %[[GLOBAL_ADDR:.*]] = fir.address_of(@_QMacc_declare_allocatable_testEdata1) {acc.declare = #acc.declare<dataClause =  acc_create>} : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
 ! CHECK:         %[[DEVICEPTR:.*]] = acc.getdeviceptr varPtr(%[[GLOBAL_ADDR]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>)   -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>> {dataClause = #acc<data_clause acc_create>, name = "data1", structured = false}
 ! CHECK:         acc.declare_exit dataOperands(%[[DEVICEPTR]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>)
 ! CHECK:         acc.delete accPtr(%[[DEVICEPTR]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>) {dataClause = #acc<data_clause acc_create>, name = "data1", structured = false}
 ! CHECK:         acc.terminator
 ! CHECK:       }
+
+! Test that the pre/post alloc/dealloc attributes are set when the
+! allocate/deallocate statement are in a different module.
+module acc_declare_allocatable_test2
+contains
+  subroutine init()
+    use acc_declare_allocatable_test
+    allocate(data1(100))
+! CHECK: fir.store %{{.*}} to %{{.*}} {acc.declare_action = #acc.declare_action<postAlloc = @_QMacc_declare_allocatable_testEdata1_acc_declare_update_desc_post_alloc>} : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
+  end subroutine
+
+  subroutine finalize()
+    use acc_declare_allocatable_test
+    deallocate(data1)
+! CHECK: %{{.*}} = fir.box_addr %{{.*}} {acc.declare_action = #acc.declare_action<preDealloc = @_QMacc_declare_allocatable_testEdata1_acc_declare_update_desc_pre_dealloc>} : (!fir.box<!fir.heap<!fir.array<?xi32>>>) -> !fir.heap<!fir.array<?xi32>>
+! CHECK: fir.store %{{.*}} to %{{.*}} {acc.declare_action = #acc.declare_action<postDealloc = @_QMacc_declare_allocatable_testEdata1_acc_declare_update_desc_post_alloc>} : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
+  end subroutine
+end module