diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPConstants.h b/llvm/include/llvm/Frontend/OpenMP/OMPConstants.h --- a/llvm/include/llvm/Frontend/OpenMP/OMPConstants.h +++ b/llvm/include/llvm/Frontend/OpenMP/OMPConstants.h @@ -120,6 +120,10 @@ Runtime = 37, Auto = 38, // auto + StaticBalancedChunked = 45, // static with chunk adjustment (e.g., simd) + GuidedSimd = 46, // guided with chunk adjustment + RuntimeSimd = 47, // runtime with chunk adjustment + ModifierMonotonic = (1 << 29), // Set if the monotonic schedule modifier was present ModifierNonmonotonic = diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td --- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td +++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td @@ -137,11 +137,13 @@ def OMP_SCHEDULE_MOD_None : StrEnumAttrCase<"none", 0>; def OMP_SCHEDULE_MOD_Monotonic : StrEnumAttrCase<"monotonic", 1>; def OMP_SCHEDULE_MOD_Nonmonotonic : StrEnumAttrCase<"nonmonotonic", 2>; +def OMP_SCHEDULE_MOD_SIMD : StrEnumAttrCase<"simd", 3>; def ScheduleModifier : StrEnumAttr<"ScheduleModifier", "OpenMP Schedule Modifier", [OMP_SCHEDULE_MOD_None, OMP_SCHEDULE_MOD_Monotonic, - OMP_SCHEDULE_MOD_Nonmonotonic]> + OMP_SCHEDULE_MOD_Nonmonotonic, + OMP_SCHEDULE_MOD_SIMD]> { let cppNamespace = "::mlir::omp"; } @@ -226,6 +228,7 @@ OptionalAttr:$schedule_val, Optional:$schedule_chunk_var, OptionalAttr:$schedule_modifiers, + OptionalAttr:$simd_modifier, Confined, [IntMinValue<0>]>:$collapse_val, UnitAttr:$nowait, Confined, [IntMinValue<0>]>:$ordered_val, diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp --- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp +++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp @@ -413,7 +413,7 @@ /// sched-wo-chunk ::= `auto` | `runtime` static ParseResult parseScheduleClause(OpAsmParser &parser, SmallString<8> &schedule, - SmallVector> &modifiers, + SmallVectorImpl> &modifiers, Optional &chunkSize) { if (parser.parseLParen()) return failure(); @@ -438,7 +438,7 @@ } // If there is a comma, we have one or more modifiers.. - if (succeeded(parser.parseOptionalComma())) { + while (succeeded(parser.parseOptionalComma())) { StringRef mod; if (parser.parseKeyword(&mod)) return failure(); @@ -679,10 +679,13 @@ schedule[0] = llvm::toUpper(schedule[0]); auto attr = parser.getBuilder().getStringAttr(schedule); result.addAttribute("schedule_val", attr); - if (modifiers.size() > 0) - { - auto mod = parser.getBuilder().getStringAttr(modifiers[0]); - result.addAttribute("schedule_modifiers", mod); + if (modifiers.size() > 0) { + auto mod = parser.getBuilder().getStringAttr(modifiers[0]); + result.addAttribute("schedule_modifiers", mod); + if (modifiers.size() > 1) { + mod = parser.getBuilder().getStringAttr(modifiers[1]); + result.addAttribute("simd_modifier", mod); + } } if (scheduleChunkSize) { auto chunkSizeType = parser.getBuilder().getI32Type(); @@ -745,6 +748,10 @@ if (auto modifier = op.schedule_modifiers()) { p << ", " << modifier; } + auto simd = op.simd_modifier(); + if (simd.hasValue() && *simd != "none") { + p << ", " << simd; + } p << ")"; } diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp --- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp +++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp @@ -612,6 +612,15 @@ ompBuilder->collapseLoops(diLoc, loopInfos, {}); allocaIP = findAllocaInsertPoint(builder, moduleTranslation); + llvm::OpenMPIRBuilder::InsertPointTy afterIP; + llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder(); + + bool isSimd = false; + if (auto simd = loop.simd_modifier()) { + omp::ScheduleModifier modifier = *omp::symbolizeScheduleModifier(*simd); + isSimd = (modifier == omp::ScheduleModifier::simd); + } + if (schedule == omp::ClauseScheduleKind::Static) { ompBuilder->applyStaticWorkshareLoop(ompLoc.DL, loopInfo, allocaIP, !loop.nowait(), chunk); @@ -622,13 +631,19 @@ schedType = llvm::omp::OMPScheduleType::DynamicChunked; break; case omp::ClauseScheduleKind::Guided: - schedType = llvm::omp::OMPScheduleType::GuidedChunked; + if (isSimd) + schedType = llvm::omp::OMPScheduleType::GuidedSimd; + else + schedType = llvm::omp::OMPScheduleType::GuidedChunked; break; case omp::ClauseScheduleKind::Auto: schedType = llvm::omp::OMPScheduleType::Auto; break; case omp::ClauseScheduleKind::Runtime: - schedType = llvm::omp::OMPScheduleType::Runtime; + if (isSimd) + schedType = llvm::omp::OMPScheduleType::RuntimeSimd; + else + schedType = llvm::omp::OMPScheduleType::Runtime; break; default: llvm_unreachable("Unknown schedule value");