[Flang][OpenMP] Improve entry block argument creation and binding (llvm#110267)

The main purpose of this patch is to centralize the logic for creating
MLIR operation entry blocks and for binding them to the corresponding
symbols. This minimizes the chances of mixing arguments up for
operations having multiple entry block argument-generating clauses and
prevents divergence while binding arguments.

Some changes implemented to this end are:
- Split into two functions the creation of the entry block, and the
binding of its arguments and the corresponding Fortran symbol. This
enabled a significant simplification of the lowering of composite
constructs, where it's no longer necessary to manually ensure the lists
of arguments and symbols refer to the same variables in the same order
and also match the expected order by the `BlockArgOpenMPOpInterface`.
- Removed redundant and error-prone passing of types and locations from
`ClauseProcessor` methods. Instead, these are obtained from the values
in the appropriate clause operands structure. This also simplifies
argument lists of several lowering functions.
- Access block arguments of already created MLIR operations through the
`BlockArgOpenMPOpInterface` instead of directly indexing the argument
list of the operation, which is not scalable as more entry block
argument-generating clauses are added to an operation.
- Simplified the implementation of `genParallelOp` to no longer need to
define different callbacks depending on whether delayed privatization is
enabled.

Author: skatrak

flang/lib/Lower/OpenMP/ClauseProcessor.cpp

@@ -166,15 +166,11 @@ getIfClauseOperand(lower::AbstractConverter &converter,
 static void addUseDeviceClause(
     lower::AbstractConverter &converter, const omp::ObjectList &objects,
     llvm::SmallVectorImpl<mlir::Value> &operands,
-    llvm::SmallVectorImpl<mlir::Type> &useDeviceTypes,
-    llvm::SmallVectorImpl<mlir::Location> &useDeviceLocs,
     llvm::SmallVectorImpl<const semantics::Symbol *> &useDeviceSyms) {
   genObjectList(objects, converter, operands);
-  for (mlir::Value &operand : operands) {
+  for (mlir::Value &operand : operands)
     checkMapType(operand.getLoc(), operand.getType());
-    useDeviceTypes.push_back(operand.getType());
-    useDeviceLocs.push_back(operand.getLoc());
-  }
+
   for (const omp::Object &object : objects)
     useDeviceSyms.push_back(object.sym());
 }
@@ -832,14 +828,12 @@ bool ClauseProcessor::processDepend(mlir::omp::DependClauseOps &result) const {
 
 bool ClauseProcessor::processHasDeviceAddr(
     mlir::omp::HasDeviceAddrClauseOps &result,
-    llvm::SmallVectorImpl<mlir::Type> &isDeviceTypes,
-    llvm::SmallVectorImpl<mlir::Location> &isDeviceLocs,
-    llvm::SmallVectorImpl<const semantics::Symbol *> &isDeviceSymbols) const {
+    llvm::SmallVectorImpl<const semantics::Symbol *> &isDeviceSyms) const {
   return findRepeatableClause<omp::clause::HasDeviceAddr>(
       [&](const omp::clause::HasDeviceAddr &devAddrClause,
           const parser::CharBlock &) {
         addUseDeviceClause(converter, devAddrClause.v, result.hasDeviceAddrVars,
-                           isDeviceTypes, isDeviceLocs, isDeviceSymbols);
+                           isDeviceSyms);
       });
 }
 
@@ -864,14 +858,12 @@ bool ClauseProcessor::processIf(
 
 bool ClauseProcessor::processIsDevicePtr(
     mlir::omp::IsDevicePtrClauseOps &result,
-    llvm::SmallVectorImpl<mlir::Type> &isDeviceTypes,
-    llvm::SmallVectorImpl<mlir::Location> &isDeviceLocs,
-    llvm::SmallVectorImpl<const semantics::Symbol *> &isDeviceSymbols) const {
+    llvm::SmallVectorImpl<const semantics::Symbol *> &isDeviceSyms) const {
   return findRepeatableClause<omp::clause::IsDevicePtr>(
       [&](const omp::clause::IsDevicePtr &devPtrClause,
           const parser::CharBlock &) {
         addUseDeviceClause(converter, devPtrClause.v, result.isDevicePtrVars,
-                           isDeviceTypes, isDeviceLocs, isDeviceSymbols);
+                           isDeviceSyms);
       });
 }
 
@@ -892,9 +884,7 @@ void ClauseProcessor::processMapObjects(
     std::map<const semantics::Symbol *,
              llvm::SmallVector<OmpMapMemberIndicesData>> &parentMemberIndices,
     llvm::SmallVectorImpl<mlir::Value> &mapVars,
-    llvm::SmallVectorImpl<const semantics::Symbol *> *mapSyms,
-    llvm::SmallVectorImpl<mlir::Location> *mapSymLocs,
-    llvm::SmallVectorImpl<mlir::Type> *mapSymTypes) const {
+    llvm::SmallVectorImpl<const semantics::Symbol *> &mapSyms) const {
   fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
   for (const omp::Object &object : objects) {
     llvm::SmallVector<mlir::Value> bounds;
@@ -927,22 +917,15 @@ void ClauseProcessor::processMapObjects(
       addChildIndexAndMapToParent(object, parentMemberIndices, mapOp, semaCtx);
     } else {
       mapVars.push_back(mapOp);
-      if (mapSyms)
-        mapSyms->push_back(object.sym());
-      if (mapSymTypes)
-        mapSymTypes->push_back(baseOp.getType());
-      if (mapSymLocs)
-        mapSymLocs->push_back(baseOp.getLoc());
+      mapSyms.push_back(object.sym());
     }
   }
 }
 
 bool ClauseProcessor::processMap(
     mlir::Location currentLocation, lower::StatementContext &stmtCtx,
     mlir::omp::MapClauseOps &result,
-    llvm::SmallVectorImpl<const semantics::Symbol *> *mapSyms,
-    llvm::SmallVectorImpl<mlir::Location> *mapSymLocs,
-    llvm::SmallVectorImpl<mlir::Type> *mapSymTypes) const {
+    llvm::SmallVectorImpl<const semantics::Symbol *> *mapSyms) const {
   // We always require tracking of symbols, even if the caller does not,
   // so we create an optionally used local set of symbols when the mapSyms
   // argument is not present.
@@ -999,12 +982,11 @@ bool ClauseProcessor::processMap(
         }
         processMapObjects(stmtCtx, clauseLocation,
                           std::get<omp::ObjectList>(clause.t), mapTypeBits,
-                          parentMemberIndices, result.mapVars, ptrMapSyms,
-                          mapSymLocs, mapSymTypes);
+                          parentMemberIndices, result.mapVars, *ptrMapSyms);
       });
 
   insertChildMapInfoIntoParent(converter, parentMemberIndices, result.mapVars,
-                               *ptrMapSyms, mapSymTypes, mapSymLocs);
+                               *ptrMapSyms);
 
   return clauseFound;
 }
@@ -1027,16 +1009,15 @@ bool ClauseProcessor::processMotionClauses(lower::StatementContext &stmtCtx,
 
     processMapObjects(stmtCtx, clauseLocation, std::get<ObjectList>(clause.t),
                       mapTypeBits, parentMemberIndices, result.mapVars,
-                      &mapSymbols);
+                      mapSymbols);
   };
 
   bool clauseFound = findRepeatableClause<omp::clause::To>(callbackFn);
   clauseFound =
       findRepeatableClause<omp::clause::From>(callbackFn) || clauseFound;
 
   insertChildMapInfoIntoParent(converter, parentMemberIndices, result.mapVars,
-                               mapSymbols,
-                               /*mapSymTypes=*/nullptr, /*mapSymLocs=*/nullptr);
+                               mapSymbols);
   return clauseFound;
 }
 
@@ -1054,34 +1035,24 @@ bool ClauseProcessor::processNontemporal(
 
 bool ClauseProcessor::processReduction(
     mlir::Location currentLocation, mlir::omp::ReductionClauseOps &result,
-    llvm::SmallVectorImpl<mlir::Type> *outReductionTypes,
-    llvm::SmallVectorImpl<const semantics::Symbol *> *outReductionSyms) const {
+    llvm::SmallVectorImpl<const semantics::Symbol *> &outReductionSyms) const {
   return findRepeatableClause<omp::clause::Reduction>(
       [&](const omp::clause::Reduction &clause, const parser::CharBlock &) {
         llvm::SmallVector<mlir::Value> reductionVars;
         llvm::SmallVector<bool> reduceVarByRef;
         llvm::SmallVector<mlir::Attribute> reductionDeclSymbols;
         llvm::SmallVector<const semantics::Symbol *> reductionSyms;
         ReductionProcessor rp;
-        rp.addDeclareReduction(
-            currentLocation, converter, clause, reductionVars, reduceVarByRef,
-            reductionDeclSymbols, outReductionSyms ? &reductionSyms : nullptr);
+        rp.addDeclareReduction(currentLocation, converter, clause,
+                               reductionVars, reduceVarByRef,
+                               reductionDeclSymbols, reductionSyms);
 
         // Copy local lists into the output.
         llvm::copy(reductionVars, std::back_inserter(result.reductionVars));
         llvm::copy(reduceVarByRef, std::back_inserter(result.reductionByref));
         llvm::copy(reductionDeclSymbols,
                    std::back_inserter(result.reductionSyms));
-
-        if (outReductionTypes) {
-          outReductionTypes->reserve(outReductionTypes->size() +
-                                     reductionVars.size());
-          llvm::transform(reductionVars, std::back_inserter(*outReductionTypes),
-                          [](mlir::Value v) { return v.getType(); });
-        }
-
-        if (outReductionSyms)
-          llvm::copy(reductionSyms, std::back_inserter(*outReductionSyms));
+        llvm::copy(reductionSyms, std::back_inserter(outReductionSyms));
       });
 }
 
@@ -1107,8 +1078,6 @@ bool ClauseProcessor::processEnter(
 
 bool ClauseProcessor::processUseDeviceAddr(
     lower::StatementContext &stmtCtx, mlir::omp::UseDeviceAddrClauseOps &result,
-    llvm::SmallVectorImpl<mlir::Type> &useDeviceTypes,
-    llvm::SmallVectorImpl<mlir::Location> &useDeviceLocs,
     llvm::SmallVectorImpl<const semantics::Symbol *> &useDeviceSyms) const {
   std::map<const semantics::Symbol *,
            llvm::SmallVector<OmpMapMemberIndicesData>>
@@ -1122,19 +1091,16 @@ bool ClauseProcessor::processUseDeviceAddr(
             llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_FROM;
         processMapObjects(stmtCtx, location, clause.v, mapTypeBits,
                           parentMemberIndices, result.useDeviceAddrVars,
-                          &useDeviceSyms, &useDeviceLocs, &useDeviceTypes);
+                          useDeviceSyms);
       });
 
   insertChildMapInfoIntoParent(converter, parentMemberIndices,
-                               result.useDeviceAddrVars, useDeviceSyms,
-                               &useDeviceTypes, &useDeviceLocs);
+                               result.useDeviceAddrVars, useDeviceSyms);
   return clauseFound;
 }
 
 bool ClauseProcessor::processUseDevicePtr(
     lower::StatementContext &stmtCtx, mlir::omp::UseDevicePtrClauseOps &result,
-    llvm::SmallVectorImpl<mlir::Type> &useDeviceTypes,
-    llvm::SmallVectorImpl<mlir::Location> &useDeviceLocs,
     llvm::SmallVectorImpl<const semantics::Symbol *> &useDeviceSyms) const {
   std::map<const semantics::Symbol *,
            llvm::SmallVector<OmpMapMemberIndicesData>>
@@ -1148,12 +1114,11 @@ bool ClauseProcessor::processUseDevicePtr(
             llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_FROM;
         processMapObjects(stmtCtx, location, clause.v, mapTypeBits,
                           parentMemberIndices, result.useDevicePtrVars,
-                          &useDeviceSyms, &useDeviceLocs, &useDeviceTypes);
+                          useDeviceSyms);
       });
 
   insertChildMapInfoIntoParent(converter, parentMemberIndices,
-                               result.useDevicePtrVars, useDeviceSyms,
-                               &useDeviceTypes, &useDeviceLocs);
+                               result.useDevicePtrVars, useDeviceSyms);
   return clauseFound;
 }
 

flang/lib/Lower/OpenMP/ClauseProcessor.h

@@ -68,9 +68,7 @@ class ClauseProcessor {
                     mlir::omp::FinalClauseOps &result) const;
   bool processHasDeviceAddr(
       mlir::omp::HasDeviceAddrClauseOps &result,
-      llvm::SmallVectorImpl<mlir::Type> &isDeviceTypes,
-      llvm::SmallVectorImpl<mlir::Location> &isDeviceLocs,
-      llvm::SmallVectorImpl<const semantics::Symbol *> &isDeviceSymbols) const;
+      llvm::SmallVectorImpl<const semantics::Symbol *> &isDeviceSyms) const;
   bool processHint(mlir::omp::HintClauseOps &result) const;
   bool processMergeable(mlir::omp::MergeableClauseOps &result) const;
   bool processNowait(mlir::omp::NowaitClauseOps &result) const;
@@ -104,43 +102,33 @@ class ClauseProcessor {
                  mlir::omp::IfClauseOps &result) const;
   bool processIsDevicePtr(
       mlir::omp::IsDevicePtrClauseOps &result,
-      llvm::SmallVectorImpl<mlir::Type> &isDeviceTypes,
-      llvm::SmallVectorImpl<mlir::Location> &isDeviceLocs,
-      llvm::SmallVectorImpl<const semantics::Symbol *> &isDeviceSymbols) const;
+      llvm::SmallVectorImpl<const semantics::Symbol *> &isDeviceSyms) const;
   bool
   processLink(llvm::SmallVectorImpl<DeclareTargetCapturePair> &result) const;
 
   // This method is used to process a map clause.
-  // The optional parameters - mapSymTypes, mapSymLocs & mapSyms are used to
-  // store the original type, location and Fortran symbol for the map operands.
-  // They may be used later on to create the block_arguments for some of the
-  // target directives that require it.
-  bool processMap(
-      mlir::Location currentLocation, lower::StatementContext &stmtCtx,
-      mlir::omp::MapClauseOps &result,
-      llvm::SmallVectorImpl<const semantics::Symbol *> *mapSyms = nullptr,
-      llvm::SmallVectorImpl<mlir::Location> *mapSymLocs = nullptr,
-      llvm::SmallVectorImpl<mlir::Type> *mapSymTypes = nullptr) const;
+  // The optional parameter mapSyms is used to store the original Fortran symbol
+  // for the map operands. It may be used later on to create the block_arguments
+  // for some of the directives that require it.
+  bool processMap(mlir::Location currentLocation,
+                  lower::StatementContext &stmtCtx,
+                  mlir::omp::MapClauseOps &result,
+                  llvm::SmallVectorImpl<const semantics::Symbol *> *mapSyms =
+                      nullptr) const;
   bool processMotionClauses(lower::StatementContext &stmtCtx,
                             mlir::omp::MapClauseOps &result);
   bool processNontemporal(mlir::omp::NontemporalClauseOps &result) const;
   bool processReduction(
       mlir::Location currentLocation, mlir::omp::ReductionClauseOps &result,
-      llvm::SmallVectorImpl<mlir::Type> *reductionTypes = nullptr,
-      llvm::SmallVectorImpl<const semantics::Symbol *> *reductionSyms =
-          nullptr) const;
+      llvm::SmallVectorImpl<const semantics::Symbol *> &reductionSyms) const;
   bool processTo(llvm::SmallVectorImpl<DeclareTargetCapturePair> &result) const;
   bool processUseDeviceAddr(
       lower::StatementContext &stmtCtx,
       mlir::omp::UseDeviceAddrClauseOps &result,
-      llvm::SmallVectorImpl<mlir::Type> &useDeviceTypes,
-      llvm::SmallVectorImpl<mlir::Location> &useDeviceLocs,
       llvm::SmallVectorImpl<const semantics::Symbol *> &useDeviceSyms) const;
   bool processUseDevicePtr(
       lower::StatementContext &stmtCtx,
       mlir::omp::UseDevicePtrClauseOps &result,
-      llvm::SmallVectorImpl<mlir::Type> &useDeviceTypes,
-      llvm::SmallVectorImpl<mlir::Location> &useDeviceLocs,
       llvm::SmallVectorImpl<const semantics::Symbol *> &useDeviceSyms) const;
 
   // Call this method for these clauses that should be supported but are not
@@ -181,9 +169,7 @@ class ClauseProcessor {
       std::map<const semantics::Symbol *,
                llvm::SmallVector<OmpMapMemberIndicesData>> &parentMemberIndices,
       llvm::SmallVectorImpl<mlir::Value> &mapVars,
-      llvm::SmallVectorImpl<const semantics::Symbol *> *mapSyms,
-      llvm::SmallVectorImpl<mlir::Location> *mapSymLocs = nullptr,
-      llvm::SmallVectorImpl<mlir::Type> *mapSymTypes = nullptr) const;
+      llvm::SmallVectorImpl<const semantics::Symbol *> &mapSyms) const;
 
   lower::AbstractConverter &converter;
   semantics::SemanticsContext &semaCtx;