diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td index fab8c2cd0de45..5556411cb878a 100644 --- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td +++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td @@ -463,6 +463,91 @@ def YieldOp : OpenMP_Op<"yield", let assemblyFormat = [{ ( `(` $results^ `:` type($results) `)` )? attr-dict}]; } +//===----------------------------------------------------------------------===// +// 2.10.1 task Construct +//===----------------------------------------------------------------------===// + +def TaskOp : OpenMP_Op<"task", [AttrSizedOperandSegments, + OutlineableOpenMPOpInterface, AutomaticAllocationScope, + ReductionClauseInterface]> { + let summary = "task construct"; + let description = [{ + The task construct defines an explicit task. + + For definitions of "undeferred task", "included task", "final task" and + "mergeable task", please check OpenMP Specification. + + When an `if` clause is present on a task construct, and the value of + `if_expr` evaluates to `false`, an "undeferred task" is generated, and the + encountering thread must suspend the current task region, for which + execution cannot be resumed until execution of the structured block that is + associated with the generated task is completed. + + When a `final` clause is present on a task construct and the `final_expr` + evaluates to `true`, the generated task will be a "final task". All task + constructs encountered during execution of a final task will generate final + and included tasks. + + If the `untied` clause is present on a task construct, any thread in the + team can resume the task region after a suspension. The `untied` clause is + ignored if a `final` clause is present on the same task construct and the + `final_expr` evaluates to `true`, or if a task is an included task. + + When the `mergeable` clause is present on a task construct, the generated + task is a "mergeable task". + + The `in_reduction` clause specifies that this particular task (among all the + tasks in current taskgroup, if any) participates in a reduction. + + The `priority` clause is a hint for the priority of the generated task. + The `priority` is a non-negative integer expression that provides a hint for + task execution order. Among all tasks ready to be executed, higher priority + tasks (those with a higher numerical value in the priority clause + expression) are recommended to execute before lower priority ones. The + default priority-value when no priority clause is specified should be + assumed to be zero (the lowest priority). + + The `allocators_vars` and `allocate_vars` parameters are a variadic list of + values that specify the memory allocator to be used to obtain storage for + private values. + + }]; + + // TODO: depend, affinity and detach clauses + let arguments = (ins Optional:$if_expr, + Optional:$final_expr, + UnitAttr:$untied, + UnitAttr:$mergeable, + Variadic:$in_reduction_vars, + OptionalAttr:$in_reductions, + Optional:$priority, + Variadic:$allocate_vars, + Variadic:$allocators_vars); + let regions = (region AnyRegion:$region); + let assemblyFormat = [{ + oilist(`if` `(` $if_expr `)` + |`final` `(` $final_expr `)` + |`untied` $untied + |`mergeable` $mergeable + |`in_reduction` `(` + custom( + $in_reduction_vars, type($in_reduction_vars), $in_reductions + ) `)` + |`priority` `(` $priority `)` + |`allocate` `(` + custom( + $allocate_vars, type($allocate_vars), + $allocators_vars, type($allocators_vars) + ) `)` + ) $region attr-dict + }]; + let extraClassDeclaration = [{ + /// Returns the reduction variables + operand_range getReductionVars() { return in_reduction_vars(); } + }]; + let hasVerifier = 1; +} + //===----------------------------------------------------------------------===// // 2.10.4 taskyield Construct //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp index 176c27b132164..5118a8954688c 100644 --- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp +++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp @@ -725,6 +725,13 @@ LogicalResult ReductionOp::verify() { return emitOpError() << "the accumulator is not used by the parent"; } +//===----------------------------------------------------------------------===// +// TaskOp +//===----------------------------------------------------------------------===// +LogicalResult TaskOp::verify() { + return verifyReductionVarList(*this, in_reductions(), in_reduction_vars()); +} + //===----------------------------------------------------------------------===// // WsLoopOp //===----------------------------------------------------------------------===// diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir index 8c42052211dbc..dd684dda83a26 100644 --- a/mlir/test/Dialect/OpenMP/invalid.mlir +++ b/mlir/test/Dialect/OpenMP/invalid.mlir @@ -980,3 +980,70 @@ func @omp_single(%data_var : memref) -> () { }) {operand_segment_sizes = dense<[1,0]> : vector<2xi32>} : (memref) -> () return } + +// ----- + +func @omp_task(%ptr: !llvm.ptr) { + // expected-error @below {{op expected symbol reference @add_f32 to point to a reduction declaration}} + omp.task in_reduction(@add_f32 -> %ptr : !llvm.ptr) { + // CHECK: "test.foo"() : () -> () + "test.foo"() : () -> () + // CHECK: omp.terminator + omp.terminator + } +} + +// ----- + +omp.reduction.declare @add_f32 : f32 +init { +^bb0(%arg: f32): + %0 = arith.constant 0.0 : f32 + omp.yield (%0 : f32) +} +combiner { +^bb1(%arg0: f32, %arg1: f32): + %1 = arith.addf %arg0, %arg1 : f32 + omp.yield (%1 : f32) +} + +func @omp_task(%ptr: !llvm.ptr) { + // expected-error @below {{op accumulator variable used more than once}} + omp.task in_reduction(@add_f32 -> %ptr : !llvm.ptr, @add_f32 -> %ptr : !llvm.ptr) { + // CHECK: "test.foo"() : () -> () + "test.foo"() : () -> () + // CHECK: omp.terminator + omp.terminator + } +} + +// ----- + +omp.reduction.declare @add_i32 : i32 +init { +^bb0(%arg: i32): + %0 = arith.constant 0 : i32 + omp.yield (%0 : i32) +} +combiner { +^bb1(%arg0: i32, %arg1: i32): + %1 = arith.addi %arg0, %arg1 : i32 + omp.yield (%1 : i32) +} +atomic { +^bb2(%arg2: !llvm.ptr, %arg3: !llvm.ptr): + %2 = llvm.load %arg3 : !llvm.ptr + llvm.atomicrmw add %arg2, %2 monotonic : i32 + omp.yield +} + +func @omp_task(%mem: memref<1xf32>) { + // expected-error @below {{op expected accumulator ('memref<1xf32>') to be the same type as reduction declaration ('!llvm.ptr')}} + omp.task in_reduction(@add_i32 -> %mem : memref<1xf32>) { + // CHECK: "test.foo"() : () -> () + "test.foo"() : () -> () + // CHECK: omp.terminator + omp.terminator + } + return +} diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir index 7dcbaba39ba29..b620afb9c6a89 100644 --- a/mlir/test/Dialect/OpenMP/ops.mlir +++ b/mlir/test/Dialect/OpenMP/ops.mlir @@ -897,6 +897,96 @@ func @omp_single_allocate_nowait(%data_var: memref) { return } +// CHECK-LABEL: @omp_task +// CHECK-SAME: (%[[bool_var:.*]]: i1, %[[i64_var:.*]]: i64, %[[i32_var:.*]]: i32, %[[data_var:.*]]: memref) +func @omp_task(%bool_var: i1, %i64_var: i64, %i32_var: i32, %data_var: memref) { + + // Checking simple task + // CHECK: omp.task { + omp.task { + // CHECK: "test.foo"() : () -> () + "test.foo"() : () -> () + // CHECK: omp.terminator + omp.terminator + } + + // Checking `if` clause + // CHECK: omp.task if(%[[bool_var]]) { + omp.task if(%bool_var) { + // CHECK: "test.foo"() : () -> () + "test.foo"() : () -> () + // CHECK: omp.terminator + omp.terminator + } + + // Checking `final` clause + // CHECK: omp.task final(%[[bool_var]]) { + omp.task final(%bool_var) { + // CHECK: "test.foo"() : () -> () + "test.foo"() : () -> () + // CHECK: omp.terminator + omp.terminator + } + + // Checking `untied` clause + // CHECK: omp.task untied { + omp.task untied { + // CHECK: "test.foo"() : () -> () + "test.foo"() : () -> () + // CHECK: omp.terminator + omp.terminator + } + + // Checking `in_reduction` clause + %c1 = arith.constant 1 : i32 + // CHECK: %[[redn_var1:.*]] = llvm.alloca %{{.*}} x f32 : (i32) -> !llvm.ptr + %0 = llvm.alloca %c1 x f32 : (i32) -> !llvm.ptr + // CHECK: %[[redn_var2:.*]] = llvm.alloca %{{.*}} x f32 : (i32) -> !llvm.ptr + %1 = llvm.alloca %c1 x f32 : (i32) -> !llvm.ptr + // CHECK: omp.task in_reduction(@add_f32 -> %[[redn_var1]] : !llvm.ptr, @add_f32 -> %[[redn_var2]] : !llvm.ptr) { + omp.task in_reduction(@add_f32 -> %0 : !llvm.ptr, @add_f32 -> %1 : !llvm.ptr) { + // CHECK: "test.foo"() : () -> () + "test.foo"() : () -> () + // CHECK: omp.terminator + omp.terminator + } + + // Checking priority clause + // CHECK: omp.task priority(%[[i32_var]]) { + omp.task priority(%i32_var) { + // CHECK: "test.foo"() : () -> () + "test.foo"() : () -> () + // CHECK: omp.terminator + omp.terminator + } + + // Checking allocate clause + // CHECK: omp.task allocate(%[[data_var]] : memref -> %[[data_var]] : memref) { + omp.task allocate(%data_var : memref -> %data_var : memref) { + // CHECK: "test.foo"() : () -> () + "test.foo"() : () -> () + // CHECK: omp.terminator + omp.terminator + } + + // Checking multiple clauses + // CHECK: omp.task if(%[[bool_var]]) final(%[[bool_var]]) untied + omp.task if(%bool_var) final(%bool_var) untied + // CHECK-SAME: in_reduction(@add_f32 -> %[[redn_var1]] : !llvm.ptr, @add_f32 -> %[[redn_var2]] : !llvm.ptr) + in_reduction(@add_f32 -> %0 : !llvm.ptr, @add_f32 -> %1 : !llvm.ptr) + // CHECK-SAME: priority(%[[i32_var]]) + priority(%i32_var) + // CHECK-SAME: allocate(%[[data_var]] : memref -> %[[data_var]] : memref) + allocate(%data_var : memref -> %data_var : memref) { + // CHECK: "test.foo"() : () -> () + "test.foo"() : () -> () + // CHECK: omp.terminator + omp.terminator + } + + return +} + // ----- func @omp_threadprivate() {