Skip to content

Commit 98f5ec0

Browse files
skc7skc7
committed
[flang] Add support for workdistribute construct in flang frontend
1 parent 8f60ac8 commit 98f5ec0

File tree

7 files changed

+187
-3
lines changed

7 files changed

+187
-3
lines changed

‎flang/include/flang/Semantics/openmp-directive-sets.h‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,7 @@ static const OmpDirectiveSet topTargetSet{
143143
Directive::OMPD_target_teams_distribute_parallel_do_simd,
144144
Directive::OMPD_target_teams_distribute_simd,
145145
Directive::OMPD_target_teams_loop,
146+
Directive::OMPD_target_teams_workdistribute,
146147
};
147148

148149
static const OmpDirectiveSet allTargetSet{topTargetSet};
@@ -172,6 +173,7 @@ static const OmpDirectiveSet topTeamsSet{
172173
Directive::OMPD_teams_distribute_parallel_do_simd,
173174
Directive::OMPD_teams_distribute_simd,
174175
Directive::OMPD_teams_loop,
176+
Directive::OMPD_teams_workdistribute,
175177
};
176178

177179
static const OmpDirectiveSet bottomTeamsSet{
@@ -187,9 +189,16 @@ static const OmpDirectiveSet allTeamsSet{
187189
Directive::OMPD_target_teams_distribute_parallel_do_simd,
188190
Directive::OMPD_target_teams_distribute_simd,
189191
Directive::OMPD_target_teams_loop,
192+
Directive::OMPD_target_teams_workdistribute,
190193
} | topTeamsSet,
191194
};
192195

196+
static const OmpDirectiveSet allWorkdistributeSet{
197+
Directive::OMPD_workdistribute,
198+
Directive::OMPD_teams_workdistribute,
199+
Directive::OMPD_target_teams_workdistribute,
200+
};
201+
193202
//===----------------------------------------------------------------------===//
194203
// Directive sets for groups of multiple directives
195204
//===----------------------------------------------------------------------===//
@@ -230,6 +239,9 @@ static const OmpDirectiveSet blockConstructSet{
230239
Directive::OMPD_taskgroup,
231240
Directive::OMPD_teams,
232241
Directive::OMPD_workshare,
242+
Directive::OMPD_target_teams_workdistribute,
243+
Directive::OMPD_teams_workdistribute,
244+
Directive::OMPD_workdistribute,
233245
};
234246

235247
static const OmpDirectiveSet loopConstructSet{
@@ -294,6 +306,7 @@ static const OmpDirectiveSet workShareSet{
294306
Directive::OMPD_scope,
295307
Directive::OMPD_sections,
296308
Directive::OMPD_single,
309+
Directive::OMPD_workdistribute,
297310
} | allDoSet,
298311
};
299312

@@ -376,6 +389,7 @@ static const OmpDirectiveSet nestedReduceWorkshareAllowedSet{
376389
};
377390

378391
static const OmpDirectiveSet nestedTeamsAllowedSet{
392+
Directive::OMPD_workdistribute,
379393
Directive::OMPD_distribute,
380394
Directive::OMPD_distribute_parallel_do,
381395
Directive::OMPD_distribute_parallel_do_simd,

‎flang/lib/Lower/OpenMP/OpenMP.cpp‎

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -535,6 +535,16 @@ static void processHostEvalClauses(lower::AbstractConverter &converter,
535535
cp.processCollapse(loc, eval, hostInfo->ops, hostInfo->iv);
536536
break;
537537

538+
case OMPD_teams_workdistribute:
539+
cp.processThreadLimit(stmtCtx, hostInfo.ops);
540+
[[fallthrough]];
541+
case OMPD_target_teams_workdistribute:
542+
cp.processNumTeams(stmtCtx, hostInfo.ops);
543+
processSingleNestedIf([](Directive nestedDir) {
544+
return topDistributeSet.test(nestedDir) || topLoopSet.test(nestedDir);
545+
});
546+
break;
547+
538548
// Standalone 'target' case.
539549
case OMPD_target: {
540550
processSingleNestedIf(
@@ -2794,6 +2804,17 @@ genTeamsOp(lower::AbstractConverter &converter, lower::SymMap &symTable,
27942804
queue, item, clauseOps);
27952805
}
27962806

2807+
static mlir::omp::WorkdistributeOp genWorkdistributeOp(
2808+
lower::AbstractConverter &converter, lower::SymMap &symTable,
2809+
semantics::SemanticsContext &semaCtx, lower::pft::Evaluation &eval,
2810+
mlir::Location loc, const ConstructQueue &queue,
2811+
ConstructQueue::const_iterator item) {
2812+
return genOpWithBody<mlir::omp::WorkdistributeOp>(
2813+
OpWithBodyGenInfo(converter, symTable, semaCtx, loc, eval,
2814+
llvm::omp::Directive::OMPD_workdistribute),
2815+
queue, item);
2816+
}
2817+
27972818
//===----------------------------------------------------------------------===//
27982819
// Code generation functions for the standalone version of constructs that can
27992820
// also be a leaf of a composite construct
@@ -3416,7 +3437,10 @@ static void genOMPDispatch(lower::AbstractConverter &converter,
34163437
case llvm::omp::Directive::OMPD_unroll:
34173438
genUnrollOp(converter, symTable, stmtCtx, semaCtx, eval, loc, queue, item);
34183439
break;
3419-
// case llvm::omp::Directive::OMPD_workdistribute:
3440+
case llvm::omp::Directive::OMPD_workdistribute:
3441+
newOp = genWorkdistributeOp(converter, symTable, semaCtx, eval, loc, queue,
3442+
item);
3443+
break;
34203444
case llvm::omp::Directive::OMPD_workshare:
34213445
newOp = genWorkshareOp(converter, symTable, stmtCtx, semaCtx, eval, loc,
34223446
queue, item);

‎flang/lib/Parser/openmp-parsers.cpp‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1590,12 +1590,17 @@ TYPE_PARSER(
15901590
"TARGET DATA" >> pure(llvm::omp::Directive::OMPD_target_data),
15911591
"TARGET_DATA" >> pure(llvm::omp::Directive::OMPD_target_data),
15921592
"TARGET PARALLEL" >> pure(llvm::omp::Directive::OMPD_target_parallel),
1593+
"TARGET TEAMS WORKDISTRIBUTE" >>
1594+
pure(llvm::omp::Directive::OMPD_target_teams_workdistribute),
15931595
"TARGET TEAMS" >> pure(llvm::omp::Directive::OMPD_target_teams),
15941596
"TARGET" >> pure(llvm::omp::Directive::OMPD_target),
15951597
"TASK"_id >> pure(llvm::omp::Directive::OMPD_task),
15961598
"TASKGROUP" >> pure(llvm::omp::Directive::OMPD_taskgroup),
1599+
"TEAMS WORKDISTRIBUTE" >>
1600+
pure(llvm::omp::Directive::OMPD_teams_workdistribute),
15971601
"TEAMS" >> pure(llvm::omp::Directive::OMPD_teams),
1598-
"WORKSHARE" >> pure(llvm::omp::Directive::OMPD_workshare))))
1602+
"WORKSHARE" >> pure(llvm::omp::Directive::OMPD_workshare),
1603+
"WORKDISTRIBUTE" >> pure(llvm::omp::Directive::OMPD_workdistribute))))
15991604

16001605
TYPE_PARSER(sourced(construct<OmpBeginBlockDirective>(
16011606
sourced(Parser<OmpBlockDirective>{}), Parser<OmpClauseList>{})))

‎flang/lib/Semantics/resolve-directives.cpp‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1714,10 +1714,13 @@ bool OmpAttributeVisitor::Pre(const parser::OpenMPBlockConstruct &x) {
17141714
case llvm::omp::Directive::OMPD_task:
17151715
case llvm::omp::Directive::OMPD_taskgroup:
17161716
case llvm::omp::Directive::OMPD_teams:
1717+
case llvm::omp::Directive::OMPD_workdistribute:
17171718
case llvm::omp::Directive::OMPD_workshare:
17181719
case llvm::omp::Directive::OMPD_parallel_workshare:
17191720
case llvm::omp::Directive::OMPD_target_teams:
1721+
case llvm::omp::Directive::OMPD_target_teams_workdistribute:
17201722
case llvm::omp::Directive::OMPD_target_parallel:
1723+
case llvm::omp::Directive::OMPD_teams_workdistribute:
17211724
PushContext(beginDir.source, beginDir.v);
17221725
break;
17231726
default:
@@ -1747,9 +1750,12 @@ void OmpAttributeVisitor::Post(const parser::OpenMPBlockConstruct &x) {
17471750
case llvm::omp::Directive::OMPD_target:
17481751
case llvm::omp::Directive::OMPD_task:
17491752
case llvm::omp::Directive::OMPD_teams:
1753+
case llvm::omp::Directive::OMPD_workdistribute:
17501754
case llvm::omp::Directive::OMPD_parallel_workshare:
17511755
case llvm::omp::Directive::OMPD_target_teams:
1752-
case llvm::omp::Directive::OMPD_target_parallel: {
1756+
case llvm::omp::Directive::OMPD_target_parallel:
1757+
case llvm::omp::Directive::OMPD_target_teams_workdistribute:
1758+
case llvm::omp::Directive::OMPD_teams_workdistribute: {
17531759
bool hasPrivate;
17541760
for (const auto *allocName : allocateNames_) {
17551761
hasPrivate = false;

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

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
! RUN: %flang_fc1 -emit-hlfir -fopenmp %s -o - | FileCheck %s
2+
3+
! CHECK-LABEL: func @_QPtarget_teams_workdistribute
4+
subroutine target_teams_workdistribute()
5+
! CHECK: omp.target
6+
! CHECK: omp.teams
7+
! CHECK: omp.workdistribute
8+
!$omp target teams workdistribute
9+
! CHECK: fir.call
10+
call f1()
11+
! CHECK: omp.terminator
12+
! CHECK: omp.terminator
13+
! CHECK: omp.terminator
14+
!$omp end target teams workdistribute
15+
end subroutine target_teams_workdistribute
16+
17+
! CHECK-LABEL: func @_QPteams_workdistribute
18+
subroutine teams_workdistribute()
19+
! CHECK: omp.teams
20+
! CHECK: omp.workdistribute
21+
!$omp teams workdistribute
22+
! CHECK: fir.call
23+
call f1()
24+
! CHECK: omp.terminator
25+
! CHECK: omp.terminator
26+
!$omp end teams workdistribute
27+
end subroutine teams_workdistribute
28+
29+
! CHECK-LABEL: func @_QPtarget_teams_workdistribute_m
30+
subroutine target_teams_workdistribute_m()
31+
! CHECK: omp.target
32+
! CHECK: omp.teams
33+
! CHECK: omp.workdistribute
34+
!$omp target
35+
!$omp teams
36+
!$omp workdistribute
37+
! CHECK: fir.call
38+
call f1()
39+
! CHECK: omp.terminator
40+
! CHECK: omp.terminator
41+
! CHECK: omp.terminator
42+
!$omp end workdistribute
43+
!$omp end teams
44+
!$omp end target
45+
end subroutine target_teams_workdistribute_m
46+
47+
! CHECK-LABEL: func @_QPteams_workdistribute_m
48+
subroutine teams_workdistribute_m()
49+
! CHECK: omp.teams
50+
! CHECK: omp.workdistribute
51+
!$omp teams
52+
!$omp workdistribute
53+
! CHECK: fir.call
54+
call f1()
55+
! CHECK: omp.terminator
56+
! CHECK: omp.terminator
57+
!$omp end workdistribute
58+
!$omp end teams
59+
end subroutine teams_workdistribute_m

‎llvm/include/llvm/Frontend/OpenMP/OMP.td‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1309,6 +1309,15 @@ def OMP_EndWorkshare : Directive<[Spelling<"end workshare">]> {
13091309
let category = OMP_Workshare.category;
13101310
let languages = [L_Fortran];
13111311
}
1312+
def OMP_Workdistribute : Directive<[Spelling<"workdistribute">]> {
1313+
let association = AS_Block;
1314+
let category = CA_Executable;
1315+
}
1316+
def OMP_EndWorkdistribute : Directive<[Spelling<"end workdistribute">]> {
1317+
let leafConstructs = OMP_Workdistribute.leafConstructs;
1318+
let association = OMP_Workdistribute.association;
1319+
let category = OMP_Workdistribute.category;
1320+
}
13121321

13131322
//===----------------------------------------------------------------------===//
13141323
// Definitions of OpenMP compound directives
@@ -2452,6 +2461,34 @@ def OMP_TargetTeamsDistributeSimd
24522461
let leafConstructs = [OMP_Target, OMP_Teams, OMP_Distribute, OMP_Simd];
24532462
let category = CA_Executable;
24542463
}
2464+
def OMP_TargetTeamsWorkdistribute : Directive<[Spelling<"target teams workdistribute">]> {
2465+
let allowedClauses = [
2466+
VersionedClause<OMPC_Allocate>,
2467+
VersionedClause<OMPC_Depend>,
2468+
VersionedClause<OMPC_FirstPrivate>,
2469+
VersionedClause<OMPC_HasDeviceAddr, 51>,
2470+
VersionedClause<OMPC_If>,
2471+
VersionedClause<OMPC_IsDevicePtr>,
2472+
VersionedClause<OMPC_Map>,
2473+
VersionedClause<OMPC_OMPX_Attribute>,
2474+
VersionedClause<OMPC_Private>,
2475+
VersionedClause<OMPC_Reduction>,
2476+
VersionedClause<OMPC_Shared>,
2477+
VersionedClause<OMPC_UsesAllocators, 50>,
2478+
];
2479+
let allowedOnceClauses = [
2480+
VersionedClause<OMPC_Default>,
2481+
VersionedClause<OMPC_DefaultMap>,
2482+
VersionedClause<OMPC_Device>,
2483+
VersionedClause<OMPC_NoWait>,
2484+
VersionedClause<OMPC_NumTeams>,
2485+
VersionedClause<OMPC_OMPX_DynCGroupMem>,
2486+
VersionedClause<OMPC_OMPX_Bare>,
2487+
VersionedClause<OMPC_ThreadLimit>,
2488+
];
2489+
let leafConstructs = [OMP_Target, OMP_Teams, OMP_Workdistribute];
2490+
let category = CA_Executable;
2491+
}
24552492
def OMP_target_teams_loop : Directive<[Spelling<"target teams loop">]> {
24562493
let allowedClauses = [
24572494
VersionedClause<OMPC_Allocate>,
@@ -2682,6 +2719,24 @@ def OMP_TeamsDistributeSimd : Directive<[Spelling<"teams distribute simd">]> {
26822719
let leafConstructs = [OMP_Teams, OMP_Distribute, OMP_Simd];
26832720
let category = CA_Executable;
26842721
}
2722+
def OMP_TeamsWorkdistribute : Directive<[Spelling<"teams workdistribute">]> {
2723+
let allowedClauses = [
2724+
VersionedClause<OMPC_Allocate>,
2725+
VersionedClause<OMPC_FirstPrivate>,
2726+
VersionedClause<OMPC_OMPX_Attribute>,
2727+
VersionedClause<OMPC_Private>,
2728+
VersionedClause<OMPC_Reduction>,
2729+
VersionedClause<OMPC_Shared>,
2730+
];
2731+
let allowedOnceClauses = [
2732+
VersionedClause<OMPC_Default>,
2733+
VersionedClause<OMPC_If, 52>,
2734+
VersionedClause<OMPC_NumTeams>,
2735+
VersionedClause<OMPC_ThreadLimit>,
2736+
];
2737+
let leafConstructs = [OMP_Teams, OMP_Workdistribute];
2738+
let category = CA_Executable;
2739+
}
26852740
def OMP_teams_loop : Directive<[Spelling<"teams loop">]> {
26862741
let allowedClauses = [
26872742
VersionedClause<OMPC_Allocate>,

‎mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2205,6 +2205,27 @@ def TargetFreeMemOp : OpenMP_Op<"target_freemem",
22052205
Arg<I64, "", [MemFree]>:$heapref
22062206
);
22072207
let assemblyFormat = "$device `,` $heapref attr-dict `:` type($device) `,` qualified(type($heapref))";
2208+
2209+
//===----------------------------------------------------------------------===//
2210+
// workdistribute Construct
2211+
//===----------------------------------------------------------------------===//
2212+
2213+
def WorkdistributeOp : OpenMP_Op<"workdistribute"> {
2214+
let summary = "workdistribute directive";
2215+
let description = [{
2216+
workdistribute divides execution of the enclosed structured block into
2217+
separate units of work, each executed only once by each
2218+
initial thread in the league.
2219+
```
2220+
!$omp target teams
2221+
!$omp workdistribute
2222+
y = a * x + y
2223+
!$omp end workdistribute
2224+
!$omp end target teams
2225+
```
2226+
}];
2227+
let regions = (region AnyRegion:$region);
2228+
let assemblyFormat = "$region attr-dict";
22082229
}
22092230

22102231
#endif // OPENMP_OPS

0 commit comments

Comments
 (0)