diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td index b41c53ef282047..f8d7e934edbc61 100644 --- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td +++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td @@ -302,6 +302,10 @@ def WorkshareOp : OpenMP_Op<"workshare", traits = [ block into separate units of work, and causes the threads of the team to share the work such that each unit is executed only once by one thread, in the context of its implicit task + + This operation is used for the intermediate representation of the workshare + block before the work gets divided between the threads. See the flang + LowerWorkshare pass for details. }] # clausesDescription; let builders = [ @@ -314,10 +318,16 @@ def WorkshareLoopWrapperOp : OpenMP_Op<"workshare_loop_wrapper", traits = [ RecursiveMemoryEffects, SingleBlock ], singleRegion = true> { let summary = "contains loop nests to be parallelized by workshare"; + let description = [{ + This operation wraps a loop nest that is marked for dividing into units of + work by an encompassing omp.workshare operation. + }]; + let builders = [ OpBuilder<(ins), [{ build($_builder, $_state, {}); }]> ]; let assemblyFormat = "$region attr-dict"; + let hasVerifier = 1; } //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp index 230af836c1531e..f4acbd97ca6d1a 100644 --- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp +++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp @@ -1698,6 +1698,20 @@ void WorkshareOp::build(OpBuilder &builder, OperationState &state, WorkshareOp::build(builder, state, clauses.nowait); } +//===----------------------------------------------------------------------===// +// WorkshareLoopWrapperOp +//===----------------------------------------------------------------------===// + +LogicalResult WorkshareLoopWrapperOp::verify() { + if (!isWrapper()) + return emitOpError() << "must be a loop wrapper"; + if (getNestedWrapper()) + return emitError() << "nested wrappers not supported"; + if (!(*this)->getParentOfType()) + return emitError() << "must be nested in an omp.workshare"; + return success(); +} + //===----------------------------------------------------------------------===// // WsloopOp //===----------------------------------------------------------------------===// diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir index c76b07ec94a597..a67dcb34427062 100644 --- a/mlir/test/Dialect/OpenMP/invalid.mlir +++ b/mlir/test/Dialect/OpenMP/invalid.mlir @@ -2545,3 +2545,45 @@ func.func @omp_taskloop_invalid_composite(%lb: index, %ub: index, %step: index) } {omp.composite} return } + +// ----- +func.func @nested_wrapper(%idx : index) { + omp.workshare { + // expected-error @below {{nested wrappers not supported}} + omp.workshare_loop_wrapper { + omp.simd { + omp.loop_nest (%iv) : index = (%idx) to (%idx) step (%idx) { + omp.yield + } + omp.terminator + } + omp.terminator + } + omp.terminator + } + return +} + +// ----- +func.func @not_wrapper() { + omp.workshare { + // expected-error @below {{must be a loop wrapper}} + omp.workshare_loop_wrapper { + omp.terminator + } + omp.terminator + } + return +} + +// ----- +func.func @missing_workshare(%idx : index) { + // expected-error @below {{must be nested in an omp.workshare}} + omp.workshare_loop_wrapper { + omp.loop_nest (%iv) : index = (%idx) to (%idx) step (%idx) { + omp.yield + } + omp.terminator + } + return +} diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir index 8ef49ddc807b6a..1e0ccfb79ed54e 100644 --- a/mlir/test/Dialect/OpenMP/ops.mlir +++ b/mlir/test/Dialect/OpenMP/ops.mlir @@ -2847,11 +2847,15 @@ func.func @omp_workshare_multiple_blocks() { // CHECK-LABEL: func @omp_workshare_loop_wrapper func.func @omp_workshare_loop_wrapper(%idx : index) { - // CHECK-NEXT: omp.workshare_loop_wrapper - omp.workshare_loop_wrapper { - // CHECK-NEXT: omp.loop_nest - omp.loop_nest (%iv) : index = (%idx) to (%idx) step (%idx) { - omp.yield + // CHECK-NEXT: omp.workshare { + omp.workshare { + // CHECK-NEXT: omp.workshare_loop_wrapper + omp.workshare_loop_wrapper { + // CHECK-NEXT: omp.loop_nest + omp.loop_nest (%iv) : index = (%idx) to (%idx) step (%idx) { + omp.yield + } + omp.terminator } omp.terminator } @@ -2860,14 +2864,18 @@ func.func @omp_workshare_loop_wrapper(%idx : index) { // CHECK-LABEL: func @omp_workshare_loop_wrapper_attrs func.func @omp_workshare_loop_wrapper_attrs(%idx : index) { - // CHECK-NEXT: omp.workshare_loop_wrapper { - omp.workshare_loop_wrapper { - // CHECK-NEXT: omp.loop_nest - omp.loop_nest (%iv) : index = (%idx) to (%idx) step (%idx) { - omp.yield - } + // CHECK-NEXT: omp.workshare { + omp.workshare { + // CHECK-NEXT: omp.workshare_loop_wrapper { + omp.workshare_loop_wrapper { + // CHECK-NEXT: omp.loop_nest + omp.loop_nest (%iv) : index = (%idx) to (%idx) step (%idx) { + omp.yield + } + omp.terminator + // CHECK: } {attr_in_dict} + } {attr_in_dict} omp.terminator - // CHECK: } {attr_in_dict} - } {attr_in_dict} + } return }