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 @@ -228,6 +228,55 @@ let assemblyFormat = [{ ( `(` $results^ `:` type($results) `)` )? attr-dict}]; } +//===---------------------------------------------------------------------===// +// 2.9.4 distribute loop 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. + + + `private_var`, `firstprivate_var`, and `lastprivate_var` arguments are variadic + list of operands that specify the data sharing attributes of the list of values + + 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, + Variadic:$private_var, + Variadic:$firstprivate_var, + Variadic:$lastprivate_var, + OptionalAttr:$dist_schedule_var, + Optional:$schedule_chunk, + Confined, [IntMinValue<1>]>:$collapse, + Variadic:$allocate, + Variadic:$allocators, + UnitAttr:$inclusive); + + let regions = (region AnyRegion:$region); +} + //===----------------------------------------------------------------------===// // 2.10.4 taskyield Construct //===----------------------------------------------------------------------===// 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 @@ -297,3 +297,17 @@ return } + +// CHECK-LABEL: omp_distribute +func @omp_distribute(%lb : index, %ub : index, %step : index, %data_var : memref, %chunk_var : si32) -> () { + + // CHECK: omp.distribute + "omp.distribute" (%lb, %ub, %step, %data_var,%data_var, %data_var, %data_var, %data_var, %chunk_var) ({ + }) {operand_segment_sizes = dense<[1,1,1,1,1,1,1,1,1]> : vector<9xi32>, dist_schedule_val = "Static", collapse_val = 3} : (index, index, index, memref, memref, memref, memref, memref, si32) -> () + + // CHECK: omp.barrier + omp.barrier + + return + +}