[mlir][OpenMP] Added omp.task

[shraiysh]

Summary

This patch adds tasking construct according to Section 2.10.1 of OpenMP 5.0

[peixin]

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

63

1. depend is not one construct. It is one clause.

2. dependsource and dependsink are only used in ordered construct. The depend clause in task, taskwait, and some other offloading constructs use another format: depend(in | out | ...), which also has modifier and locator-list. Maybe we should split these two formats into two definitions to clearly differentiate their usage. BTW, it seems this depend clause format is not supported in parser currently. If you want to delay the depend clause support, it is also ok, but it's better not move this for now.

521

nit: I32 -> I64

[abidmalikwaterloo]

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

515

It is a boolean value.

516

Same here. Can we define it as a Boolean value? I think we have an option in MLIR.

[abidmalikwaterloo]

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

63

The depend clause for the task construct is going through a revision in the language committee. One should go through the proposed revision before defining it.

[shraiysh]

Addressed comments.

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

63

Yes, thank you for pointing this out. I was planning to do this later but forgot to undo the move.

515

A boolean value is an I1. Are you referring to UnitAttr? It cannot be used here because the value of UnitAttr is a compile-time constant but here it is a value at runtime.

521

I kept it I32 because the task related runtime functions accept i32 as priority values. Allowing an I64 value as priority here and having a INT64_MAX as priority might lead to unexpected behavior due to data loss during translation.

[peixin]

LGTM. Please wait for at least one more approve befire land.

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

63

Right. OpenMP 5.2 made great changes to depend clause. But I think we should support OpenMP 5.0/5.1 since there are some old workloads using them. We can delay the support of new features from OpenMP 5.2. Finally, we should support both of new depend clause in OpenMP 5.2 and old depend clause before OpenMP 5.1.

521

OK. I didn't check the runtime argument type. I32 is right.

mlir/test/Dialect/OpenMP/ops.mlir

1015

Nit: split the test case with "// -----"

[abidmalikwaterloo]

LGTM

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

515

thanks. Got it

[shraiysh]

Rebase and added a verifier check. Hence requesting review again.

[shraiysh]

mlir/test/Dialect/OpenMP/ops.mlir

1015

With the verifier, I needed the symbol @add_f32 in scope. So, I moved the testcase above which gives me access to the symbol.

[kiranchandramohan]

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

470

Can you add the OutlineableAopenMPOpInterface and the AutomaticAllocationScope? What about the ReductionClauseInterface?

[shraiysh]

Added interfaces.

[shraiysh]

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

470

I do not see anything using the OutlineableOpenMPOpInterface. Plus its description also seems to be incomplete. I understand what kind of operations will have this interface, but is something actually using it?

[shraiysh]

Added OutlineableOpenMPOpInterface. Missed it last time.

[kiranchandramohan]

LGTM.

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

470

At the moment, it is used to find a place to insert allocas in flang/lib/Optimizer/Builder/FIRBuilder.cpp and also to check whether allocas should be marked as pinned.
There are potentially other uses as you suspect.

[shraiysh]

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

470

Alright. Thank you for the details.

[shraiysh]

This has been merged but the commit is somehow not being picked up by phabricator. Similar thing happened with D123562 (the commit right before this one)

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

def OpenMP_PointerLikeTypeInterface : TypeInterface<"PointerLikeType"> {
   ];
 }
 
 def OpenMP_PointerLikeType : Type<
   CPred<"$_self.isa<::mlir::omp::PointerLikeType>()">,
   "OpenMP-compatible variable type", "::mlir::omp::PointerLikeType">;
 
 //===----------------------------------------------------------------------===//
 // 2.6 parallel Construct
 //===----------------------------------------------------------------------===//
 
 def ParallelOp : OpenMP_Op<"parallel", [
                  AutomaticAllocationScope, AttrSizedOperandSegments,
                  DeclareOpInterfaceMethods<OutlineableOpenMPOpInterface>,
                  RecursiveSideEffects, ReductionClauseInterface]> {
   let summary = "parallel construct";
   let description = [{
def YieldOp : OpenMP_Op<"yield",
   let builders = [
     OpBuilder<(ins), [{ build($_builder, $_state, {}); }]>
   ];
 
   let assemblyFormat = [{ ( `(` $results^ `:` type($results) `)` )? attr-dict}];
 }
 
 //===----------------------------------------------------------------------===//
+// 2.10.1 task Construct
+//===----------------------------------------------------------------------===//
+
+def TaskOp : OpenMP_Op<"task", [AttrSizedOperandSegments,
+                       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<I1>:$if_expr,
+                       Optional<I1>:$final_expr,
+                       UnitAttr:$untied,
+                       UnitAttr:$mergeable,
+                       Variadic<OpenMP_PointerLikeType>:$in_reduction_vars,
+                       OptionalAttr<SymbolRefArrayAttr>:$in_reductions,
+                       Optional<I32>:$priority,
+                       Variadic<AnyType>:$allocate_vars,
+                       Variadic<AnyType>:$allocators_vars);
+  let regions = (region AnyRegion:$region);
+  let assemblyFormat = [{
+    oilist(`if` `(` $if_expr `)`
+          |`final` `(` $final_expr `)`
+          |`untied` $untied
+          |`mergeable` $mergeable
+          |`in_reduction` `(`
+              custom<ReductionVarList>(
+                $in_reduction_vars, type($in_reduction_vars), $in_reductions
+              ) `)`
+          |`priority` `(` $priority `)`
+          |`allocate` `(`
+              custom<AllocateAndAllocator>(
+                $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
 //===----------------------------------------------------------------------===//
 
 def TaskyieldOp : OpenMP_Op<"taskyield"> {
   let summary = "taskyield construct";
   let description = [{
     The taskyield construct specifies that the current task can be suspended
     in favor of execution of a different task.

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

for (const auto &var :
       if (var == accumulator())
         return success();
     op = op->getParentWithTrait<ReductionClauseInterface::Trait>();
   }
   return emitOpError() << "the accumulator is not used by the parent";
 }
 
 //===----------------------------------------------------------------------===//
+// TaskOp
+//===----------------------------------------------------------------------===//
+LogicalResult TaskOp::verify() {
+  return verifyReductionVarList(*this, in_reductions(), in_reduction_vars());
+}
+
+//===----------------------------------------------------------------------===//
 // WsLoopOp
 //===----------------------------------------------------------------------===//
 
 void WsLoopOp::build(OpBuilder &builder, OperationState &state,
                      ValueRange lowerBound, ValueRange upperBound,
                      ValueRange step, ArrayRef<NamedAttribute> attributes) {
   build(builder, state, lowerBound, upperBound, step,
         /*linear_vars=*/ValueRange(),

mlir/test/Dialect/OpenMP/invalid.mlir

 
 func @omp_single(%data_var : memref<i32>) -> () {
   // expected-error @below {{expected equal sizes for allocate and allocator variables}}
   "omp.single" (%data_var) ({
     omp.barrier
   }) {operand_segment_sizes = dense<[1,0]> : vector<2xi32>} : (memref<i32>) -> ()
   return
 }
+
+// -----
+
+func @omp_task(%ptr: !llvm.ptr<f32>) {
+  // expected-error @below {{op expected symbol reference @add_f32 to point to a reduction declaration}}
+  omp.task in_reduction(@add_f32 -> %ptr : !llvm.ptr<f32>) {
+    // 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<f32>) {
+  // expected-error @below {{op accumulator variable used more than once}}
+  omp.task in_reduction(@add_f32 -> %ptr : !llvm.ptr<f32>, @add_f32 -> %ptr : !llvm.ptr<f32>) {
+    // 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<i32>, %arg3: !llvm.ptr<i32>):
+  %2 = llvm.load %arg3 : !llvm.ptr<i32>
+  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<i32>')}}
+  omp.task in_reduction(@add_i32 -> %mem : memref<1xf32>) {
+    // CHECK: "test.foo"() : () -> ()
+    "test.foo"() : () -> ()
+    // CHECK: omp.terminator
+    omp.terminator
+  }
+  return
+}

mlir/test/Dialect/OpenMP/ops.mlir

omp.single allocate(%data_var : memref<i32> -> %data_var : memref<i32>) nowait {
       omp.terminator
     }
     // CHECK: omp.terminator
     omp.terminator
   }
   return
 }
 
+// CHECK-LABEL: @omp_task
+// CHECK-SAME: (%[[bool_var:.*]]: i1, %[[i64_var:.*]]: i64, %[[i32_var:.*]]: i32, %[[data_var:.*]]: memref<i32>)
+func @omp_task(%bool_var: i1, %i64_var: i64, %i32_var: i32, %data_var: memref<i32>) {
+
+  // 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<f32>
+  %0 = llvm.alloca %c1 x f32 : (i32) -> !llvm.ptr<f32>
+  // CHECK: %[[redn_var2:.*]] = llvm.alloca %{{.*}} x f32 : (i32) -> !llvm.ptr<f32>
+  %1 = llvm.alloca %c1 x f32 : (i32) -> !llvm.ptr<f32>
+  // CHECK: omp.task in_reduction(@add_f32 -> %[[redn_var1]] : !llvm.ptr<f32>, @add_f32 -> %[[redn_var2]] : !llvm.ptr<f32>) {
+  omp.task in_reduction(@add_f32 -> %0 : !llvm.ptr<f32>, @add_f32 -> %1 : !llvm.ptr<f32>) {
+    // 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<i32> -> %[[data_var]] : memref<i32>) {
+  omp.task allocate(%data_var : memref<i32> -> %data_var : memref<i32>) {
+    // 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<f32>, @add_f32 -> %[[redn_var2]] : !llvm.ptr<f32>)
+      in_reduction(@add_f32 -> %0 : !llvm.ptr<f32>, @add_f32 -> %1 : !llvm.ptr<f32>)
+      // CHECK-SAME: priority(%[[i32_var]])
+      priority(%i32_var)
+      // CHECK-SAME: allocate(%[[data_var]] : memref<i32> -> %[[data_var]] : memref<i32>)
+      allocate(%data_var : memref<i32> -> %data_var : memref<i32>) {
+    // CHECK: "test.foo"() : () -> ()
+    "test.foo"() : () -> ()
+    // CHECK: omp.terminator
+    omp.terminator
+  }
+
+  return
+}
+
 // -----
 
 func @omp_threadprivate() {
   %0 = arith.constant 1 : i32
   %1 = arith.constant 2 : i32
   %2 = arith.constant 3 : i32
 
   // CHECK: [[ARG0:%.*]] = llvm.mlir.addressof @_QFsubEx : !llvm.ptr<i32>
   // CHECK: {{.*}} = omp.threadprivate [[ARG0]] : !llvm.ptr<i32> -> !llvm.ptr<i32>
   %3 = llvm.mlir.addressof @_QFsubEx : !llvm.ptr<i32>
   %4 = omp.threadprivate %3 : !llvm.ptr<i32> -> !llvm.ptr<i32>
   llvm.store %0, %4 : !llvm.ptr<i32>
 
   // CHECK:  omp.parallel
   // CHECK:    {{.*}} = omp.threadprivate [[ARG0]] : !llvm.ptr<i32> -> !llvm.ptr<i32>
   omp.parallel  {
     %5 = omp.threadprivate %3 : !llvm.ptr<i32> -> !llvm.ptr<i32>
     llvm.store %1, %5 : !llvm.ptr<i32>
     omp.terminator
   }
   llvm.store %2, %4 : !llvm.ptr<i32>
   return
 }
 
 llvm.mlir.global internal @_QFsubEx() : i32