[mlir][OpenMP][flang] make private variable allocation implicit in omp.private (#124019)

The intention of this work is to give MLIR->LLVMIR conversion freedom to
control how the private variable is allocated so that it can be
allocated on the stack in ordinary cases or as part of a structure used
to give closure context for tasks which might outlive the current stack
frame. See RFC:

https://discourse.llvm.org/t/rfc-openmp-supporting-delayed-task-execution-with-firstprivate-variables/83084

For example, a privatizer for an integer used to look like
```mlir
  omp.private {type = private} @x.privatizer : !fir.ref<i32> alloc {
  ^bb0(%arg0: !fir.ref<i32>):
    %0 = ... allocate proper memory for the private clone ...
    omp.yield(%0 : !fir.ref<i32>)
  }
```

After this change, allocation become implicit in the operation:
```mlir
  omp.private {type = private} @x.privatizer : i32
```

For more complex types that require initialization after allocation, an
init region can be used:
``` mlir
  omp.private {type = private} @x.privatizer : !some.type init {
  ^bb0(%arg0: !some.pointer<!some.type>, %arg1: !some.pointer<!some.type>):
    // initialize %arg1, using %arg0 as a mold for allocations
    omp.yield(%arg1 : !some.pointer<!some.type>)
  } dealloc {
    ^bb0(%arg0: !some.pointer<!some.type>):
    ... deallocate memory allocated by the init region ...
    omp.yield
  }
```

This patch lays the groundwork for delayed task execution but is not
enough on its own.

After this patch all gfortran tests which previously passed still pass.
There
are the following changes to the Fujitsu test suite:
- 0380_0009 and 0435_0009 are fixed
- 0688_0041 now fails at runtime. This patch is testing firstprivate
variables with tasks. Previously we got lucky with the undefined
behavior and won the race. After these changes we no longer get lucky.
This patch lays the groundwork for a proper fix for this issue.

In flang the lowering re-uses the existing lowering used for reduction
init and dealloc regions.

In flang, before this patch we hit a TODO with the same wording when
generating the copy region for firstprivate polymorphic variables. After
this patch the box-like fir.class is passed by reference into the copy
region, leading to a different path that didn't hit that old TODO but
the generated code still didn't work so I added a new TODO in
DataSharingProcessor.

Author: tblah

flang/lib/Lower/OpenMP/DataSharingProcessor.cpp

@@ -12,13 +12,17 @@
 
 #include "DataSharingProcessor.h"
 
+#include "PrivateReductionUtils.h"
 #include "Utils.h"
 #include "flang/Lower/ConvertVariable.h"
 #include "flang/Lower/PFTBuilder.h"
 #include "flang/Lower/SymbolMap.h"
+#include "flang/Optimizer/Builder/BoxValue.h"
 #include "flang/Optimizer/Builder/HLFIRTools.h"
 #include "flang/Optimizer/Builder/Todo.h"
+#include "flang/Optimizer/HLFIR/HLFIRDialect.h"
 #include "flang/Optimizer/HLFIR/HLFIROps.h"
+#include "flang/Semantics/attr.h"
 #include "flang/Semantics/tools.h"
 
 namespace Fortran {
@@ -85,35 +89,65 @@ void DataSharingProcessor::insertDeallocs() {
         converter.createHostAssociateVarCloneDealloc(*sym);
         continue;
       }
-
-      lower::SymbolBox hsb = converter.lookupOneLevelUpSymbol(*sym);
-      assert(hsb && "Host symbol box not found");
-      mlir::Type symType = hsb.getAddr().getType();
-      mlir::Location symLoc = hsb.getAddr().getLoc();
-      fir::ExtendedValue symExV = converter.getSymbolExtendedValue(*sym);
-      mlir::omp::PrivateClauseOp privatizer = symToPrivatizer.at(sym);
-
-      lower::SymMapScope scope(symTable);
-      mlir::OpBuilder::InsertionGuard guard(firOpBuilder);
-
-      mlir::Region &deallocRegion = privatizer.getDeallocRegion();
-      fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
-      mlir::Block *deallocEntryBlock = firOpBuilder.createBlock(
-          &deallocRegion, /*insertPt=*/{}, symType, symLoc);
-
-      firOpBuilder.setInsertionPointToEnd(deallocEntryBlock);
-      symTable.addSymbol(*sym,
-                         fir::substBase(symExV, deallocRegion.getArgument(0)));
-
-      converter.createHostAssociateVarCloneDealloc(*sym);
-      firOpBuilder.create<mlir::omp::YieldOp>(hsb.getAddr().getLoc());
+      // For delayed privatization deallocs are created by
+      // populateByRefInitAndCleanupRegions
     }
 }
 
 void DataSharingProcessor::cloneSymbol(const semantics::Symbol *sym) {
   bool isFirstPrivate = sym->test(semantics::Symbol::Flag::OmpFirstPrivate);
-  bool success = converter.createHostAssociateVarClone(
-      *sym, /*skipDefaultInit=*/isFirstPrivate);
+
+  // If we are doing eager-privatization on a symbol created using delayed
+  // privatization there could be incompatible types here e.g.
+  // fir.ref<fir.box<fir.array<>>>
+  bool success = [&]() -> bool {
+    const auto *details =
+        sym->detailsIf<Fortran::semantics::HostAssocDetails>();
+    assert(details && "No host-association found");
+    const Fortran::semantics::Symbol &hsym = details->symbol();
+    mlir::Value addr = converter.getSymbolAddress(hsym);
+
+    if (auto refTy = mlir::dyn_cast<fir::ReferenceType>(addr.getType())) {
+      if (auto boxTy = mlir::dyn_cast<fir::BoxType>(refTy.getElementType())) {
+        if (auto arrayTy =
+                mlir::dyn_cast<fir::SequenceType>(boxTy.getElementType())) {
+          // FirConverter/fir::ExtendedValue considers all references to boxes
+          // as mutable boxes. Outside of OpenMP it doesn't make sense to have a
+          // mutable box of an array. Work around this here by loading the
+          // reference so it is a normal boxed array.
+          fir::FirOpBuilder &builder = converter.getFirOpBuilder();
+          mlir::Location loc = converter.genLocation(hsym.name());
+          fir::ExtendedValue hexv = converter.getSymbolExtendedValue(hsym);
+
+          llvm::SmallVector<mlir::Value> extents =
+              fir::factory::getExtents(loc, builder, hexv);
+
+          // TODO: uniqName, name
+          mlir::Value allocVal =
+              builder.allocateLocal(loc, arrayTy, /*uniqName=*/"",
+                                    /*name=*/"", extents, /*typeParams=*/{},
+                                    sym->GetUltimate().attrs().test(
+                                        Fortran::semantics::Attr::TARGET));
+          mlir::Value shape = builder.genShape(loc, extents);
+          mlir::Value box = builder.createBox(loc, boxTy, allocVal, shape,
+                                              nullptr, {}, nullptr);
+
+          // This can't be a CharArrayBoxValue because otherwise
+          // boxTy.getElementType() would be a character type.
+          // Assume the array element type isn't polymorphic because we are
+          // privatizing.
+          fir::ExtendedValue newExv = fir::ArrayBoxValue{box, extents};
+
+          converter.bindSymbol(*sym, newExv);
+          return true;
+        }
+      }
+    }
+
+    // Normal case:
+    return converter.createHostAssociateVarClone(
+        *sym, /*skipDefaultInit=*/isFirstPrivate);
+  }();
   (void)success;
   assert(success && "Privatization failed due to existing binding");
 
@@ -132,7 +166,7 @@ void DataSharingProcessor::cloneSymbol(const semantics::Symbol *sym) {
 
   if (needInitClone()) {
     Fortran::lower::initializeCloneAtRuntime(converter, *sym, symTable);
-    callsInitClone = true;
+    mightHaveReadHostSym = true;
   }
 }
 
@@ -184,7 +218,8 @@ bool DataSharingProcessor::needBarrier() {
   // Emit implicit barrier for linear clause. Maybe on somewhere else.
   for (const semantics::Symbol *sym : allPrivatizedSymbols) {
     if (sym->test(semantics::Symbol::Flag::OmpLastPrivate) &&
-        (sym->test(semantics::Symbol::Flag::OmpFirstPrivate) || callsInitClone))
+        (sym->test(semantics::Symbol::Flag::OmpFirstPrivate) ||
+         mightHaveReadHostSym))
       return true;
   }
   return false;
@@ -468,15 +503,47 @@ void DataSharingProcessor::doPrivatize(const semantics::Symbol *sym,
   lower::SymbolBox hsb = converter.lookupOneLevelUpSymbol(*sym);
   assert(hsb && "Host symbol box not found");
 
-  mlir::Type symType = hsb.getAddr().getType();
   mlir::Location symLoc = hsb.getAddr().getLoc();
   std::string privatizerName = sym->name().ToString() + ".privatizer";
   bool isFirstPrivate = sym->test(semantics::Symbol::Flag::OmpFirstPrivate);
 
+  mlir::Value privVal = hsb.getAddr();
+  mlir::Type allocType = privVal.getType();
+  if (!mlir::isa<fir::PointerType>(privVal.getType()))
+    allocType = fir::unwrapRefType(privVal.getType());
+
+  if (auto poly = mlir::dyn_cast<fir::ClassType>(allocType)) {
+    if (!mlir::isa<fir::PointerType>(poly.getEleTy()) && isFirstPrivate)
+      TODO(symLoc, "create polymorphic host associated copy");
+  }
+
+  // fir.array<> cannot be converted to any single llvm type and fir helpers
+  // are not available in openmp to llvmir translation so we cannot generate
+  // an alloca for a fir.array type there. Get around this by boxing all
+  // arrays.
+  if (mlir::isa<fir::SequenceType>(allocType)) {
+    hlfir::Entity entity{hsb.getAddr()};
+    entity = genVariableBox(symLoc, firOpBuilder, entity);
+    privVal = entity.getBase();
+    allocType = privVal.getType();
+  }
+
+  if (mlir::isa<fir::BaseBoxType>(privVal.getType())) {
+    // Boxes should be passed by reference into nested regions:
+    auto oldIP = firOpBuilder.saveInsertionPoint();
+    firOpBuilder.setInsertionPointToStart(firOpBuilder.getAllocaBlock());
+    auto alloca = firOpBuilder.create<fir::AllocaOp>(symLoc, privVal.getType());
+    firOpBuilder.restoreInsertionPoint(oldIP);
+    firOpBuilder.create<fir::StoreOp>(symLoc, privVal, alloca);
+    privVal = alloca;
+  }
+
+  mlir::Type argType = privVal.getType();
+
   mlir::omp::PrivateClauseOp privatizerOp = [&]() {
     auto moduleOp = firOpBuilder.getModule();
     auto uniquePrivatizerName = fir::getTypeAsString(
-        symType, converter.getKindMap(),
+        allocType, converter.getKindMap(),
         converter.mangleName(*sym) +
             (isFirstPrivate ? "_firstprivate" : "_private"));
 
@@ -488,44 +555,40 @@ void DataSharingProcessor::doPrivatize(const semantics::Symbol *sym,
     mlir::OpBuilder::InsertionGuard guard(firOpBuilder);
     firOpBuilder.setInsertionPointToStart(moduleOp.getBody());
     auto result = firOpBuilder.create<mlir::omp::PrivateClauseOp>(
-        symLoc, uniquePrivatizerName, symType,
+        symLoc, uniquePrivatizerName, allocType,
         isFirstPrivate ? mlir::omp::DataSharingClauseType::FirstPrivate
                        : mlir::omp::DataSharingClauseType::Private);
     fir::ExtendedValue symExV = converter.getSymbolExtendedValue(*sym);
     lower::SymMapScope outerScope(symTable);
 
-    // Populate the `alloc` region.
-    {
-      mlir::Region &allocRegion = result.getAllocRegion();
-      mlir::Block *allocEntryBlock = firOpBuilder.createBlock(
-          &allocRegion, /*insertPt=*/{}, symType, symLoc);
-
-      firOpBuilder.setInsertionPointToEnd(allocEntryBlock);
-
-      fir::ExtendedValue localExV =
-          hlfir::translateToExtendedValue(
-              symLoc, firOpBuilder, hlfir::Entity{allocRegion.getArgument(0)},
-              /*contiguousHint=*/
-              evaluate::IsSimplyContiguous(*sym, converter.getFoldingContext()))
-              .first;
-
-      symTable.addSymbol(*sym, localExV);
-      lower::SymMapScope innerScope(symTable);
-      cloneSymbol(sym);
-      mlir::Value cloneAddr = symTable.shallowLookupSymbol(*sym).getAddr();
-      mlir::Type cloneType = cloneAddr.getType();
-
-      // A `convert` op is required for variables that are storage associated
-      // via `equivalence`. The problem is that these variables are declared as
-      // `fir.ptr`s while their privatized storage is declared as `fir.ref`,
-      // therefore we convert to proper symbol type.
-      mlir::Value yieldedValue =
-          (symType == cloneType) ? cloneAddr
-                                 : firOpBuilder.createConvert(
-                                       cloneAddr.getLoc(), symType, cloneAddr);
-
-      firOpBuilder.create<mlir::omp::YieldOp>(hsb.getAddr().getLoc(),
-                                              yieldedValue);
+    // Populate the `init` region.
+    // We need to initialize in the following cases:
+    // 1. The allocation was for a derived type which requires initialization
+    //    (this can be skipped if it will be initialized anyway by the copy
+    //    region, unless the derived type has allocatable components)
+    // 2. The allocation was for any kind of box
+    // 3. The allocation was for a boxed character
+    const bool needsInitialization =
+        (Fortran::lower::hasDefaultInitialization(sym->GetUltimate()) &&
+         (!isFirstPrivate || hlfir::mayHaveAllocatableComponent(allocType))) ||
+        mlir::isa<fir::BaseBoxType>(allocType) ||
+        mlir::isa<fir::BoxCharType>(allocType);
+    if (needsInitialization) {
+      mlir::Region &initRegion = result.getInitRegion();
+      mlir::Block *initBlock = firOpBuilder.createBlock(
+          &initRegion, /*insertPt=*/{}, {argType, argType}, {symLoc, symLoc});
+
+      populateByRefInitAndCleanupRegions(
+          converter, symLoc, argType, /*scalarInitValue=*/nullptr, initBlock,
+          result.getInitPrivateArg(), result.getInitMoldArg(),
+          result.getDeallocRegion(),
+          isFirstPrivate ? DeclOperationKind::FirstPrivate
+                         : DeclOperationKind::Private,
+          sym);
+      // TODO: currently there are false positives from dead uses of the mold
+      // arg
+      if (!result.getInitMoldArg().getUses().empty())
+        mightHaveReadHostSym = true;
     }
 
     // Populate the `copy` region if this is a `firstprivate`.
@@ -534,7 +597,7 @@ void DataSharingProcessor::doPrivatize(const semantics::Symbol *sym,
       // First block argument corresponding to the original/host value while
       // second block argument corresponding to the privatized value.
       mlir::Block *copyEntryBlock = firOpBuilder.createBlock(
-          &copyRegion, /*insertPt=*/{}, {symType, symType}, {symLoc, symLoc});
+          &copyRegion, /*insertPt=*/{}, {argType, argType}, {symLoc, symLoc});
       firOpBuilder.setInsertionPointToEnd(copyEntryBlock);
 
       auto addSymbol = [&](unsigned argIdx, bool force = false) {
@@ -565,7 +628,7 @@ void DataSharingProcessor::doPrivatize(const semantics::Symbol *sym,
 
   if (clauseOps) {
     clauseOps->privateSyms.push_back(mlir::SymbolRefAttr::get(privatizerOp));
-    clauseOps->privateVars.push_back(hsb.getAddr());
+    clauseOps->privateVars.push_back(privVal);
   }
 
   symToPrivatizer[sym] = privatizerOp;

flang/lib/Lower/OpenMP/DataSharingProcessor.h

@@ -86,7 +86,7 @@ class DataSharingProcessor {
   lower::pft::Evaluation &eval;
   bool shouldCollectPreDeterminedSymbols;
   bool useDelayedPrivatization;
-  bool callsInitClone = false;
+  bool mightHaveReadHostSym = false;
   lower::SymMap &symTable;
   OMPConstructSymbolVisitor visitor;