[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.

539

nit: I32 -> I64

[abidmalikwaterloo]

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

533

It is a boolean value.

534

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.

533

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.

539

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.

539

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

mlir/test/Dialect/OpenMP/ops.mlir

899

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

[abidmalikwaterloo]

LGTM

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

533

thanks. Got it

[shraiysh]

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

[shraiysh]

mlir/test/Dialect/OpenMP/ops.mlir

899

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

488

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

[shraiysh]

Added interfaces.

[shraiysh]

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

488

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

488

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

488

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.17.11 depend Construct
+//===----------------------------------------------------------------------===//
+
+def ClauseDependSource : I32EnumAttrCase<"dependsource", 0>;
+def ClauseDependSink   : I32EnumAttrCase<"dependsink",   1>;
+
+def ClauseDepend : I32EnumAttr<
+    "ClauseDepend",
+    "depend clause",
+    [ClauseDependSource, ClauseDependSink]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::omp";
+}
+def ClauseDependAttr : EnumAttr<OpenMP_Dialect, ClauseDepend, "clause_depend"> {
+  let assemblyFormat = "`(` $value `)`";
+}
+
+//===----------------------------------------------------------------------===//
 // 2.6 parallel Construct
 //===----------------------------------------------------------------------===//
 
 def ParallelOp : OpenMP_Op<"parallel", [
                  AutomaticAllocationScope, AttrSizedOperandSegments,
                  DeclareOpInterfaceMethods<OutlineableOpenMPOpInterface>,
                  RecursiveSideEffects, ReductionClauseInterface]> {
   let summary = "parallel construct";
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]> {
+  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
+  }];
+}
+
+
+//===----------------------------------------------------------------------===//
 // 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.
def BarrierOp : OpenMP_Op<"barrier"> {
 
   let assemblyFormat = "attr-dict";
 }
 
 //===----------------------------------------------------------------------===//
 // [5.1] 2.19.9 ordered Construct
 //===----------------------------------------------------------------------===//
 
-def ClauseDependSource : I32EnumAttrCase<"dependsource", 0>;
-def ClauseDependSink   : I32EnumAttrCase<"dependsink",   1>;
-
-def ClauseDepend : I32EnumAttr<
-    "ClauseDepend",
-    "depend clause",
-    [ClauseDependSource, ClauseDependSink]> {
-  let genSpecializedAttr = 0;
-  let cppNamespace = "::mlir::omp";
-}
-def ClauseDependAttr : EnumAttr<OpenMP_Dialect, ClauseDepend, "clause_depend"> {
-  let assemblyFormat = "`(` $value `)`";
-}
-
 def OrderedOp : OpenMP_Op<"ordered"> {
   let summary = "ordered construct without region";
   let description = [{
     The ordered construct without region is a stand-alone directive that
     specifies cross-iteration dependences in a doacross loop nest.
 
     The `depend_type_val` attribute refers to either the DEPEND(SOURCE) clause
     or the DEPEND(SINK: vec) clause.

mlir/test/Dialect/OpenMP/ops.mlir

omp.single allocate(%data_var : memref<i32> -> %data_var : memref<i32>) nowait {
       // CHECK: omp.terminator
       omp.terminator
     }
     // CHECK: omp.terminator
     omp.terminator
   }
   return
 }
+
+func @foo() { 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: call @foo() : () -> ()
+    call @foo() : () -> ()
+    // CHECK: omp.terminator
+    omp.terminator
+  }
+
+  // Checking `if` clause
+  // CHECK: omp.task if(%[[bool_var]]) {
+  omp.task if(%bool_var) {
+    // CHECK: call @foo() : () -> ()
+    call @foo() : () -> ()
+    // CHECK: omp.terminator
+    omp.terminator
+  }
+
+  // Checking `final` clause
+  // CHECK: omp.task final(%[[bool_var]]) {
+  omp.task final(%bool_var) {
+    // CHECK: call @foo() : () -> ()
+    call @foo() : () -> ()
+    // CHECK: omp.terminator
+    omp.terminator
+  }
+
+  // Checking `untied` clause
+  // CHECK: omp.task untied {
+  omp.task untied {
+    // CHECK: call @foo() : () -> ()
+    call @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: call @foo() : () -> ()
+    call @foo() : () -> ()
+    // CHECK: omp.terminator
+    omp.terminator
+  }
+
+  // Checking priority clause
+  // CHECK: omp.task priority(%[[i32_var]]) {
+  omp.task priority(%i32_var) {
+    // CHECK: call @foo() : () -> ()
+    call @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: call @foo() : () -> ()
+    call @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: call @foo() : () -> ()
+    call @foo() : () -> ()
+    // CHECK: omp.terminator
+    omp.terminator
+  }
+
+  return
+}