Skip to content

[MLIR][OpenMP] Add MLIR Lowering Support for dist_schedule #152736

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion llvm/include/llvm/Frontend/OpenMP/OMP.td
Original file line number Diff line number Diff line change
Expand Up @@ -458,7 +458,8 @@ def OMP_SCHEDULE_Dynamic : EnumVal<"dynamic", 3, 1> {}
def OMP_SCHEDULE_Guided : EnumVal<"guided", 4, 1> {}
def OMP_SCHEDULE_Auto : EnumVal<"auto", 5, 1> {}
def OMP_SCHEDULE_Runtime : EnumVal<"runtime", 6, 1> {}
def OMP_SCHEDULE_Default : EnumVal<"default", 7, 0> { let isDefault = 1; }
def OMP_SCHEDULE_Distribute : EnumVal<"distribute", 7, 1> {}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this the change to ensure the schedule type passed to the openmp runtime matches what clang does?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes

def OMP_SCHEDULE_Default : EnumVal<"default", 8, 0> { let isDefault = 1; }
def OMPC_Schedule : Clause<[Spelling<"schedule">]> {
let clangClass = "OMPScheduleClause";
let flangClass = "OmpScheduleClause";
Expand All @@ -469,6 +470,7 @@ def OMPC_Schedule : Clause<[Spelling<"schedule">]> {
OMP_SCHEDULE_Guided,
OMP_SCHEDULE_Auto,
OMP_SCHEDULE_Runtime,
OMP_SCHEDULE_Distribute,
OMP_SCHEDULE_Default
];
}
Expand Down
47 changes: 34 additions & 13 deletions llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
Original file line number Diff line number Diff line change
Expand Up @@ -1096,11 +1096,17 @@ class OpenMPIRBuilder {
/// \param NeedsBarrier Indicates whether a barrier must be inserted after
/// the loop.
/// \param LoopType Type of workshare loop.
/// \param HasDistSchedule Defines if the clause being lowered is
/// dist_schedule as this is handled slightly differently
/// \param DistScheduleSchedType Defines the Schedule Type for the Distribute
/// loop. Defaults to None if no Distribute loop is present.
///
/// \returns Point where to insert code after the workshare construct.
InsertPointOrErrorTy applyStaticWorkshareLoop(
DebugLoc DL, CanonicalLoopInfo *CLI, InsertPointTy AllocaIP,
omp::WorksharingLoopType LoopType, bool NeedsBarrier);
omp::WorksharingLoopType LoopType, bool NeedsBarrier,
bool HasDistSchedule = false,
omp::OMPScheduleType DistScheduleSchedType = omp::OMPScheduleType::None);

/// Modifies the canonical loop a statically-scheduled workshare loop with a
/// user-specified chunk size.
Expand All @@ -1113,13 +1119,22 @@ class OpenMPIRBuilder {
/// \param NeedsBarrier Indicates whether a barrier must be inserted after the
/// loop.
/// \param ChunkSize The user-specified chunk size.
/// \param SchedType Optional type of scheduling to be passed to the init
/// function.
/// \param DistScheduleChunkSize The size of dist_shcedule chunk considered
/// as a unit when
/// scheduling. If \p nullptr, defaults to 1.
/// \param DistScheduleSchedType Defines the Schedule Type for the Distribute
/// loop. Defaults to None if no Distribute loop is present.
///
/// \returns Point where to insert code after the workshare construct.
InsertPointOrErrorTy applyStaticChunkedWorkshareLoop(DebugLoc DL,
CanonicalLoopInfo *CLI,
InsertPointTy AllocaIP,
bool NeedsBarrier,
Value *ChunkSize);
InsertPointOrErrorTy applyStaticChunkedWorkshareLoop(
DebugLoc DL, CanonicalLoopInfo *CLI, InsertPointTy AllocaIP,
bool NeedsBarrier, Value *ChunkSize,
omp::OMPScheduleType SchedType =
omp::OMPScheduleType::UnorderedStaticChunked,
Value *DistScheduleChunkSize = nullptr,
omp::OMPScheduleType DistScheduleSchedType = omp::OMPScheduleType::None);

/// Modifies the canonical loop to be a dynamically-scheduled workshare loop.
///
Expand All @@ -1139,14 +1154,15 @@ class OpenMPIRBuilder {
/// the loop.
/// \param Chunk The size of loop chunk considered as a unit when
/// scheduling. If \p nullptr, defaults to 1.
/// \param DistScheduleChunk The size of dist_shcedule chunk considered as
/// a unit when
/// scheduling. If \p nullptr, defaults to 1.
///
/// \returns Point where to insert code after the workshare construct.
InsertPointOrErrorTy applyDynamicWorkshareLoop(DebugLoc DL,
CanonicalLoopInfo *CLI,
InsertPointTy AllocaIP,
omp::OMPScheduleType SchedType,
bool NeedsBarrier,
Value *Chunk = nullptr);
InsertPointOrErrorTy applyDynamicWorkshareLoop(
DebugLoc DL, CanonicalLoopInfo *CLI, InsertPointTy AllocaIP,
omp::OMPScheduleType SchedType, bool NeedsBarrier, Value *Chunk = nullptr,
Value *DistScheduleChunk = nullptr);

/// Create alternative version of the loop to support if clause
///
Expand Down Expand Up @@ -1197,6 +1213,10 @@ class OpenMPIRBuilder {
/// present.
/// \param LoopType Information about type of loop worksharing.
/// It corresponds to type of loop workshare OpenMP pragma.
/// \param HasDistSchedule Defines if the clause being lowered is
/// dist_schedule as this is handled slightly differently
///
/// \param ChunkSize The chunk size for dist_schedule loop
///
/// \returns Point where to insert code after the workshare construct.
LLVM_ABI InsertPointOrErrorTy applyWorkshareLoop(
Expand All @@ -1207,7 +1227,8 @@ class OpenMPIRBuilder {
bool HasMonotonicModifier = false, bool HasNonmonotonicModifier = false,
bool HasOrderedClause = false,
omp::WorksharingLoopType LoopType =
omp::WorksharingLoopType::ForStaticLoop);
omp::WorksharingLoopType::ForStaticLoop,
bool HasDistSchedule = false, Value *DistScheduleChunkSize = nullptr);

/// Tile a loop nest.
///
Expand Down
Loading