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 @@ -97,6 +97,7 @@ ]; let parser = [{ return parseParallelOp(parser, result); }]; let printer = [{ return printParallelOp(p, *this); }]; + let verifier = [{ return ::verifyParallelOp(*this); }]; } @@ -227,6 +228,53 @@ 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<0>]>:$collapse, + Variadic:$allocate, + Variadic:$allocators, + UnitAttr:$inclusive); + + let regions = (region AnyRegion:$region); +} + //===----------------------------------------------------------------------===// // 2.10.4 taskyield Construct //===----------------------------------------------------------------------===// @@ -283,7 +331,7 @@ }]; - let arguments = (ins Optional:$if_expr, + let arguments = (ins Optional:$if_expr, Optional:$device, Optional:$thread_limit, UnitAttr:$nowait); @@ -296,9 +344,6 @@ "IntegerAttr":$thread_limit, "UnitAttr":$nowait)> ]; - let parser = [{ return parseTargetOp(parser, result); }]; - let printer = [{ return printTargetOp(p, *this); }]; - // let verifier = [{ return ::verifyTargetOp(*this); }]; } 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 @@ -791,133 +791,5 @@ } } - - - - - -//===----------------------------------------------------------------------===// -// TargetOp -//===----------------------------------------------------------------------===// - -void TargetOp::build(OpBuilder &builder, OperationState &state, IntegerAttr if_expr, - IntegerAttr device, IntegerAttr thread_limit, - UnitAttr nowait ) { - if (if_expr) - state.addAttribute("if_expr", if_expr); - if (device) - state.addAttribute("device", device); - if (thread_limit) - state.addAttribute("thread_limit", thread_limit); - if (nowait) - state.addAttribute("nowait", nowait); - -} - -/// Parse a target operation. -/// -/// operation ::= `omp.target' clause-list -/// clause-list ::= clause | clause clause-list -/// clause ::= if | thread_limit | device | nowait -/// if ::= `if` `(` ssa-id `)` -/// thread_limit ::= `thread_limit` `(` ssa-id-and-type`)` -/// device ::= `device` `(`ssa-id-and-type`)` -/// nowait ::= `nowait` -/// TODO: other clauses -/// TODO: Will update as soon as the other clauses will be finalized - -static ParseResult parseTargetOp(OpAsmParser &parser, - OperationState &result){ - - std::pair ifCond; - std::pair deviceNum; - std::pair threadLimit; - - - std::array segments{0, 0, 0}; - - StringRef keyword; - const StringRef opName = result.name.getStringRef(); - - const int ifClausePos = 0; - const int deviceNumPos = 1; - const int threadLimitPos = 2; - - - - - while (succeeded(parser.parseOptionalKeyword(&keyword))) { - if (keyword == "if") { - // Fail if there was already another if condition. - if (segments[ifClausePos]) - return allowedOnce(parser, "if", opName); - if (parser.parseLParen() || parser.parseOperand(ifCond.first) || - parser.parseColonType(ifCond.second) || parser.parseRParen()) - return failure(); - segments[ifClausePos] = 1; - } else if (keyword == "thread_limit") { - // Fail if there was already another thread_limit clause. - if (segments[threadLimitPos]) - return allowedOnce(parser, "thread_limit", opName); - if (parser.parseLParen() || parser.parseOperand(threadLimit.first) || - parser.parseColonType(threadLimit.second) || parser.parseRParen()) - return failure(); - segments[threadLimitPos] = 1; - } else if (keyword == "device") { - // Fail if there was already another device clause. - if (segments[deviceNumPos]) - return allowedOnce(parser, "device", opName); - if (parser.parseLParen() || parser.parseOperand(deviceNum.first) || - parser.parseColonType(deviceNum.second) || parser.parseRParen()) - return failure(); - segments[deviceNumPos] = 1; - } else if (keyword == "nowait") { - auto attr = UnitAttr::get(parser.getBuilder().getContext()); - result.addAttribute("nowait", attr); - } - } - - if (segments[ifClausePos] && - parser.resolveOperand(ifCond.first, ifCond.second, result.operands)) - return failure(); - - if (segments[threadLimitPos] && - parser.resolveOperand(threadLimit.first, threadLimit.second, result.operands)) - return failure(); - - if (segments[deviceNumPos] && - parser.resolveOperand(deviceNum.first, deviceNum.second, result.operands)) - return failure(); - - result.addAttribute("operand_segment_sizes", - parser.getBuilder().getI32VectorAttr(segments)); - - Region *body = result.addRegion(); - SmallVector regionArgs; - SmallVector regionArgTypes; - if (parser.parseRegion(*body, regionArgs, regionArgTypes)) - return failure(); - - return success(); -} - -static void printTargetOp(OpAsmPrinter &p, TargetOp op) { - p << "omp.target"; - if ( auto ifCond = op.if_expr()) - p << " if (" << ifCond << " : " << ifCond.getType() << ")" ; - if ( auto threads = op.thread_limit()) - p << " thread_limit (" << threads << " : " << threads.getType() << ")"; - if ( auto devices = op.device()) - p << " device (" << devices << " : " << devices.getType() << ")"; - if (op.nowait()) - p << " nowait"; - - p.printRegion(op.getRegion()); -} - - - - - #define GET_OP_CLASSES #include "mlir/Dialect/OpenMP/OpenMPOps.cpp.inc" 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 @@ -283,14 +283,14 @@ // CHECK-LABEL: omp_target -func @omp_target(%if_cond : si32, %device : si32, %num_threads : si32) -> () { +func @omp_target(%if_cond : i1, %device : si32, %num_threads : si32) -> () { // Test with optional operands; if_expr, device, thread_limit, and nowait. // CHECK: omp.target "omp.target"(%if_cond, %device, %num_threads) ({ // CHECK: omp.terminator omp.terminator - }) {operand_segment_sizes = dense<[1,1,1]>: vector<3xi32>, nowait } : ( si32, si32, si32 ) -> () + }) {operand_segment_sizes = dense<[1,1,1]>: vector<3xi32>, nowait } : ( i1, si32, si32 ) -> () // CHECK: omp.barrier omp.barrier @@ -299,20 +299,14 @@ } +// CHECK-LABEL: omp_distribute +func @omp_distribute(%lb : index, %ub : index, %step : index, %data_var : memref, %chunk_var : si32) -> () { -func @omp_target_pretty( %if_cond : si32, %device : si32, %thread_limit : si32) -> () { - // CHECK: omp.target - omp.target { - omp.terminator - } - - // CHECK: omp.target device - omp.target device(%device : si32) { - omp.terminator - } + // 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 + return }