[Flang][OpenMP][Lower] Program level implicit SAVE variable handling for declare target

[agozillon]

Summary

This is an attempt at mimicing the method in which
threadprivate handles the following type of variables:

program main
  integer :: i
  !$omp declare target to(i)
end

Which essentially generates a GlobalOp for the variable (which
would normally only be an alloca) when it's instantiated. The
main difference is there is no operation generated within the
function, instead the declare target attribute is appended
later within handleDeclareTarget.

[agozillon]

• missing newline as always, sorry for the spam.

[agozillon]

Will need to rebase this when https://reviews.llvm.org/D150329 is landed to get it to pass CI, I tried to do some arcanist trickery to base it on that patch, but alas it didn't work.

This is one method of handling the implicit SAVE attributes that is similar to what threadprivate currently does (it also does some other fun edge cases too, but I don't believe they're needed for declare target, at least at the moment), I believe it should be viable to generate working LLVM-IR from this, but it is possible it may need further tweaking after the fact as this isn't a case I've currently had the chance to look into (and I've just begun to look into the mapping now that we have our initial test off the ground).

[agozillon]

Rebase on upstream now that dependent patch has been commited, this will hopefully satisfiy CI

[agozillon]

Rebase

[agozillon]

The rebase seems to have fixed the buildbot failure. If someone could please give this a look over to see if following threadprivate's approach is a reasonable direction for fixing this issue with declare target, thank you very much

[agozillon]

A small ping on this for some reviewer attention if at all possible please.

[kiranchandramohan]

flang/lib/Lower/OpenMP.cpp

2743–2770

Can this code be shared with the threadprivate usage?

[agozillon]

• Share declare target behaviour with current Firstprivate function

[agozillon]

flang/lib/Lower/OpenMP.cpp

2743–2770

Done, we only use a little segment of the behavior of the function (perhaps more later if there's other issues that the function fixes that declare target needs) but it's shareable!

[kiranchandramohan]

I was hoping you will factor the common portion into a separate function and call that function from both threadprivate and declare target. Could you give that and try and if it is not possible you can revert to the previous version that duplicated code and submit it since it is more readable.

LG.

[agozillon]

In D152037#4493826, @kiranchandramohan wrote:

I was hoping you will factor the common portion into a separate function and call that function from both threadprivate and declare target. Could you give that and try and if it is not possible you can revert to the previous version that duplicated code and submit it since it is more readable.

LG.

Sure, that shouldn't be a problem I imagine, the separated function would still have to be called from Bridge.cpp for declare target and then internally inside of genThreadprivateOp for threadprivate.

[agozillon]

• Revert "Share declare target behaviour with current Firstprivate function"
• Tidy code try two

[agozillon]

Made an attempt at tidying it up again as requested.

I'll upstream this tomorrow as it's a little late in the day here for me to risk angering the buildbots and I'd like to give the patch CI a chance to pass! So if you do get some time please do skim over the changes to see if they're what you were hoping for.

flang/include/flang/Lower/OpenMP.h

void genOpenMPTerminator(fir::FirOpBuilder &, mlir::Operation *,
                          mlir::Location);
 
 void genOpenMPConstruct(AbstractConverter &, pft::Evaluation &,
                         const parser::OpenMPConstruct &);
 void genOpenMPDeclarativeConstruct(AbstractConverter &, pft::Evaluation &,
                                    const parser::OpenMPDeclarativeConstruct &);
 int64_t getCollapseValue(const Fortran::parser::OmpClauseList &clauseList);
 void genThreadprivateOp(AbstractConverter &, const pft::Variable &);
+void genDeclareTargetIntGlobal(AbstractConverter &, const pft::Variable &);
 void genOpenMPReduction(AbstractConverter &,
                         const Fortran::parser::OmpClauseList &clauseList);
 
 mlir::Operation *findReductionChain(mlir::Value, mlir::Value * = nullptr);
 fir::ConvertOp getConvertFromReductionOp(mlir::Operation *, mlir::Value);
 void updateReduction(mlir::Operation *, fir::FirOpBuilder &, mlir::Value,
                      mlir::Value, fir::ConvertOp * = nullptr);
 void removeStoreOp(mlir::Operation *, mlir::Value);
 } // namespace lower
 } // namespace Fortran
 
 #endif // FORTRAN_LOWER_OPENMP_H

flang/lib/Lower/Bridge.cpp

void mapDummiesAndResults(Fortran::lower::pft::FunctionLikeUnit &funit,
     }
   }
 
   /// Instantiate variable \p var and add it to the symbol map.
   /// See ConvertVariable.cpp.
   void instantiateVar(const Fortran::lower::pft::Variable &var,
                       Fortran::lower::AggregateStoreMap &storeMap) {
     Fortran::lower::instantiateVariable(*this, var, localSymbols, storeMap);
-    if (var.hasSymbol() &&
-        var.getSymbol().test(
+    if (var.hasSymbol()) {
+      if (var.getSymbol().test(
               Fortran::semantics::Symbol::Flag::OmpThreadprivate))
         Fortran::lower::genThreadprivateOp(*this, var);
+
+      if (var.getSymbol().test(
+              Fortran::semantics::Symbol::Flag::OmpDeclareTarget))
+        Fortran::lower::genDeclareTargetIntGlobal(*this, var);
+    }
   }
 
   /// Start translation of a function.
   void startNewFunction(Fortran::lower::pft::FunctionLikeUnit &funit) {
     assert(!builder && "expected nullptr");
     bridge.fctCtx().pushScope();
     const Fortran::semantics::Scope &scope = funit.getScope();
     LLVM_DEBUG(llvm::dbgs() << "\n[bridge - startNewFunction]";

flang/lib/Lower/OpenMP.cpp

std::visit(
           [&](const Fortran::parser::OpenMPCriticalConstruct
                   &criticalConstruct) {
             genOMP(converter, eval, criticalConstruct);
           },
       },
       ompConstruct.u);
 }
 
+fir::GlobalOp globalInitialization(Fortran::lower::AbstractConverter &converter,
+                                   fir::FirOpBuilder &firOpBuilder,
+                                   const Fortran::semantics::Symbol &sym,
+                                   const Fortran::lower::pft::Variable &var,
+                                   mlir::Location currentLocation) {
+  mlir::Type ty = converter.genType(sym);
+  std::string globalName = converter.mangleName(sym);
+  mlir::StringAttr linkage = firOpBuilder.createInternalLinkage();
+  fir::GlobalOp global =
+      firOpBuilder.createGlobal(currentLocation, ty, globalName, linkage);
+
+  // Create default initialization for non-character scalar.
+  if (Fortran::semantics::IsAllocatableOrPointer(sym)) {
+    mlir::Type baseAddrType = ty.dyn_cast<fir::BoxType>().getEleTy();
+    Fortran::lower::createGlobalInitialization(
+        firOpBuilder, global, [&](fir::FirOpBuilder &b) {
+          mlir::Value nullAddr =
+              b.createNullConstant(currentLocation, baseAddrType);
+          mlir::Value box =
+              b.create<fir::EmboxOp>(currentLocation, ty, nullAddr);
+          b.create<fir::HasValueOp>(currentLocation, box);
+        });
+  } else {
+    Fortran::lower::createGlobalInitialization(
+        firOpBuilder, global, [&](fir::FirOpBuilder &b) {
+          mlir::Value undef = b.create<fir::UndefOp>(currentLocation, ty);
+          b.create<fir::HasValueOp>(currentLocation, undef);
+        });
+  }
+
+  return global;
+}
+
 void Fortran::lower::genThreadprivateOp(
     Fortran::lower::AbstractConverter &converter,
     const Fortran::lower::pft::Variable &var) {
   fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
   mlir::Location currentLocation = converter.getCurrentLocation();
 
   const Fortran::semantics::Symbol &sym = var.getSymbol();
   mlir::Value symThreadprivateValue;
if (const Fortran::semantics::Symbol *common =
     // Generate the threadprivate value for the common block member.
     symThreadprivateValue =
         genCommonBlockMember(converter, sym, commonThreadprivateValue);
   } else if (!var.isGlobal()) {
     // Non-global variable which can be in threadprivate directive must be one
     // variable in main program, and it has implicit SAVE attribute. Take it as
     // with SAVE attribute, so to create GlobalOp for it to simplify the
     // translation to LLVM IR.
-    mlir::Type ty = converter.genType(sym);
-    std::string globalName = converter.mangleName(sym);
-    mlir::StringAttr linkage = firOpBuilder.createInternalLinkage();
-    fir::GlobalOp global =
-        firOpBuilder.createGlobal(currentLocation, ty, globalName, linkage);
+    fir::GlobalOp global = globalInitialization(converter, firOpBuilder, sym,
+                                                var, currentLocation);
 
-    // Create default initialization for non-character scalar.
-    if (Fortran::semantics::IsAllocatableOrPointer(sym)) {
-      mlir::Type baseAddrType = ty.dyn_cast<fir::BoxType>().getEleTy();
-      Fortran::lower::createGlobalInitialization(
-          firOpBuilder, global, [&](fir::FirOpBuilder &b) {
-            mlir::Value nullAddr =
-                b.createNullConstant(currentLocation, baseAddrType);
-            mlir::Value box =
-                b.create<fir::EmboxOp>(currentLocation, ty, nullAddr);
-            b.create<fir::HasValueOp>(currentLocation, box);
-          });
-    } else {
-      Fortran::lower::createGlobalInitialization(
-          firOpBuilder, global, [&](fir::FirOpBuilder &b) {
-            mlir::Value undef = b.create<fir::UndefOp>(currentLocation, ty);
-            b.create<fir::HasValueOp>(currentLocation, undef);
-          });
-    }
     mlir::Value symValue = firOpBuilder.create<fir::AddrOfOp>(
         currentLocation, global.resultType(), global.getSymbol());
     symThreadprivateValue = firOpBuilder.create<mlir::omp::ThreadprivateOp>(
         currentLocation, symValue.getType(), symValue);
   } else {
     mlir::Value symValue = converter.getSymbolAddress(sym);
     mlir::Operation *op = symValue.getDefiningOp();
     // The symbol may be use-associated multiple times, and nothing needs to be
     // done after the original symbol is mapped to the threadprivatized value
     // for the first time. Use the threadprivatized value directly.
     if (mlir::isa<mlir::omp::ThreadprivateOp>(op))
       return;
     symThreadprivateValue = firOpBuilder.create<mlir::omp::ThreadprivateOp>(
         currentLocation, symValue.getType(), symValue);
   }
 
   fir::ExtendedValue sexv = converter.getSymbolExtendedValue(sym);
   fir::ExtendedValue symThreadprivateExv =
       getExtendedValue(sexv, symThreadprivateValue);
   converter.bindSymbol(sym, symThreadprivateExv);
 }
 
+// This function replicates threadprivate's behaviour of generating
+// an internal fir.GlobalOp for non-global variables in the main program
+// that have the implicit SAVE attribute, to simplifiy LLVM-IR and MLIR
+// generation.
+void Fortran::lower::genDeclareTargetIntGlobal(
+    Fortran::lower::AbstractConverter &converter,
+    const Fortran::lower::pft::Variable &var) {
+  if (!var.isGlobal()) {
+    // A non-global variable which can be in a declare target directive must
+    // be a variable in the main program, and it has the implicit SAVE
+    // attribute. We create a GlobalOp for it to simplify the translation to
+    // LLVM IR.
+    globalInitialization(converter, converter.getFirOpBuilder(),
+                         var.getSymbol(), var, converter.getCurrentLocation());
+  }
+}
+
 void handleDeclareTarget(Fortran::lower::AbstractConverter &converter,
                          Fortran::lower::pft::Evaluation &eval,
                          const Fortran::parser::OpenMPDeclareTargetConstruct
                              &declareTargetConstruct) {
   llvm::SmallVector<std::pair<mlir::omp::DeclareTargetCaptureClause,
                               Fortran::semantics::Symbol>,
                     0>
       symbolAndClause;
   mlir::ModuleOp mod = converter.getFirOpBuilder().getModule();
 
   auto findFuncAndVarSyms = [&](const Fortran::parser::OmpObjectList &objList,
                                 mlir::omp::DeclareTargetCaptureClause clause) {
     for (const Fortran::parser::OmpObject &ompObject : objList.v) {
       Fortran::common::visit(
           Fortran::common::visitors{
               [&](const Fortran::parser::Designator &designator) {
                 if (const Fortran::parser::Name *name =
                         getDesignatorNameIfDataRef(designator)) {
                   symbolAndClause.push_back(
                       std::make_pair(clause, *name->symbol));
                 }
               },
               [&](const Fortran::parser::Name &name) {
                 symbolAndClause.push_back(std::make_pair(clause, *name.symbol));
               }},
           ompObject.u);
     }
   };
 

flang/test/Lower/OpenMP/omp-declare-target-program-var.f90

+!RUN: %flang_fc1 -emit-fir -fopenmp %s -o - | FileCheck %s 
+!RUN: %flang_fc1 -emit-fir -fopenmp -fopenmp-is-device %s -o - | FileCheck %s
+
+PROGRAM main
+    ! CHECK-DAG: %0 = fir.alloca f32 {bindc_name = "i", uniq_name = "_QFEi"}
+    REAL :: I
+    ! CHECK-DAG: fir.global internal @_QFEi {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>} : f32 {
+    ! CHECK-DAG: %0 = fir.undefined f32
+    ! CHECK-DAG: fir.has_value %0 : f32
+    ! CHECK-DAG: }
+    !$omp declare target(I)
+END