Skip to content

Commit 47ae3ea

Browse files
Stylie777Stylie777
authored
[MLIR][OpenMP] Add MLIR Lowering Support for dist_schedule (#152736)
`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 0e5633f commit 47ae3ea

File tree

10 files changed

+464
-93
lines changed

10 files changed

+464
-93
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
@@ -490,7 +490,8 @@ def OMP_SCHEDULE_Dynamic : EnumVal<"dynamic", 3, 1> {}
490490
def OMP_SCHEDULE_Guided : EnumVal<"guided", 4, 1> {}
491491
def OMP_SCHEDULE_Auto : EnumVal<"auto", 5, 1> {}
492492
def OMP_SCHEDULE_Runtime : EnumVal<"runtime", 6, 1> {}
493-
def OMP_SCHEDULE_Default : EnumVal<"default", 7, 0> { let isDefault = 1; }
493+
def OMP_SCHEDULE_Distribute : EnumVal<"distribute", 7, 1> {}
494+
def OMP_SCHEDULE_Default : EnumVal<"default", 8, 0> { let isDefault = 1; }
494495
def OMPC_Schedule : Clause<[Spelling<"schedule">]> {
495496
let clangClass = "OMPScheduleClause";
496497
let flangClass = "OmpScheduleClause";
@@ -501,6 +502,7 @@ def OMPC_Schedule : Clause<[Spelling<"schedule">]> {
501502
OMP_SCHEDULE_Guided,
502503
OMP_SCHEDULE_Auto,
503504
OMP_SCHEDULE_Runtime,
505+
OMP_SCHEDULE_Distribute,
504506
OMP_SCHEDULE_Default
505507
];
506508
}

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

Lines changed: 27 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1133,11 +1133,17 @@ class OpenMPIRBuilder {
11331133
/// \param NeedsBarrier Indicates whether a barrier must be inserted after
11341134
/// the loop.
11351135
/// \param LoopType Type of workshare loop.
1136+
/// \param HasDistSchedule Defines if the clause being lowered is
1137+
/// dist_schedule as this is handled slightly differently
1138+
/// \param DistScheduleSchedType Defines the Schedule Type for the Distribute
1139+
/// loop. Defaults to None if no Distribute loop is present.
11361140
///
11371141
/// \returns Point where to insert code after the workshare construct.
11381142
InsertPointOrErrorTy applyStaticWorkshareLoop(
11391143
DebugLoc DL, CanonicalLoopInfo *CLI, InsertPointTy AllocaIP,
1140-
omp::WorksharingLoopType LoopType, bool NeedsBarrier);
1144+
omp::WorksharingLoopType LoopType, bool NeedsBarrier,
1145+
bool HasDistSchedule = false,
1146+
omp::OMPScheduleType DistScheduleSchedType = omp::OMPScheduleType::None);
11411147

11421148
/// Modifies the canonical loop a statically-scheduled workshare loop with a
11431149
/// user-specified chunk size.
@@ -1150,13 +1156,22 @@ class OpenMPIRBuilder {
11501156
/// \param NeedsBarrier Indicates whether a barrier must be inserted after the
11511157
/// loop.
11521158
/// \param ChunkSize The user-specified chunk size.
1159+
/// \param SchedType Optional type of scheduling to be passed to the init
1160+
/// function.
1161+
/// \param DistScheduleChunkSize The size of dist_shcedule chunk considered
1162+
/// as a unit when
1163+
/// scheduling. If \p nullptr, defaults to 1.
1164+
/// \param DistScheduleSchedType Defines the Schedule Type for the Distribute
1165+
/// loop. Defaults to None if no Distribute loop is present.
11531166
///
11541167
/// \returns Point where to insert code after the workshare construct.
1155-
InsertPointOrErrorTy applyStaticChunkedWorkshareLoop(DebugLoc DL,
1156-
CanonicalLoopInfo *CLI,
1157-
InsertPointTy AllocaIP,
1158-
bool NeedsBarrier,
1159-
Value *ChunkSize);
1168+
InsertPointOrErrorTy applyStaticChunkedWorkshareLoop(
1169+
DebugLoc DL, CanonicalLoopInfo *CLI, InsertPointTy AllocaIP,
1170+
bool NeedsBarrier, Value *ChunkSize,
1171+
omp::OMPScheduleType SchedType =
1172+
omp::OMPScheduleType::UnorderedStaticChunked,
1173+
Value *DistScheduleChunkSize = nullptr,
1174+
omp::OMPScheduleType DistScheduleSchedType = omp::OMPScheduleType::None);
11601175

11611176
/// Modifies the canonical loop to be a dynamically-scheduled workshare loop.
11621177
///
@@ -1235,6 +1250,10 @@ class OpenMPIRBuilder {
12351250
/// \param LoopType Information about type of loop worksharing.
12361251
/// It corresponds to type of loop workshare OpenMP pragma.
12371252
/// \param NoLoop If true, no-loop code is generated.
1253+
/// \param HasDistSchedule Defines if the clause being lowered is
1254+
/// dist_schedule as this is handled slightly differently
1255+
///
1256+
/// \param DistScheduleChunkSize The chunk size for dist_schedule loop
12381257
///
12391258
/// \returns Point where to insert code after the workshare construct.
12401259
LLVM_ABI InsertPointOrErrorTy applyWorkshareLoop(
@@ -1246,7 +1265,8 @@ class OpenMPIRBuilder {
12461265
bool HasOrderedClause = false,
12471266
omp::WorksharingLoopType LoopType =
12481267
omp::WorksharingLoopType::ForStaticLoop,
1249-
bool NoLoop = false);
1268+
bool NoLoop = false, bool HasDistSchedule = false,
1269+
Value *DistScheduleChunkSize = nullptr);
12501270

12511271
/// Tile a loop nest.
12521272
///

0 commit comments

Comments
 (0)