Skip to content

Commit 3f00e10

Browse files
LeporacanthicusLeporacanthicus
committed
[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
1 parent 286e98b commit 3f00e10

File tree

5 files changed

+89
-5
lines changed

5 files changed

+89
-5
lines changed

‎mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,18 @@ def TerminatorOp : OpenMP_Op<"terminator", [Terminator]> {
136136
let assemblyFormat = "attr-dict";
137137
}
138138

139+
def OMP_SCHEDULE_MOD_None : StrEnumAttrCase<"none", 0>;
140+
def OMP_SCHEDULE_MOD_Monotonic : StrEnumAttrCase<"monotonic", 1>;
141+
def OMP_SCHEDULE_MOD_Nonmonotonic : StrEnumAttrCase<"nonmonotonic", 2>;
142+
143+
def ScheduleModifier : StrEnumAttr<"ScheduleModifier", "OpenMP Schedule Modifier",
144+
[OMP_SCHEDULE_MOD_None,
145+
OMP_SCHEDULE_MOD_Monotonic,
146+
OMP_SCHEDULE_MOD_Nonmonotonic]>
147+
{
148+
let cppNamespace = "::mlir::omp";
149+
}
150+
139151
//===----------------------------------------------------------------------===//
140152
// 2.9.2 Workshare Loop Construct
141153
//===----------------------------------------------------------------------===//
@@ -214,6 +226,7 @@ def WsLoopOp : OpenMP_Op<"wsloop", [AttrSizedOperandSegments,
214226
OptionalAttr<SymbolRefArrayAttr>:$reductions,
215227
OptionalAttr<ScheduleKind>:$schedule_val,
216228
Optional<AnyType>:$schedule_chunk_var,
229+
OptionalAttr<ScheduleModifier>:$schedule_modifier,
217230
Confined<OptionalAttr<I64Attr>, [IntMinValue<0>]>:$collapse_val,
218231
UnitAttr:$nowait,
219232
Confined<OptionalAttr<I64Attr>, [IntMinValue<0>]>:$ordered_val,

‎mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -253,6 +253,7 @@ static void printLinearClause(OpAsmPrinter &p, OperandRange linearVars,
253253
/// sched-wo-chunk ::= `auto` | `runtime`
254254
static ParseResult
255255
parseScheduleClause(OpAsmParser &parser, SmallString<8> &schedule,
256+
SmallVectorImpl<SmallString<12>> &modifiers,
256257
Optional<OpAsmParser::OperandType> &chunkSize) {
257258
if (parser.parseLParen())
258259
return failure();
@@ -276,6 +277,14 @@ parseScheduleClause(OpAsmParser &parser, SmallString<8> &schedule,
276277
return parser.emitError(parser.getNameLoc()) << " expected schedule kind";
277278
}
278279

280+
// If there is a comma, we have one or more modifiers..
281+
if (succeeded(parser.parseOptionalComma())) {
282+
StringRef mod;
283+
if (parser.parseKeyword(&mod))
284+
return failure();
285+
modifiers.push_back(mod);
286+
}
287+
279288
if (parser.parseRParen())
280289
return failure();
281290

@@ -284,11 +293,14 @@ parseScheduleClause(OpAsmParser &parser, SmallString<8> &schedule,
284293

285294
/// Print schedule clause
286295
static void printScheduleClause(OpAsmPrinter &p, StringRef &sched,
296+
llvm::Optional<StringRef> modifier,
287297
Value scheduleChunkVar) {
288298
std::string schedLower = sched.lower();
289299
p << "(" << schedLower;
290300
if (scheduleChunkVar)
291301
p << " = " << scheduleChunkVar;
302+
if (modifier && modifier.getValue() != "none")
303+
p << ", " << modifier;
292304
p << ") ";
293305
}
294306

@@ -551,6 +563,7 @@ static ParseResult parseClauses(OpAsmParser &parser, OperationState &result,
551563
SmallVector<OpAsmParser::OperandType> linearSteps;
552564

553565
SmallString<8> schedule;
566+
SmallVector<SmallString<12>> modifiers;
554567
Optional<OpAsmParser::OperandType> scheduleChunkSize;
555568

556569
// Compute the position of clauses in operand segments
@@ -670,7 +683,7 @@ static ParseResult parseClauses(OpAsmParser &parser, OperationState &result,
670683
clauseSegments[pos[linearClause] + 1] = linearSteps.size();
671684
} else if (clauseKeyword == "schedule") {
672685
if (checkAllowed(scheduleClause) ||
673-
parseScheduleClause(parser, schedule, scheduleChunkSize))
686+
parseScheduleClause(parser, schedule, modifiers, scheduleChunkSize))
674687
return failure();
675688
if (scheduleChunkSize) {
676689
clauseSegments[pos[scheduleClause]] = 1;
@@ -797,6 +810,10 @@ static ParseResult parseClauses(OpAsmParser &parser, OperationState &result,
797810
schedule[0] = llvm::toUpper(schedule[0]);
798811
auto attr = parser.getBuilder().getStringAttr(schedule);
799812
result.addAttribute("schedule_val", attr);
813+
if (modifiers.size() > 0) {
814+
auto mod = parser.getBuilder().getStringAttr(modifiers[0]);
815+
result.addAttribute("schedule_modifier", mod);
816+
}
800817
if (scheduleChunkSize) {
801818
auto chunkSizeType = parser.getBuilder().getI32Type();
802819
parser.resolveOperand(*scheduleChunkSize, chunkSizeType, result.operands);
@@ -916,7 +933,8 @@ static void printWsLoopOp(OpAsmPrinter &p, WsLoopOp op) {
916933

917934
if (auto sched = op.schedule_val()) {
918935
p << "schedule";
919-
printScheduleClause(p, sched.getValue(), op.schedule_chunk_var());
936+
printScheduleClause(p, sched.getValue(), op.schedule_modifier(),
937+
op.schedule_chunk_var());
920938
}
921939

922940
if (auto collapse = op.collapse_val())

‎mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -707,8 +707,23 @@ convertOmpWsLoop(Operation &opInst, llvm::IRBuilderBase &builder,
707707
break;
708708
}
709709

710-
ompBuilder->applyDynamicWorkshareLoop(ompLoc.DL, loopInfo, allocaIP,
711-
schedType, !loop.nowait(), chunk);
710+
if (loop.schedule_modifier().hasValue()) {
711+
omp::ScheduleModifier modifier =
712+
*omp::symbolizeScheduleModifier(loop.schedule_modifier().getValue());
713+
switch (modifier) {
714+
case omp::ScheduleModifier::monotonic:
715+
schedType |= llvm::omp::OMPScheduleType::ModifierMonotonic;
716+
break;
717+
case omp::ScheduleModifier::nonmonotonic:
718+
schedType |= llvm::omp::OMPScheduleType::ModifierNonmonotonic;
719+
break;
720+
default:
721+
// Nothing to do here.
722+
break;
723+
}
724+
}
725+
afterIP = ompBuilder->applyDynamicWorkshareLoop(
726+
ompLoc.DL, loopInfo, allocaIP, schedType, !loop.nowait(), chunk);
712727
}
713728

714729
// Continue building IR after the loop. Note that the LoopInfo returned by

‎mlir/test/Dialect/OpenMP/ops.mlir

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -177,7 +177,7 @@ func @omp_wsloop_pretty(%lb : index, %ub : index, %step : index,
177177
}
178178

179179
// CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(static)
180-
omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) schedule(static) lastprivate(%data_var : memref<i32>) linear(%data_var = %linear_var : memref<i32>) {
180+
omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) schedule(static, none) lastprivate(%data_var : memref<i32>) linear(%data_var = %linear_var : memref<i32>) {
181181
omp.yield
182182
}
183183

@@ -188,6 +188,20 @@ func @omp_wsloop_pretty(%lb : index, %ub : index, %step : index,
188188
omp.yield
189189
}
190190

191+
// 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)
192+
omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) ordered(2) private(%data_var : memref<i32>)
193+
firstprivate(%data_var : memref<i32>) lastprivate(%data_var : memref<i32>) linear(%data_var = %linear_var : memref<i32>)
194+
schedule(dynamic = %chunk_var, nonmonotonic) collapse(3) {
195+
omp.yield
196+
}
197+
198+
// 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)
199+
omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) ordered(2) private(%data_var : memref<i32>)
200+
firstprivate(%data_var : memref<i32>) lastprivate(%data_var : memref<i32>) linear(%data_var = %linear_var : memref<i32>)
201+
schedule(dynamic = %chunk_var, monotonic) collapse(3) {
202+
omp.yield
203+
}
204+
191205
// CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private({{.*}} : memref<i32>)
192206
omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) private(%data_var : memref<i32>) {
193207
omp.yield

‎mlir/test/Target/LLVMIR/openmp-llvm.mlir

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -467,6 +467,30 @@ llvm.func @test_omp_wsloop_guided(%lb : i64, %ub : i64, %step : i64) -> () {
467467
llvm.return
468468
}
469469

470+
llvm.func @test_omp_wsloop_dynamic_nonmonotonic(%lb : i64, %ub : i64, %step : i64) -> () {
471+
omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(dynamic, nonmonotonic) {
472+
// CHECK: call void @__kmpc_dispatch_init_8u(%struct.ident_t* @{{.*}}, i32 %{{.*}}, i32 1073741859
473+
// CHECK: %[[continue:.*]] = call i32 @__kmpc_dispatch_next_8u
474+
// CHECK: %[[cond:.*]] = icmp ne i32 %[[continue]], 0
475+
// CHECK br i1 %[[cond]], label %omp_loop.header{{.*}}, label %omp_loop.exit{{.*}}
476+
llvm.call @body(%iv) : (i64) -> ()
477+
omp.yield
478+
}
479+
llvm.return
480+
}
481+
482+
llvm.func @test_omp_wsloop_dynamic_monotonic(%lb : i64, %ub : i64, %step : i64) -> () {
483+
omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(dynamic, monotonic) {
484+
// CHECK: call void @__kmpc_dispatch_init_8u(%struct.ident_t* @{{.*}}, i32 %{{.*}}, i32 536870947
485+
// CHECK: %[[continue:.*]] = call i32 @__kmpc_dispatch_next_8u
486+
// CHECK: %[[cond:.*]] = icmp ne i32 %[[continue]], 0
487+
// CHECK br i1 %[[cond]], label %omp_loop.header{{.*}}, label %omp_loop.exit{{.*}}
488+
llvm.call @body(%iv) : (i64) -> ()
489+
omp.yield
490+
}
491+
llvm.return
492+
}
493+
470494
// -----
471495

472496
omp.critical.declare @mutex hint(contended)

0 commit comments

Comments
 (0)