[flang-commits] [flang] [flang][OpenMP] Support lowering of metadirective (part 1) (PR #193664)

[Abid Qadeer via flang-commits]

================
@@ -4413,11 +4413,258 @@ static void genOMP(lower::AbstractConverter &converter, lower::SymMap &symTable,
   // support the case of threadprivate variable declared in module.
 }
 
+namespace {
+struct TargetOMPContext final : public llvm::omp::OMPContext {
+  TargetOMPContext(mlir::ModuleOp module,
+                   llvm::ArrayRef<llvm::omp::TraitProperty> constructTraits)
+      // DeviceNum is set to -1 (unknown) because the
+      // target_device={device_num()} selector is not yet supported. OMPContext
+      // uses DeviceNum > -1 to activate target_device traits from the offload
+      // triple; without a concrete device number those traits are left
+      // inactive.
+      : OMPContext(isDeviceCompilation(module), fir::getTargetTriple(module),
+                   getOffloadTargetTriple(module),
+                   /*DeviceNum=*/-1),
+        targetFeatures(fir::getTargetFeatures(module)) {
+    for (llvm::omp::TraitProperty trait : constructTraits)
+      addTrait(trait);
+  }
+
+  bool matchesISATrait(llvm::StringRef rawString) const override {
+    if (!targetFeatures || targetFeatures.nullOrEmpty())
+      return false;
+    return targetFeatures.contains(("+" + rawString).str());
+  }
+
+private:
+  static bool isDeviceCompilation(mlir::ModuleOp module) {
+    return llvm::cast<mlir::omp::OffloadModuleInterface>(*module.getOperation())
+        .getIsTargetDevice();
+  }
+
+  static llvm::Triple getOffloadTargetTriple(mlir::ModuleOp module) {
+    auto offloadMod =
+        llvm::cast<mlir::omp::OffloadModuleInterface>(*module.getOperation());
+    auto targetTriples = offloadMod.getTargetTriples();
+
+    if (!targetTriples.empty())
+      if (auto tripleAttr =
+              llvm::dyn_cast<mlir::StringAttr>(targetTriples.front()))
+        return llvm::Triple(tripleAttr.getValue());
+
+    return llvm::Triple();
+  }
+
+  mlir::LLVM::TargetFeaturesAttr targetFeatures;
+};
+
+struct MetadirectiveCandidate {
+  MetadirectiveCandidate(const parser::OmpDirectiveSpecification *spec,
+                         llvm::omp::VariantMatchInfo vmi, bool isExplicit)
+      : spec(spec), vmi(vmi), isExplicit(isExplicit) {}
+
+  const parser::OmpDirectiveSpecification *spec = nullptr;
+  llvm::omp::VariantMatchInfo vmi;
+  bool isExplicit = false;
+};
+} // namespace
+
+static void appendConstructTraits(
+    llvm::omp::Directive dir,
+    llvm::SmallVectorImpl<llvm::omp::TraitProperty> &constructTraits) {
+  // Record compound directives inside-out so callers can reverse the final
+  // sequence into the outermost-to-innermost order required by OMPContext.
+  if (llvm::omp::allSimdSet.test(dir))
+    constructTraits.push_back(llvm::omp::TraitProperty::construct_simd_simd);
+  if (llvm::omp::allDoSet.test(dir))
+    constructTraits.push_back(llvm::omp::TraitProperty::construct_for_for);
+  if (llvm::omp::allParallelSet.test(dir))
+    constructTraits.push_back(
+        llvm::omp::TraitProperty::construct_parallel_parallel);
+  if (llvm::omp::allTeamsSet.test(dir))
+    constructTraits.push_back(llvm::omp::TraitProperty::construct_teams_teams);
+  if (llvm::omp::allTargetSet.test(dir))
+    constructTraits.push_back(
+        llvm::omp::TraitProperty::construct_target_target);
+}
+
+static void appendConstructTraits(
+    mlir::Operation *op,
+    llvm::SmallVectorImpl<llvm::omp::TraitProperty> &constructTraits) {
+  if (mlir::isa<mlir::omp::SimdOp>(op))
+    constructTraits.push_back(llvm::omp::TraitProperty::construct_simd_simd);
+  if (mlir::isa<mlir::omp::WsloopOp>(op))
+    constructTraits.push_back(llvm::omp::TraitProperty::construct_for_for);
+  if (mlir::isa<mlir::omp::ParallelOp>(op))
+    constructTraits.push_back(
+        llvm::omp::TraitProperty::construct_parallel_parallel);
+  if (mlir::isa<mlir::omp::TeamsOp>(op))
+    constructTraits.push_back(llvm::omp::TraitProperty::construct_teams_teams);
+  if (mlir::isa<mlir::omp::TargetOp>(op))
+    constructTraits.push_back(
+        llvm::omp::TraitProperty::construct_target_target);
+}
+static void genMetadirective(lower::AbstractConverter &converter,
+                             lower::SymMap &symTable,
+                             semantics::SemanticsContext &semaCtx,
+                             lower::pft::Evaluation &eval,
+                             const parser::OmpClauseList &clauseList) {
+  fir::FirOpBuilder &builder = converter.getFirOpBuilder();
+
+  llvm::SmallVector<llvm::omp::TraitProperty, 8> constructTraits;
+  for (auto *parentEval = eval.parentConstruct; parentEval;
+       parentEval = parentEval->parentConstruct) {
+    const auto *ompConstruct = parentEval->getIf<parser::OpenMPConstruct>();
+    if (!ompConstruct)
+      continue;
+    llvm::omp::Directive dir =
+        parser::omp::GetOmpDirectiveName(*ompConstruct).v;
+    appendConstructTraits(dir, constructTraits);
+  }
+  if (constructTraits.empty()) {
+    for (mlir::Operation *op = builder.getInsertionBlock()->getParentOp(); op;
+         op = op->getParentOp())
+      appendConstructTraits(op, constructTraits);
+  }
+  std::reverse(constructTraits.begin(), constructTraits.end());
+  TargetOMPContext ompCtx(builder.getModule(), constructTraits);
+
+  llvm::SmallVector<MetadirectiveCandidate, 4> candidates;
+  // A null directive specification represents either the implicit `nothing`
+  // variant or the absence of an explicit otherwise/default clause.
+  const parser::OmpDirectiveSpecification *fallback = nullptr;
+
+  auto getContextSelector = [](const parser::OmpClause::When &whenClause)
+      -> const parser::modifier::OmpContextSelector * {
+    const auto &modifiers = std::get<0>(whenClause.v.t);
+    if (!modifiers)
+      return nullptr;
+    for (const auto &mod : *modifiers) {
+      if (const auto *ctxSel =
+              std::get_if<parser::modifier::OmpContextSelector>(&mod.u))
+        return ctxSel;
+    }
+    return nullptr;
+  };
+
+  // Extract the directive variant spec from a when clause.
+  // Returns {spec_ptr, isExplicit}. A null spec means "nothing".
+  auto getDirectiveVariant = [](const parser::OmpClause::When &whenClause)
+      -> std::pair<const parser::OmpDirectiveSpecification *, bool> {
+    const auto &opt = std::get<1>(whenClause.v.t);
+    if (!opt)
+      return {nullptr, false};
+    if (opt->value().DirId() == llvm::omp::Directive::OMPD_nothing)
+      return {nullptr, true};
+    return {&opt->value(), true};
+  };
+
+  // Return the directive spec pointer, or nullptr for "nothing".
+  auto getFallbackVariant = [](const parser::OmpDirectiveSpecification &spec)
+      -> const parser::OmpDirectiveSpecification * {
+    if (spec.DirId() == llvm::omp::Directive::OMPD_nothing)
+      return nullptr;
+    return &spec;
+  };
+
+  for (const auto &clause : clauseList.v) {
+    if (const auto *whenClause =
+            std::get_if<parser::OmpClause::When>(&clause.u)) {
+      const auto *ctxSel = getContextSelector(*whenClause);
+      auto [spec, isExplicit] = getDirectiveVariant(*whenClause);
+
+      // Always match when there is no context selector.
+      if (!ctxSel) {
+        candidates.emplace_back(spec, llvm::omp::VariantMatchInfo(),
+                                isExplicit);
+        continue;
+      }
----------------
abidh wrote:

Is there a testcase for this condition?

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