[Flang][OpenMP] Lowering of host-evaluated clauses (llvm#116219)

This patch adds support for lowering OpenMP clauses and expressions
attached to constructs nested inside of a target region that need to be
evaluated in the host device. This is done through the use of the
`OpenMP_HostEvalClause` `omp.target` set of operands and entry block
arguments.

When lowering clauses for a target construct, a more involved
`processHostEvalClauses()` function is called, which looks at the
current and potentially other nested constructs in order to find and
lower clauses that need to be processed outside of the `omp.target`
operation under construction. This populates an instance of a global
structure with the resulting MLIR values.

The resulting list of host-evaluated values is used to initialize the
`host_eval` operands when constructing the `omp.target` operation, and
then replaced with the corresponding block arguments after creating that
operation's region.

Afterwards, while lowering nested operations, those that might
potentially be evaluated on the host (i.e. `num_teams`, `thread_limit`,
`num_threads` and `collapse`) check first whether there is an active
global host-evaluated information structure and whether it holds values
referring to these clauses. If that is the case, the stored values
(`omp.target` entry block arguments at that stage) are used instead of
lowering these clauses again.

Author: skatrak

flang/include/flang/Common/OpenMP-utils.h

@@ -34,6 +34,7 @@ struct EntryBlockArgsEntry {
 /// Structure holding the information needed to create and bind entry block
 /// arguments associated to all clauses that can define them.
 struct EntryBlockArgs {
+  llvm::ArrayRef<mlir::Value> hostEvalVars;
   EntryBlockArgsEntry inReduction;
   EntryBlockArgsEntry map;
   EntryBlockArgsEntry priv;
@@ -49,18 +50,25 @@ struct EntryBlockArgs {
   }
 
   auto getSyms() const {
-    return llvm::concat<const Fortran::semantics::Symbol *const>(
-        inReduction.syms, map.syms, priv.syms, reduction.syms,
-        taskReduction.syms, useDeviceAddr.syms, useDevicePtr.syms);
+    return llvm::concat<const semantics::Symbol *const>(inReduction.syms,
+        map.syms, priv.syms, reduction.syms, taskReduction.syms,
+        useDeviceAddr.syms, useDevicePtr.syms);
   }
 
   auto getVars() const {
-    return llvm::concat<const mlir::Value>(inReduction.vars, map.vars,
-        priv.vars, reduction.vars, taskReduction.vars, useDeviceAddr.vars,
-        useDevicePtr.vars);
+    return llvm::concat<const mlir::Value>(hostEvalVars, inReduction.vars,
+        map.vars, priv.vars, reduction.vars, taskReduction.vars,
+        useDeviceAddr.vars, useDevicePtr.vars);
   }
 };
 
+/// Create an entry block for the given region, including the clause-defined
+/// arguments specified.
+///
+/// \param [in] builder - MLIR operation builder.
+/// \param [in]    args - entry block arguments information for the given
+///                       operation.
+/// \param [in]  region - Empty region in which to create the entry block.
 mlir::Block *genEntryBlock(
     mlir::OpBuilder &builder, const EntryBlockArgs &args, mlir::Region &region);
 } // namespace Fortran::common::openmp

flang/lib/Common/OpenMP-utils.cpp

@@ -18,10 +18,10 @@ mlir::Block *genEntryBlock(mlir::OpBuilder &builder, const EntryBlockArgs &args,
 
   llvm::SmallVector<mlir::Type> types;
   llvm::SmallVector<mlir::Location> locs;
-  unsigned numVars = args.inReduction.vars.size() + args.map.vars.size() +
-      args.priv.vars.size() + args.reduction.vars.size() +
-      args.taskReduction.vars.size() + args.useDeviceAddr.vars.size() +
-      args.useDevicePtr.vars.size();
+  unsigned numVars = args.hostEvalVars.size() + args.inReduction.vars.size() +
+      args.map.vars.size() + args.priv.vars.size() +
+      args.reduction.vars.size() + args.taskReduction.vars.size() +
+      args.useDeviceAddr.vars.size() + args.useDevicePtr.vars.size();
   types.reserve(numVars);
   locs.reserve(numVars);
 
@@ -34,6 +34,7 @@ mlir::Block *genEntryBlock(mlir::OpBuilder &builder, const EntryBlockArgs &args,
 
   // Populate block arguments in clause name alphabetical order to match
   // expected order by the BlockArgOpenMPOpInterface.
+  extractTypeLoc(args.hostEvalVars);
   extractTypeLoc(args.inReduction.vars);
   extractTypeLoc(args.map.vars);
   extractTypeLoc(args.priv.vars);