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 @@ -139,11 +139,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"; } @@ -289,6 +291,7 @@ OptionalAttr:$schedule_val, Optional:$schedule_chunk_var, OptionalAttr:$schedule_modifier, + UnitAttr:$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 @@ -244,12 +244,51 @@ // Parser and printer for Schedule Clause //===----------------------------------------------------------------------===// +static ParseResult +verifyScheduleModifiers(OpAsmParser &parser, + SmallVectorImpl> &modifiers) { + if (modifiers.size() > 2) + return parser.emitError(parser.getNameLoc()) << " unexpected modifier(s)"; + for (auto mod : modifiers) { + // Translate the string. If it has no value, then it was not a valid + // modifier! + auto symbol = symbolizeScheduleModifier(mod); + if (!symbol.hasValue()) + return parser.emitError(parser.getNameLoc()) + << " unknown modifier type: " << mod; + } + + // If we have one modifier that is "simd", then stick a "none" modiifer in + // index 0. + if (modifiers.size() == 1) { + if (symbolizeScheduleModifier(modifiers[0]) == + mlir::omp::ScheduleModifier::simd) { + modifiers.push_back(modifiers[0]); + modifiers[0] = + stringifyScheduleModifier(mlir::omp::ScheduleModifier::none); + } + } else if (modifiers.size() == 2) { + // If there are two modifier: + // First modifier should not be simd, second one should be simd + if (symbolizeScheduleModifier(modifiers[0]) == + mlir::omp::ScheduleModifier::simd || + symbolizeScheduleModifier(modifiers[1]) != + mlir::omp::ScheduleModifier::simd) + return parser.emitError(parser.getNameLoc()) + << " incorrect modifier order"; + } + return success(); +} + /// schedule ::= `schedule` `(` sched-list `)` -/// sched-list ::= sched-val | sched-val sched-list +/// sched-list ::= sched-val | sched-val sched-list | +/// sched-val `,` sched-modifier /// sched-val ::= sched-with-chunk | sched-wo-chunk /// sched-with-chunk ::= sched-with-chunk-types (`=` ssa-id-and-type)? /// sched-with-chunk-types ::= `static` | `dynamic` | `guided` /// sched-wo-chunk ::= `auto` | `runtime` +/// sched-modifier ::= sched-mod-val | sched-mod-val `,` sched-mod-val +/// sched-mod-val ::= `monotonic` | `nonmonotonic` | `simd` | `none` static ParseResult parseScheduleClause(OpAsmParser &parser, SmallString<8> &schedule, SmallVectorImpl> &modifiers, @@ -277,7 +316,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(); @@ -287,19 +326,24 @@ if (parser.parseRParen()) return failure(); + if (verifyScheduleModifiers(parser, modifiers)) + return failure(); + return success(); } /// Print schedule clause static void printScheduleClause(OpAsmPrinter &p, StringRef &sched, - llvm::Optional modifier, + llvm::Optional modifier, bool simd, Value scheduleChunkVar) { std::string schedLower = sched.lower(); p << "schedule(" << schedLower; if (scheduleChunkVar) p << " = " << scheduleChunkVar; - if (modifier && modifier.getValue() != "none") + if (modifier && modifier.hasValue()) p << ", " << modifier; + if (simd) + p << ", simd"; p << ") "; } @@ -822,6 +866,13 @@ if (modifiers.size() > 0) { auto mod = parser.getBuilder().getStringAttr(modifiers[0]); result.addAttribute("schedule_modifier", mod); + // Only SIMD attribute is allowed here! + if (modifiers.size() > 1) { + assert(symbolizeScheduleModifier(modifiers[1]) == + mlir::omp::ScheduleModifier::simd); + auto attr = UnitAttr::get(parser.getBuilder().getContext()); + result.addAttribute("simd_modifier", attr); + } } if (scheduleChunkSize) { auto chunkSizeType = parser.getBuilder().getI32Type(); @@ -1026,7 +1077,7 @@ if (auto sched = op.schedule_val()) printScheduleClause(p, sched.getValue(), op.schedule_modifier(), - op.schedule_chunk_var()); + op.simd_modifier(), op.schedule_chunk_var()); if (auto collapse = op.collapse_val()) p << "collapse(" << collapse << ") "; 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 @@ -684,6 +684,9 @@ ompBuilder->collapseLoops(diLoc, loopInfos, {}); allocaIP = findAllocaInsertPoint(builder, moduleTranslation); + + bool isSimd = loop.simd_modifier(); + if (schedule == omp::ClauseScheduleKind::Static) { ompBuilder->applyStaticWorkshareLoop(ompLoc.DL, loopInfo, allocaIP, !loop.nowait(), chunk); @@ -694,13 +697,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"); diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir --- a/mlir/test/Dialect/OpenMP/ops.mlir +++ b/mlir/test/Dialect/OpenMP/ops.mlir @@ -197,7 +197,7 @@ } // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) linear(%{{.*}} = %{{.*}} : memref) schedule(static) - omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) schedule(static, none) lastprivate(%data_var : memref) linear(%data_var = %linear_var : memref) { + omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) schedule(static) lastprivate(%data_var : memref) linear(%data_var = %linear_var : memref) { omp.yield } diff --git a/mlir/test/Target/LLVMIR/openmp-llvm-bad-schedule-modifier.mlir b/mlir/test/Target/LLVMIR/openmp-llvm-bad-schedule-modifier.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Target/LLVMIR/openmp-llvm-bad-schedule-modifier.mlir @@ -0,0 +1,49 @@ +// RUN: not mlir-translate -mlir-to-llvmir -split-input-file %s 2>&1 | FileCheck %s + +llvm.func @test_omp_wsloop_dynamic_bad_modifier(%lb : i64, %ub : i64, %step : i64) -> () { + omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(dynamic, ginandtonic) { + // CHECK: unknown modifier type: ginandtonic + omp.yield + } + llvm.return +} + +// ----- + +llvm.func @test_omp_wsloop_dynamic_many_modifier(%lb : i64, %ub : i64, %step : i64) -> () { + omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(dynamic, monotonic, monotonic, monotonic) { + // CHECK: unexpected modifier(s) + omp.yield + } + llvm.return +} + +// ----- + +llvm.func @test_omp_wsloop_dynamic_wrong_modifier(%lb : i64, %ub : i64, %step : i64) -> () { + omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(dynamic, simd, monotonic) { + // CHECK: incorrect modifier order + omp.yield + } + llvm.return +} + +// ----- + +llvm.func @test_omp_wsloop_dynamic_wrong_modifier2(%lb : i64, %ub : i64, %step : i64) -> () { + omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(dynamic, monotonic, monotonic) { + // CHECK: incorrect modifier order + omp.yield + } + llvm.return +} + +// ----- + +llvm.func @test_omp_wsloop_dynamic_wrong_modifier3(%lb : i64, %ub : i64, %step : i64) -> () { + omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(dynamic, simd, simd) { + // CHECK: incorrect modifier order + omp.yield + } + llvm.return +} diff --git a/mlir/test/Target/LLVMIR/openmp-llvm.mlir b/mlir/test/Target/LLVMIR/openmp-llvm.mlir --- a/mlir/test/Target/LLVMIR/openmp-llvm.mlir +++ b/mlir/test/Target/LLVMIR/openmp-llvm.mlir @@ -458,15 +458,15 @@ llvm.func @body(i64) llvm.func @test_omp_wsloop_runtime(%lb : i64, %ub : i64, %step : i64) -> () { - omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(runtime) { - // CHECK: call void @__kmpc_dispatch_init_8u - // CHECK: %[[continue:.*]] = call i32 @__kmpc_dispatch_next_8u - // CHECK: %[[cond:.*]] = icmp ne i32 %[[continue]], 0 - // CHECK br i1 %[[cond]], label %omp_loop.header{{.*}}, label %omp_loop.exit{{.*}} - llvm.call @body(%iv) : (i64) -> () - omp.yield - } - llvm.return + omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(runtime) { + // CHECK: call void @__kmpc_dispatch_init_8u + // CHECK: %[[continue:.*]] = call i32 @__kmpc_dispatch_next_8u + // CHECK: %[[cond:.*]] = icmp ne i32 %[[continue]], 0 + // CHECK br i1 %[[cond]], label %omp_loop.header{{.*}}, label %omp_loop.exit{{.*}} + llvm.call @body(%iv) : (i64) -> () + omp.yield + } + llvm.return } // ----- @@ -474,15 +474,15 @@ llvm.func @body(i64) llvm.func @test_omp_wsloop_guided(%lb : i64, %ub : i64, %step : i64) -> () { - omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(guided) { - // CHECK: call void @__kmpc_dispatch_init_8u - // CHECK: %[[continue:.*]] = call i32 @__kmpc_dispatch_next_8u - // CHECK: %[[cond:.*]] = icmp ne i32 %[[continue]], 0 - // CHECK br i1 %[[cond]], label %omp_loop.header{{.*}}, label %omp_loop.exit{{.*}} - llvm.call @body(%iv) : (i64) -> () - omp.yield - } - llvm.return + omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(guided) { + // CHECK: call void @__kmpc_dispatch_init_8u + // CHECK: %[[continue:.*]] = call i32 @__kmpc_dispatch_next_8u + // CHECK: %[[cond:.*]] = icmp ne i32 %[[continue]], 0 + // CHECK br i1 %[[cond]], label %omp_loop.header{{.*}}, label %omp_loop.exit{{.*}} + llvm.call @body(%iv) : (i64) -> () + omp.yield + } + llvm.return } // ----- @@ -490,15 +490,15 @@ llvm.func @body(i64) llvm.func @test_omp_wsloop_dynamic_nonmonotonic(%lb : i64, %ub : i64, %step : i64) -> () { - omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(dynamic, nonmonotonic) { - // CHECK: call void @__kmpc_dispatch_init_8u(%struct.ident_t* @{{.*}}, i32 %{{.*}}, i32 1073741859 - // CHECK: %[[continue:.*]] = call i32 @__kmpc_dispatch_next_8u - // CHECK: %[[cond:.*]] = icmp ne i32 %[[continue]], 0 - // CHECK br i1 %[[cond]], label %omp_loop.header{{.*}}, label %omp_loop.exit{{.*}} - llvm.call @body(%iv) : (i64) -> () - omp.yield - } - llvm.return + omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(dynamic, nonmonotonic) { + // CHECK: call void @__kmpc_dispatch_init_8u(%struct.ident_t* @{{.*}}, i32 %{{.*}}, i32 1073741859 + // CHECK: %[[continue:.*]] = call i32 @__kmpc_dispatch_next_8u + // CHECK: %[[cond:.*]] = icmp ne i32 %[[continue]], 0 + // CHECK br i1 %[[cond]], label %omp_loop.header{{.*}}, label %omp_loop.exit{{.*}} + llvm.call @body(%iv) : (i64) -> () + omp.yield + } + llvm.return } // ----- @@ -506,15 +506,39 @@ llvm.func @body(i64) llvm.func @test_omp_wsloop_dynamic_monotonic(%lb : i64, %ub : i64, %step : i64) -> () { - omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(dynamic, monotonic) { - // CHECK: call void @__kmpc_dispatch_init_8u(%struct.ident_t* @{{.*}}, i32 %{{.*}}, i32 536870947 - // CHECK: %[[continue:.*]] = call i32 @__kmpc_dispatch_next_8u - // CHECK: %[[cond:.*]] = icmp ne i32 %[[continue]], 0 - // CHECK br i1 %[[cond]], label %omp_loop.header{{.*}}, label %omp_loop.exit{{.*}} - llvm.call @body(%iv) : (i64) -> () - omp.yield - } - llvm.return + omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(dynamic, monotonic) { + // CHECK: call void @__kmpc_dispatch_init_8u(%struct.ident_t* @{{.*}}, i32 %{{.*}}, i32 536870947 + // CHECK: %[[continue:.*]] = call i32 @__kmpc_dispatch_next_8u + // CHECK: %[[cond:.*]] = icmp ne i32 %[[continue]], 0 + // CHECK br i1 %[[cond]], label %omp_loop.header{{.*}}, label %omp_loop.exit{{.*}} + llvm.call @body(%iv) : (i64) -> () + omp.yield + } + llvm.return +} + +llvm.func @test_omp_wsloop_runtime_simd(%lb : i64, %ub : i64, %step : i64) -> () { + omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(runtime, simd) { + // CHECK: call void @__kmpc_dispatch_init_8u(%struct.ident_t* @{{.*}}, i32 %{{.*}}, i32 47 + // CHECK: %[[continue:.*]] = call i32 @__kmpc_dispatch_next_8u + // CHECK: %[[cond:.*]] = icmp ne i32 %[[continue]], 0 + // CHECK br i1 %[[cond]], label %omp_loop.header{{.*}}, label %omp_loop.exit{{.*}} + llvm.call @body(%iv) : (i64) -> () + omp.yield + } + llvm.return +} + +llvm.func @test_omp_wsloop_guided_simd(%lb : i64, %ub : i64, %step : i64) -> () { + omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(guided, simd) { + // CHECK: call void @__kmpc_dispatch_init_8u(%struct.ident_t* @{{.*}}, i32 %{{.*}}, i32 46 + // CHECK: %[[continue:.*]] = call i32 @__kmpc_dispatch_next_8u + // CHECK: %[[cond:.*]] = icmp ne i32 %[[continue]], 0 + // CHECK br i1 %[[cond]], label %omp_loop.header{{.*}}, label %omp_loop.exit{{.*}} + llvm.call @body(%iv) : (i64) -> () + omp.yield + } + llvm.return } // -----