[flang][OpenMP] Lower OpenMP `taskloop` construct

[NimishMishra]

Summary

Lower Fortran OpenMP taskloop to FIR + OpenMP Dialect.

[kiranchandramohan]

Can we share some code for creating the Operands from the Clauses? It can be completing the work started by @shraiysh.

[NimishMishra]

Rebase with main. Fixed failing test.

[NimishMishra]

In D136182#3865953, @kiranchandramohan wrote:

Can we share some code for creating the Operands from the Clauses? It can be completing the work started by @shraiysh.

Hello @kiranchandramohan,

I am taking care of this revision now. Can you explain a little on this? Do we mean we should refactor code that extracts clauses into a separate function?

[kiranchandramohan]

In D136182#4386726, @NimishMishra wrote:

In D136182#3865953, @kiranchandramohan wrote:

Can we share some code for creating the Operands from the Clauses? It can be completing the work started by @shraiysh.

Hello @kiranchandramohan,

I am taking care of this revision now. Can you explain a little on this? Do we mean we should refactor code that extracts clauses into a separate function?

The suggestion was to use D121583 to refactor clause handling code. This can be done separately.

Changes in convert-to-llvm-openmp-and-fir.fir seem to be removing some unrelated code. Is this really intended?

Please submit this only after the MLIR to LLVM IR changes are submitted, otherwise we will get some crashes when code with taskloop is run.

flang/test/Lower/OpenMP/taskloop.f90

2

Please add more tests to cover handling of all the clauses that are lowered.

flang/lib/Lower/OpenMP.cpp

 static void genOMP(Fortran::lower::AbstractConverter &converter,
                    Fortran::lower::pft::Evaluation &eval,
                    const Fortran::parser::OpenMPLoopConstruct &loopConstruct) {
 
   fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
   mlir::Location currentLocation = converter.getCurrentLocation();
   llvm::SmallVector<mlir::Value> lowerBound, upperBound, step, linearVars,
       linearStepVars, reductionVars;
-  mlir::Value scheduleChunkClauseOperand, ifClauseOperand;
+  mlir::Value scheduleChunkClauseOperand, ifClauseOperand,
+      grainSizeClauseOperand, numTasksClauseOperand, finalClauseOperand,
+      priorityClauseOperand;
   mlir::Attribute scheduleClauseOperand, noWaitClauseOperand,
       orderedClauseOperand, orderClauseOperand;
+  mlir::UnitAttr untiedAttr, mergeableAttr, nogroupAttr;
   mlir::IntegerAttr simdlenClauseOperand, safelenClauseOperand;
   SmallVector<Attribute> reductionDeclSymbols;
   Fortran::lower::StatementContext stmtCtx;
   const auto &loopOpClauseList = std::get<Fortran::parser::OmpClauseList>(
       std::get<Fortran::parser::OmpBeginLoopDirective>(loopConstruct.t).t);
 
   const auto ompDirective =
       std::get<Fortran::parser::OmpLoopDirective>(
           std::get<Fortran::parser::OmpBeginLoopDirective>(loopConstruct.t).t)
           .v;
   if (llvm::omp::OMPD_parallel_do == ompDirective) {
     createCombinedParallelOp<Fortran::parser::OmpBeginLoopDirective>(
         converter, eval,
         std::get<Fortran::parser::OmpBeginLoopDirective>(loopConstruct.t));
   } else if (llvm::omp::OMPD_do != ompDirective &&
+             llvm::omp::OMPD_taskloop != ompDirective &&
              llvm::omp::OMPD_simd != ompDirective) {
     TODO(converter.getCurrentLocation(), "Construct enclosing do loop");
   }
 
   // Collect the loops to collapse.
   auto *doConstructEval = &eval.getFirstNestedEvaluation();
 
   std::int64_t collapseValue =
if (const auto &scheduleClause =
       simdlenClauseOperand = firOpBuilder.getI64IntegerAttr(*simdlenVal);
     } else if (const auto &safelenClause =
                    std::get_if<Fortran::parser::OmpClause::Safelen>(
                        &clause.u)) {
       const auto *expr = Fortran::semantics::GetExpr(safelenClause->v);
       const std::optional<std::int64_t> safelenVal =
           Fortran::evaluate::ToInt64(*expr);
       safelenClauseOperand = firOpBuilder.getI64IntegerAttr(*safelenVal);
+    } // `TaskLoop` construct clauses.
+    else if (const auto &grainSizeClause =
+                 std::get_if<Fortran::parser::OmpClause::Grainsize>(
+                     &clause.u)) {
+      const auto *expr = Fortran::semantics::GetExpr(grainSizeClause->v);
+      grainSizeClauseOperand =
+          fir::getBase(converter.genExprValue(*expr, stmtCtx));
+    } else if (const auto &numTasksClause =
+                   std::get_if<Fortran::parser::OmpClause::NumTasks>(
+                       &clause.u)) {
+      const auto *expr = Fortran::semantics::GetExpr(numTasksClause->v);
+      numTasksClauseOperand =
+          fir::getBase(converter.genExprValue(*expr, stmtCtx));
+    } else if (std::get_if<Fortran::parser::OmpClause::Mergeable>(&clause.u)) {
+      mergeableAttr = firOpBuilder.getUnitAttr();
+    } else if (std::get_if<Fortran::parser::OmpClause::Untied>(&clause.u)) {
+      untiedAttr = firOpBuilder.getUnitAttr();
+    } else if (std::get_if<Fortran::parser::OmpClause::Nogroup>(&clause.u)) {
+      nogroupAttr = firOpBuilder.getUnitAttr();
+    } 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 (const auto &priorityClause =
+                   std::get_if<Fortran::parser::OmpClause::Priority>(
+                       &clause.u)) {
+      priorityClauseOperand = fir::getBase(converter.genExprValue(
+          *Fortran::semantics::GetExpr(priorityClause->v), stmtCtx));
     }
   }
 
   // The types of lower bound, upper bound, and step are converted into the
   // type of the loop variable if necessary.
   mlir::Type loopVarType = getLoopVarType(converter, loopVarTypeSize);
   for (unsigned it = 0; it < (unsigned)lowerBound.size(); it++) {
     lowerBound[it] = firOpBuilder.createConvert(currentLocation, loopVarType,
auto SimdLoopOp = firOpBuilder.create<mlir::omp::SimdLoopOp>(
         currentLocation, resultType, lowerBound, upperBound, step,
         ifClauseOperand, simdlenClauseOperand, safelenClauseOperand,
         /*inclusive=*/firOpBuilder.getUnitAttr());
     createBodyOfOp<omp::SimdLoopOp>(SimdLoopOp, converter, currentLocation,
                                     eval, &loopOpClauseList, iv);
     return;
   }
 
+  // TODO: Support all the clauses
+  if (llvm::omp::OMPD_taskloop == ompDirective) {
+    auto TaskLoopOp = firOpBuilder.create<mlir::omp::TaskLoopOp>(
+        currentLocation, lowerBound, upperBound, step,
+        /*inclusive=*/firOpBuilder.getUnitAttr(), ifClauseOperand,
+        finalClauseOperand, untiedAttr, mergeableAttr,
+        /*in_reduction_vars*/ ValueRange(), /*in_reductions*/ nullptr,
+        /*reduction_vars*/ ValueRange(),
+        /*reductions*/ nullptr, priorityClauseOperand,
+        /*allocate_vars*/ ValueRange(), /*allocators_vars*/ ValueRange(),
+        grainSizeClauseOperand, numTasksClauseOperand, nogroupAttr);
+    createBodyOfOp<omp::TaskLoopOp>(TaskLoopOp, converter, currentLocation,
+                                    eval, &loopOpClauseList, iv);
+    return;
+  }
   // FIXME: Add support for following clauses:
   // 1. linear
   // 2. order
   auto wsLoopOp = firOpBuilder.create<mlir::omp::WsLoopOp>(
       currentLocation, lowerBound, upperBound, step, linearVars, linearStepVars,
       reductionVars,
       reductionDeclSymbols.empty()
           ? nullptr

flang/test/Fir/convert-to-llvm-openmp-and-fir.fir

 // CHECK:              omp.terminator
 // CHECK:            }
 // CHECK:            omp.terminator
 // CHECK:          }
 // CHECK:          llvm.return
 // CHECK:        }
 // CHECK:        llvm.func @_QPwork() attributes {sym_visibility = "private"}
 // CHECK:      }
+
+// -----
+
+func.func @_QPomp_taskloop() {
+    %0 = fir.alloca i32 {adapt.valuebyref}
+    %1 = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFomp_taskloopEi"}
+    %c1_i32 = arith.constant 1 : i32
+    %c10000_i32 = arith.constant 10000 : i32
+    %c1_i32_0 = arith.constant 1 : i32
+    %c2_i32 = arith.constant 2 : i32
+    omp.taskloop   num_tasks(%c2_i32 : i32) nogroup for  (%arg0) : i32 = (%c1_i32) to (%c10000_i32) inclusive step (%c1_i32_0) {
+      fir.store %arg0 to %0 : !fir.ref<i32>
+      fir.call @_QPwork(%0) : (!fir.ref<i32>) -> ()
+      omp.terminator
+    }
+    return
+  }
+  func.func private @_QPwork(!fir.ref<i32>)
+
+// CHECK-LABEL: llvm.func @_QPomp_taskloop() {
+// CHECK:    %0 = llvm.mlir.constant(1 : i64) : i64
+// CHECK:    %1 = llvm.alloca %0 x i32 {adapt.valuebyref, in_type = i32, operand_segment_sizes = array<i32: 0, 0>} : (i64) -> !llvm.ptr<i32>
+// CHECK:    %2 = llvm.mlir.constant(1 : i64) : i64
+// CHECK:    %3 = llvm.alloca %2 x i32 {bindc_name = "i", in_type = i32, operand_segment_sizes = array<i32: 0, 0>, uniq_name = "_QFomp_taskloopEi"} : (i64) -> !llvm.ptr<i32>
+// CHECK:    %[[LOWER:.*]] = llvm.mlir.constant(1 : i32) : i32
+// CHECK:    %[[UPPER:.*]] = llvm.mlir.constant(10000 : i32) : i32
+// CHECK:    %[[STEP:.*]] = llvm.mlir.constant(1 : i32) : i32
+// CHECK:    %7 = llvm.mlir.constant(2 : i32) : i32
+// CHECK:    omp.taskloop   num_tasks(%7 : i32) nogroup for  (%arg0) : i32 = (%[[LOWER]]) to (%[[UPPER]]) inclusive step (%[[STEP]]) {
+// CHECK:      llvm.store %arg0, %1 : !llvm.ptr<i32>
+// CHECK:      llvm.call @_QPwork(%1) : (!llvm.ptr<i32>) -> ()
+// CHECK:      omp.terminator
+// CHECK:    }
+// CHECK:    llvm.return
+// CHECK:  }
+// CHECK:  llvm.func @_QPwork(!llvm.ptr<i32>) attributes {sym_visibility = "private"}

flang/test/Lower/OpenMP/taskloop.f90

+!RUN: %flang_fc1 -emit-fir -fopenmp %s -o - | FileCheck %s
+
+!CHECK-LABEL: @_QPomp_taskloop
+subroutine omp_taskloop
+integer i
+
+!CHECK: omp.taskloop   num_tasks(%{{.*}} : i32) nogroup for  (%{{.*}}) : i32 = (%c1_i32) to ({{.*}}) inclusive step (%{{.*}}) {
+ !$omp taskloop private(i) num_tasks(2) nogroup
+ do i=1,10000
+!CHECK: fir.store %{{.*}} to %{{.*}} : !fir.ref<i32>
+!CHECK: fir.call @_QPwork(%{{.*}}) : (!fir.ref<i32>) -> ()
+ call work(i)
+ end do
+!CHECK: omp.terminator
+ !$omp end taskloop
+
+ end subroutine