diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td index f79a3eb88e4b5..156e6eb371b85 100644 --- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td +++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td @@ -313,6 +313,49 @@ def SingleOp : OpenMP_Op<"single", traits = [ let hasVerifier = 1; } +//===----------------------------------------------------------------------===// +// 2.8.3 Workshare Construct +//===----------------------------------------------------------------------===// + +def WorkshareOp : OpenMP_Op<"workshare", traits = [ + RecursiveMemoryEffects, + ], clauses = [ + OpenMP_NowaitClause, + ], singleRegion = true> { + let summary = "workshare directive"; + let description = [{ + The workshare construct divides the execution of the enclosed structured + 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 = [ + OpBuilder<(ins CArg<"const WorkshareOperands &">:$clauses)> + ]; +} + +def WorkshareLoopWrapperOp : OpenMP_Op<"workshare.loop_wrapper", traits = [ + DeclareOpInterfaceMethods, NoTerminator, + 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; +} + //===----------------------------------------------------------------------===// // Loop Nest //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp index 19e0fa30a7571..94e71e089d4b1 100644 --- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp +++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp @@ -1897,6 +1897,27 @@ LogicalResult SingleOp::verify() { getCopyprivateSyms()); } +//===----------------------------------------------------------------------===// +// WorkshareOp +//===----------------------------------------------------------------------===// + +void WorkshareOp::build(OpBuilder &builder, OperationState &state, + const WorkshareOperands &clauses) { + WorkshareOp::build(builder, state, clauses.nowait); +} + +//===----------------------------------------------------------------------===// +// WorkshareLoopWrapperOp +//===----------------------------------------------------------------------===// + +LogicalResult WorkshareLoopWrapperOp::verify() { + if (!(*this)->getParentOfType()) + return emitError() << "must be nested in an omp.workshare"; + if (getNestedWrapper()) + return emitError() << "cannot be composite"; + return success(); +} + //===----------------------------------------------------------------------===// // LoopWrapperInterface //===----------------------------------------------------------------------===// diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir index aa41eea44f3ef..2a19e4837f550 100644 --- a/mlir/test/Dialect/OpenMP/invalid.mlir +++ b/mlir/test/Dialect/OpenMP/invalid.mlir @@ -2620,6 +2620,44 @@ func.func @omp_loop_invalid_binding(%lb : index, %ub : index, %step : index) { omp.yield } } + return +} + +// ----- +func.func @nested_wrapper(%idx : index) { + omp.workshare { + // expected-error @below {{cannot be composite}} + omp.workshare.loop_wrapper { + omp.simd { + omp.loop_nest (%iv) : index = (%idx) to (%idx) step (%idx) { + omp.yield + } + } {omp.composite} + } + omp.terminator + } + return +} + +// ----- +func.func @not_wrapper() { + omp.workshare { + // expected-error @below {{op nested in loop wrapper is not another loop wrapper or `omp.loop_nest`}} + omp.workshare.loop_wrapper { + %0 = arith.constant 0 : index + } + 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 + } + } return } diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir index 69c53d1f77e84..c25a6ef4b4849 100644 --- a/mlir/test/Dialect/OpenMP/ops.mlir +++ b/mlir/test/Dialect/OpenMP/ops.mlir @@ -2789,3 +2789,70 @@ func.func @omp_loop(%lb : index, %ub : index, %step : index) { return } + +// CHECK-LABEL: func @omp_workshare +func.func @omp_workshare() { + // CHECK: omp.workshare { + omp.workshare { + "test.payload"() : () -> () + // CHECK: omp.terminator + omp.terminator + } + return +} + +// CHECK-LABEL: func @omp_workshare_nowait +func.func @omp_workshare_nowait() { + // CHECK: omp.workshare nowait { + omp.workshare nowait { + "test.payload"() : () -> () + // CHECK: omp.terminator + omp.terminator + } + return +} + +// CHECK-LABEL: func @omp_workshare_multiple_blocks +func.func @omp_workshare_multiple_blocks() { + // CHECK: omp.workshare { + omp.workshare { + cf.br ^bb2 + ^bb2: + // CHECK: omp.terminator + omp.terminator + } + return +} + +// CHECK-LABEL: func @omp_workshare_loop_wrapper +func.func @omp_workshare_loop_wrapper(%idx : index) { + // 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 + } + return +} + +// CHECK-LABEL: func @omp_workshare_loop_wrapper_attrs +func.func @omp_workshare_loop_wrapper_attrs(%idx : index) { + // 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 + } + // CHECK: } {attr_in_dict} + } {attr_in_dict} + omp.terminator + } + return +}