[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`.
2024-11-26 16:10:26 +05:30
parent c0192a008c
commit b9e3a769b9
8 changed files with 42 additions and 44 deletions
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/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,
--- a/flang/test/Lower/OpenMP/Todo/task_mergeable.f90
+++ b/flang/test/Lower/OpenMP/Todo/task_mergeable.f90
@@ -1,13 +0,0 @@
-! RUN: %not_todo_cmd bbc -emit-fir -fopenmp -o - %s 2>&1 | FileCheck %s
-! RUN: %not_todo_cmd %flang_fc1 -emit-fir -fopenmp -o - %s 2>&1 | FileCheck %s
-
-!===============================================================================
-! `mergeable` clause
-!===============================================================================
-
-! CHECK: not yet implemented: Unhandled clause MERGEABLE in TASK construct
-subroutine omp_task_mergeable()
-  !$omp task mergeable
-  call foo()
-  !$omp end task
-end subroutine omp_task_mergeable
--- a/flang/test/Lower/OpenMP/task.f90
+++ b/flang/test/Lower/OpenMP/task.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
--- a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
+++ b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
@@ -1262,12 +1262,12 @@ public:
  ///                    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
  ///
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -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.
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.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();
--- a/mlir/test/Target/LLVMIR/openmp-llvm.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-llvm.mlir
@@ -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() -> ()
--- a/mlir/test/Target/LLVMIR/openmp-todo.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-todo.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}}