Skip to content

Commit 211ed03

Browse files
skatrakskatrak
committed
[Flang][OpenMP][Lower] Support lowering of teams directive to MLIR
This patch adds support for translating `teams` OpenMP directives to MLIR, when appearing as either loop or block constructs and as part of combined constructs or on its own. The current Fortran parser does not allow the specification of the optional lower bound for the "num_teams" clause, so only the `num_teams_upper` MLIR argument is set by this patch. Depends on D156809 Differential Revision: https://reviews.llvm.org/D156884
1 parent 8c177ae commit 211ed03

File tree

4 files changed

+251
-6
lines changed

4 files changed

+251
-6
lines changed

‎flang/lib/Lower/OpenMP.cpp

Lines changed: 55 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -496,6 +496,8 @@ class ClauseProcessor {
496496
bool processHint(mlir::IntegerAttr &result) const;
497497
bool processMergeable(mlir::UnitAttr &result) const;
498498
bool processNowait(mlir::UnitAttr &result) const;
499+
bool processNumTeams(Fortran::lower::StatementContext &stmtCtx,
500+
mlir::Value &result) const;
499501
bool processNumThreads(Fortran::lower::StatementContext &stmtCtx,
500502
mlir::Value &result) const;
501503
bool processOrdered(mlir::IntegerAttr &result) const;
@@ -1347,6 +1349,18 @@ bool ClauseProcessor::processNowait(mlir::UnitAttr &result) const {
13471349
return markClauseOccurrence<ClauseTy::Nowait>(result);
13481350
}
13491351

1352+
bool ClauseProcessor::processNumTeams(Fortran::lower::StatementContext &stmtCtx,
1353+
mlir::Value &result) const {
1354+
// TODO Get lower and upper bounds for num_teams when parser is updated to
1355+
// accept both.
1356+
if (auto *numTeamsClause = findUniqueClause<ClauseTy::NumTeams>()) {
1357+
result = fir::getBase(converter.genExprValue(
1358+
*Fortran::semantics::GetExpr(numTeamsClause->v), stmtCtx));
1359+
return true;
1360+
}
1361+
return false;
1362+
}
1363+
13501364
bool ClauseProcessor::processNumThreads(
13511365
Fortran::lower::StatementContext &stmtCtx, mlir::Value &result) const {
13521366
if (auto *numThreadsClause = findUniqueClause<ClauseTy::NumThreads>()) {
@@ -2359,6 +2373,40 @@ genTargetOp(Fortran::lower::AbstractConverter &converter,
23592373
mapOperands, mapTypesArrayAttr);
23602374
}
23612375

2376+
static mlir::omp::TeamsOp
2377+
genTeamsOp(Fortran::lower::AbstractConverter &converter,
2378+
Fortran::lower::pft::Evaluation &eval,
2379+
mlir::Location currentLocation,
2380+
const Fortran::parser::OmpClauseList &clauseList,
2381+
bool outerCombined = false) {
2382+
Fortran::lower::StatementContext stmtCtx;
2383+
mlir::Value numTeamsClauseOperand, ifClauseOperand, threadLimitClauseOperand;
2384+
llvm::SmallVector<mlir::Value> allocateOperands, allocatorOperands,
2385+
reductionVars;
2386+
llvm::SmallVector<mlir::Attribute> reductionDeclSymbols;
2387+
2388+
ClauseProcessor cp(converter, clauseList);
2389+
cp.processIf(stmtCtx,
2390+
Fortran::parser::OmpIfClause::DirectiveNameModifier::Teams,
2391+
ifClauseOperand);
2392+
cp.processAllocate(allocatorOperands, allocateOperands);
2393+
cp.processDefault();
2394+
cp.processNumTeams(stmtCtx, numTeamsClauseOperand);
2395+
cp.processThreadLimit(stmtCtx, threadLimitClauseOperand);
2396+
if (cp.processReduction(currentLocation, reductionVars, reductionDeclSymbols))
2397+
TODO(currentLocation, "Reduction of TEAMS directive");
2398+
2399+
return genOpWithBody<mlir::omp::TeamsOp>(
2400+
converter, eval, currentLocation, outerCombined, &clauseList,
2401+
/*num_teams_lower=*/nullptr, numTeamsClauseOperand, ifClauseOperand,
2402+
threadLimitClauseOperand, allocateOperands, allocatorOperands,
2403+
reductionVars,
2404+
reductionDeclSymbols.empty()
2405+
? nullptr
2406+
: mlir::ArrayAttr::get(converter.getFirOpBuilder().getContext(),
2407+
reductionDeclSymbols));
2408+
}
2409+
23622410
//===----------------------------------------------------------------------===//
23632411
// genOMP() Code generation helper functions
23642412
//===----------------------------------------------------------------------===//
@@ -2483,7 +2531,8 @@ static void genOMP(Fortran::lower::AbstractConverter &converter,
24832531
if ((llvm::omp::allTeamsSet & llvm::omp::loopConstructSet)
24842532
.test(ompDirective)) {
24852533
validDirective = true;
2486-
TODO(currentLocation, "Teams construct");
2534+
genTeamsOp(converter, eval, currentLocation, loopOpClauseList,
2535+
/*outerCombined=*/true);
24872536
}
24882537
if (llvm::omp::allDistributeSet.test(ompDirective)) {
24892538
validDirective = true;
@@ -2628,7 +2677,8 @@ genOMP(Fortran::lower::AbstractConverter &converter,
26282677
!std::get_if<Fortran::parser::OmpClause::Map>(&clause.u) &&
26292678
!std::get_if<Fortran::parser::OmpClause::UseDevicePtr>(&clause.u) &&
26302679
!std::get_if<Fortran::parser::OmpClause::UseDeviceAddr>(&clause.u) &&
2631-
!std::get_if<Fortran::parser::OmpClause::ThreadLimit>(&clause.u)) {
2680+
!std::get_if<Fortran::parser::OmpClause::ThreadLimit>(&clause.u) &&
2681+
!std::get_if<Fortran::parser::OmpClause::NumTeams>(&clause.u)) {
26322682
TODO(clauseLocation, "OpenMP Block construct clause");
26332683
}
26342684
}
@@ -2667,7 +2717,8 @@ genOMP(Fortran::lower::AbstractConverter &converter,
26672717
genTaskGroupOp(converter, eval, currentLocation, beginClauseList);
26682718
break;
26692719
case llvm::omp::Directive::OMPD_teams:
2670-
TODO(currentLocation, "Teams construct");
2720+
genTeamsOp(converter, eval, currentLocation, beginClauseList,
2721+
/*outerCombined=*/false);
26712722
break;
26722723
case llvm::omp::Directive::OMPD_workshare:
26732724
TODO(currentLocation, "Workshare construct");
@@ -2683,7 +2734,7 @@ genOMP(Fortran::lower::AbstractConverter &converter,
26832734
}
26842735
if ((llvm::omp::allTeamsSet & llvm::omp::blockConstructSet)
26852736
.test(directive.v)) {
2686-
TODO(currentLocation, "Teams construct");
2737+
genTeamsOp(converter, eval, currentLocation, beginClauseList);
26872738
combinedDirective = true;
26882739
}
26892740
if ((llvm::omp::allParallelSet & llvm::omp::blockConstructSet)

‎flang/test/Lower/OpenMP/Todo/reduction-teams.f90

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
! RUN: %not_todo_cmd bbc -emit-fir -fopenmp -o - %s 2>&1 | FileCheck %s
2+
! RUN: %not_todo_cmd %flang_fc1 -emit-fir -fopenmp -o - %s 2>&1 | FileCheck %s
3+
4+
! CHECK: not yet implemented: Reduction of TEAMS directive
5+
subroutine reduction_teams()
6+
integer :: i
7+
i = 0
8+
9+
!$omp teams reduction(+:i)
10+
i = i + 1
11+
!$omp end teams
12+
end subroutine reduction_teams

‎flang/test/Lower/OpenMP/if-clause.f90

Lines changed: 70 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,15 +13,13 @@ program main
1313
! - PARALLEL SECTIONS
1414
! - PARALLEL WORKSHARE
1515
! - TARGET PARALLEL
16-
! - TARGET TEAMS
1716
! - TARGET TEAMS DISTRIBUTE
1817
! - TARGET TEAMS DISTRIBUTE PARALLEL DO
1918
! - TARGET TEAMS DISTRIBUTE PARALLEL DO SIMD
2019
! - TARGET TEAMS DISTRIBUTE SIMD
2120
! - TARGET UPDATE
2221
! - TASKLOOP
2322
! - TASKLOOP SIMD
24-
! - TEAMS
2523
! - TEAMS DISTRIBUTE
2624
! - TEAMS DISTRIBUTE PARALLEL DO
2725
! - TEAMS DISTRIBUTE PARALLEL DO SIMD
@@ -416,6 +414,54 @@ program main
416414
end do
417415
!$omp end target simd
418416

417+
! ----------------------------------------------------------------------------
418+
! TARGET TEAMS
419+
! ----------------------------------------------------------------------------
420+
421+
! CHECK: omp.target
422+
! CHECK-NOT: if({{.*}})
423+
! CHECK-SAME: {
424+
! CHECK: omp.teams
425+
! CHECK-NOT: if({{.*}})
426+
! CHECK-SAME: {
427+
!$omp target teams
428+
i = 1
429+
!$omp end target teams
430+
431+
! CHECK: omp.target
432+
! CHECK-SAME: if({{.*}})
433+
! CHECK: omp.teams
434+
! CHECK-SAME: if({{.*}})
435+
!$omp target teams if(.true.)
436+
i = 1
437+
!$omp end target teams
438+
439+
! CHECK: omp.target
440+
! CHECK-SAME: if({{.*}})
441+
! CHECK: omp.teams
442+
! CHECK-SAME: if({{.*}})
443+
!$omp target teams if(target: .true.) if(teams: .false.)
444+
i = 1
445+
!$omp end target teams
446+
447+
! CHECK: omp.target
448+
! CHECK-SAME: if({{.*}})
449+
! CHECK: omp.teams
450+
! CHECK-NOT: if({{.*}})
451+
! CHECK-SAME: {
452+
!$omp target teams if(target: .true.)
453+
i = 1
454+
!$omp end target teams
455+
456+
! CHECK: omp.target
457+
! CHECK-NOT: if({{.*}})
458+
! CHECK-SAME: {
459+
! CHECK: omp.teams
460+
! CHECK-SAME: if({{.*}})
461+
!$omp target teams if(teams: .true.)
462+
i = 1
463+
!$omp end target teams
464+
419465
! ----------------------------------------------------------------------------
420466
! TASK
421467
! ----------------------------------------------------------------------------
@@ -434,4 +480,26 @@ program main
434480
! CHECK-SAME: if({{.*}})
435481
!$omp task if(task: .true.)
436482
!$omp end task
483+
484+
! ----------------------------------------------------------------------------
485+
! TEAMS
486+
! ----------------------------------------------------------------------------
487+
! CHECK: omp.teams
488+
! CHECK-NOT: if({{.*}})
489+
! CHECK-SAME: {
490+
!$omp teams
491+
i = 1
492+
!$omp end teams
493+
494+
! CHECK: omp.teams
495+
! CHECK-SAME: if({{.*}})
496+
!$omp teams if(.true.)
497+
i = 1
498+
!$omp end teams
499+
500+
! CHECK: omp.teams
501+
! CHECK-SAME: if({{.*}})
502+
!$omp teams if(teams: .true.)
503+
i = 1
504+
!$omp end teams
437505
end program main

‎flang/test/Lower/OpenMP/teams.f90

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
! RUN: %flang_fc1 -emit-fir -fopenmp %s -o - | FileCheck %s
2+
3+
! CHECK-LABEL: func @_QPteams_simple
4+
subroutine teams_simple()
5+
! CHECK: omp.teams
6+
!$omp teams
7+
! CHECK: fir.call
8+
call f1()
9+
! CHECK: omp.terminator
10+
!$omp end teams
11+
end subroutine teams_simple
12+
13+
!===============================================================================
14+
! `num_teams` clause
15+
!===============================================================================
16+
17+
! CHECK-LABEL: func @_QPteams_numteams
18+
subroutine teams_numteams(num_teams)
19+
integer, intent(inout) :: num_teams
20+
21+
! CHECK: omp.teams
22+
! CHECK-SAME: num_teams( to %{{.*}}: i32)
23+
!$omp teams num_teams(4)
24+
! CHECK: fir.call
25+
call f1()
26+
! CHECK: omp.terminator
27+
!$omp end teams
28+
29+
! CHECK: omp.teams
30+
! CHECK-SAME: num_teams( to %{{.*}}: i32)
31+
!$omp teams num_teams(num_teams)
32+
! CHECK: fir.call
33+
call f2()
34+
! CHECK: omp.terminator
35+
!$omp end teams
36+
37+
end subroutine teams_numteams
38+
39+
!===============================================================================
40+
! `if` clause
41+
!===============================================================================
42+
43+
! CHECK-LABEL: func @_QPteams_if
44+
subroutine teams_if(alpha)
45+
integer, intent(in) :: alpha
46+
logical :: condition
47+
48+
! CHECK: omp.teams
49+
! CHECK-SAME: if(%{{.*}})
50+
!$omp teams if(.false.)
51+
! CHECK: fir.call
52+
call f1()
53+
! CHECK: omp.terminator
54+
!$omp end teams
55+
56+
! CHECK: omp.teams
57+
! CHECK-SAME: if(%{{.*}})
58+
!$omp teams if(alpha .le. 0)
59+
! CHECK: fir.call
60+
call f2()
61+
! CHECK: omp.terminator
62+
!$omp end teams
63+
64+
! CHECK: omp.teams
65+
! CHECK-SAME: if(%{{.*}})
66+
!$omp teams if(condition)
67+
! CHECK: fir.call
68+
call f3()
69+
! CHECK: omp.terminator
70+
!$omp end teams
71+
end subroutine teams_if
72+
73+
!===============================================================================
74+
! `thread_limit` clause
75+
!===============================================================================
76+
77+
! CHECK-LABEL: func @_QPteams_threadlimit
78+
subroutine teams_threadlimit(thread_limit)
79+
integer, intent(inout) :: thread_limit
80+
81+
! CHECK: omp.teams
82+
! CHECK-SAME: thread_limit(%{{.*}}: i32)
83+
!$omp teams thread_limit(4)
84+
! CHECK: fir.call
85+
call f1()
86+
! CHECK: omp.terminator
87+
!$omp end teams
88+
89+
! CHECK: omp.teams
90+
! CHECK-SAME: thread_limit(%{{.*}}: i32)
91+
!$omp teams thread_limit(thread_limit)
92+
! CHECK: fir.call
93+
call f2()
94+
! CHECK: omp.terminator
95+
!$omp end teams
96+
97+
end subroutine teams_threadlimit
98+
99+
!===============================================================================
100+
! `allocate` clause
101+
!===============================================================================
102+
103+
! CHECK-LABEL: func @_QPteams_allocate
104+
subroutine teams_allocate()
105+
use omp_lib
106+
integer :: x
107+
! CHECK: omp.teams
108+
! CHECK-SAME: allocate(%{{.+}} : i32 -> %{{.+}} : !fir.ref<i32>)
109+
!$omp teams allocate(omp_high_bw_mem_alloc: x) private(x)
110+
! CHECK: arith.addi
111+
x = x + 12
112+
! CHECK: omp.terminator
113+
!$omp end teams
114+
end subroutine teams_allocate

0 commit comments

Comments
 (0)