[flang][OpenMP] Use StylizedInstance in converted clauses (#171907)

Invent `StylizedInstance` class to store special variables together with
the instantiated expression in omp::clause::Initializer. This will
eliminate the need for visiting the original AST nodes in lowering to
MLIR.

flang/include/flang/Lower/OpenMP/Clauses.h

@@ -113,6 +113,8 @@ Object makeObject(const parser::Designator &dsg,
                   semantics::SemanticsContext &semaCtx);
 Object makeObject(const parser::StructureComponent &comp,
                   semantics::SemanticsContext &semaCtx);
+Object makeObject(const parser::EntityDecl &decl,
+                  semantics::SemanticsContext &semaCtx);
 
 inline auto makeObjectFn(semantics::SemanticsContext &semaCtx) {
   return [&](auto &&s) { return makeObject(s, semaCtx); };
@@ -172,6 +174,7 @@ std::optional<Object> getBaseObject(const Object &object,
                                     semantics::SemanticsContext &semaCtx);
 
 namespace clause {
+using StylizedInstance = tomp::type::StylizedInstanceT<IdTy, ExprTy>;
 using Range = tomp::type::RangeT<ExprTy>;
 using Mapper = tomp::type::MapperT<IdTy, ExprTy>;
 using Iterator = tomp::type::IteratorT<TypeTy, IdTy, ExprTy>;

flang/include/flang/Semantics/openmp-utils.h

@@ -97,6 +97,8 @@ const SomeExpr *HasStorageOverlap(
     const SomeExpr &base, llvm::ArrayRef<SomeExpr> exprs);
 bool IsAssignment(const parser::ActionStmt *x);
 bool IsPointerAssignment(const evaluate::Assignment &x);
+
+MaybeExpr MakeEvaluateExpr(const parser::OmpStylizedInstance &inp);
 } // namespace omp
 } // namespace Fortran::semantics
 

flang/lib/Lower/OpenMP/ClauseProcessor.cpp

@@ -383,36 +383,36 @@ bool ClauseProcessor::processInclusive(
 }
 
 bool ClauseProcessor::processInitializer(
-    lower::SymMap &symMap, const parser::OmpClause::Initializer &inp,
+    lower::SymMap &symMap,
     ReductionProcessor::GenInitValueCBTy &genInitValueCB) const {
   if (auto *clause = findUniqueClause<omp::clause::Initializer>()) {
     genInitValueCB = [&, clause](fir::FirOpBuilder &builder, mlir::Location loc,
                                  mlir::Type type, mlir::Value ompOrig) {
       lower::SymMapScope scope(symMap);
-      const parser::OmpInitializerExpression &iexpr = inp.v.v;
-      const parser::OmpStylizedInstance &styleInstance = iexpr.v.front();
-      const std::list<parser::OmpStylizedDeclaration> &declList =
-          std::get<std::list<parser::OmpStylizedDeclaration>>(styleInstance.t);
       mlir::Value ompPrivVar;
-      for (const parser::OmpStylizedDeclaration &decl : declList) {
-        auto &name = std::get<parser::ObjectName>(decl.var.t);
-        assert(name.symbol && "Name does not have a symbol");
+      const clause::StylizedInstance &inst = clause->v.front();
+
+      for (const Object &object :
+           std::get<clause::StylizedInstance::Variables>(inst.t)) {
         mlir::Value addr = builder.createTemporary(loc, ompOrig.getType());
         fir::StoreOp::create(builder, loc, ompOrig, addr);
         fir::FortranVariableFlagsEnum extraFlags = {};
         fir::FortranVariableFlagsAttr attributes =
-            Fortran::lower::translateSymbolAttributes(builder.getContext(),
-                                                      *name.symbol, extraFlags);
-        auto declareOp = hlfir::DeclareOp::create(
-            builder, loc, addr, name.ToString(), nullptr, {}, nullptr, nullptr,
-            0, attributes);
-        if (name.ToString() == "omp_priv")
+            Fortran::lower::translateSymbolAttributes(
+                builder.getContext(), *object.sym(), extraFlags);
+        std::string name = object.sym()->name().ToString();
+        auto declareOp =
+            hlfir::DeclareOp::create(builder, loc, addr, name, nullptr, {},
+                                     nullptr, nullptr, 0, attributes);
+        if (name == "omp_priv")
           ompPrivVar = declareOp.getResult(0);
-        symMap.addVariableDefinition(*name.symbol, declareOp);
+        symMap.addVariableDefinition(*object.sym(), declareOp);
       }
+
       // Lower the expression/function call
       lower::StatementContext stmtCtx;
-      const semantics::SomeExpr &initExpr = clause->v.front();
+      const semantics::SomeExpr &initExpr =
+          std::get<clause::StylizedInstance::Instance>(inst.t);
       mlir::Value result = common::visit(
           common::visitors{
               [&](const evaluate::ProcedureRef &procRef) -> mlir::Value {

flang/lib/Lower/OpenMP/ClauseProcessor.h

@@ -20,7 +20,7 @@
 #include "flang/Lower/OpenMP/Clauses.h"
 #include "flang/Lower/Support/ReductionProcessor.h"
 #include "flang/Optimizer/Builder/Todo.h"
-#include "flang/Parser/parse-tree.h"
+#include "flang/Parser/char-block.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 
 namespace fir {
@@ -89,7 +89,7 @@ public:
   bool processInclusive(mlir::Location currentLocation,
                         mlir::omp::InclusiveClauseOps &result) const;
   bool processInitializer(
-      lower::SymMap &symMap, const parser::OmpClause::Initializer &inp,
+      lower::SymMap &symMap,
       ReductionProcessor::GenInitValueCBTy &genInitValueCB) const;
   bool processMergeable(mlir::omp::MergeableClauseOps &result) const;
   bool processNogroup(mlir::omp::NogroupClauseOps &result) const;

flang/lib/Lower/OpenMP/Clauses.cpp

@@ -13,6 +13,7 @@
 #include "flang/Parser/parse-tree.h"
 #include "flang/Semantics/expression.h"
 #include "flang/Semantics/openmp-modifiers.h"
+#include "flang/Semantics/openmp-utils.h"
 #include "flang/Semantics/symbol.h"
 
 #include <list>
@@ -128,6 +129,11 @@ Object makeObject(const parser::OmpObject &object,
   return makeObject(std::get<parser::Designator>(object.u), semaCtx);
 }
 
+Object makeObject(const parser::EntityDecl &decl,
+                  semantics::SemanticsContext &semaCtx) {
+  return makeObject(std::get<parser::ObjectName>(decl.t), semaCtx);
+}
+
 ObjectList makeObjects(const parser::OmpArgumentList &objects,
                        semantics::SemanticsContext &semaCtx) {
   return makeList(objects.v, [&](const parser::OmpArgument &arg) {
@@ -275,12 +281,10 @@ makeIteratorSpecifiers(const parser::OmpIteratorSpecifier &inp,
   auto &tds = std::get<parser::TypeDeclarationStmt>(inp.t);
   auto &entities = std::get<std::list<parser::EntityDecl>>(tds.t);
   for (const parser::EntityDecl &ed : entities) {
-    auto &name = std::get<parser::ObjectName>(ed.t);
-    assert(name.symbol && "Expecting symbol for iterator variable");
-    auto *stype = name.symbol->GetType();
-    assert(stype && "Expecting symbol type");
-    IteratorSpecifier spec{{evaluate::DynamicType::From(*stype),
-                            makeObject(name, semaCtx), range}};
+    auto *symbol = std::get<parser::ObjectName>(ed.t).symbol;
+    auto *type = DEREF(symbol).GetType();
+    IteratorSpecifier spec{{evaluate::DynamicType::From(DEREF(type)),
+                            makeObject(ed, semaCtx), range}};
     specifiers.emplace_back(std::move(spec));
   }
 
@@ -983,19 +987,24 @@ Initializer make(const parser::OmpClause::Initializer &inp,
                  semantics::SemanticsContext &semaCtx) {
   const parser::OmpInitializerExpression &iexpr = inp.v.v;
   Initializer initializer;
-  for (const parser::OmpStylizedInstance &styleInstance : iexpr.v) {
-    auto &instance =
-        std::get<parser::OmpStylizedInstance::Instance>(styleInstance.t);
-    if (const auto *as = std::get_if<parser::AssignmentStmt>(&instance.u)) {
-      auto &expr = std::get<parser::Expr>(as->t);
-      initializer.v.push_back(makeExpr(expr, semaCtx));
-    } else if (const auto *call = std::get_if<parser::CallStmt>(&instance.u)) {
-      assert(call->typedCall && "Expecting typedCall");
-      const auto &procRef = *call->typedCall;
-      initializer.v.push_back(semantics::SomeExpr(procRef));
-    } else {
-      llvm_unreachable("Unexpected initializer");
-    }
+
+  for (const parser::OmpStylizedInstance &sinst : iexpr.v) {
+    ObjectList variables;
+    llvm::transform(
+        std::get<std::list<parser::OmpStylizedDeclaration>>(sinst.t),
+        std::back_inserter(variables),
+        [&](const parser::OmpStylizedDeclaration &s) {
+          return makeObject(s.var, semaCtx);
+        });
+
+    SomeExpr instance = [&]() {
+      if (auto &&expr = semantics::omp::MakeEvaluateExpr(sinst))
+        return std::move(*expr);
+      llvm_unreachable("Expecting expression instance");
+    }();
+
+    initializer.v.push_back(
+        StylizedInstance{{std::move(variables), std::move(instance)}});
   }
 
   return initializer;

flang/lib/Lower/OpenMP/OpenMP.cpp

@@ -3764,9 +3764,7 @@ static void genOMP(
     List<Clause> clauses = makeClauses(initializer, semaCtx);
     ReductionProcessor::GenInitValueCBTy genInitValueCB;
     ClauseProcessor cp(converter, semaCtx, clauses);
-    const parser::OmpClause::Initializer &iclause{
-        std::get<parser::OmpClause::Initializer>(initializer.v.front().u)};
-    cp.processInitializer(symTable, iclause, genInitValueCB);
+    cp.processInitializer(symTable, genInitValueCB);
     const auto &identifier =
         std::get<parser::OmpReductionIdentifier>(specifier.t);
     const auto &designator =

flang/lib/Semantics/openmp-utils.cpp

@@ -496,4 +496,24 @@ bool IsPointerAssignment(const evaluate::Assignment &x) {
   return std::holds_alternative<evaluate::Assignment::BoundsSpec>(x.u) ||
       std::holds_alternative<evaluate::Assignment::BoundsRemapping>(x.u);
 }
+
+MaybeExpr MakeEvaluateExpr(const parser::OmpStylizedInstance &inp) {
+  auto &instance = std::get<parser::OmpStylizedInstance::Instance>(inp.t);
+
+  return common::visit( //
+      common::visitors{
+          [&](const parser::AssignmentStmt &s) -> MaybeExpr {
+            return GetEvaluateExpr(std::get<parser::Expr>(s.t));
+          },
+          [&](const parser::CallStmt &s) -> MaybeExpr {
+            assert(s.typedCall && "Expecting typedCall");
+            const auto &procRef = *s.typedCall;
+            return SomeExpr(procRef);
+          },
+          [&](const common::Indirection<parser::Expr> &s) -> MaybeExpr {
+            return GetEvaluateExpr(s.value());
+          },
+      },
+      instance.u);
+}
 } // namespace Fortran::semantics::omp

llvm/include/llvm/Frontend/OpenMP/ClauseT.h

@@ -189,6 +189,14 @@ template <typename I, typename E> using ObjectListT = ListT<ObjectT<I, E>>;
 
 using DirectiveName = llvm::omp::Directive;
 
+template <typename I, typename E> //
+struct StylizedInstanceT {
+  using Variables = ObjectListT<I, E>;
+  using Instance = E;
+  using TupleTrait = std::true_type;
+  std::tuple<Variables, Instance> t;
+};
+
 template <typename I, typename E> //
 struct DefinedOperatorT {
   struct DefinedOpName {
@@ -762,8 +770,7 @@ struct InitT {
 // V5.2: [5.5.4] `initializer` clause
 template <typename T, typename I, typename E> //
 struct InitializerT {
-  using InitializerExpr = E;
-  using List = ListT<InitializerExpr>;
+  using List = ListT<type::StylizedInstanceT<I, E>>;
   using WrapperTrait = std::true_type;
   List v;
 };