[flang][mlir][llvm][OpenMP] Add lowering and translation support for mergeable clause on task (#114662)

Add FIR generation and LLVMIR translation support for mergeable clause
on task construct. If mergeable clause is present on a task, the
relevant flag in `ompt_task_flag_t` is set and passed to
`__kmpc_omp_task_alloc`.

Author: NimishMishra

flang/lib/Lower/OpenMP/OpenMP.cpp

@@ -1349,8 +1349,8 @@ static void genTaskClauses(lower::AbstractConverter &converter,
   cp.processUntied(clauseOps);
   // TODO Support delayed privatization.
 
-  cp.processTODO<clause::Affinity, clause::Detach, clause::InReduction,
-                 clause::Mergeable>(loc, llvm::omp::Directive::OMPD_task);
+  cp.processTODO<clause::Affinity, clause::Detach, clause::InReduction>(
+      loc, llvm::omp::Directive::OMPD_task);
 }
 
 static void genTaskgroupClauses(lower::AbstractConverter &converter,

flang/test/Lower/OpenMP/Todo/task_mergeable.f90

@@ -234,3 +234,11 @@ subroutine task_multiple_clauses()
   !CHECK: omp.terminator
   !$omp end task
 end subroutine task_multiple_clauses
+
+subroutine task_mergeable()
+!CHECK: omp.task mergeable {
+!CHECK: omp.terminator
+!CHECK: }
+ !$omp task mergeable
+ !$omp end task
+end subroutine

flang/test/Lower/OpenMP/task.f90

@@ -1262,12 +1262,12 @@ class OpenMPIRBuilder {
   ///                    cannot be resumed until execution of the structured
   ///                    block that is associated with the generated task is
   ///                    completed.
-  InsertPointOrErrorTy createTask(const LocationDescription &Loc,
-                                  InsertPointTy AllocaIP,
-                                  BodyGenCallbackTy BodyGenCB, bool Tied = true,
-                                  Value *Final = nullptr,
-                                  Value *IfCondition = nullptr,
-                                  SmallVector<DependData> Dependencies = {});
+  /// \param Mergeable	 If the given task is `mergeable`
+  InsertPointOrErrorTy
+  createTask(const LocationDescription &Loc, InsertPointTy AllocaIP,
+             BodyGenCallbackTy BodyGenCB, bool Tied = true,
+             Value *Final = nullptr, Value *IfCondition = nullptr,
+             SmallVector<DependData> Dependencies = {}, bool Mergeable = false);
 
   /// Generator for the taskgroup construct
   ///

llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h

@@ -1815,11 +1815,10 @@ static Value *emitTaskDependencies(
   return DepArray;
 }
 
-OpenMPIRBuilder::InsertPointOrErrorTy
-OpenMPIRBuilder::createTask(const LocationDescription &Loc,
-                            InsertPointTy AllocaIP, BodyGenCallbackTy BodyGenCB,
-                            bool Tied, Value *Final, Value *IfCondition,
-                            SmallVector<DependData> Dependencies) {
+OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createTask(
+    const LocationDescription &Loc, InsertPointTy AllocaIP,
+    BodyGenCallbackTy BodyGenCB, bool Tied, Value *Final, Value *IfCondition,
+    SmallVector<DependData> Dependencies, bool Mergeable) {
 
   if (!updateToLocation(Loc))
     return InsertPointTy();
@@ -1865,7 +1864,8 @@ OpenMPIRBuilder::createTask(const LocationDescription &Loc,
       Builder, AllocaIP, ToBeDeleted, TaskAllocaIP, "global.tid", false));
 
   OI.PostOutlineCB = [this, Ident, Tied, Final, IfCondition, Dependencies,
-                      TaskAllocaBB, ToBeDeleted](Function &OutlinedFn) mutable {
+                      Mergeable, TaskAllocaBB,
+                      ToBeDeleted](Function &OutlinedFn) mutable {
     // Replace the Stale CI by appropriate RTL function call.
     assert(OutlinedFn.getNumUses() == 1 &&
            "there must be a single user for the outlined function");
@@ -1890,6 +1890,8 @@ OpenMPIRBuilder::createTask(const LocationDescription &Loc,
     // Task is untied iff (Flags & 1) == 0.
     // Task is final iff (Flags & 2) == 2.
     // Task is not final iff (Flags & 2) == 0.
+    // Task is mergeable iff (Flags & 4) == 4.
+    // Task is not mergeable iff (Flags & 4) == 0.
     // TODO: Handle the other flags.
     Value *Flags = Builder.getInt32(Tied);
     if (Final) {
@@ -1898,6 +1900,9 @@ OpenMPIRBuilder::createTask(const LocationDescription &Loc,
       Flags = Builder.CreateOr(FinalFlag, Flags);
     }
 
+    if (Mergeable)
+      Flags = Builder.CreateOr(Builder.getInt32(4), Flags);
+
     // Argument - `sizeof_kmp_task_t` (TaskSize)
     // Tasksize refers to the size in bytes of kmp_task_t data structure
     // including private vars accessed in task.

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

@@ -191,10 +191,6 @@ static LogicalResult checkImplementationStatus(Operation &op) {
     if (!op.getLinearVars().empty() || !op.getLinearStepVars().empty())
       result = todo("linear");
   };
-  auto checkMergeable = [&todo](auto op, LogicalResult &result) {
-    if (op.getMergeable())
-      result = todo("mergeable");
-  };
   auto checkNontemporal = [&todo](auto op, LogicalResult &result) {
     if (!op.getNontemporalVars().empty())
       result = todo("nontemporal");
@@ -257,7 +253,6 @@ static LogicalResult checkImplementationStatus(Operation &op) {
       .Case([&](omp::TaskOp op) {
         checkAllocate(op, result);
         checkInReduction(op, result);
-        checkMergeable(op, result);
         checkPriority(op, result);
         checkUntied(op, result);
       })
@@ -1662,7 +1657,8 @@ convertOmpTaskOp(omp::TaskOp taskOp, llvm::IRBuilderBase &builder,
       moduleTranslation.getOpenMPBuilder()->createTask(
           ompLoc, allocaIP, bodyCB, !taskOp.getUntied(),
           moduleTranslation.lookupValue(taskOp.getFinal()),
-          moduleTranslation.lookupValue(taskOp.getIfExpr()), dds);
+          moduleTranslation.lookupValue(taskOp.getIfExpr()), dds,
+          taskOp.getMergeable());
 
   if (failed(handleError(afterIP, *taskOp)))
     return failure();

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

@@ -3003,6 +3003,19 @@ module attributes {omp.is_target_device = true} {
 
 // -----
 
+// Third argument is 5: essentially (4 || 1)
+// signifying this task is TIED and MERGEABLE
+
+// CHECK: {{.*}} = call ptr @__kmpc_omp_task_alloc(ptr @1, i32 %omp_global_thread_num, i32 5, i64 40, i64 0, ptr @omp_task_mergeable..omp_par)
+llvm.func @omp_task_mergeable() {
+  omp.task mergeable {
+    omp.terminator
+  }
+  llvm.return
+}
+
+// -----
+
 llvm.func external @foo_before() -> ()
 llvm.func external @foo() -> ()
 llvm.func external @foo_after() -> ()

mlir/test/Target/LLVMIR/openmp-llvm.mlir

@@ -447,17 +447,6 @@ llvm.func @task_in_reduction(%x : !llvm.ptr) {
 
 // -----
 
-llvm.func @task_mergeable() {
-  // expected-error@below {{not yet implemented: Unhandled clause mergeable in omp.task operation}}
-  // expected-error@below {{LLVM Translation failed for operation: omp.task}}
-  omp.task mergeable {
-    omp.terminator
-  }
-  llvm.return
-}
-
-// -----
-
 llvm.func @task_priority(%x : i32) {
   // expected-error@below {{not yet implemented: Unhandled clause priority in omp.task operation}}
   // expected-error@below {{LLVM Translation failed for operation: omp.task}}