[FLANG] Lower groupprivate to omp mlir

Author: skc7

flang/include/flang/Lower/OpenMP.h

@@ -82,6 +82,7 @@ void genOpenMPSymbolProperties(AbstractConverter &converter,
                                const pft::Variable &var);
 
 void genThreadprivateOp(AbstractConverter &, const pft::Variable &);
+void genGroupprivateOp(AbstractConverter &, const pft::Variable &);
 void genDeclareTargetIntGlobal(AbstractConverter &, const pft::Variable &);
 bool isOpenMPTargetConstruct(const parser::OpenMPConstruct &);
 bool isOpenMPDeviceDeclareTarget(Fortran::lower::AbstractConverter &,

flang/lib/Lower/OpenMP/OpenMP.cpp

@@ -692,6 +692,91 @@ static void threadPrivatizeVars(lower::AbstractConverter &converter,
   }
 }
 
+static void groupprivatizeVars(lower::AbstractConverter &converter,
+                               lower::pft::Evaluation &eval) {
+  fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
+  mlir::Location currentLocation = converter.getCurrentLocation();
+  mlir::OpBuilder::InsertionGuard guard(firOpBuilder);
+  firOpBuilder.setInsertionPointToStart(firOpBuilder.getAllocaBlock());
+
+  auto module = converter.getModuleOp();
+
+  // Create a groupprivate operation for the symbol.
+  // TODO: Extract device_type from the groupprivate directive.
+  auto genGroupprivateOp = [&](const semantics::Symbol &sym) -> mlir::Value {
+    std::string globalName = converter.mangleName(sym);
+    fir::GlobalOp global = module.lookupSymbol<fir::GlobalOp>(globalName);
+    if (!global) {
+      return mlir::Value();
+    }
+
+    // Generate fir.address_of to get the global address
+    mlir::Value symAddr = fir::AddrOfOp::create(
+        firOpBuilder, currentLocation, global.resultType(), global.getSymbol());
+
+    mlir::omp::DeclareTargetDeviceType deviceTypeEnum =
+        mlir::omp::DeclareTargetDeviceType::any;
+    mlir::omp::DeclareTargetDeviceTypeAttr deviceTypeAttr =
+        mlir::omp::DeclareTargetDeviceTypeAttr::get(firOpBuilder.getContext(),
+                                                    deviceTypeEnum);
+
+    return mlir::omp::GroupprivateOp::create(firOpBuilder, currentLocation,
+                                             symAddr.getType(), symAddr,
+                                             deviceTypeAttr);
+  };
+
+  llvm::SetVector<const semantics::Symbol *> groupprivateSyms;
+  converter.collectSymbolSet(eval, groupprivateSyms,
+                             semantics::Symbol::Flag::OmpGroupPrivate,
+                             /*collectSymbols=*/true,
+                             /*collectHostAssociatedSymbols=*/true);
+  std::set<semantics::SourceName> groupprivateSymNames;
+
+  // For a COMMON block, the GroupprivateOp is generated for the block itself
+  // instead of its members.
+  llvm::SetVector<const semantics::Symbol *> commonSyms;
+
+  for (std::size_t i = 0; i < groupprivateSyms.size(); i++) {
+    const semantics::Symbol *sym = groupprivateSyms[i];
+    mlir::Value symGroupprivateValue;
+    // The variable may be used more than once, and each reference has one
+    // symbol with the same name. Only do once for references of one variable.
+    if (groupprivateSymNames.find(sym->name()) != groupprivateSymNames.end())
+      continue;
+    groupprivateSymNames.insert(sym->name());
+
+    if (const semantics::Symbol *common =
+            semantics::FindCommonBlockContaining(sym->GetUltimate())) {
+      // Handle common block members: create groupprivate op for the entire
+      // common block, then compute member offset.
+      mlir::Value commonGroupprivateValue;
+      if (commonSyms.contains(common)) {
+        commonGroupprivateValue = converter.getSymbolAddress(*common);
+      } else {
+        commonGroupprivateValue = genGroupprivateOp(*common);
+        if (!commonGroupprivateValue)
+          continue;
+        converter.bindSymbol(*common, commonGroupprivateValue);
+        commonSyms.insert(common);
+      }
+      symGroupprivateValue = lower::genCommonBlockMember(
+          converter, currentLocation, sym->GetUltimate(),
+          commonGroupprivateValue, common->size());
+    } else {
+      symGroupprivateValue = genGroupprivateOp(*sym);
+    }
+
+    if (!symGroupprivateValue) {
+      continue;
+    }
+
+    fir::ExtendedValue sexv = converter.getSymbolExtendedValue(*sym);
+    fir::ExtendedValue symGroupprivateExv =
+        getExtendedValue(sexv, symGroupprivateValue);
+    converter.bindSymbol(*sym, symGroupprivateExv);
+  }
+}
+
 static mlir::Operation *setLoopVar(lower::AbstractConverter &converter,
                                    mlir::Location loc, mlir::Value indexVal,
                                    const semantics::Symbol *sym) {
@@ -1230,6 +1315,10 @@ static void createBodyOfOp(mlir::Operation &op, const OpWithBodyGenInfo &info,
     }
   }
 
+  if (info.dir == llvm::omp::Directive::OMPD_teams) {
+    groupprivatizeVars(info.converter, info.eval);
+  }
+
   if (!info.genSkeletonOnly) {
     if (ConstructQueue::const_iterator next = std::next(item);
         next != queue.end()) {
@@ -2736,6 +2825,11 @@ genTargetOp(lower::AbstractConverter &converter, lower::SymMap &symTable,
         !symbolsWithDynamicSubstring.contains(&sym.GetUltimate()))
       return;
 
+    // Skip groupprivate symbols - they don't need to be mapped because
+    // groupprivate creates its own LDS storage.
+    if (sym.GetUltimate().test(semantics::Symbol::Flag::OmpGroupPrivate))
+      return;
+
     if (!isDuplicateMappedSymbol(sym, dsp.getAllSymbolsToPrivatize(),
                                  hasDeviceAddrSyms, mapSyms, isDevicePtrSyms)) {
       if (const auto *details =
@@ -4015,7 +4109,8 @@ static void genOMP(lower::AbstractConverter &converter, lower::SymMap &symTable,
                    semantics::SemanticsContext &semaCtx,
                    lower::pft::Evaluation &eval,
                    const parser::OpenMPGroupprivate &directive) {
-  TODO(converter.getCurrentLocation(), "GROUPPRIVATE");
+  // The groupprivate directive is lowered when the variable is referenced
+  // inside target/teams regions.
 }
 
 static void genOMP(lower::AbstractConverter &converter, lower::SymMap &symTable,
@@ -4419,6 +4514,9 @@ void Fortran::lower::genOpenMPSymbolProperties(
 
   if (sym.test(semantics::Symbol::Flag::OmpDeclareTarget))
     lower::genDeclareTargetIntGlobal(converter, var);
+
+  if (sym.test(semantics::Symbol::Flag::OmpGroupPrivate))
+    lower::genGroupprivateOp(converter, var);
 }
 
 void Fortran::lower::genThreadprivateOp(lower::AbstractConverter &converter,
@@ -4487,6 +4585,42 @@ void Fortran::lower::genThreadprivateOp(lower::AbstractConverter &converter,
   converter.bindSymbol(sym, symThreadprivateExv);
 }
 
+void Fortran::lower::genGroupprivateOp(lower::AbstractConverter &converter,
+                                       const lower::pft::Variable &var) {
+  fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
+  mlir::Location currentLocation = converter.getCurrentLocation();
+
+  const semantics::Symbol &sym = var.getSymbol();
+
+  // For common block members, the groupprivate op is generated for the entire
+  // common block in groupprivatizeVars, not for individual members here.
+  // The common block already has a global, so nothing to do here.
+  if (semantics::FindCommonBlockContaining(sym.GetUltimate())) {
+    return;
+  }
+
+  fir::GlobalOp global;
+  auto module = converter.getModuleOp();
+  std::string globalName = converter.mangleName(sym);
+
+  // Handle non-global variables - create a GlobalOp for them.
+  // Local variables with SAVE attribute can be groupprivate.
+  // Promote them to fir.global so that omp.groupprivate
+  // can reference them via fir.address_of.
+  if (!var.isGlobal()) {
+    if (module.lookupSymbol<fir::GlobalOp>(globalName))
+      global = module.lookupSymbol<fir::GlobalOp>(globalName);
+    else
+      global = globalInitialization(converter, firOpBuilder, sym, var,
+                                    currentLocation);
+  } else {
+    global = module.lookupSymbol<fir::GlobalOp>(globalName);
+  }
+
+  // The actual omp.groupprivate operation is created by groupprivatizeVars
+  // when entering a teams region.
+}
+
 // 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

flang/test/Lower/OpenMP/Todo/groupprivate.f90

@@ -0,0 +1,57 @@
+!RUN: %flang_fc1 -emit-hlfir -fopenmp -fopenmp-version=60 %s -o - | FileCheck %s
+
+! Test lowering of groupprivate directive to omp.groupprivate.
+
+! CHECK-DAG: fir.global common @blk_
+! CHECK-DAG: fir.global @_QMmEx : i32
+
+! Test: basic groupprivate with single variable
+
+module m
+  implicit none
+  integer, save :: x
+  !$omp groupprivate(x)
+end module
+
+! CHECK-LABEL: func.func @_QPtest_groupprivate
+! CHECK: omp.target
+! CHECK:   omp.teams
+! CHECK:     fir.address_of(@_QMmEx)
+! CHECK:     omp.groupprivate
+subroutine test_groupprivate()
+  use m
+
+  !$omp target
+    !$omp teams
+      x = 10
+    !$omp end teams
+  !$omp end target
+end subroutine
+
+! Test: groupprivate with common block
+module m2
+  implicit none
+  integer :: cb_x, cb_y
+  real :: cb_z
+  common /blk/ cb_x, cb_y, cb_z
+  !$omp groupprivate(/blk/)
+end module
+
+! CHECK-LABEL: func.func @_QPtest_common_block_groupprivate
+! CHECK: omp.target
+! CHECK:   omp.teams
+! CHECK:     fir.address_of(@blk_)
+! CHECK:     omp.groupprivate
+! CHECK:     fir.convert
+! CHECK:     fir.coordinate_of
+subroutine test_common_block_groupprivate()
+  use m2
+
+  !$omp target
+    !$omp teams
+      cb_x = 1
+      cb_y = 2
+      cb_z = 3.0
+    !$omp end teams
+  !$omp end target
+end subroutine