[flang-commits] [flang] [flang][OpenMP] Lower DO, SIMD, and DO SIMD metadirective variants (PR #218555)

[via flang-commits]

================
@@ -6822,23 +6834,402 @@ static void genOMP(lower::AbstractConverter &converter, lower::SymMap &symTable,
 }
 
 namespace {
-struct MetadirectiveCandidate {
-  MetadirectiveCandidate(const parser::OmpDirectiveSpecification *spec,
-                         llvm::omp::VariantMatchInfo vmi, bool isExplicit,
-                         std::optional<semantics::omp::DynamicUserCondition>
-                             dynamicCond = std::nullopt,
-                         bool conditionShouldBeTrue = true)
-      : spec(spec), vmi(vmi), isExplicit(isExplicit), dynamicCond(dynamicCond),
-        conditionShouldBeTrue(conditionShouldBeTrue) {}
-
-  const parser::OmpDirectiveSpecification *spec = nullptr;
-  llvm::omp::VariantMatchInfo vmi;
-  bool isExplicit = false;
-  std::optional<semantics::omp::DynamicUserCondition> dynamicCond;
-  bool conditionShouldBeTrue = true;
+struct SplicedAssociatedEvaluations {
+  using Iterator = lower::pft::EvaluationList::iterator;
+
+  void record(lower::pft::EvaluationList &parent, Iterator evaluation) {
+    assert((!parentList || parentList == &parent) &&
+           "associated evaluations have different parents");
+    parentList = &parent;
+    evaluations.emplace_back(evaluation, std::next(evaluation));
+  }
+
+  void restore(lower::pft::EvaluationList &nested) {
+    if (evaluations.empty())
+      return;
+    assert(parentList && "missing parent evaluation list");
+    // A saved successor may also have been spliced. Restore in reverse order
+    // so every insertion point is back in the parent list before it is used.
+    for (auto &entry : llvm::reverse(evaluations)) {
+      entry.first->skipNextLowering = true;
+      parentList->splice(entry.second, nested, entry.first);
+    }
+    if (entryEvaluation) {
+      entryEvaluation->isNewBlock = true;
+      entryEvaluation->block = entryBlock;
+    }
+  }
+
+  void suppressEntryBlock(lower::pft::Evaluation &evaluation) {
+    assert(!entryEvaluation && evaluation.isNewBlock && evaluation.block &&
+           "invalid associated entry evaluation");
+    // Do not let either cloned loop arm enter a function-region block. The
+    // metadirective selection will be placed in this block for an active ENTRY.
+    entryEvaluation = &evaluation;
+    entryBlock = evaluation.block;
+    evaluation.isNewBlock = false;
+    evaluation.block = nullptr;
+  }
+
+  mlir::Block *getEntryBlock() const { return entryBlock; }
+
+private:
+  lower::pft::EvaluationList *parentList = nullptr;
+  llvm::SmallVector<std::pair<Iterator, Iterator>, 4> evaluations;
+  lower::pft::Evaluation *entryEvaluation = nullptr;
+  mlir::Block *entryBlock = nullptr;
 };
 } // namespace
 
+static bool isSupportedMetadirectiveLoopCompilerDirective(
+    const parser::CompilerDirective &directive) {
+  // Only directives known not to emit executable operations may be processed
+  // before runtime selection. New variants remain unsupported until classified.
+  using SupportedDirectives = std::tuple<
+      std::list<parser::CompilerDirective::IgnoreTKR>,
+      parser::CompilerDirective::LoopCount,
+      std::list<parser::CompilerDirective::AssumeAligned>,
+      parser::CompilerDirective::VectorAlways,
+      parser::CompilerDirective::VectorLength,
+      std::list<parser::CompilerDirective::NameValue>,
+      parser::CompilerDirective::Unroll,
+      parser::CompilerDirective::UnrollAndJam,
+      parser::CompilerDirective::Unrecognized,
+      parser::CompilerDirective::NoVector, parser::CompilerDirective::NoUnroll,
+      parser::CompilerDirective::NoUnrollAndJam,
+      parser::CompilerDirective::ForceInline, parser::CompilerDirective::Inline,
+      parser::CompilerDirective::NoInline,
+      parser::CompilerDirective::InlineAlways, parser::CompilerDirective::IVDep,
+      parser::CompilerDirective::Simd>;
+
+  return common::visit(
+      [](const auto &value) {
+        using T = std::decay_t<decltype(value)>;
+        return common::HasMember<T, SupportedDirectives>;
+      },
+      directive.u);
+}
+
+static bool
+isIgnorableMetadirectiveLoopAssociationEval(lower::pft::Evaluation &eval) {
+  const auto *directive = eval.getIf<parser::CompilerDirective>();
+  return eval.isEndStmt() ||
+         (directive &&
+          isSupportedMetadirectiveLoopCompilerDirective(*directive));
+}
+
+static bool
+isUnsupportedMetadirectiveLoopAssociationEval(lower::pft::Evaluation &eval) {
+  if (const auto *directive = eval.getIf<parser::CompilerDirective>())
+    return !isSupportedMetadirectiveLoopCompilerDirective(*directive);
+  return eval.isOtherStmt() ||
+         (eval.isDirective() && !eval.isExecutableDirective());
+}
+
+/// A loop-associated metadirective is lowered like a real loop construct, but
+/// the PFT leaves its associated loop nest as the following sibling instead of
+/// nesting it underneath. Splice that sibling into the metadirective's own
+/// nested evaluations so the shared loop-lowering path can find it. Return
+/// nullptr if no associated DO loop follows.
+static lower::pft::Evaluation *spliceAssociatedDoEval(
+    lower::pft::Evaluation &eval,
+    SplicedAssociatedEvaluations *splicedEvaluations = nullptr,
+    lower::pft::Evaluation **unsupportedInterveningEval = nullptr) {
+  if (unsupportedInterveningEval)
+    *unsupportedInterveningEval = nullptr;
+
+  if (eval.hasNestedEvaluations()) {
+    auto nestedIt =
+        llvm::find_if(eval.getNestedEvaluations(), [](auto &nested) {
+          return !isIgnorableMetadirectiveLoopAssociationEval(nested);
+        });
+    if (nestedIt != eval.getNestedEvaluations().end()) {
+      if (nestedIt->getIf<parser::DoConstruct>())
+        return &*nestedIt;
+      if (unsupportedInterveningEval &&
+          isUnsupportedMetadirectiveLoopAssociationEval(*nestedIt))
+        *unsupportedInterveningEval = &*nestedIt;
----------------
MattPD wrote:

Confirmed, the nested association cases are now covered.

https://github.com/llvm/llvm-project/pull/218555