[flang][OpenMP] Lowering for task construct

[shraiysh]

Summary

This patch adds lowering for task construct from Fortran to
omp.task operation in OpenMPDialect Dialect (mlir). Also added tests
for the same.

[kiranchandramohan]

LGTM. Please wait a day in case others have comments.

Nit: We generally call the step from parse-tree to MLIR lowering.

[peixin]

One general question: Are those clauses for task construct supported when lowering to LLVM IR? Specifically, if, final, unted, mergeable, priority, allocate clauses. If not, should we wait for those support in OMPIRBuilder and lowering from MLIR to LLVMIR. Or will there be one TODO when lowering from MLIR to LLVMIR or OMPIRBuilder?

Nit: We generally call the step from parse-tree to MLIR lowering.

+1

In addition, this reminds me that there is one drawback collecting clauses in one place for all block constructs. For those unsupported clauses, one hard TODO is changed into one unknown result, maybe one segfault. It is easy to omit some features support and get users into trouble if there is segfault/ICE. For example, this patch makes private and firstprivate clauses passed. Are those two clauses supported for task construct if using the code in https://reviews.llvm.org/D122595?

[peixin]

@kiranchandramohan What do you think?

[shraiysh]

Or will there be one TODO when lowering from MLIR to LLVMIR or OMPIRBuilder?

There will be an error reporting when this is encountered. It is already done in patch D123919.

For example, this patch makes private and firstprivate clauses passed.

I agree, thanks for pointing out. Would it better if we split the genOmpBlockConstruct into genOmpParallelConstruct, genOMPTaskConstruct etc.?

[peixin]

Would it better if we split the genOmpBlockConstruct into genOmpParallelConstruct, genOMPTaskConstruct etc.?

I think so, too. Now there are 5 constructs using the same clauses handling. I think they should be split.

[kiranchandramohan]

In D124138#3470782, @peixin wrote:

One general question: Are those clauses for task construct supported when lowering to LLVM IR? Specifically, if, final, unted, mergeable, priority, allocate clauses. If not, should we wait for those support in OMPIRBuilder and lowering from MLIR to LLVMIR. Or will there be one TODO when lowering from MLIR to LLVMIR or OMPIRBuilder?

Nit: We generally call the step from parse-tree to MLIR lowering.

+1

In addition, this reminds me that there is one drawback collecting clauses in one place for all block constructs. For those unsupported clauses, one hard TODO is changed into one unknown result, maybe one segfault. It is easy to omit some features support and get users into trouble if there is segfault/ICE. For example, this patch makes private and firstprivate clauses passed. Are those two clauses supported for task construct if using the code in https://reviews.llvm.org/D122595?

Privatisation should work the same for all the block constructs. If it is not working then it is likely a bug. If there is a difference in handling the clauses depending on where it appears then we can split or parameterise a common function. Until then we can keep it the same is my opinion. But if you feel strongly then feel free to go ahead.

[peixin]

Privatisation should work the same for all the block constructs. If it is not working then it is likely a bug. If there is a difference in handling the clauses depending on where it appears then we can split or parameterise a common function. Until then we can keep it the same is my opinion. But if you feel strongly then feel free to go ahead.

I am also OK to this patch considering all the necessary clauses have TODOs when lowering to LLVM IR. @shraiysh Can you check if private and firstprivate works ok for task construct with this patch? If ok, can you add two test cases for them? We can delay the refactor work and discussion in D121583.

[peixin]

@shraiysh Can you add one test case of basic data type such as integer for private and firstprivate? I think you can move forward with this patch without the need to support complex data type for privatization. Maybe it's better to not let it block other works for task construct. Also, can you change the title and summary to "Support lowering to ..."?

flang/test/Lower/OpenMP/task.f90

3

Can you remove the fir-opt test? Also, can you remove the --check-prefixes?

[shraiysh]

@peixin sure I will add it soon and address your comments. I was working on this but was distracted by the issue in firstprivate implementation. For the time being, I will go ahead with the current implementation of firstprivate. Also, I will edit the revision to "lowering..." for this revision, but according to MLIR terminology, "lowering" refers to changing MLIR from one dialect operation to another (input is mlir and output is also mlir), and "translation" refers to the act of "importing to mlir" or "exporting from mlir" (either input or output is not mlir). I think this terminology is confusing and we should change it in the future.

A small question though, all captured task variables are firstprivate by default. So, this will be handled like default clause was handled, right? (we need to collect all the symbols under the region and emit firstprivate code for them).

[shraiysh]

Addressed comments

[peixin]

For fortran code, I think we usually use lowering for parse-tree to FIR Dialect.

A small question though, all captured task variables are firstprivate by default.

Is this implementation defined? Does clang do this? I don't find this requirement in the standard.

So, this will be handled like default clause was handled, right? (we need to collect all the symbols under the region and emit firstprivate code for them).

If this is implementation defined, I am not sure if it should be handled in lowering stage or OMPIRBuilder.

[shraiysh]

Added private clauses in multiple clauses test

[shraiysh]

Is this implementation defined? Does clang do this? I don't find this requirement in the standard.

Ah, I must have misread this from somewhere🤦‍♂️. This whole time I was focusing on it. I rechecked the standard and this requirement is not there. Thank you for correcting me.

[peixin]

LGTM

flang/lib/Lower/OpenMP.cpp

const auto &blockDirective =
       std::get<Fortran::parser::OmpBlockDirective>(beginBlockDirective.t);
   const auto &endBlockDirective =
       std::get<Fortran::parser::OmpEndBlockDirective>(blockConstruct.t);
   fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
   mlir::Location currentLocation = converter.getCurrentLocation();
 
   Fortran::lower::StatementContext stmtCtx;
   llvm::ArrayRef<mlir::Type> argTy;
-  mlir::Value ifClauseOperand, numThreadsClauseOperand;
+  mlir::Value ifClauseOperand, numThreadsClauseOperand, finalClauseOperand,
+      priorityClauseOperand;
   mlir::omp::ClauseProcBindKindAttr procBindKindAttr;
   SmallVector<Value> allocateOperands, allocatorOperands;
-  mlir::UnitAttr nowaitAttr;
+  mlir::UnitAttr nowaitAttr, untiedAttr, mergeableAttr;
 
   const auto &opClauseList =
       std::get<Fortran::parser::OmpClauseList>(beginBlockDirective.t);
   for (const auto &clause : opClauseList.v) {
     if (const auto &ifClause =
             std::get_if<Fortran::parser::OmpClause::If>(&clause.u)) {
       auto &expr = std::get<Fortran::parser::ScalarLogicalExpr>(ifClause->v.t);
       mlir::Value ifVal = fir::getBase(
for (const auto &clause : opClauseList.v) {
     } else if (std::get_if<Fortran::parser::OmpClause::Private>(&clause.u) ||
                std::get_if<Fortran::parser::OmpClause::Firstprivate>(
                    &clause.u)) {
       // Privatisation clauses are handled elsewhere.
       continue;
     } else if (std::get_if<Fortran::parser::OmpClause::Threads>(&clause.u)) {
       // Nothing needs to be done for threads clause.
       continue;
+    } else if (const auto &finalClause =
+                   std::get_if<Fortran::parser::OmpClause::Final>(&clause.u)) {
+      mlir::Value finalVal = fir::getBase(converter.genExprValue(
+          *Fortran::semantics::GetExpr(finalClause->v), stmtCtx));
+      finalClauseOperand = firOpBuilder.createConvert(
+          currentLocation, firOpBuilder.getI1Type(), finalVal);
+    } else if (std::get_if<Fortran::parser::OmpClause::Untied>(&clause.u)) {
+      untiedAttr = firOpBuilder.getUnitAttr();
+    } else if (std::get_if<Fortran::parser::OmpClause::Mergeable>(&clause.u)) {
+      mergeableAttr = firOpBuilder.getUnitAttr();
+    } else if (const auto &priorityClause =
+                   std::get_if<Fortran::parser::OmpClause::Priority>(
+                       &clause.u)) {
+      priorityClauseOperand = fir::getBase(converter.genExprValue(
+          *Fortran::semantics::GetExpr(priorityClause->v), stmtCtx));
     } else {
       TODO(currentLocation, "OpenMP Block construct clauses");
     }
   }
 
   for (const auto &clause :
        std::get<Fortran::parser::OmpClauseList>(endBlockDirective.t).v) {
     if (std::get_if<Fortran::parser::OmpClause::Nowait>(&clause.u))
genOMP(Fortran::lower::AbstractConverter &converter,
   } else if (blockDirective.v == llvm::omp::OMPD_single) {
     auto singleOp = firOpBuilder.create<mlir::omp::SingleOp>(
         currentLocation, allocateOperands, allocatorOperands, nowaitAttr);
     createBodyOfOp<omp::SingleOp>(singleOp, converter, currentLocation);
   } else if (blockDirective.v == llvm::omp::OMPD_ordered) {
     auto orderedOp = firOpBuilder.create<mlir::omp::OrderedRegionOp>(
         currentLocation, /*simd=*/nullptr);
     createBodyOfOp<omp::OrderedRegionOp>(orderedOp, converter, currentLocation);
+  } else if (blockDirective.v == llvm::omp::OMPD_task) {
+    auto taskOp = firOpBuilder.create<mlir::omp::TaskOp>(
+        currentLocation, ifClauseOperand, finalClauseOperand, untiedAttr,
+        mergeableAttr, /*in_reduction_vars=*/ValueRange(),
+        /*in_reductions=*/nullptr, priorityClauseOperand, allocateOperands,
+        allocatorOperands);
+    createBodyOfOp(taskOp, converter, currentLocation, &opClauseList);
   } else {
     TODO(converter.getCurrentLocation(), "Unhandled block directive");
   }
 }
 
 static void genOMP(Fortran::lower::AbstractConverter &converter,
                    Fortran::lower::pft::Evaluation &eval,
                    const Fortran::parser::OpenMPLoopConstruct &loopConstruct) {

flang/test/Lower/OpenMP/task.f90

+!RUN: %flang_fc1 -emit-fir -fopenmp %s -o - | FileCheck %s
+
+!CHECK-LABEL: func @_QPomp_task_simple() {
+subroutine omp_task_simple
+  !CHECK: omp.task {
+  !$omp task
+  !CHECK: fir.call @_QPfoo() : () -> ()
+  call foo()
+  !CHECK: omp.terminator
+  !$omp end task
+end subroutine omp_task_simple
+
+!===============================================================================
+! `if` clause
+!===============================================================================
+
+!CHECK-LABEL: func @_QPomp_task_if(%{{.+}}) {
+subroutine omp_task_if(bar)
+  logical, intent(inout) :: bar
+  !CHECK: omp.task if(%{{.+}}) {
+  !$omp task if(bar)
+  !CHECK: fir.call @_QPfoo() : () -> ()
+  call foo()
+  !CHECK: omp.terminator
+  !$omp end task
+end subroutine omp_task_if
+
+!===============================================================================
+! `final` clause
+!===============================================================================
+
+!CHECK-LABEL: func @_QPomp_task_final(%{{.+}}) {
+subroutine omp_task_final(bar)
+  logical, intent(inout) :: bar
+  !CHECK: omp.task final(%{{.+}}) {
+  !$omp task final(bar)
+  !CHECK: fir.call @_QPfoo() : () -> ()
+  call foo()
+  !CHECK: omp.terminator
+  !$omp end task
+end subroutine omp_task_final
+
+!===============================================================================
+! `untied` clause
+!===============================================================================
+
+!CHECK-LABEL: func @_QPomp_task_untied() {
+subroutine omp_task_untied()
+  !CHECK: omp.task untied {
+  !$omp task untied
+  !CHECK: fir.call @_QPfoo() : () -> ()
+  call foo()
+  !CHECK: omp.terminator
+  !$omp end task
+end subroutine omp_task_untied
+
+!===============================================================================
+! `mergeable` clause
+!===============================================================================
+
+!CHECK-LABEL: func @_QPomp_task_mergeable() {
+subroutine omp_task_mergeable()
+  !CHECK: omp.task mergeable {
+  !$omp task mergeable
+  !CHECK: fir.call @_QPfoo() : () -> ()
+  call foo()
+  !CHECK: omp.terminator
+  !$omp end task
+end subroutine omp_task_mergeable
+
+!===============================================================================
+! `priority` clause
+!===============================================================================
+
+!CHECK-LABEL: func @_QPomp_task_priority(%{{.+}}) {
+subroutine omp_task_priority(bar)
+  integer, intent(inout) :: bar
+  !CHECK: omp.task priority(%{{.+}}) {
+  !$omp task priority(bar)
+  !CHECK: fir.call @_QPfoo() : () -> ()
+  call foo()
+  !CHECK: omp.terminator
+  !$omp end task
+end subroutine omp_task_priority
+
+!===============================================================================
+! `allocate` clause
+!===============================================================================
+
+!CHECK-LABEL: func @_QPtask_allocate
+subroutine task_allocate()
+  use omp_lib
+  integer :: x
+  !CHECK: omp.task allocate(%{{.+}} : i32 -> %{{.+}} : !fir.ref<i32>) {
+  !$omp task allocate(omp_high_bw_mem_alloc: x) private(x)
+  !CHECK: arith.addi
+  x = x + 12
+  !CHECK: omp.terminator
+  !$omp end task
+end subroutine task_allocate
+
+!===============================================================================
+! `private` clause
+!===============================================================================
+!CHECK-LABEL: func @_QPtask_private
+subroutine task_private
+  type mytype
+  integer :: x
+  end type mytype
+
+  !CHECK: %[[int_var:.+]] = fir.alloca i32
+  !CHECK: %[[mytype_var:.+]] = fir.alloca !fir.type<_QFtask_privateTmytype{x:i32}>
+  integer :: int_var
+  type(mytype) :: mytype_var
+
+  !CHECK: fir.call @_QPbar(%[[int_var]], %[[mytype_var]]) : (!fir.ref<i32>, !fir.ref<!fir.type<_QFtask_privateTmytype{x:i32}>>) -> ()
+  call bar(int_var, mytype_var)
+
+  !CHECK: omp.task {
+  !$omp task private(int_var, mytype_var)
+  !CHECK: %[[int_var_private:.+]] = fir.alloca i32
+  !CHECK: %[[mytype_var_private:.+]] = fir.alloca !fir.type<_QFtask_privateTmytype{x:i32}>
+
+  !CHECK: fir.call @_QPbar(%[[int_var_private]], %[[mytype_var_private]]) : (!fir.ref<i32>, !fir.ref<!fir.type<_QFtask_privateTmytype{x:i32}>>) -> ()
+  call bar(int_var, mytype_var)
+  !CHECK: omp.terminator
+  !$omp end task
+end subroutine task_private
+
+!===============================================================================
+! `firstprivate` clause
+!===============================================================================
+!CHECK-LABEL: func @_QPtask_firstprivate
+subroutine task_firstprivate
+  type mytype
+  integer :: x
+  end type mytype
+
+  !CHECK: %[[int_var:.+]] = fir.alloca i32
+  !CHECK: %[[mytype_var:.+]] = fir.alloca !fir.type<_QFtask_firstprivateTmytype{x:i32}>
+  integer :: int_var
+  type(mytype) :: mytype_var
+
+  !CHECK: fir.call @_QPbaz(%[[int_var]], %[[mytype_var]]) : (!fir.ref<i32>, !fir.ref<!fir.type<_QFtask_firstprivateTmytype{x:i32}>>) -> ()
+  call baz(int_var, mytype_var)
+
+  !CHECK: omp.task {
+  !$omp task firstprivate(int_var, mytype_var)
+  !CHECK: %[[int_var_firstprivate:.+]] = fir.alloca i32
+  !CHECK: %[[int_var_load:.+]] = fir.load %[[int_var]] : !fir.ref<i32>
+  !CHECK: fir.store %[[int_var_load]] to %[[int_var_firstprivate]] : !fir.ref<i32>
+  !CHECK: %[[mytype_var_firstprivate:.+]] = fir.alloca !fir.type<_QFtask_firstprivateTmytype{x:i32}>
+  !CHECK: %[[mytype_var_load:.+]] = fir.load %[[mytype_var]] : !fir.ref<!fir.type<_QFtask_firstprivateTmytype{x:i32}>>
+  !CHECK: fir.store %[[mytype_var_load]] to %[[mytype_var_firstprivate]]
+  !CHECK: fir.call @_QPbaz(%[[int_var_firstprivate]], %[[mytype_var_firstprivate]]) : (!fir.ref<i32>, !fir.ref<!fir.type<_QFtask_firstprivateTmytype{x:i32}>>) -> ()
+  call baz(int_var, mytype_var)
+  !CHECK: omp.terminator
+  !$omp end task
+end subroutine task_firstprivate
+
+!===============================================================================
+! Multiple clauses
+!===============================================================================
+
+!CHECK-LABEL: func @_QPtask_multiple_clauses
+subroutine task_multiple_clauses()
+  use omp_lib
+
+  !CHECK: %[[x:.+]] = fir.alloca i32 {bindc_name = "x", uniq_name = "_QFtask_multiple_clausesEx"}
+  !CHECK: %[[y:.+]] = fir.alloca i32 {bindc_name = "y", uniq_name = "_QFtask_multiple_clausesEy"}
+  !CHECK: %[[z:.+]] = fir.alloca i32 {bindc_name = "z", uniq_name = "_QFtask_multiple_clausesEz"}
+  integer :: x, y, z
+  logical :: buzz
+
+  !CHECK: omp.task if(%{{.+}}) final(%{{.+}}) untied mergeable priority(%{{.+}}) allocate(%{{.+}} : i32 -> %{{.+}} : !fir.ref<i32>) {
+  !$omp task if(buzz) final(buzz) untied mergeable priority(z) allocate(omp_high_bw_mem_alloc: x) private(x) firstprivate(y)
+
+  !CHECK: %[[x_priv:.+]] = fir.alloca i32
+  !CHECK: %[[y_priv:.+]] = fir.alloca i32
+  !CHECK: %[[y_load:.+]] = fir.load %[[y]] : !fir.ref<i32>
+  !CHECK: fir.store %[[y_load]] to %[[y_priv]] : !fir.ref<i32>
+
+  !CHECK: arith.addi
+  x = x + 12
+  !CHECK: arith.subi
+  y = y - 12
+
+  !CHECK: omp.terminator
+  !$omp end task
+end subroutine task_multiple_clauses