[flang-commits] [flang] [llvm] [flang][OpenMP] Handle multiple ranges in `num_teams` clause (PR #102535)

[Krzysztof Parzyszek via flang-commits]

https://github.com/kparzysz created https://github.com/llvm/llvm-project/pull/102535

Commit cee594cf36 added support to clang for multiple expressions in `num_teams` clause. Add follow-up changes to flang.

>From b905da1600109454a8c3893a209bbba3842a0f8d Mon Sep 17 00:00:00 2001
From: Krzysztof Parzyszek <Krzysztof.Parzyszek at amd.com>
Date: Thu, 8 Aug 2024 15:46:01 -0500
Subject: [PATCH] [flang][OpenMP] Handle multiple ranges in `num_teams` clause

Commit cee594cf36 added support to clang for multiple expressions
in `num_teams` clause. Add follow-up changes to flang.
---
 flang/lib/Lower/OpenMP/ClauseProcessor.cpp  |  8 ++++++--
 flang/lib/Lower/OpenMP/Clauses.cpp          |  5 +++--
 llvm/include/llvm/Frontend/OpenMP/ClauseT.h | 14 ++++++++++++--
 3 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
index 310c0b0b5fb636..c30844f40b7e04 100644
--- a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
+++ b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
@@ -382,8 +382,12 @@ bool ClauseProcessor::processNumTeams(
   // TODO Get lower and upper bounds for num_teams when parser is updated to
   // accept both.
   if (auto *clause = findUniqueClause<omp::clause::NumTeams>()) {
-    // auto lowerBound = std::get<std::optional<ExprTy>>(clause->t);
-    auto &upperBound = std::get<ExprTy>(clause->t);
+    // The num_teams directive accepts a list of team lower/upper bounds.
+    // This is an extension to support grid specification for ompx_bare.
+    // Here, only expect a single element in the list.
+    assert(clause->v.size() == 1);
+    // auto lowerBound = std::get<std::optional<ExprTy>>(clause->v[0]->t);
+    auto &upperBound = std::get<ExprTy>(clause->v[0].t);
     result.numTeamsUpper =
         fir::getBase(converter.genExprValue(upperBound, stmtCtx));
     return true;
diff --git a/flang/lib/Lower/OpenMP/Clauses.cpp b/flang/lib/Lower/OpenMP/Clauses.cpp
index 256a603014e496..307be994644ed0 100644
--- a/flang/lib/Lower/OpenMP/Clauses.cpp
+++ b/flang/lib/Lower/OpenMP/Clauses.cpp
@@ -935,8 +935,9 @@ NumTasks make(const parser::OmpClause::NumTasks &inp,
 NumTeams make(const parser::OmpClause::NumTeams &inp,
               semantics::SemanticsContext &semaCtx) {
   // inp.v -> parser::ScalarIntExpr
-  return NumTeams{{/*LowerBound=*/std::nullopt,
-                   /*UpperBound=*/makeExpr(inp.v, semaCtx)}};
+  List<NumTeams::Range> v{{{/*LowerBound=*/std::nullopt,
+                           /*UpperBound=*/makeExpr(inp.v, semaCtx)}}};
+  return NumTeams{/*List=*/v};
 }
 
 NumThreads make(const parser::OmpClause::NumThreads &inp,
diff --git a/llvm/include/llvm/Frontend/OpenMP/ClauseT.h b/llvm/include/llvm/Frontend/OpenMP/ClauseT.h
index a16673f82aa9c7..43b2fb558e0ead 100644
--- a/llvm/include/llvm/Frontend/OpenMP/ClauseT.h
+++ b/llvm/include/llvm/Frontend/OpenMP/ClauseT.h
@@ -885,10 +885,20 @@ struct NumTasksT {
 // V5.2: [10.2.1] `num_teams` clause
 template <typename T, typename I, typename E> //
 struct NumTeamsT {
-  using TupleTrait = std::true_type;
   using LowerBound = E;
   using UpperBound = E;
-  std::tuple<OPT(LowerBound), UpperBound> t;
+
+  // The name Range is not a spec name.
+  struct Range {
+    using TupleTrait = std::true_type;
+    std::tuple<OPT(LowerBound), UpperBound> t;
+  };
+
+  // The name List is not a spec name. The list is an extension to allow
+  // specifying a grid with connection with the ompx_bare clause.
+  using List = ListT<Range>;
+  using WrapperTrait = std::true_type;
+  List v;
 };
 
 // V5.2: [10.1.2] `num_threads` clause