diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td index 5da05476a68372..89d2ac342281cc 100644 --- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td +++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td @@ -636,6 +636,57 @@ def YieldOp : OpenMP_Op<"yield", let assemblyFormat = [{ ( `(` $results^ `:` type($results) `)` )? attr-dict}]; } +//===----------------------------------------------------------------------===// +// Distribute construct [2.9.4.1] +//===----------------------------------------------------------------------===// +def DistributeOp : OpenMP_Op<"distribute", [AttrSizedOperandSegments, + RecursiveMemoryEffects]> { + let summary = "distribute 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. + + The `dist_schedule_static` 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: private_var, firstprivate_var, lastprivate_var, collapse + }]; + let arguments = (ins + UnitAttr:$dist_schedule_static, + Optional:$chunk_size, + Variadic:$allocate_vars, + Variadic:$allocators_vars, + OptionalAttr:$order_val); + + let regions = (region AnyRegion:$region); + + let assemblyFormat = [{ + oilist(`dist_schedule_static` $dist_schedule_static + |`chunk_size` `(` $chunk_size `:` type($chunk_size) `)` + |`order` `(` custom($order_val) `)` + |`allocate` `(` + custom( + $allocate_vars, type($allocate_vars), + $allocators_vars, type($allocators_vars) + ) `)` + ) $region attr-dict + }]; + + 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 index 2bf9355ed62676..3c3fbe2981b018 100644 --- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp +++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp @@ -1053,6 +1053,22 @@ LogicalResult SimdLoopOp::verify() { return success(); } +//===----------------------------------------------------------------------===// +// Verifier for Distribute construct [2.9.4.1] +//===----------------------------------------------------------------------===// + +LogicalResult DistributeOp::verify() { + if (this->getChunkSize() && !this->getDistScheduleStatic()) + return emitOpError() << "chunk size set without " + "dist_schedule_static being present"; + + if (getAllocateVars().size() != getAllocatorsVars().size()) + return emitError( + "expected equal sizes for allocate and allocator variables"); + + return success(); +} + //===----------------------------------------------------------------------===// // ReductionOp //===----------------------------------------------------------------------===// diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir index 13cbea6c9923c2..a509c6cb162667 100644 --- a/mlir/test/Dialect/OpenMP/ops.mlir +++ b/mlir/test/Dialect/OpenMP/ops.mlir @@ -479,6 +479,36 @@ func.func @omp_simdloop_pretty_multiple(%lb1 : index, %ub1 : index, %step1 : ind return } +// CHECK-LABEL: omp_distribute +func.func @omp_distribute(%chunk_size : i32, %data_var : memref) -> () { + // CHECK: omp.distribute + "omp.distribute" () ({ + omp.terminator + }) {} : () -> () + // CHECK: omp.distribute + omp.distribute { + omp.terminator + } + // CHECK: omp.distribute dist_schedule_static + omp.distribute dist_schedule_static { + omp.terminator + } + // CHECK: omp.distribute dist_schedule_static chunk_size(%{{.+}} : i32) + omp.distribute dist_schedule_static chunk_size(%chunk_size : i32) { + omp.terminator + } + // CHECK: omp.distribute order(concurrent) + omp.distribute order(concurrent) { + omp.terminator + } + // CHECK: omp.distribute allocate(%{{.+}} : memref -> %{{.+}} : memref) + omp.distribute allocate(%data_var : memref -> %data_var : memref) { + omp.terminator + } +return +} + + // CHECK-LABEL: omp_target func.func @omp_target(%if_cond : i1, %device : si32, %num_threads : i32, %map1: memref, %map2: memref) -> () {