Skip to content

Commit 3064646

Browse files
authored
[flang][OpenMP] Handle multiple ranges in num_teams clause (#102535)
Commit cee594c added support to clang for multiple expressions in `num_teams` clause. Add follow-up changes to flang.
1 parent 24be4d5 commit 3064646

File tree

3 files changed

+21
-6
lines changed

3 files changed

+21
-6
lines changed

flang/lib/Lower/OpenMP/ClauseProcessor.cpp

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -382,8 +382,12 @@ bool ClauseProcessor::processNumTeams(
382382
// TODO Get lower and upper bounds for num_teams when parser is updated to
383383
// accept both.
384384
if (auto *clause = findUniqueClause<omp::clause::NumTeams>()) {
385-
// auto lowerBound = std::get<std::optional<ExprTy>>(clause->t);
386-
auto &upperBound = std::get<ExprTy>(clause->t);
385+
// The num_teams directive accepts a list of team lower/upper bounds.
386+
// This is an extension to support grid specification for ompx_bare.
387+
// Here, only expect a single element in the list.
388+
assert(clause->v.size() == 1);
389+
// auto lowerBound = std::get<std::optional<ExprTy>>(clause->v[0]->t);
390+
auto &upperBound = std::get<ExprTy>(clause->v[0].t);
387391
result.numTeamsUpper =
388392
fir::getBase(converter.genExprValue(upperBound, stmtCtx));
389393
return true;

flang/lib/Lower/OpenMP/Clauses.cpp

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -935,8 +935,9 @@ NumTasks make(const parser::OmpClause::NumTasks &inp,
935935
NumTeams make(const parser::OmpClause::NumTeams &inp,
936936
semantics::SemanticsContext &semaCtx) {
937937
// inp.v -> parser::ScalarIntExpr
938-
return NumTeams{{/*LowerBound=*/std::nullopt,
939-
/*UpperBound=*/makeExpr(inp.v, semaCtx)}};
938+
List<NumTeams::Range> v{{{/*LowerBound=*/std::nullopt,
939+
/*UpperBound=*/makeExpr(inp.v, semaCtx)}}};
940+
return NumTeams{/*List=*/v};
940941
}
941942

942943
NumThreads make(const parser::OmpClause::NumThreads &inp,

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

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -885,10 +885,20 @@ struct NumTasksT {
885885
// V5.2: [10.2.1] `num_teams` clause
886886
template <typename T, typename I, typename E> //
887887
struct NumTeamsT {
888-
using TupleTrait = std::true_type;
889888
using LowerBound = E;
890889
using UpperBound = E;
891-
std::tuple<OPT(LowerBound), UpperBound> t;
890+
891+
// The name Range is not a spec name.
892+
struct Range {
893+
using TupleTrait = std::true_type;
894+
std::tuple<OPT(LowerBound), UpperBound> t;
895+
};
896+
897+
// The name List is not a spec name. The list is an extension to allow
898+
// specifying a grid with connection with the ompx_bare clause.
899+
using List = ListT<Range>;
900+
using WrapperTrait = std::true_type;
901+
List v;
892902
};
893903

894904
// V5.2: [10.1.2] `num_threads` clause

0 commit comments

Comments
 (0)