Skip to content

Commit dd42112

Browse files
authored
[Flang][OpenMP] MLIR lowering support for grainsize and num_tasks clause (#128490)
- Added MLIR lowering for grainsize and num_tasks clauses of taskloop construct.
1 parent 7618ce8 commit dd42112

File tree

5 files changed

+162
-12
lines changed

5 files changed

+162
-12
lines changed

flang/lib/Lower/OpenMP/ClauseProcessor.cpp

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -388,6 +388,27 @@ bool ClauseProcessor::processNowait(mlir::omp::NowaitClauseOps &result) const {
388388
return markClauseOccurrence<omp::clause::Nowait>(result.nowait);
389389
}
390390

391+
bool ClauseProcessor::processNumTasks(
392+
lower::StatementContext &stmtCtx,
393+
mlir::omp::NumTasksClauseOps &result) const {
394+
using NumTasks = omp::clause::NumTasks;
395+
if (auto *clause = findUniqueClause<NumTasks>()) {
396+
fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
397+
mlir::MLIRContext *context = firOpBuilder.getContext();
398+
const auto &modifier =
399+
std::get<std::optional<NumTasks::Prescriptiveness>>(clause->t);
400+
if (modifier && *modifier == NumTasks::Prescriptiveness::Strict) {
401+
result.numTasksMod = mlir::omp::ClauseNumTasksTypeAttr::get(
402+
context, mlir::omp::ClauseNumTasksType::Strict);
403+
}
404+
const auto &numtasksExpr = std::get<omp::SomeExpr>(clause->t);
405+
result.numTasks =
406+
fir::getBase(converter.genExprValue(numtasksExpr, stmtCtx));
407+
return true;
408+
}
409+
return false;
410+
}
411+
391412
bool ClauseProcessor::processNumTeams(
392413
lower::StatementContext &stmtCtx,
393414
mlir::omp::NumTeamsClauseOps &result) const {
@@ -934,6 +955,27 @@ bool ClauseProcessor::processDepend(lower::SymMap &symMap,
934955
return findRepeatableClause<omp::clause::Depend>(process);
935956
}
936957

958+
bool ClauseProcessor::processGrainsize(
959+
lower::StatementContext &stmtCtx,
960+
mlir::omp::GrainsizeClauseOps &result) const {
961+
using Grainsize = omp::clause::Grainsize;
962+
if (auto *clause = findUniqueClause<Grainsize>()) {
963+
fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
964+
mlir::MLIRContext *context = firOpBuilder.getContext();
965+
const auto &modifier =
966+
std::get<std::optional<Grainsize::Prescriptiveness>>(clause->t);
967+
if (modifier && *modifier == Grainsize::Prescriptiveness::Strict) {
968+
result.grainsizeMod = mlir::omp::ClauseGrainsizeTypeAttr::get(
969+
context, mlir::omp::ClauseGrainsizeType::Strict);
970+
}
971+
const auto &grainsizeExpr = std::get<omp::SomeExpr>(clause->t);
972+
result.grainsize =
973+
fir::getBase(converter.genExprValue(grainsizeExpr, stmtCtx));
974+
return true;
975+
}
976+
return false;
977+
}
978+
937979
bool ClauseProcessor::processHasDeviceAddr(
938980
lower::StatementContext &stmtCtx, mlir::omp::HasDeviceAddrClauseOps &result,
939981
llvm::SmallVectorImpl<const semantics::Symbol *> &hasDeviceSyms) const {

flang/lib/Lower/OpenMP/ClauseProcessor.h

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,8 @@ class ClauseProcessor {
7373
mlir::omp::FilterClauseOps &result) const;
7474
bool processFinal(lower::StatementContext &stmtCtx,
7575
mlir::omp::FinalClauseOps &result) const;
76+
bool processGrainsize(lower::StatementContext &stmtCtx,
77+
mlir::omp::GrainsizeClauseOps &result) const;
7678
bool processHasDeviceAddr(
7779
lower::StatementContext &stmtCtx,
7880
mlir::omp::HasDeviceAddrClauseOps &result,
@@ -82,6 +84,8 @@ class ClauseProcessor {
8284
mlir::omp::InclusiveClauseOps &result) const;
8385
bool processMergeable(mlir::omp::MergeableClauseOps &result) const;
8486
bool processNowait(mlir::omp::NowaitClauseOps &result) const;
87+
bool processNumTasks(lower::StatementContext &stmtCtx,
88+
mlir::omp::NumTasksClauseOps &result) const;
8589
bool processNumTeams(lower::StatementContext &stmtCtx,
8690
mlir::omp::NumTeamsClauseOps &result) const;
8791
bool processNumThreads(lower::StatementContext &stmtCtx,

flang/lib/Lower/OpenMP/OpenMP.cpp

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1783,17 +1783,19 @@ static void genTaskgroupClauses(
17831783

17841784
static void genTaskloopClauses(lower::AbstractConverter &converter,
17851785
semantics::SemanticsContext &semaCtx,
1786+
lower::StatementContext &stmtCtx,
17861787
const List<Clause> &clauses, mlir::Location loc,
17871788
mlir::omp::TaskloopOperands &clauseOps) {
17881789

17891790
ClauseProcessor cp(converter, semaCtx, clauses);
1791+
cp.processGrainsize(stmtCtx, clauseOps);
1792+
cp.processNumTasks(stmtCtx, clauseOps);
17901793

17911794
cp.processTODO<clause::Allocate, clause::Collapse, clause::Default,
1792-
clause::Final, clause::Grainsize, clause::If,
1793-
clause::InReduction, clause::Lastprivate, clause::Mergeable,
1794-
clause::Nogroup, clause::NumTasks, clause::Priority,
1795-
clause::Reduction, clause::Shared, clause::Untied>(
1796-
loc, llvm::omp::Directive::OMPD_taskloop);
1795+
clause::Final, clause::If, clause::InReduction,
1796+
clause::Lastprivate, clause::Mergeable, clause::Nogroup,
1797+
clause::Priority, clause::Reduction, clause::Shared,
1798+
clause::Untied>(loc, llvm::omp::Directive::OMPD_taskloop);
17971799
}
17981800

17991801
static void genTaskwaitClauses(lower::AbstractConverter &converter,
@@ -3270,12 +3272,12 @@ genStandaloneSimd(lower::AbstractConverter &converter, lower::SymMap &symTable,
32703272

32713273
static mlir::omp::TaskloopOp genStandaloneTaskloop(
32723274
lower::AbstractConverter &converter, lower::SymMap &symTable,
3273-
semantics::SemanticsContext &semaCtx, lower::pft::Evaluation &eval,
3274-
mlir::Location loc, const ConstructQueue &queue,
3275-
ConstructQueue::const_iterator item) {
3275+
lower::StatementContext &stmtCtx, semantics::SemanticsContext &semaCtx,
3276+
lower::pft::Evaluation &eval, mlir::Location loc,
3277+
const ConstructQueue &queue, ConstructQueue::const_iterator item) {
32763278
mlir::omp::TaskloopOperands taskloopClauseOps;
3277-
genTaskloopClauses(converter, semaCtx, item->clauses, loc, taskloopClauseOps);
3278-
3279+
genTaskloopClauses(converter, semaCtx, stmtCtx, item->clauses, loc,
3280+
taskloopClauseOps);
32793281
DataSharingProcessor dsp(converter, semaCtx, item->clauses, eval,
32803282
/*shouldCollectPreDeterminedSymbols=*/true,
32813283
enableDelayedPrivatization, symTable);
@@ -3736,8 +3738,8 @@ static void genOMPDispatch(lower::AbstractConverter &converter,
37363738
genTaskgroupOp(converter, symTable, semaCtx, eval, loc, queue, item);
37373739
break;
37383740
case llvm::omp::Directive::OMPD_taskloop:
3739-
newOp = genStandaloneTaskloop(converter, symTable, semaCtx, eval, loc,
3740-
queue, item);
3741+
newOp = genStandaloneTaskloop(converter, symTable, stmtCtx, semaCtx, eval,
3742+
loc, queue, item);
37413743
break;
37423744
case llvm::omp::Directive::OMPD_taskwait:
37433745
newOp = genTaskwaitOp(converter, symTable, semaCtx, eval, loc, queue, item);
Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
! This test checks lowering of grainsize clause in taskloop directive.
2+
3+
! RUN: bbc -emit-hlfir -fopenmp -fopenmp-version=50 -o - %s 2>&1 | FileCheck %s
4+
! RUN: %flang_fc1 -emit-hlfir -fopenmp -fopenmp-version=50 -o - %s 2>&1 | FileCheck %s
5+
6+
! CHECK-LABEL: omp.private
7+
! CHECK-SAME: {type = private} @[[I_PRIVATE_TEST2:.*]] : i32
8+
9+
! CHECK-LABEL: omp.private
10+
! CHECK-SAME: {type = firstprivate} @[[X_FIRSTPRIVATE_TEST2:.*]] : i32
11+
! CHECK-SAME: copy {
12+
! CHECK: hlfir.assign
13+
14+
! CHECK-LABEL: omp.private
15+
! CHECK-SAME: {type = private} @[[I_PRIVATE:.*]] : i32
16+
17+
! CHECK-LABEL: omp.private
18+
! CHECK-SAME: {type = firstprivate} @[[X_FIRSTPRIVATE:.*]] : i32
19+
! CHECK-SAME: copy {
20+
! CHECK: hlfir.assign
21+
22+
! CHECK-LABEL: func.func @_QPtest_grainsize
23+
! CHECK: %[[ALLOCA_I:.*]] = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFtest_grainsizeEi"}
24+
! CHECK: %[[DECL_I:.*]]:2 = hlfir.declare %[[ALLOCA_I]] {uniq_name = "_QFtest_grainsizeEi"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
25+
! CHECK: %[[ALLOCA_X:.*]] = fir.alloca i32 {bindc_name = "x", uniq_name = "_QFtest_grainsizeEx"}
26+
! CHECK: %[[DECL_X:.*]]:2 = hlfir.declare %[[ALLOCA_X]] {uniq_name = "_QFtest_grainsizeEx"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
27+
! CHECK: %[[GRAINSIZE:.*]] = arith.constant 10 : i32
28+
subroutine test_grainsize
29+
integer :: i, x
30+
! CHECK: omp.taskloop grainsize(%[[GRAINSIZE]]: i32)
31+
! CHECK-SAME: private(@[[X_FIRSTPRIVATE]] %[[DECL_X]]#0 -> %[[ARG0:.*]], @[[I_PRIVATE]] %[[DECL_I]]#0 -> %[[ARG1:.*]] : !fir.ref<i32>, !fir.ref<i32>) {
32+
! CHECK: omp.loop_nest (%[[ARG2:.*]]) : i32 = (%{{.*}}) to (%{{.*}}) inclusive step (%{{.*}}) {
33+
!$omp taskloop grainsize(10)
34+
do i = 1, 1000
35+
x = x + 1
36+
end do
37+
!$omp end taskloop
38+
end subroutine test_grainsize
39+
40+
!CHECK-LABEL: func.func @_QPtest_grainsize_strict()
41+
subroutine test_grainsize_strict
42+
integer :: i, x
43+
! CHECK: %[[GRAINSIZE:.*]] = arith.constant 10 : i32
44+
! CHECK: omp.taskloop grainsize(strict, %[[GRAINSIZE]]: i32)
45+
!$omp taskloop grainsize(strict:10)
46+
do i = 1, 1000
47+
!CHECK: arith.addi
48+
x = x + 1
49+
end do
50+
!$omp end taskloop
51+
end subroutine
Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
! This test checks lowering of num_tasks clause in taskloop directive.
2+
3+
! RUN: bbc -emit-hlfir -fopenmp -fopenmp-version=50 -o - %s 2>&1 | FileCheck %s
4+
! RUN: %flang_fc1 -emit-hlfir -fopenmp -fopenmp-version=50 -o - %s 2>&1 | FileCheck %s
5+
6+
! CHECK-LABEL: omp.private
7+
! CHECK-SAME: {type = private} @[[I_PRIVATE_TEST2:.*]] : i32
8+
9+
! CHECK-LABEL: omp.private
10+
! CHECK-SAME: {type = firstprivate} @[[X_FIRSTPRIVATE_TEST2:.*]] : i32
11+
! CHECK-SAME: copy {
12+
! CHECK: hlfir.assign
13+
14+
! CHECK-LABEL: omp.private
15+
! CHECK-SAME: {type = private} @[[I_PRIVATE:.*]] : i32
16+
17+
! CHECK-LABEL: omp.private
18+
! CHECK-SAME: {type = firstprivate} @[[X_FIRSTPRIVATE:.*]] : i32
19+
! CHECK-SAME: copy {
20+
! CHECK: hlfir.assign
21+
22+
! CHECK-LABEL: func.func @_QPtest_num_tasks
23+
! CHECK: %[[ALLOCA_I:.*]] = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFtest_num_tasksEi"}
24+
! CHECK: %[[DECL_I:.*]]:2 = hlfir.declare %[[ALLOCA_I]] {uniq_name = "_QFtest_num_tasksEi"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
25+
! CHECK: %[[ALLOCA_X:.*]] = fir.alloca i32 {bindc_name = "x", uniq_name = "_QFtest_num_tasksEx"}
26+
! CHECK: %[[DECL_X:.*]]:2 = hlfir.declare %[[ALLOCA_X]] {uniq_name = "_QFtest_num_tasksEx"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
27+
! CHECK: %[[VAL_NUMTASKS:.*]] = arith.constant 10 : i32
28+
subroutine test_num_tasks
29+
integer :: i, x
30+
! CHECK: omp.taskloop num_tasks(%[[VAL_NUMTASKS]]: i32)
31+
! CHECK-SAME: private(@[[X_FIRSTPRIVATE]] %[[DECL_X]]#0 -> %[[ARG0:.*]], @[[I_PRIVATE]] %[[DECL_I]]#0 -> %[[ARG1:.*]] : !fir.ref<i32>, !fir.ref<i32>) {
32+
! CHECK: omp.loop_nest (%[[ARG2:.*]]) : i32 = (%{{.*}}) to (%{{.*}}) inclusive step (%{{.*}}) {
33+
!$omp taskloop num_tasks(10)
34+
do i = 1, 1000
35+
x = x + 1
36+
end do
37+
!$omp end taskloop
38+
end subroutine test_num_tasks
39+
40+
! CHECK-LABEL: func.func @_QPtest_num_tasks_strict
41+
subroutine test_num_tasks_strict
42+
integer :: x, i
43+
! CHECK: %[[NUM_TASKS:.*]] = arith.constant 10 : i32
44+
! CHECK: omp.taskloop num_tasks(strict, %[[NUM_TASKS]]: i32)
45+
!$omp taskloop num_tasks(strict:10)
46+
do i = 1, 100
47+
!CHECK: arith.addi
48+
x = x + 1
49+
end do
50+
!$omp end taskloop
51+
end subroutine

0 commit comments

Comments
 (0)