[mlir][OpenMP] Change the definition of omp.private

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

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.

---

Please read mlir changes first and then flang changes.

In flang lowering I box up all arrays and pass the boxes by reference so
that the existing code for reduction init and dealloc regions can be
shared.

The TODOs for pointers, derived types, characters etc will be resolved
in later patches in this same series (to be squashed into this one). I
separated it to make it easier to review.

Assumed rank was already broken before this patch.
I can't find any mention of "assumed rank" in the openmp standard so
I guess it is not prohibited.

Other than the omp.private operation definition changes, the test changes
are mostly down to slightly different codegen from re-using the reduction
init region. That code is already well tested so I didn't want to change it.

Author: tblah

flang/lib/Lower/OpenMP/DataSharingProcessor.cpp

@@ -12,13 +12,15 @@
 
 #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/HLFIRTools.h"
 #include "flang/Optimizer/Builder/Todo.h"
 #include "flang/Optimizer/HLFIR/HLFIROps.h"
+#include "flang/Semantics/attr.h"
 #include "flang/Semantics/tools.h"
 
 namespace Fortran {
@@ -85,28 +87,8 @@ 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
     }
 }
 
@@ -468,15 +450,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::Value privVal = hsb.getAddr();
+  mlir::Type allocType = fir::unwrapRefType(privVal.getType());
   mlir::Location symLoc = hsb.getAddr().getLoc();
   std::string privatizerName = sym->name().ToString() + ".privatizer";
   bool isFirstPrivate = sym->test(semantics::Symbol::Flag::OmpFirstPrivate);
 
+  if (mlir::isa<fir::PointerType>(hsb.getAddr().getType()))
+    TODO(symLoc, "Privatization of pointers");
+
+  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 +502,33 @@ 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.
+    const bool needsInitialization =
+        !fir::isa_trivial(allocType) ||
+        Fortran::lower::hasDefaultInitialization(sym->GetUltimate();
+    if (needsInitialization) {
+      mlir::Region &initRegion = result.getInitRegion();
+      mlir::Block *initBlock = firOpBuilder.createBlock(
+          &initRegion, /*insertPt=*/{}, {argType, argType}, {symLoc, symLoc});
+
+      if (fir::isa_char(allocType))
+        TODO(symLoc, "Privatization init of characters");
+      if (fir::isa_derived(allocType))
+        TODO(symLoc, "Privatization init of derived types");
+      if (Fortran::lower::hasDefaultInitialization(sym->GetUltimate()))
+        TODO(symLoc,
+             "Privatization init of symbol with default initialization");
+
+      populateByRefInitAndCleanupRegions(
+          firOpBuilder, symLoc, argType, /*scalarInitValue=*/nullptr, initBlock,
+          result.getInitPrivateArg(), result.getInitMoldArg(),
+          result.getDeallocRegion());
     }
 
     // Populate the `copy` region if this is a `firstprivate`.
@@ -534,7 +537,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 +568,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/Optimizer/CodeGen/CodeGenOpenMP.cpp

@@ -90,9 +90,45 @@ struct MapInfoOpConversion
     return mlir::success();
   }
 };
+
+// FIR op specific conversion for PrivateClauseOp that overwrites the default
+// OpenMP Dialect lowering, this allows FIR-aware lowering of types, required
+// for boxes because the OpenMP dialect conversion doesn't know anything about
+// FIR types.
+struct PrivateClauseOpConversion
+    : public OpenMPFIROpConversion<mlir::omp::PrivateClauseOp> {
+  using OpenMPFIROpConversion::OpenMPFIROpConversion;
+
+  llvm::LogicalResult
+  matchAndRewrite(mlir::omp::PrivateClauseOp curOp, OpAdaptor adaptor,
+                  mlir::ConversionPatternRewriter &rewriter) const override {
+    const fir::LLVMTypeConverter &converter = lowerTy();
+    mlir::Type convertedAllocType;
+    if (auto box = mlir::dyn_cast<fir::BaseBoxType>(curOp.getType())) {
+      // In LLVM codegen fir.box<> == fir.ref<fir.box<>> == llvm.ptr
+      // Here we really do want the actual structure
+      if (box.isAssumedRank())
+        TODO(curOp->getLoc(), "Privatize an assumed rank array");
+      unsigned rank = 0;
+      if (auto seqTy = mlir::dyn_cast<fir::SequenceType>(
+              fir::unwrapRefType(box.getEleTy())))
+        rank = seqTy.getShape().size();
+      convertedAllocType = converter.convertBoxTypeAsStruct(box, rank);
+    } else {
+      convertedAllocType = converter.convertType(adaptor.getType());
+    }
+    if (!convertedAllocType)
+      return mlir::failure();
+    rewriter.startOpModification(curOp);
+    curOp.setType(convertedAllocType);
+    rewriter.finalizeOpModification(curOp);
+    return mlir::success();
+  }
+};
 } // namespace
 
 void fir::populateOpenMPFIRToLLVMConversionPatterns(
     const LLVMTypeConverter &converter, mlir::RewritePatternSet &patterns) {
   patterns.add<MapInfoOpConversion>(converter);
+  patterns.add<PrivateClauseOpConversion>(converter);
 }

flang/test/Analysis/AliasAnalysis/alias-analysis-omp-private-allocatable.mlir

@@ -20,15 +20,7 @@
 // CHECK: ar2#0 <-> ar1#1: NoAlias
 // CHECK: ar2#1 <-> ar1#1: NoAlias
 
-omp.private {type = private} @_QFmysubEar1_private_ref_box_heap_Uxf64 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>> alloc {
-^bb0(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>):
-  %0 = fir.alloca !fir.box<!fir.heap<!fir.array<?xf64>>> {bindc_name = "ar1", pinned, uniq_name = "_QFmysubEar1"}
-  %5:2 = hlfir.declare %0 {fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFmysubEar1"} : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>) -> (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>)
-  omp.yield(%5#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>)
-} dealloc {
-^bb0(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>):
-  omp.yield
-}
+omp.private {type = private} @_QFmysubEar1_private_ref_box_heap_Uxf64 : !fir.box<!fir.heap<!fir.array<?xf64>>>
 func.func @testPrivateAllocatable(%arg0: !fir.ref<i32> {fir.bindc_name = "ns"}, %arg1: !fir.ref<i32> {fir.bindc_name = "ne"}) {
   %0 = fir.dummy_scope : !fir.dscope
   %1 = fir.alloca !fir.box<!fir.heap<!fir.array<?xf64>>> {bindc_name = "ar1", uniq_name = "_QFmysubEar1"}

flang/test/Analysis/AliasAnalysis/alias-analysis-omp-teams-distribute-private-ptr.mlir

@@ -17,18 +17,8 @@
 // CHECK-LABEL: Testing : "_QQmain"
 // CHECK-DAG:   ptrA#0 <-> ArrayA#0: MayAlias
 
-omp.private {type = private} @_QFEi_private_ref_i32 : !fir.ref<i32> alloc {
-^bb0(%arg0: !fir.ref<i32>):
-  %0 = fir.alloca i32 {bindc_name = "i", pinned, uniq_name = "_QFEi"}
-  %1:2 = hlfir.declare %0 {uniq_name = "_QFEi"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
-  omp.yield(%1#0 : !fir.ref<i32>)
-}
-omp.private {type = firstprivate} @_QFEptra_firstprivate_ref_box_ptr_Uxi32 : !fir.ref<!fir.box<!fir.ptr<!fir.array<?xi32>>>> alloc {
-^bb0(%arg0: !fir.ref<!fir.box<!fir.ptr<!fir.array<?xi32>>>>):
-  %0 = fir.alloca !fir.box<!fir.ptr<!fir.array<?xi32>>> {bindc_name = "ptra", pinned, uniq_name = "_QFEptra"}
-  %1:2 = hlfir.declare %0 {fortran_attrs = #fir.var_attrs<pointer>, uniq_name = "_QFEptra"} : (!fir.ref<!fir.box<!fir.ptr<!fir.array<?xi32>>>>) -> (!fir.ref<!fir.box<!fir.ptr<!fir.array<?xi32>>>>, !fir.ref<!fir.box<!fir.ptr<!fir.array<?xi32>>>>)
-  omp.yield(%1#0 : !fir.ref<!fir.box<!fir.ptr<!fir.array<?xi32>>>>)
-} copy {
+omp.private {type = private} @_QFEi_private_ref_i32 : i32
+omp.private {type = firstprivate} @_QFEptra_firstprivate_ref_box_ptr_Uxi32 : !fir.box<!fir.ptr<!fir.array<?xi32>>> copy {
 ^bb0(%arg0: !fir.ref<!fir.box<!fir.ptr<!fir.array<?xi32>>>>, %arg1: !fir.ref<!fir.box<!fir.ptr<!fir.array<?xi32>>>>):
   %0 = fir.load %arg0 : !fir.ref<!fir.box<!fir.ptr<!fir.array<?xi32>>>>
   fir.store %0 to %arg1 : !fir.ref<!fir.box<!fir.ptr<!fir.array<?xi32>>>>

flang/test/Analysis/AliasAnalysis/alias-analysis-omp-teams-distribute-private.mlir

@@ -21,26 +21,10 @@
 // CHECK-DAG: tmp_private_array#0 <-> unnamed_array#0: NoAlias
 // CHECK-DAG: tmp_private_array#1 <-> unnamed_array#0: NoAlias
 
-omp.private {type = private} @_QFEi_private_ref_i32 : !fir.ref<i32> alloc {
-^bb0(%arg0: !fir.ref<i32>):
-  %0 = fir.alloca i32 {bindc_name = "i", pinned, uniq_name = "_QFEi"}
-  %1:2 = hlfir.declare %0 {uniq_name = "_QFEi"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
-  omp.yield(%1#0 : !fir.ref<i32>)
-}
-omp.private {type = private} @_QFEj_private_ref_i32 : !fir.ref<i32> alloc {
-^bb0(%arg0: !fir.ref<i32>):
-  %0 = fir.alloca i32 {bindc_name = "j", pinned, uniq_name = "_QFEj"}
-  %1:2 = hlfir.declare %0 {uniq_name = "_QFEj"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
-  omp.yield(%1#0 : !fir.ref<i32>)
-}
-omp.private {type = private} @_QFEtmp_private_ref_2xi32 : !fir.ref<!fir.array<2xi32>> alloc {
-^bb0(%arg0: !fir.ref<!fir.array<2xi32>>):
-  %c2 = arith.constant 2 : index
-  %0 = fir.alloca !fir.array<2xi32> {bindc_name = "tmp", pinned, uniq_name = "_QFEtmp"}
-  %1 = fir.shape %c2 : (index) -> !fir.shape<1>
-  %2:2 = hlfir.declare %0(%1) {uniq_name = "_QFEtmp"} : (!fir.ref<!fir.array<2xi32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<2xi32>>, !fir.ref<!fir.array<2xi32>>)
-  omp.yield(%2#0 : !fir.ref<!fir.array<2xi32>>)
-}
+omp.private {type = private} @_QFEi_private_ref_i32 : i32
+omp.private {type = private} @_QFEj_private_ref_i32 : i32
+omp.private {type = private} @_QFEtmp_private_ref_2xi32 : !fir.array<2xi32>
+
 func.func @_QQmain() attributes {fir.bindc_name = "main"} {
   %0 = fir.address_of(@_QFEarraya) : !fir.ref<!fir.array<10x10xi32>>
   %c10 = arith.constant 10 : index