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 @@ -394,7 +394,7 @@ custom( $reduction_vars, type($reduction_vars), $reductions ) `)` - ) `for` custom($region, $lowerBound, $upperBound, $step, + ) `for` custom($region, $lowerBound, $upperBound, $step, type($step), $inclusive) attr-dict }]; let hasVerifier = 1; @@ -446,7 +446,7 @@ def YieldOp : OpenMP_Op<"yield", [NoSideEffect, 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 @@ -463,6 +463,68 @@ 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 optional `dist_schedule_var` 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. + + The optional `collapse` attribute specifies the number of loops which + are collapsed to form the distribute loop. + }]; + + let arguments = (ins Variadic:$lowerBound, + Variadic:$upperBound, + Variadic:$step, + UnitAttr:$static_dist_schedule, + Optional:$schedule_chunk, + Confined, [IntMinValue<1>]>:$collapse_val, + Variadic:$allocate_vars, + Variadic:$allocators_vars, + UnitAttr:$inclusive); + + let assemblyFormat = [{ + oilist(`dist_schedule` `(` `static` + ($static_dist_schedule^)? + ( $schedule_chunk^ `:` type($schedule_chunk))? `)` + |`collapse` `(` $collapse_val `)` + | `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 @@ -207,6 +207,20 @@ return success(); } +//===---------------------------------------------------------------------===// +// Verifier for Dristribute Op +//===---------------------------------------------------------------------===// + +LogicalResult DistributeOp::verify(){ + if (this->lowerBound().empty()) + return emitOpError() << "empty lowerbound for distribute loop operation"; + + if (this->upperBound().empty()) + return emitOpError() << "empty upperbound for distribute loop operation"; + + return success(); +} + /// schedule ::= `schedule` `(` sched-list `)` /// sched-list ::= sched-val | sched-val sched-list | /// sched-val `,` sched-modifier @@ -510,14 +524,14 @@ } //===----------------------------------------------------------------------===// -// WsLoopOp +// LoopControl //===----------------------------------------------------------------------===// /// loop-control ::= `(` ssa-id-list `)` `:` type `=` loop-bounds /// loop-bounds := `(` ssa-id-list `)` to `(` ssa-id-list `)` inclusive? steps /// steps := `step` `(`ssa-id-list`)` ParseResult -parseWsLoopControl(OpAsmParser &parser, Region ®ion, +parseLoopControl(OpAsmParser &parser, Region ®ion, SmallVectorImpl &lowerBound, SmallVectorImpl &upperBound, SmallVectorImpl &steps, @@ -560,7 +574,7 @@ return success(); } -void printWsLoopControl(OpAsmPrinter &p, Operation *op, Region ®ion, +void printLoopControl(OpAsmPrinter &p, Operation *op, Region ®ion, ValueRange lowerBound, ValueRange upperBound, ValueRange steps, TypeRange loopVarTypes, UnitAttr inclusive) { 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 @@ -121,6 +121,19 @@ return } +// CHECK-LABEL: omp_DistributeOp +func.func @omp_DistributeOp(%lb : index, %ub : index, %step : index, %data_var : memref, %chunk_var : i32) -> () { + // CHECK: omp.distribute collapse(2) + // CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) + "omp.distribute" (%lb, %ub, %step) ({ + ^bb0(%iv: index): + omp.yield + }) {operand_segment_sizes = dense<[1,1,1,0,0,0]> : vector<6xi32>, collapse_val = 2} : + (index, index, index) -> () + + return + } + // CHECK-LABEL: omp_wsloop func.func @omp_wsloop(%lb : index, %ub : index, %step : index, %data_var : memref, %linear_var : i32, %chunk_var : i32) -> () {