[flang] Extend localization support for do concurrent (init regions) (#142564)

Extends support for locality specifiers in `do concurrent` by supporting
data types that need `init` regions.

This further unifies the paths taken by the compiler for OpenMP
privatization clauses and `do concurrent` locality specifiers.

Author: ergawy

flang/lib/Lower/OpenMP/PrivateReductionUtils.h renamed to flang/include/flang/Lower/Support/PrivateReductionUtils.h

@@ -35,12 +35,14 @@ namespace Fortran {
 namespace lower {
 class AbstractConverter;
 
-namespace omp {
-
-enum class DeclOperationKind { Private, FirstPrivate, Reduction };
+enum class DeclOperationKind {
+  PrivateOrLocal,
+  FirstPrivateOrLocalInit,
+  Reduction
+};
 inline bool isPrivatization(DeclOperationKind kind) {
-  return (kind == DeclOperationKind::FirstPrivate) ||
-         (kind == DeclOperationKind::Private);
+  return (kind == DeclOperationKind::FirstPrivateOrLocalInit) ||
+         (kind == DeclOperationKind::PrivateOrLocal);
 }
 inline bool isReduction(DeclOperationKind kind) {
   return kind == DeclOperationKind::Reduction;
@@ -56,7 +58,7 @@ void populateByRefInitAndCleanupRegions(
     mlir::Value allocatedPrivVarArg, mlir::Value moldArg,
     mlir::Region &cleanupRegion, DeclOperationKind kind,
     const Fortran::semantics::Symbol *sym = nullptr,
-    bool cannotHaveNonDefaultLowerBounds = false);
+    bool cannotHaveNonDefaultLowerBounds = false, bool isDoConcurrent = false);
 
 /// Generate a fir::ShapeShift op describing the provided boxed array.
 /// `cannotHaveNonDefaultLowerBounds` should be set if `box` is known to have
@@ -69,7 +71,6 @@ fir::ShapeShiftOp getShapeShift(fir::FirOpBuilder &builder, mlir::Location loc,
                                 bool cannotHaveNonDefaultLowerBounds = false,
                                 bool useDefaultLowerBounds = false);
 
-} // namespace omp
 } // namespace lower
 } // namespace Fortran
 

flang/include/flang/Lower/Support/Utils.h

@@ -20,6 +20,7 @@
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/BuiltinAttributes.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/StringRef.h"
 
 namespace Fortran::lower {
@@ -98,8 +99,9 @@ bool isEqual(const Fortran::lower::ExplicitIterSpace::ArrayBases &x,
 template <typename OpType, typename OperandsStructType>
 void privatizeSymbol(
     lower::AbstractConverter &converter, fir::FirOpBuilder &firOpBuilder,
-    lower::SymMap &symTable, std::function<void(OpType, mlir::Type)> initGen,
+    lower::SymMap &symTable,
     llvm::SetVector<const semantics::Symbol *> &allPrivatizedSymbols,
+    llvm::SmallSet<const semantics::Symbol *, 16> &mightHaveReadHostSym,
     const semantics::Symbol *symToPrivatize, OperandsStructType *clauseOps);
 
 } // end namespace Fortran::lower

flang/lib/Lower/Bridge.cpp

@@ -12,7 +12,6 @@
 
 #include "flang/Lower/Bridge.h"
 
-#include "OpenMP/DataSharingProcessor.h"
 #include "flang/Lower/Allocatable.h"
 #include "flang/Lower/CallInterface.h"
 #include "flang/Lower/Coarray.h"
@@ -2040,44 +2039,38 @@ class FirConverter : public Fortran::lower::AbstractConverter {
     bool useDelayedPriv =
         enableDelayedPrivatizationStaging && doConcurrentLoopOp;
     llvm::SetVector<const Fortran::semantics::Symbol *> allPrivatizedSymbols;
+    llvm::SmallSet<const Fortran::semantics::Symbol *, 16> mightHaveReadHostSym;
 
-    for (const Fortran::semantics::Symbol *sym : info.localSymList) {
+    for (const Fortran::semantics::Symbol *symToPrivatize : info.localSymList) {
       if (useDelayedPriv) {
         Fortran::lower::privatizeSymbol<fir::LocalitySpecifierOp>(
-            *this, this->getFirOpBuilder(), localSymbols,
-            [this](fir::LocalitySpecifierOp result, mlir::Type argType) {
-              TODO(this->toLocation(),
-                   "Localizers that need init regions are not supported yet.");
-            },
-            allPrivatizedSymbols, sym, &privateClauseOps);
+            *this, this->getFirOpBuilder(), localSymbols, allPrivatizedSymbols,
+            mightHaveReadHostSym, symToPrivatize, &privateClauseOps);
         continue;
       }
 
-      createHostAssociateVarClone(*sym, /*skipDefaultInit=*/false);
+      createHostAssociateVarClone(*symToPrivatize, /*skipDefaultInit=*/false);
     }
 
-    for (const Fortran::semantics::Symbol *sym : info.localInitSymList) {
+    for (const Fortran::semantics::Symbol *symToPrivatize :
+         info.localInitSymList) {
       if (useDelayedPriv) {
         Fortran::lower::privatizeSymbol<fir::LocalitySpecifierOp>(
-            *this, this->getFirOpBuilder(), localSymbols,
-            [this](fir::LocalitySpecifierOp result, mlir::Type argType) {
-              TODO(this->toLocation(),
-                   "Localizers that need init regions are not supported yet.");
-            },
-            allPrivatizedSymbols, sym, &privateClauseOps);
+            *this, this->getFirOpBuilder(), localSymbols, allPrivatizedSymbols,
+            mightHaveReadHostSym, symToPrivatize, &privateClauseOps);
         continue;
       }
 
-      createHostAssociateVarClone(*sym, /*skipDefaultInit=*/true);
+      createHostAssociateVarClone(*symToPrivatize, /*skipDefaultInit=*/true);
       const auto *hostDetails =
-          sym->detailsIf<Fortran::semantics::HostAssocDetails>();
+          symToPrivatize->detailsIf<Fortran::semantics::HostAssocDetails>();
       assert(hostDetails && "missing locality spec host symbol");
       const Fortran::semantics::Symbol *hostSym = &hostDetails->symbol();
       Fortran::evaluate::ExpressionAnalyzer ea{semanticsContext};
       Fortran::evaluate::Assignment assign{
-          ea.Designate(Fortran::evaluate::DataRef{*sym}).value(),
+          ea.Designate(Fortran::evaluate::DataRef{*symToPrivatize}).value(),
           ea.Designate(Fortran::evaluate::DataRef{*hostSym}).value()};
-      if (Fortran::semantics::IsPointer(*sym))
+      if (Fortran::semantics::IsPointer(*symToPrivatize))
         assign.u = Fortran::evaluate::Assignment::BoundsSpec{};
       genAssignment(assign);
     }

flang/lib/Lower/CMakeLists.txt

@@ -28,11 +28,11 @@ add_flang_library(FortranLower
   OpenMP/DataSharingProcessor.cpp
   OpenMP/Decomposer.cpp
   OpenMP/OpenMP.cpp
-  OpenMP/PrivateReductionUtils.cpp
   OpenMP/ReductionProcessor.cpp
   OpenMP/Utils.cpp
   PFTBuilder.cpp
   Runtime.cpp
+  Support/PrivateReductionUtils.cpp
   Support/Utils.cpp
   SymbolMap.cpp
   VectorSubscripts.cpp

flang/lib/Lower/OpenMP/DataSharingProcessor.cpp

@@ -12,10 +12,10 @@
 
 #include "DataSharingProcessor.h"
 
-#include "PrivateReductionUtils.h"
 #include "Utils.h"
 #include "flang/Lower/ConvertVariable.h"
 #include "flang/Lower/PFTBuilder.h"
+#include "flang/Lower/Support/PrivateReductionUtils.h"
 #include "flang/Lower/Support/Utils.h"
 #include "flang/Lower/SymbolMap.h"
 #include "flang/Optimizer/Builder/BoxValue.h"
@@ -537,38 +537,10 @@ void DataSharingProcessor::privatizeSymbol(
     return;
   }
 
-  auto initGen = [&](mlir::omp::PrivateClauseOp result, mlir::Type argType) {
-    lower::SymbolBox hsb = converter.lookupOneLevelUpSymbol(*symToPrivatize);
-    assert(hsb && "Host symbol box not found");
-    hlfir::Entity entity{hsb.getAddr()};
-    bool cannotHaveNonDefaultLowerBounds =
-        !entity.mayHaveNonDefaultLowerBounds();
-
-    mlir::Region &initRegion = result.getInitRegion();
-    mlir::Location symLoc = hsb.getAddr().getLoc();
-    mlir::Block *initBlock = firOpBuilder.createBlock(
-        &initRegion, /*insertPt=*/{}, {argType, argType}, {symLoc, symLoc});
-
-    bool emitCopyRegion =
-        symToPrivatize->test(semantics::Symbol::Flag::OmpFirstPrivate);
-
-    populateByRefInitAndCleanupRegions(
-        converter, symLoc, argType, /*scalarInitValue=*/nullptr, initBlock,
-        result.getInitPrivateArg(), result.getInitMoldArg(),
-        result.getDeallocRegion(),
-        emitCopyRegion ? omp::DeclOperationKind::FirstPrivate
-                       : omp::DeclOperationKind::Private,
-        symToPrivatize, cannotHaveNonDefaultLowerBounds);
-    // TODO: currently there are false positives from dead uses of the mold
-    // arg
-    if (result.initReadsFromMold())
-      mightHaveReadHostSym.insert(symToPrivatize);
-  };
-
   Fortran::lower::privatizeSymbol<mlir::omp::PrivateClauseOp,
                                   mlir::omp::PrivateClauseOps>(
-      converter, firOpBuilder, symTable, initGen, allPrivatizedSymbols,
-      symToPrivatize, clauseOps);
+      converter, firOpBuilder, symTable, allPrivatizedSymbols,
+      mightHaveReadHostSym, symToPrivatize, clauseOps);
 }
 } // namespace omp
 } // namespace lower

flang/lib/Lower/OpenMP/ReductionProcessor.cpp

@@ -12,9 +12,9 @@
 
 #include "ReductionProcessor.h"
 
-#include "PrivateReductionUtils.h"
 #include "flang/Lower/AbstractConverter.h"
 #include "flang/Lower/ConvertType.h"
+#include "flang/Lower/Support/PrivateReductionUtils.h"
 #include "flang/Lower/SymbolMap.h"
 #include "flang/Optimizer/Builder/Complex.h"
 #include "flang/Optimizer/Builder/HLFIRTools.h"

flang/lib/Lower/OpenMP/PrivateReductionUtils.cpp renamed to flang/lib/Lower/Support/PrivateReductionUtils.cpp

@@ -10,7 +10,7 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "PrivateReductionUtils.h"
+#include "flang/Lower/Support/PrivateReductionUtils.h"
 
 #include "flang/Lower/AbstractConverter.h"
 #include "flang/Lower/Allocatable.h"
@@ -42,7 +42,8 @@ static bool hasFinalization(const Fortran::semantics::Symbol &sym) {
 static void createCleanupRegion(Fortran::lower::AbstractConverter &converter,
                                 mlir::Location loc, mlir::Type argType,
                                 mlir::Region &cleanupRegion,
-                                const Fortran::semantics::Symbol *sym) {
+                                const Fortran::semantics::Symbol *sym,
+                                bool isDoConcurrent) {
   fir::FirOpBuilder &builder = converter.getFirOpBuilder();
   assert(cleanupRegion.empty());
   mlir::Block *block = builder.createBlock(&cleanupRegion, cleanupRegion.end(),
@@ -72,7 +73,10 @@ static void createCleanupRegion(Fortran::lower::AbstractConverter &converter,
         fir::MutableBoxValue mutableBox{converted, /*lenParameters=*/{},
                                         /*mutableProperties=*/{}};
         Fortran::lower::genDeallocateIfAllocated(converter, mutableBox, loc);
-        builder.create<mlir::omp::YieldOp>(loc);
+        if (isDoConcurrent)
+          builder.create<fir::YieldOp>(loc);
+        else
+          builder.create<mlir::omp::YieldOp>(loc);
         return;
       }
     }
@@ -100,7 +104,10 @@ static void createCleanupRegion(Fortran::lower::AbstractConverter &converter,
     builder.create<fir::FreeMemOp>(loc, cast);
 
     builder.setInsertionPointAfter(ifOp);
-    builder.create<mlir::omp::YieldOp>(loc);
+    if (isDoConcurrent)
+      builder.create<fir::YieldOp>(loc);
+    else
+      builder.create<mlir::omp::YieldOp>(loc);
     return;
   }
 
@@ -115,14 +122,18 @@ static void createCleanupRegion(Fortran::lower::AbstractConverter &converter,
     addr = builder.createConvert(loc, heapTy, addr);
 
     builder.create<fir::FreeMemOp>(loc, addr);
-    builder.create<mlir::omp::YieldOp>(loc);
+    if (isDoConcurrent)
+      builder.create<fir::YieldOp>(loc);
+    else
+      builder.create<mlir::omp::YieldOp>(loc);
+
     return;
   }
 
   typeError();
 }
 
-fir::ShapeShiftOp Fortran::lower::omp::getShapeShift(
+fir::ShapeShiftOp Fortran::lower::getShapeShift(
     fir::FirOpBuilder &builder, mlir::Location loc, mlir::Value box,
     bool cannotHaveNonDefaultLowerBounds, bool useDefaultLowerBounds) {
   fir::SequenceType sequenceType = mlir::cast<fir::SequenceType>(
@@ -262,7 +273,7 @@ static mlir::Value generateZeroShapeForRank(fir::FirOpBuilder &builder,
 }
 
 namespace {
-using namespace Fortran::lower::omp;
+using namespace Fortran::lower;
 /// Class to store shared data so we don't have to maintain so many function
 /// arguments
 class PopulateInitAndCleanupRegionsHelper {
@@ -273,12 +284,13 @@ class PopulateInitAndCleanupRegionsHelper {
       mlir::Value allocatedPrivVarArg, mlir::Value moldArg,
       mlir::Block *initBlock, mlir::Region &cleanupRegion,
       DeclOperationKind kind, const Fortran::semantics::Symbol *sym,
-      bool cannotHaveLowerBounds)
+      bool cannotHaveLowerBounds, bool isDoConcurrent)
       : converter{converter}, builder{converter.getFirOpBuilder()}, loc{loc},
         argType{argType}, scalarInitValue{scalarInitValue},
         allocatedPrivVarArg{allocatedPrivVarArg}, moldArg{moldArg},
         initBlock{initBlock}, cleanupRegion{cleanupRegion}, kind{kind},
-        sym{sym}, cannotHaveNonDefaultLowerBounds{cannotHaveLowerBounds} {
+        sym{sym}, cannotHaveNonDefaultLowerBounds{cannotHaveLowerBounds},
+        isDoConcurrent{isDoConcurrent} {
     valType = fir::unwrapRefType(argType);
   }
 
@@ -324,8 +336,13 @@ class PopulateInitAndCleanupRegionsHelper {
   /// lower bounds then we don't need to generate code to read them.
   bool cannotHaveNonDefaultLowerBounds;
 
+  bool isDoConcurrent;
+
   void createYield(mlir::Value ret) {
-    builder.create<mlir::omp::YieldOp>(loc, ret);
+    if (isDoConcurrent)
+      builder.create<fir::YieldOp>(loc, ret);
+    else
+      builder.create<mlir::omp::YieldOp>(loc, ret);
   }
 
   void initTrivialType() {
@@ -429,11 +446,12 @@ void PopulateInitAndCleanupRegionsHelper::initAndCleanupBoxedScalar(
       /*slice=*/mlir::Value{}, lenParams);
   initializeIfDerivedTypeBox(
       builder, loc, box, getLoadedMoldArg(), needsInitialization,
-      /*isFirstPrivate=*/kind == DeclOperationKind::FirstPrivate);
+      /*isFirstPrivate=*/kind == DeclOperationKind::FirstPrivateOrLocalInit);
   fir::StoreOp lastOp =
       builder.create<fir::StoreOp>(loc, box, allocatedPrivVarArg);
 
-  createCleanupRegion(converter, loc, argType, cleanupRegion, sym);
+  createCleanupRegion(converter, loc, argType, cleanupRegion, sym,
+                      isDoConcurrent);
 
   if (ifUnallocated)
     builder.setInsertionPointAfter(ifUnallocated);
@@ -470,13 +488,14 @@ void PopulateInitAndCleanupRegionsHelper::initAndCleanupBoxedArray(
                                                         allocatedArray, shape);
     initializeIfDerivedTypeBox(
         builder, loc, firClass, source, needsInitialization,
-        /*isFirstprivate=*/kind == DeclOperationKind::FirstPrivate);
+        /*isFirstprivate=*/kind == DeclOperationKind::FirstPrivateOrLocalInit);
     builder.create<fir::StoreOp>(loc, firClass, allocatedPrivVarArg);
     if (ifUnallocated)
       builder.setInsertionPointAfter(ifUnallocated);
     createYield(allocatedPrivVarArg);
     mlir::OpBuilder::InsertionGuard guard(builder);
-    createCleanupRegion(converter, loc, argType, cleanupRegion, sym);
+    createCleanupRegion(converter, loc, argType, cleanupRegion, sym,
+                        isDoConcurrent);
     return;
   }
 
@@ -492,7 +511,8 @@ void PopulateInitAndCleanupRegionsHelper::initAndCleanupBoxedArray(
          "createTempFromMold decides this statically");
   if (cstNeedsDealloc.has_value() && *cstNeedsDealloc != false) {
     mlir::OpBuilder::InsertionGuard guard(builder);
-    createCleanupRegion(converter, loc, argType, cleanupRegion, sym);
+    createCleanupRegion(converter, loc, argType, cleanupRegion, sym,
+                        isDoConcurrent);
   } else {
     assert(!isAllocatableOrPointer &&
            "Pointer-like arrays must be heap allocated");
@@ -520,7 +540,7 @@ void PopulateInitAndCleanupRegionsHelper::initAndCleanupBoxedArray(
 
   initializeIfDerivedTypeBox(
       builder, loc, box, getLoadedMoldArg(), needsInitialization,
-      /*isFirstPrivate=*/kind == DeclOperationKind::FirstPrivate);
+      /*isFirstPrivate=*/kind == DeclOperationKind::FirstPrivateOrLocalInit);
 
   builder.create<fir::StoreOp>(loc, box, allocatedPrivVarArg);
   if (ifUnallocated)
@@ -548,7 +568,8 @@ void PopulateInitAndCleanupRegionsHelper::initAndCleanupBoxchar(
       loc, eleTy, /*name=*/{}, /*shape=*/{}, /*lenParams=*/len);
   mlir::Value boxChar = charExprHelper.createEmboxChar(privateAddr, len);
 
-  createCleanupRegion(converter, loc, argType, cleanupRegion, sym);
+  createCleanupRegion(converter, loc, argType, cleanupRegion, sym,
+                      isDoConcurrent);
 
   builder.setInsertionPointToEnd(initBlock);
   createYield(boxChar);
@@ -563,10 +584,11 @@ void PopulateInitAndCleanupRegionsHelper::initAndCleanupUnboxedDerivedType(
   mlir::Value moldBox = builder.create<fir::EmboxOp>(loc, boxedTy, moldArg);
   initializeIfDerivedTypeBox(builder, loc, newBox, moldBox, needsInitialization,
                              /*isFirstPrivate=*/kind ==
-                                 DeclOperationKind::FirstPrivate);
+                                 DeclOperationKind::FirstPrivateOrLocalInit);
 
   if (sym && hasFinalization(*sym))
-    createCleanupRegion(converter, loc, argType, cleanupRegion, sym);
+    createCleanupRegion(converter, loc, argType, cleanupRegion, sym,
+                        isDoConcurrent);
 
   builder.setInsertionPointToEnd(initBlock);
   createYield(allocatedPrivVarArg);
@@ -632,15 +654,17 @@ void PopulateInitAndCleanupRegionsHelper::populateByRefInitAndCleanupRegions() {
        "creating reduction/privatization init region for unsupported type");
 }
 
-void Fortran::lower::omp::populateByRefInitAndCleanupRegions(
+void Fortran::lower::populateByRefInitAndCleanupRegions(
     Fortran::lower::AbstractConverter &converter, mlir::Location loc,
     mlir::Type argType, mlir::Value scalarInitValue, mlir::Block *initBlock,
     mlir::Value allocatedPrivVarArg, mlir::Value moldArg,
     mlir::Region &cleanupRegion, DeclOperationKind kind,
-    const Fortran::semantics::Symbol *sym, bool cannotHaveLowerBounds) {
+    const Fortran::semantics::Symbol *sym, bool cannotHaveLowerBounds,
+    bool isDoConcurrent) {
   PopulateInitAndCleanupRegionsHelper helper(
       converter, loc, argType, scalarInitValue, allocatedPrivVarArg, moldArg,
-      initBlock, cleanupRegion, kind, sym, cannotHaveLowerBounds);
+      initBlock, cleanupRegion, kind, sym, cannotHaveLowerBounds,
+      isDoConcurrent);
   helper.populateByRefInitAndCleanupRegions();
 
   // Often we load moldArg to check something (e.g. length parameters, shape)