[flang][openacc] Lower acc routine with function name

The routine directive can appear in the specification part of
a subroutine, function or module and therefore appear before the
function or subroutine is lowered. We keep track of the created
routine info attribute and attach them to the function at the end
of the lowering if the directive appeared before the function was
lowered.

Reviewed By: razvanlupusoru

Differential Revision: https://reviews.llvm.org/D158204

Author: clementval

flang/include/flang/Lower/OpenACC.h

@@ -24,6 +24,7 @@ class StringRef;
 namespace mlir {
 class Location;
 class Type;
+class ModuleOp;
 class OpBuilder;
 class Value;
 } // namespace mlir
@@ -52,6 +53,9 @@ namespace pft {
 struct Evaluation;
 } // namespace pft
 
+using AccRoutineInfoMappingList =
+    llvm::SmallVector<std::pair<std::string, mlir::SymbolRefAttr>>;
+
 static constexpr llvm::StringRef declarePostAllocSuffix =
     "_acc_declare_update_desc_post_alloc";
 static constexpr llvm::StringRef declarePreDeallocSuffix =
@@ -62,10 +66,14 @@ static constexpr llvm::StringRef declarePostDeallocSuffix =
 void genOpenACCConstruct(AbstractConverter &,
                          Fortran::semantics::SemanticsContext &,
                          pft::Evaluation &, const parser::OpenACCConstruct &);
-void genOpenACCDeclarativeConstruct(
-    AbstractConverter &, Fortran::semantics::SemanticsContext &,
-    StatementContext &, pft::Evaluation &,
-    const parser::OpenACCDeclarativeConstruct &);
+void genOpenACCDeclarativeConstruct(AbstractConverter &,
+                                    Fortran::semantics::SemanticsContext &,
+                                    StatementContext &, pft::Evaluation &,
+                                    const parser::OpenACCDeclarativeConstruct &,
+                                    AccRoutineInfoMappingList &);
+
+void finalizeOpenACCRoutineAttachment(mlir::ModuleOp &,
+                                      AccRoutineInfoMappingList &);
 
 /// Get a acc.private.recipe op for the given type or create it if it does not
 /// exist yet.

flang/lib/Lower/Bridge.cpp

@@ -350,6 +350,8 @@ class FirConverter : public Fortran::lower::AbstractConverter {
         fir::runtime::genEnvironmentDefaults(*builder, toLocation(),
                                              bridge.getEnvironmentDefaults());
       });
+
+    finalizeOpenACCLowering();
   }
 
   /// Declare a function.
@@ -2248,7 +2250,8 @@ class FirConverter : public Fortran::lower::AbstractConverter {
   void genFIR(const Fortran::parser::OpenACCDeclarativeConstruct &accDecl) {
     mlir::OpBuilder::InsertPoint insertPt = builder->saveInsertionPoint();
     genOpenACCDeclarativeConstruct(*this, bridge.getSemanticsContext(),
-                                   bridge.fctCtx(), getEval(), accDecl);
+                                   bridge.fctCtx(), getEval(), accDecl,
+                                   accRoutineInfos);
     for (Fortran::lower::pft::Evaluation &e : getEval().getNestedEvaluations())
       genFIR(e);
     builder->restoreInsertionPoint(insertPt);
@@ -4683,6 +4686,13 @@ class FirConverter : public Fortran::lower::AbstractConverter {
         expr.u);
   }
 
+  /// Performing OpenACC lowering action that were deferred to the end of
+  /// lowering.
+  void finalizeOpenACCLowering() {
+    Fortran::lower::finalizeOpenACCRoutineAttachment(getModuleOp(),
+                                                     accRoutineInfos);
+  }
+
   //===--------------------------------------------------------------------===//
 
   Fortran::lower::LoweringBridge &bridge;
@@ -4726,6 +4736,9 @@ class FirConverter : public Fortran::lower::AbstractConverter {
 
   /// A counter for uniquing names in `literalNamesMap`.
   std::uint64_t uniqueLitId = 0;
+
+  /// Deferred OpenACC routine attachment.
+  Fortran::lower::AccRoutineInfoMappingList accRoutineInfos;
 };
 
 } // namespace

flang/lib/Lower/OpenACC.cpp

@@ -2925,29 +2925,53 @@ GetConstExpr(Fortran::semantics::SemanticsContext &semanticsContext,
   return std::nullopt;
 }
 
+static void attachRoutineInfo(mlir::func::FuncOp func,
+                              mlir::SymbolRefAttr routineAttr) {
+  llvm::SmallVector<mlir::SymbolRefAttr> routines;
+  if (func.getOperation()->hasAttr(mlir::acc::getRoutineInfoAttrName())) {
+    auto routineInfo =
+        func.getOperation()->getAttrOfType<mlir::acc::RoutineInfoAttr>(
+            mlir::acc::getRoutineInfoAttrName());
+    routines.append(routineInfo.getAccRoutines().begin(),
+                    routineInfo.getAccRoutines().end());
+  }
+  routines.push_back(routineAttr);
+  func.getOperation()->setAttr(
+      mlir::acc::getRoutineInfoAttrName(),
+      mlir::acc::RoutineInfoAttr::get(func.getContext(), routines));
+}
+
 static void
 genACC(Fortran::lower::AbstractConverter &converter,
        Fortran::semantics::SemanticsContext &semanticsContext,
        Fortran::lower::pft::Evaluation &eval,
-       const Fortran::parser::OpenACCRoutineConstruct &routineConstruct) {
+       const Fortran::parser::OpenACCRoutineConstruct &routineConstruct,
+       Fortran::lower::AccRoutineInfoMappingList &accRoutineInfos) {
   fir::FirOpBuilder &builder = converter.getFirOpBuilder();
   mlir::Location loc = converter.genLocation(routineConstruct.source);
   std::optional<Fortran::parser::Name> name =
       std::get<std::optional<Fortran::parser::Name>>(routineConstruct.t);
   const auto &clauses =
       std::get<Fortran::parser::AccClauseList>(routineConstruct.t);
-  if (name)
-    TODO(loc, "acc routine with name");
 
-  mlir::func::FuncOp func = builder.getFunction();
   mlir::ModuleOp mod = builder.getModule();
+  mlir::func::FuncOp funcOp;
+  std::string funcName;
+  if (name) {
+    funcName = converter.mangleName(*name->symbol);
+    funcOp = builder.getNamedFunction(funcName);
+  } else {
+    funcOp = builder.getFunction();
+    funcName = funcOp.getName();
+  }
+
   mlir::OpBuilder modBuilder(mod.getBodyRegion());
   std::stringstream routineOpName;
   routineOpName << accRoutinePrefix.str() << routineCounter++;
   auto routineOp = modBuilder.create<mlir::acc::RoutineOp>(
-      loc, routineOpName.str(), func.getName(), mlir::StringAttr{},
+      loc, routineOpName.str(), funcName, mlir::StringAttr{}, mlir::UnitAttr{},
       mlir::UnitAttr{}, mlir::UnitAttr{}, mlir::UnitAttr{}, mlir::UnitAttr{},
-      mlir::UnitAttr{}, mlir::UnitAttr{}, mlir::IntegerAttr{});
+      mlir::UnitAttr{}, mlir::IntegerAttr{});
 
   for (const Fortran::parser::AccClause &clause : clauses.v) {
     if (std::get_if<Fortran::parser::AccClause::Seq>(&clause.u)) {
@@ -2994,18 +3018,27 @@ genACC(Fortran::lower::AbstractConverter &converter,
     }
   }
 
-  llvm::SmallVector<mlir::SymbolRefAttr> routines;
-  if (func.getOperation()->hasAttr(mlir::acc::getRoutineInfoAttrName())) {
-    auto routineInfo =
-        func.getOperation()->getAttrOfType<mlir::acc::RoutineInfoAttr>(
-            mlir::acc::getRoutineInfoAttrName());
-    routines.append(routineInfo.getAccRoutines().begin(),
-                    routineInfo.getAccRoutines().end());
+  if (funcOp)
+    attachRoutineInfo(funcOp, builder.getSymbolRefAttr(routineOpName.str()));
+  else
+    // FuncOp is not lowered yet. Keep the information so the routine info
+    // can be attached later to the funcOp.
+    accRoutineInfos.push_back(std::make_pair(
+        funcName, builder.getSymbolRefAttr(routineOpName.str())));
+}
+
+void Fortran::lower::finalizeOpenACCRoutineAttachment(
+    mlir::ModuleOp &mod,
+    Fortran::lower::AccRoutineInfoMappingList &accRoutineInfos) {
+  for (auto &mapping : accRoutineInfos) {
+    mlir::func::FuncOp funcOp =
+        mod.lookupSymbol<mlir::func::FuncOp>(mapping.first);
+    if (!funcOp)
+      llvm::report_fatal_error(
+          "could not find function to attach OpenACC routine information.");
+    attachRoutineInfo(funcOp, mapping.second);
   }
-  routines.push_back(builder.getSymbolRefAttr(routineOpName.str()));
-  func.getOperation()->setAttr(
-      mlir::acc::getRoutineInfoAttrName(),
-      mlir::acc::RoutineInfoAttr::get(builder.getContext(), routines));
+  accRoutineInfos.clear();
 }
 
 void Fortran::lower::genOpenACCConstruct(
@@ -3050,7 +3083,8 @@ void Fortran::lower::genOpenACCDeclarativeConstruct(
     Fortran::semantics::SemanticsContext &semanticsContext,
     Fortran::lower::StatementContext &fctCtx,
     Fortran::lower::pft::Evaluation &eval,
-    const Fortran::parser::OpenACCDeclarativeConstruct &accDeclConstruct) {
+    const Fortran::parser::OpenACCDeclarativeConstruct &accDeclConstruct,
+    Fortran::lower::AccRoutineInfoMappingList &accRoutineInfos) {
 
   std::visit(
       common::visitors{
@@ -3061,7 +3095,8 @@ void Fortran::lower::genOpenACCDeclarativeConstruct(
           },
           [&](const Fortran::parser::OpenACCRoutineConstruct
                   &routineConstruct) {
-            genACC(converter, semanticsContext, eval, routineConstruct);
+            genACC(converter, semanticsContext, eval, routineConstruct,
+                   accRoutineInfos);
           },
       },
       accDeclConstruct.u);

flang/test/Lower/OpenACC/acc-routine-named.f90

@@ -0,0 +1,25 @@
+! This test checks lowering of OpenACC routine directive.
+
+! RUN: bbc -fopenacc -emit-fir %s -o - | FileCheck %s
+
+module acc_routines
+
+! CHECK: acc.routine @acc_routine_1 func(@_QMacc_routinesPacc2)
+! CHECK: acc.routine @acc_routine_0 func(@_QMacc_routinesPacc1) seq
+
+!$acc routine(acc1) seq
+
+contains
+
+  subroutine acc1()
+  end subroutine
+
+! CHECK-LABEL: func.func @_QMacc_routinesPacc1() attributes {acc.routine_info = #acc.routine_info<[@acc_routine_0]>}
+
+  subroutine acc2()
+    !$acc routine(acc2)
+  end subroutine
+
+! CHECK-LABEL: func.func @_QMacc_routinesPacc2() attributes {acc.routine_info = #acc.routine_info<[@acc_routine_1]>}
+
+end module