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

[clementval]

Summary

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.

[clementval]

Clear the list at the end

[razvanlupusoru]

Looks great!

[clementval]

Rebase

flang/include/flang/Lower/OpenACC.h

 template <typename T, unsigned N>
 class SmallVector;
 class StringRef;
 }
 
 namespace mlir {
 class Location;
 class Type;
+class ModuleOp;
 class OpBuilder;
 class Value;
 } // namespace mlir
 
 namespace fir {
 class FirOpBuilder;
 }
 

 
 class AbstractConverter;
 class StatementContext;
 
 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 =
     "_acc_declare_update_desc_pre_dealloc";
 static constexpr llvm::StringRef declarePostDeallocSuffix =
     "_acc_declare_update_desc_post_dealloc";
 
 void genOpenACCConstruct(AbstractConverter &,
                          Fortran::semantics::SemanticsContext &,
                          pft::Evaluation &, const parser::OpenACCConstruct &);
-void genOpenACCDeclarativeConstruct(
-    AbstractConverter &, Fortran::semantics::SemanticsContext &,
+void genOpenACCDeclarativeConstruct(AbstractConverter &,
+                                    Fortran::semantics::SemanticsContext &,
                                     StatementContext &, pft::Evaluation &,
-    const parser::OpenACCDeclarativeConstruct &);
+                                    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.
 mlir::acc::PrivateRecipeOp createOrGetPrivateRecipe(mlir::OpBuilder &,
                                                     llvm::StringRef,
                                                     mlir::Location, mlir::Type);
 
 /// Get a acc.reduction.recipe op for the given type or create it if it does not

flang/lib/Lower/Bridge.cpp

if (hasMainProgram)
         // accepting the list of environment defaults. That way, we would not
         // need to add an extern pointer to the runtime and said pointer would
         // not need to be generated even if no defaults are specified.
         // However, generating main or changing when the runtime reads
         // environment variables is required to do so.
         fir::runtime::genEnvironmentDefaults(*builder, toLocation(),
                                              bridge.getEnvironmentDefaults());
       });
+
+    finalizeOpenACCLowering();
   }
 
   /// Declare a function.
   void declareFunction(Fortran::lower::pft::FunctionLikeUnit &funit) {
     setCurrentPosition(funit.getStartingSourceLoc());
     for (int entryIndex = 0, last = funit.entryPointList.size();
          entryIndex < last; ++entryIndex) {
       funit.setActiveEntry(entryIndex);
void genFIR(const Fortran::parser::OpenACCConstruct &acc) {
     for (Fortran::lower::pft::Evaluation &e : getEval().getNestedEvaluations())
       genFIR(e);
     builder->restoreInsertionPoint(insertPt);
   }
 
   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);
   }
 
   void genFIR(const Fortran::parser::OpenMPConstruct &omp) {
     mlir::OpBuilder::InsertPoint insertPt = builder->saveInsertionPoint();
     localSymbols.pushScope();
return std::visit(
             }
           } else {
             fir::emitFatalError(loc, "unexpected expression");
           }
         },
         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;
   Fortran::evaluate::FoldingContext foldingContext;
   fir::FirOpBuilder *builder = nullptr;
   Fortran::lower::pft::Evaluation *evalPtr = nullptr;
   Fortran::lower::SymMap localSymbols;
   Fortran::parser::CharBlock currentPosition;

   llvm::DenseMap<const Fortran::lower::SomeExpr *, std::string> literalNamesMap;
 
   /// Storage for Constant expressions used as keys for literalNamesMap.
   llvm::SmallVector<std::unique_ptr<Fortran::lower::SomeExpr>>
       literalExprsStorage;
 
   /// A counter for uniquing names in `literalNamesMap`.
   std::uint64_t uniqueLitId = 0;
+
+  /// Deferred OpenACC routine attachment.
+  Fortran::lower::AccRoutineInfoMappingList accRoutineInfos;
 };
 
 } // namespace
 
 Fortran::evaluate::FoldingContext
 Fortran::lower::LoweringBridge::createFoldingContext() const {
   return {getDefaultKinds(), getIntrinsicTable(), getTargetCharacteristics()};
 }

flang/lib/Lower/OpenACC.cpp

if (funcOp)
                            directiveLocation, accClauseList);
     else if (moduleOp)
       genDeclareInModule(converter, moduleOp, accClauseList);
     return;
   }
   llvm_unreachable("unsupported declarative directive");
 }
 
+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::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)) {
       routineOp.setSeqAttr(builder.getUnitAttr());
     } else if (const auto *gangClause =
                    std::get_if<Fortran::parser::AccClause::Gang>(&clause.u)) {
       routineOp.setGangAttr(builder.getUnitAttr());
       if (gangClause->v) {
if (std::get_if<Fortran::parser::AccClause::Seq>(&clause.u)) {
       routineOp.setVectorAttr(builder.getUnitAttr());
     } else if (std::get_if<Fortran::parser::AccClause::Worker>(&clause.u)) {
       routineOp.setWorkerAttr(builder.getUnitAttr());
     } else if (std::get_if<Fortran::parser::AccClause::Nohost>(&clause.u)) {
       routineOp.setNohostAttr(builder.getUnitAttr());
     }
   }
 
-  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(builder.getSymbolRefAttr(routineOpName.str()));
-  func.getOperation()->setAttr(
-      mlir::acc::getRoutineInfoAttrName(),
-      mlir::acc::RoutineInfoAttr::get(builder.getContext(), routines));
+  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);
+  }
+  accRoutineInfos.clear();
 }
 
 void Fortran::lower::genOpenACCConstruct(
     Fortran::lower::AbstractConverter &converter,
     Fortran::semantics::SemanticsContext &semanticsContext,
     Fortran::lower::pft::Evaluation &eval,
     const Fortran::parser::OpenACCConstruct &accConstruct) {
 
std::visit(
       accConstruct.u);
 }
 
 void Fortran::lower::genOpenACCDeclarativeConstruct(
     Fortran::lower::AbstractConverter &converter,
     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{
           [&](const Fortran::parser::OpenACCStandaloneDeclarativeConstruct
                   &standaloneDeclarativeConstruct) {
             genACC(converter, semanticsContext, fctCtx, eval,
                    standaloneDeclarativeConstruct);
           },
           [&](const Fortran::parser::OpenACCRoutineConstruct
                   &routineConstruct) {
-            genACC(converter, eval, routineConstruct);
+            genACC(converter, eval, routineConstruct, accRoutineInfos);
           },
       },
       accDeclConstruct.u);
 }
 
 void Fortran::lower::attachDeclarePostAllocAction(
     AbstractConverter &converter, fir::FirOpBuilder &builder,
     const Fortran::semantics::Symbol &sym) {

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

+! 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