Skip to content

Commit b0d76b0

Browse files
ivanradanovivanradanov
committed
[MLIR][omp] Add omp.workshare op
Add custom omp loop wrapper Add recursive memory effects trait to workshare Remove stray include Remove omp.workshare verifier Add assembly format for wrapper and add test Add verification and descriptions
1 parent 0f0a96b commit b0d76b0

File tree

5 files changed

+246
-4
lines changed

5 files changed

+246
-4
lines changed

‎mlir/include/mlir/Dialect/OpenMP/OpenMPClauseOperands.h

Lines changed: 69 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -44,10 +44,75 @@ struct DeviceTypeClauseOps {
4444
// TODO: Add `indirect` clause.
4545
using DeclareTargetOperands = detail::Clauses<DeviceTypeClauseOps>;
4646

47-
/// omp.target_enter_data, omp.target_exit_data and omp.target_update take the
48-
/// same clauses, so we give the structure to be shared by all of them a
49-
/// representative name.
50-
using TargetEnterExitUpdateDataOperands = TargetEnterDataOperands;
47+
using DistributeOperands =
48+
detail::Clauses<AllocateClauseOps, DistScheduleClauseOps, OrderClauseOps,
49+
PrivateClauseOps>;
50+
51+
using LoopNestOperands = detail::Clauses<LoopRelatedOps>;
52+
53+
using MaskedOperands = detail::Clauses<FilterClauseOps>;
54+
55+
using OrderedOperands = detail::Clauses<DoacrossClauseOps>;
56+
57+
using OrderedRegionOperands = detail::Clauses<ParallelizationLevelClauseOps>;
58+
59+
using ParallelOperands =
60+
detail::Clauses<AllocateClauseOps, IfClauseOps, NumThreadsClauseOps,
61+
PrivateClauseOps, ProcBindClauseOps, ReductionClauseOps>;
62+
63+
using SectionsOperands = detail::Clauses<AllocateClauseOps, NowaitClauseOps,
64+
PrivateClauseOps, ReductionClauseOps>;
65+
66+
using SimdOperands =
67+
detail::Clauses<AlignedClauseOps, IfClauseOps, LinearClauseOps,
68+
NontemporalClauseOps, OrderClauseOps, PrivateClauseOps,
69+
ReductionClauseOps, SafelenClauseOps, SimdlenClauseOps>;
70+
71+
using SingleOperands = detail::Clauses<AllocateClauseOps, CopyprivateClauseOps,
72+
NowaitClauseOps, PrivateClauseOps>;
73+
74+
// TODO `defaultmap`, `uses_allocators` clauses.
75+
using TargetOperands =
76+
detail::Clauses<AllocateClauseOps, DependClauseOps, DeviceClauseOps,
77+
HasDeviceAddrClauseOps, IfClauseOps, InReductionClauseOps,
78+
IsDevicePtrClauseOps, MapClauseOps, NowaitClauseOps,
79+
PrivateClauseOps, ThreadLimitClauseOps>;
80+
81+
using TargetDataOperands =
82+
detail::Clauses<DeviceClauseOps, IfClauseOps, MapClauseOps,
83+
UseDeviceAddrClauseOps, UseDevicePtrClauseOps>;
84+
85+
using TargetEnterExitUpdateDataOperands =
86+
detail::Clauses<DependClauseOps, DeviceClauseOps, IfClauseOps, MapClauseOps,
87+
NowaitClauseOps>;
88+
89+
// TODO `affinity`, `detach` clauses.
90+
using TaskOperands =
91+
detail::Clauses<AllocateClauseOps, DependClauseOps, FinalClauseOps,
92+
IfClauseOps, InReductionClauseOps, MergeableClauseOps,
93+
PriorityClauseOps, PrivateClauseOps, UntiedClauseOps>;
94+
95+
using TaskgroupOperands =
96+
detail::Clauses<AllocateClauseOps, TaskReductionClauseOps>;
97+
98+
using TaskloopOperands =
99+
detail::Clauses<AllocateClauseOps, FinalClauseOps, GrainsizeClauseOps,
100+
IfClauseOps, InReductionClauseOps, MergeableClauseOps,
101+
NogroupClauseOps, NumTasksClauseOps, PriorityClauseOps,
102+
PrivateClauseOps, ReductionClauseOps, UntiedClauseOps>;
103+
104+
using TaskwaitOperands = detail::Clauses<DependClauseOps, NowaitClauseOps>;
105+
106+
using TeamsOperands =
107+
detail::Clauses<AllocateClauseOps, IfClauseOps, NumTeamsClauseOps,
108+
PrivateClauseOps, ReductionClauseOps, ThreadLimitClauseOps>;
109+
110+
using WorkshareOperands = detail::Clauses<NowaitClauseOps>;
111+
112+
using WsloopOperands =
113+
detail::Clauses<AllocateClauseOps, LinearClauseOps, NowaitClauseOps,
114+
OrderClauseOps, OrderedClauseOps, PrivateClauseOps,
115+
ReductionClauseOps, ScheduleClauseOps>;
51116

52117
} // namespace omp
53118
} // namespace mlir

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

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -313,6 +313,49 @@ def SingleOp : OpenMP_Op<"single", traits = [
313313
let hasVerifier = 1;
314314
}
315315

316+
//===----------------------------------------------------------------------===//
317+
// 2.8.3 Workshare Construct
318+
//===----------------------------------------------------------------------===//
319+
320+
def WorkshareOp : OpenMP_Op<"workshare", traits = [
321+
RecursiveMemoryEffects,
322+
], clauses = [
323+
OpenMP_NowaitClause,
324+
], singleRegion = true> {
325+
let summary = "workshare directive";
326+
let description = [{
327+
The workshare construct divides the execution of the enclosed structured
328+
block into separate units of work, and causes the threads of the team to
329+
share the work such that each unit is executed only once by one thread, in
330+
the context of its implicit task
331+
332+
This operation is used for the intermediate representation of the workshare
333+
block before the work gets divided between the threads. See the flang
334+
LowerWorkshare pass for details.
335+
}] # clausesDescription;
336+
337+
let builders = [
338+
OpBuilder<(ins CArg<"const WorkshareOperands &">:$clauses)>
339+
];
340+
}
341+
342+
def WorkshareLoopWrapperOp : OpenMP_Op<"workshare.loop_wrapper", traits = [
343+
DeclareOpInterfaceMethods<LoopWrapperInterface>,
344+
RecursiveMemoryEffects, SingleBlock
345+
], singleRegion = true> {
346+
let summary = "contains loop nests to be parallelized by workshare";
347+
let description = [{
348+
This operation wraps a loop nest that is marked for dividing into units of
349+
work by an encompassing omp.workshare operation.
350+
}];
351+
352+
let builders = [
353+
OpBuilder<(ins), [{ build($_builder, $_state, {}); }]>
354+
];
355+
let assemblyFormat = "$region attr-dict";
356+
let hasVerifier = 1;
357+
}
358+
316359
//===----------------------------------------------------------------------===//
317360
// Loop Nest
318361
//===----------------------------------------------------------------------===//

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1921,6 +1921,29 @@ LogicalResult SingleOp::verify() {
19211921
getCopyprivateSyms());
19221922
}
19231923

1924+
//===----------------------------------------------------------------------===//
1925+
// WorkshareOp
1926+
//===----------------------------------------------------------------------===//
1927+
1928+
void WorkshareOp::build(OpBuilder &builder, OperationState &state,
1929+
const WorkshareOperands &clauses) {
1930+
WorkshareOp::build(builder, state, clauses.nowait);
1931+
}
1932+
1933+
//===----------------------------------------------------------------------===//
1934+
// WorkshareLoopWrapperOp
1935+
//===----------------------------------------------------------------------===//
1936+
1937+
LogicalResult WorkshareLoopWrapperOp::verify() {
1938+
if (!isWrapper())
1939+
return emitOpError() << "must be a loop wrapper";
1940+
if (getNestedWrapper())
1941+
return emitError() << "nested wrappers not supported";
1942+
if (!(*this)->getParentOfType<WorkshareOp>())
1943+
return emitError() << "must be nested in an omp.workshare";
1944+
return success();
1945+
}
1946+
19241947
//===----------------------------------------------------------------------===//
19251948
// LoopWrapperInterface
19261949
//===----------------------------------------------------------------------===//

‎mlir/test/Dialect/OpenMP/invalid.mlir

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2577,3 +2577,45 @@ func.func @omp_taskloop_invalid_composite(%lb: index, %ub: index, %step: index)
25772577
} {omp.composite}
25782578
return
25792579
}
2580+
2581+
// -----
2582+
func.func @nested_wrapper(%idx : index) {
2583+
omp.workshare {
2584+
// expected-error @below {{nested wrappers not supported}}
2585+
omp.workshare.loop_wrapper {
2586+
omp.simd {
2587+
omp.loop_nest (%iv) : index = (%idx) to (%idx) step (%idx) {
2588+
omp.yield
2589+
}
2590+
omp.terminator
2591+
}
2592+
omp.terminator
2593+
}
2594+
omp.terminator
2595+
}
2596+
return
2597+
}
2598+
2599+
// -----
2600+
func.func @not_wrapper() {
2601+
omp.workshare {
2602+
// expected-error @below {{must be a loop wrapper}}
2603+
omp.workshare.loop_wrapper {
2604+
omp.terminator
2605+
}
2606+
omp.terminator
2607+
}
2608+
return
2609+
}
2610+
2611+
// -----
2612+
func.func @missing_workshare(%idx : index) {
2613+
// expected-error @below {{must be nested in an omp.workshare}}
2614+
omp.workshare.loop_wrapper {
2615+
omp.loop_nest (%iv) : index = (%idx) to (%idx) step (%idx) {
2616+
omp.yield
2617+
}
2618+
omp.terminator
2619+
}
2620+
return
2621+
}

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

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2749,3 +2749,72 @@ func.func @omp_target_private(%map1: memref<?xi32>, %map2: memref<?xi32>, %priv_
27492749

27502750
return
27512751
}
2752+
2753+
// CHECK-LABEL: func @omp_workshare
2754+
func.func @omp_workshare() {
2755+
// CHECK: omp.workshare {
2756+
omp.workshare {
2757+
"test.payload"() : () -> ()
2758+
// CHECK: omp.terminator
2759+
omp.terminator
2760+
}
2761+
return
2762+
}
2763+
2764+
// CHECK-LABEL: func @omp_workshare_nowait
2765+
func.func @omp_workshare_nowait() {
2766+
// CHECK: omp.workshare nowait {
2767+
omp.workshare nowait {
2768+
"test.payload"() : () -> ()
2769+
// CHECK: omp.terminator
2770+
omp.terminator
2771+
}
2772+
return
2773+
}
2774+
2775+
// CHECK-LABEL: func @omp_workshare_multiple_blocks
2776+
func.func @omp_workshare_multiple_blocks() {
2777+
// CHECK: omp.workshare {
2778+
omp.workshare {
2779+
cf.br ^bb2
2780+
^bb2:
2781+
// CHECK: omp.terminator
2782+
omp.terminator
2783+
}
2784+
return
2785+
}
2786+
2787+
// CHECK-LABEL: func @omp_workshare.loop_wrapper
2788+
func.func @omp_workshare.loop_wrapper(%idx : index) {
2789+
// CHECK-NEXT: omp.workshare {
2790+
omp.workshare {
2791+
// CHECK-NEXT: omp.workshare.loop_wrapper
2792+
omp.workshare.loop_wrapper {
2793+
// CHECK-NEXT: omp.loop_nest
2794+
omp.loop_nest (%iv) : index = (%idx) to (%idx) step (%idx) {
2795+
omp.yield
2796+
}
2797+
omp.terminator
2798+
}
2799+
omp.terminator
2800+
}
2801+
return
2802+
}
2803+
2804+
// CHECK-LABEL: func @omp_workshare.loop_wrapper_attrs
2805+
func.func @omp_workshare.loop_wrapper_attrs(%idx : index) {
2806+
// CHECK-NEXT: omp.workshare {
2807+
omp.workshare {
2808+
// CHECK-NEXT: omp.workshare.loop_wrapper {
2809+
omp.workshare.loop_wrapper {
2810+
// CHECK-NEXT: omp.loop_nest
2811+
omp.loop_nest (%iv) : index = (%idx) to (%idx) step (%idx) {
2812+
omp.yield
2813+
}
2814+
omp.terminator
2815+
// CHECK: } {attr_in_dict}
2816+
} {attr_in_dict}
2817+
omp.terminator
2818+
}
2819+
return
2820+
}

0 commit comments

Comments
 (0)