[mlir][OpenMP]Support for modifiers in workshare loops

Pass the modifiers from the Flang parser to FIR/MLIR workshare
loop operation.

Not yet supporting the SIMD modifier, which is a bit more work
than just adding it to the list of modifiers, so will go in a
separate patch.

This adds a new field to the WsLoopOp.

Also add test for dynamic WSLoop, checking that dynamic schedule calls
the init and next functions as expected.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D111053
This commit is contained in:
Mats Petersson 2021-04-06 11:20:49 +01:00
parent 286e98b97e
commit 3f00e10bdd
5 changed files with 89 additions and 5 deletions

View File

@ -136,6 +136,18 @@ def TerminatorOp : OpenMP_Op<"terminator", [Terminator]> {
let assemblyFormat = "attr-dict";
}
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 ScheduleModifier : StrEnumAttr<"ScheduleModifier", "OpenMP Schedule Modifier",
[OMP_SCHEDULE_MOD_None,
OMP_SCHEDULE_MOD_Monotonic,
OMP_SCHEDULE_MOD_Nonmonotonic]>
{
let cppNamespace = "::mlir::omp";
}
//===----------------------------------------------------------------------===//
// 2.9.2 Workshare Loop Construct
//===----------------------------------------------------------------------===//
@ -214,6 +226,7 @@ def WsLoopOp : OpenMP_Op<"wsloop", [AttrSizedOperandSegments,
OptionalAttr<SymbolRefArrayAttr>:$reductions,
OptionalAttr<ScheduleKind>:$schedule_val,
Optional<AnyType>:$schedule_chunk_var,
OptionalAttr<ScheduleModifier>:$schedule_modifier,
Confined<OptionalAttr<I64Attr>, [IntMinValue<0>]>:$collapse_val,
UnitAttr:$nowait,
Confined<OptionalAttr<I64Attr>, [IntMinValue<0>]>:$ordered_val,

View File

@ -253,6 +253,7 @@ static void printLinearClause(OpAsmPrinter &p, OperandRange linearVars,
/// sched-wo-chunk ::= `auto` | `runtime`
static ParseResult
parseScheduleClause(OpAsmParser &parser, SmallString<8> &schedule,
SmallVectorImpl<SmallString<12>> &modifiers,
Optional<OpAsmParser::OperandType> &chunkSize) {
if (parser.parseLParen())
return failure();
@ -276,6 +277,14 @@ parseScheduleClause(OpAsmParser &parser, SmallString<8> &schedule,
return parser.emitError(parser.getNameLoc()) << " expected schedule kind";
}
// If there is a comma, we have one or more modifiers..
if (succeeded(parser.parseOptionalComma())) {
StringRef mod;
if (parser.parseKeyword(&mod))
return failure();
modifiers.push_back(mod);
}
if (parser.parseRParen())
return failure();
@ -284,11 +293,14 @@ parseScheduleClause(OpAsmParser &parser, SmallString<8> &schedule,
/// Print schedule clause
static void printScheduleClause(OpAsmPrinter &p, StringRef &sched,
llvm::Optional<StringRef> modifier,
Value scheduleChunkVar) {
std::string schedLower = sched.lower();
p << "(" << schedLower;
if (scheduleChunkVar)
p << " = " << scheduleChunkVar;
if (modifier && modifier.getValue() != "none")
p << ", " << modifier;
p << ") ";
}
@ -551,6 +563,7 @@ static ParseResult parseClauses(OpAsmParser &parser, OperationState &result,
SmallVector<OpAsmParser::OperandType> linearSteps;
SmallString<8> schedule;
SmallVector<SmallString<12>> modifiers;
Optional<OpAsmParser::OperandType> scheduleChunkSize;
// Compute the position of clauses in operand segments
@ -670,7 +683,7 @@ static ParseResult parseClauses(OpAsmParser &parser, OperationState &result,
clauseSegments[pos[linearClause] + 1] = linearSteps.size();
} else if (clauseKeyword == "schedule") {
if (checkAllowed(scheduleClause) ||
parseScheduleClause(parser, schedule, scheduleChunkSize))
parseScheduleClause(parser, schedule, modifiers, scheduleChunkSize))
return failure();
if (scheduleChunkSize) {
clauseSegments[pos[scheduleClause]] = 1;
@ -797,6 +810,10 @@ static ParseResult parseClauses(OpAsmParser &parser, OperationState &result,
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_modifier", mod);
}
if (scheduleChunkSize) {
auto chunkSizeType = parser.getBuilder().getI32Type();
parser.resolveOperand(*scheduleChunkSize, chunkSizeType, result.operands);
@ -916,7 +933,8 @@ static void printWsLoopOp(OpAsmPrinter &p, WsLoopOp op) {
if (auto sched = op.schedule_val()) {
p << "schedule";
printScheduleClause(p, sched.getValue(), op.schedule_chunk_var());
printScheduleClause(p, sched.getValue(), op.schedule_modifier(),
op.schedule_chunk_var());
}
if (auto collapse = op.collapse_val())

View File

@ -707,8 +707,23 @@ convertOmpWsLoop(Operation &opInst, llvm::IRBuilderBase &builder,
break;
}
ompBuilder->applyDynamicWorkshareLoop(ompLoc.DL, loopInfo, allocaIP,
schedType, !loop.nowait(), chunk);
if (loop.schedule_modifier().hasValue()) {
omp::ScheduleModifier modifier =
*omp::symbolizeScheduleModifier(loop.schedule_modifier().getValue());
switch (modifier) {
case omp::ScheduleModifier::monotonic:
schedType |= llvm::omp::OMPScheduleType::ModifierMonotonic;
break;
case omp::ScheduleModifier::nonmonotonic:
schedType |= llvm::omp::OMPScheduleType::ModifierNonmonotonic;
break;
default:
// Nothing to do here.
break;
}
}
afterIP = ompBuilder->applyDynamicWorkshareLoop(
ompLoc.DL, loopInfo, allocaIP, schedType, !loop.nowait(), chunk);
}
// Continue building IR after the loop. Note that the LoopInfo returned by

View File

@ -177,7 +177,7 @@ func @omp_wsloop_pretty(%lb : index, %ub : index, %step : index,
}
// CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(static)
omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) schedule(static) lastprivate(%data_var : memref<i32>) linear(%data_var = %linear_var : memref<i32>) {
omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) schedule(static, none) lastprivate(%data_var : memref<i32>) linear(%data_var = %linear_var : memref<i32>) {
omp.yield
}
@ -188,6 +188,20 @@ func @omp_wsloop_pretty(%lb : index, %ub : index, %step : index,
omp.yield
}
// CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) lastprivate(%{{.*}} : memref<i32>) linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(dynamic = %{{.*}}, nonmonotonic) collapse(3) ordered(2)
omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) ordered(2) private(%data_var : memref<i32>)
firstprivate(%data_var : memref<i32>) lastprivate(%data_var : memref<i32>) linear(%data_var = %linear_var : memref<i32>)
schedule(dynamic = %chunk_var, nonmonotonic) collapse(3) {
omp.yield
}
// CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) lastprivate(%{{.*}} : memref<i32>) linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(dynamic = %{{.*}}, monotonic) collapse(3) ordered(2)
omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) ordered(2) private(%data_var : memref<i32>)
firstprivate(%data_var : memref<i32>) lastprivate(%data_var : memref<i32>) linear(%data_var = %linear_var : memref<i32>)
schedule(dynamic = %chunk_var, monotonic) collapse(3) {
omp.yield
}
// CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private({{.*}} : memref<i32>)
omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) private(%data_var : memref<i32>) {
omp.yield

View File

@ -467,6 +467,30 @@ llvm.func @test_omp_wsloop_guided(%lb : i64, %ub : i64, %step : i64) -> () {
llvm.return
}
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
}
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.critical.declare @mutex hint(contended)