Skip to content

Commit 96fdc04

Browse files
Stylie777Stylie777
committed
[MLIR][OpenMP] Add support for dist_schedule
`dist_schedule` was previously supported in Flang/Clang but was not implemented in MLIR, instead a user would get a "not yet implemented" error. This patch adds support for the `dist_schedule` clause to be lowered to LLVM IR when used in an `omp.distribute` or `omp.wsloop` section. There has needed to be some rework required to ensure that MLIR/LLVM emits the correct Schedule Type for the clause, as it uses a different schedule type to other OpenMP directives/clauses in the runtime library. This patch also ensures that when using dist_schedule or a chunked schedule clause, the correct llvm loop parallel accesses details are added.
1 parent 538c850 commit 96fdc04

File tree

10 files changed

+470
-100
lines changed

10 files changed

+470
-100
lines changed

‎flang/docs/OpenMPSupport.md‎

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -42,24 +42,24 @@ Note : No distinction is made between the support in Parser/Semantics, MLIR, Low
4242
| target update construct | P | device clause not supported |
4343
| declare target directive | P | |
4444
| teams construct | Y | |
45-
| distribute construct | P | dist_schedule clause not supported |
46-
| distribute simd construct | P | dist_schedule and linear clauses are not supported |
47-
| distribute parallel loop construct | P | dist_schedule clause not supported |
48-
| distribute parallel loop simd construct | P | dist_schedule and linear clauses are not supported |
45+
| distribute construct | P | |
46+
| distribute simd construct | P | linear clauses are not supported |
47+
| distribute parallel loop construct | P | |
48+
| distribute parallel loop simd construct | P | linear clauses are not supported |
4949
| depend clause | Y | |
5050
| declare reduction construct | N | |
5151
| atomic construct extensions | Y | |
5252
| cancel construct | Y | |
5353
| cancellation point construct | Y | |
5454
| parallel do simd construct | P | linear clause not supported |
5555
| target teams construct | P | device clause not supported |
56-
| teams distribute construct | P | dist_schedule clause not supported |
57-
| teams distribute simd construct | P | dist_schedule and linear clauses are not supported |
58-
| target teams distribute construct | P | device and dist_schedule clauses are not supported |
59-
| teams distribute parallel loop construct | P | dist_schedule clause not supported |
60-
| target teams distribute parallel loop construct | P | device and dist_schedule clauses are not supported |
61-
| teams distribute parallel loop simd construct | P | dist_schedule and linear clauses are not supported |
62-
| target teams distribute parallel loop simd construct | P | device, dist_schedule and linear clauses are not supported |
56+
| teams distribute construct | P | |
57+
| teams distribute simd construct | P | linear clause is not supported |
58+
| target teams distribute construct | P | device clause is not supported |
59+
| teams distribute parallel loop construct | P | |
60+
| target teams distribute parallel loop construct | P | device clause is not supported |
61+
| teams distribute parallel loop simd construct | P | linear clause is not supported |
62+
| target teams distribute parallel loop simd construct | P | device and linear clauses are not supported |
6363

6464
## Extensions
6565
### ATOMIC construct

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -485,7 +485,8 @@ def OMP_SCHEDULE_Dynamic : EnumVal<"dynamic", 3, 1> {}
485485
def OMP_SCHEDULE_Guided : EnumVal<"guided", 4, 1> {}
486486
def OMP_SCHEDULE_Auto : EnumVal<"auto", 5, 1> {}
487487
def OMP_SCHEDULE_Runtime : EnumVal<"runtime", 6, 1> {}
488-
def OMP_SCHEDULE_Default : EnumVal<"default", 7, 0> { let isDefault = 1; }
488+
def OMP_SCHEDULE_Distribute : EnumVal<"distribute", 7, 1> {}
489+
def OMP_SCHEDULE_Default : EnumVal<"default", 8, 0> { let isDefault = 1; }
489490
def OMPC_Schedule : Clause<[Spelling<"schedule">]> {
490491
let clangClass = "OMPScheduleClause";
491492
let flangClass = "OmpScheduleClause";
@@ -496,6 +497,7 @@ def OMPC_Schedule : Clause<[Spelling<"schedule">]> {
496497
OMP_SCHEDULE_Guided,
497498
OMP_SCHEDULE_Auto,
498499
OMP_SCHEDULE_Runtime,
500+
OMP_SCHEDULE_Distribute,
499501
OMP_SCHEDULE_Default
500502
];
501503
}

‎llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h‎

Lines changed: 34 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1110,11 +1110,17 @@ class OpenMPIRBuilder {
11101110
/// \param NeedsBarrier Indicates whether a barrier must be inserted after
11111111
/// the loop.
11121112
/// \param LoopType Type of workshare loop.
1113+
/// \param HasDistSchedule Defines if the clause being lowered is
1114+
/// dist_schedule as this is handled slightly differently
1115+
/// \param DistScheduleSchedType Defines the Schedule Type for the Distribute
1116+
/// loop. Defaults to None if no Distribute loop is present.
11131117
///
11141118
/// \returns Point where to insert code after the workshare construct.
11151119
InsertPointOrErrorTy applyStaticWorkshareLoop(
11161120
DebugLoc DL, CanonicalLoopInfo *CLI, InsertPointTy AllocaIP,
1117-
omp::WorksharingLoopType LoopType, bool NeedsBarrier);
1121+
omp::WorksharingLoopType LoopType, bool NeedsBarrier,
1122+
bool HasDistSchedule = false,
1123+
omp::OMPScheduleType DistScheduleSchedType = omp::OMPScheduleType::None);
11181124

11191125
/// Modifies the canonical loop a statically-scheduled workshare loop with a
11201126
/// user-specified chunk size.
@@ -1127,13 +1133,22 @@ class OpenMPIRBuilder {
11271133
/// \param NeedsBarrier Indicates whether a barrier must be inserted after the
11281134
/// loop.
11291135
/// \param ChunkSize The user-specified chunk size.
1136+
/// \param SchedType Optional type of scheduling to be passed to the init
1137+
/// function.
1138+
/// \param DistScheduleChunkSize The size of dist_shcedule chunk considered
1139+
/// as a unit when
1140+
/// scheduling. If \p nullptr, defaults to 1.
1141+
/// \param DistScheduleSchedType Defines the Schedule Type for the Distribute
1142+
/// loop. Defaults to None if no Distribute loop is present.
11301143
///
11311144
/// \returns Point where to insert code after the workshare construct.
1132-
InsertPointOrErrorTy applyStaticChunkedWorkshareLoop(DebugLoc DL,
1133-
CanonicalLoopInfo *CLI,
1134-
InsertPointTy AllocaIP,
1135-
bool NeedsBarrier,
1136-
Value *ChunkSize);
1145+
InsertPointOrErrorTy applyStaticChunkedWorkshareLoop(
1146+
DebugLoc DL, CanonicalLoopInfo *CLI, InsertPointTy AllocaIP,
1147+
bool NeedsBarrier, Value *ChunkSize,
1148+
omp::OMPScheduleType SchedType =
1149+
omp::OMPScheduleType::UnorderedStaticChunked,
1150+
Value *DistScheduleChunkSize = nullptr,
1151+
omp::OMPScheduleType DistScheduleSchedType = omp::OMPScheduleType::None);
11371152

11381153
/// Modifies the canonical loop to be a dynamically-scheduled workshare loop.
11391154
///
@@ -1153,14 +1168,15 @@ class OpenMPIRBuilder {
11531168
/// the loop.
11541169
/// \param Chunk The size of loop chunk considered as a unit when
11551170
/// scheduling. If \p nullptr, defaults to 1.
1171+
/// \param DistScheduleChunk The size of dist_shcedule chunk considered as
1172+
/// a unit when
1173+
/// scheduling. If \p nullptr, defaults to 1.
11561174
///
11571175
/// \returns Point where to insert code after the workshare construct.
1158-
InsertPointOrErrorTy applyDynamicWorkshareLoop(DebugLoc DL,
1159-
CanonicalLoopInfo *CLI,
1160-
InsertPointTy AllocaIP,
1161-
omp::OMPScheduleType SchedType,
1162-
bool NeedsBarrier,
1163-
Value *Chunk = nullptr);
1176+
InsertPointOrErrorTy applyDynamicWorkshareLoop(
1177+
DebugLoc DL, CanonicalLoopInfo *CLI, InsertPointTy AllocaIP,
1178+
omp::OMPScheduleType SchedType, bool NeedsBarrier, Value *Chunk = nullptr,
1179+
Value *DistScheduleChunk = nullptr);
11641180

11651181
/// Create alternative version of the loop to support if clause
11661182
///
@@ -1212,6 +1228,10 @@ class OpenMPIRBuilder {
12121228
/// \param LoopType Information about type of loop worksharing.
12131229
/// It corresponds to type of loop workshare OpenMP pragma.
12141230
/// \param NoLoop If true, no-loop code is generated.
1231+
/// \param HasDistSchedule Defines if the clause being lowered is
1232+
/// dist_schedule as this is handled slightly differently
1233+
///
1234+
/// \param ChunkSize The chunk size for dist_schedule loop
12151235
///
12161236
/// \returns Point where to insert code after the workshare construct.
12171237
LLVM_ABI InsertPointOrErrorTy applyWorkshareLoop(
@@ -1223,7 +1243,8 @@ class OpenMPIRBuilder {
12231243
bool HasOrderedClause = false,
12241244
omp::WorksharingLoopType LoopType =
12251245
omp::WorksharingLoopType::ForStaticLoop,
1226-
bool NoLoop = false);
1246+
bool NoLoop = false,
1247+
bool HasDistSchedule = false, Value *DistScheduleChunkSize = nullptr);
12271248

12281249
/// Tile a loop nest.
12291250
///

0 commit comments

Comments
 (0)