[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:
parent
286e98b97e
commit
3f00e10bdd
|
@ -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,
|
||||
|
|
|
@ -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())
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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)
|
||||
|
|
Loading…
Reference in New Issue