diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td --- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td +++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td @@ -452,7 +452,7 @@ def YieldOp : OpenMP_Op<"yield", [Pure, ReturnLike, Terminator, ParentOneOf<["WsLoopOp", "ReductionDeclareOp", - "AtomicUpdateOp", "SimdLoopOp"]>]> { + "AtomicUpdateOp", "SimdLoopOp", "DistributeOp"]>]> { let summary = "loop yield and termination operation"; let description = [{ "omp.yield" yields SSA values from the OpenMP dialect op region and @@ -469,6 +469,65 @@ let assemblyFormat = [{ ( `(` $results^ `:` type($results) `)` )? attr-dict}]; } +//===----------------------------------------------------------------------===// +// 2.9.4.1 distribute Construct +//===----------------------------------------------------------------------===// + +def DistributeOp : OpenMP_Op<"distribute", [AttrSizedOperandSegments, + AllTypesMatch<["lowerBound", "upperBound", "step"]>]> { + let summary = "distribute loop construct"; + let description = [{ + The distribute construct specifies that the iterations of one or more loop + will be executed by the initial teams in the context of their implicit + tasks. The iterations are distributed across the initial threads of all + initial teams that execute the teams region to which the distribute region + binds. + + The distribute loop construct specifies that the iterations of the loop(s) + will be executed in parallel by threads in the current context. These + iterations are spread across threads that already exist in the enclosing + region. The lower and upper bounds specify a half-open range: the + range includes the lower bound but does not include the upper bound. If the + `inclusive` attribute is specified then the upper bound is also included. + + // TODO: private_var, firstprivate_var and lastprivate_var variables + // will be defined later + + The `static_dist_schedule` attribute specifies the schedule for this + loop, determining how the loop is distributed across the parallel threads. + The optional `schedule_chunk` associated with this determines further + controls this distribution. + + // TODO: collapse + }]; + + let arguments = (ins Variadic:$lowerBound, + Variadic:$upperBound, + Variadic:$step, + UnitAttr:$static_dist_schedule, + Optional:$schedule_chunk, + Variadic:$allocate_vars, + Variadic:$allocators_vars, + UnitAttr:$inclusive); + + let assemblyFormat = [{ + oilist(`dist_schedule` `(` `static` + ($static_dist_schedule^)? + ( $schedule_chunk^ `:` type($schedule_chunk))? `)` + | `allocate` `(` + custom( + $allocate_vars, type($allocate_vars), + $allocators_vars, type($allocators_vars) + ) `)` + + ) `for` custom($region, $lowerBound, $upperBound, $step, + type($step), $inclusive) attr-dict + }]; + + let regions = (region AnyRegion:$region); + let hasVerifier = 1; +} + //===----------------------------------------------------------------------===// // 2.10.1 task Construct //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp --- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp +++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp @@ -281,6 +281,26 @@ return success(); } +//===---------------------------------------------------------------------===// +// Verifier for Dristribute Op +//===---------------------------------------------------------------------===// + +LogicalResult DistributeOp::verify(){ + if (this->getLowerBound().empty()) + return emitOpError() << "empty lowerbound for distribute loop operation"; + + if (this->getUpperBound().empty()) + return emitOpError() << "empty upperbound for distribute loop operation"; + + if (this->getLowerBound().size() != this->getUpperBound().size()) + return emitOpError() << "upperbound and lowerbound sizes are not equal"; + + if (this->getLowerBound().size() != this->getStep().size()) + return emitOpError() << "upperbound, lowerbound and step sizes are not equal"; + + return success(); +} + /// schedule ::= `schedule` `(` sched-list `)` /// sched-list ::= sched-val | sched-val sched-list | /// sched-val `,` sched-modifier @@ -590,7 +610,7 @@ } //===----------------------------------------------------------------------===// -// WsLoopOp +// LoopControl //===----------------------------------------------------------------------===// /// loop-control ::= `(` ssa-id-list `)` `:` type `=` loop-bounds diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir --- a/mlir/test/Dialect/OpenMP/ops.mlir +++ b/mlir/test/Dialect/OpenMP/ops.mlir @@ -133,6 +133,30 @@ return } +// CHECK-LABEL: omp_DistributeOp +func.func @omp_DistributeOp(%lb : index, %ub : index, %step : index, %data_var : memref, %chunk_var : i32) -> () { +// CHECK: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) + "omp.distribute" (%lb, %ub, %step) ({ + ^bb0(%iv: index): + omp.yield + }) {operand_segment_sizes = array} : + (index, index, index) -> () + + return + } + +// CHECK-LABEL: omp_distribute_pretty + func.func @omp_distribute_pretty(%lb : index, %ub : index, %step : index, %data_var : memref, %chunk_var : i32) -> () { +// CHECK: omp.distribute +// CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) + omp.distribute + for (%iv) : index = (%lb) to (%ub) step (%step) { + omp.yield + } + + return + } + // CHECK-LABEL: omp_wsloop func.func @omp_wsloop(%lb : index, %ub : index, %step : index, %data_var : memref, %linear_var : i32, %chunk_var : i32) -> () { @@ -145,7 +169,7 @@ (index, index, index) -> () // CHECK: omp.wsloop linear(%{{.*}} = %{{.*}} : memref) schedule(static) - // CHECK-SAMe: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) + // CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) "omp.wsloop" (%lb, %ub, %step, %data_var, %linear_var) ({ ^bb0(%iv: index): omp.yield