[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

1

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, alignedVars, nontemporalVars;
-  mlir::Value scheduleChunkClauseOperand, ifClauseOperand;
+  mlir::Value scheduleChunkClauseOperand, ifClauseOperand,
+      grainSizeClauseOperand, numTasksClauseOperand, finalClauseOperand,
+      priorityClauseOperand;
+  mlir::UnitAttr untiedAttr, mergeableAttr, nogroupAttr;
   mlir::Attribute scheduleClauseOperand, noWaitClauseOperand,
       orderedClauseOperand, orderClauseOperand;
   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();
   if (doConstructEval->getIf<Fortran::parser::DoConstruct>()
           ->IsDoConcurrent()) {
for (const auto &clause : loopOpClauseList.v) {
     } 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,
                                                 lowerBound[it]);
auto simdLoopOp = firOpBuilder.create<mlir::omp::SimdLoopOp>(
         orderClauseOperand.dyn_cast_or_null<omp::ClauseOrderKindAttr>(),
         simdlenClauseOperand, safelenClauseOperand,
         /*inclusive=*/firOpBuilder.getUnitAttr());
     createBodyOfOp<omp::SimdLoopOp>(simdLoopOp, converter, currentLocation,
                                     eval, &loopOpClauseList, iv);
     return;
   }
 
+  // 2.10.2 taskloop construct
+  // 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);
+  }
   // 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:     llvm.store %[[N]], %[[X_REF]] : !llvm.ptr<i32>
 // CHECK:     omp.terminator
 // CHECK:   }
 // CHECK:   omp.terminator
 // CHECK: }
 // CHECK: llvm.return
 // CHECK: }
 
+
 // -----
 
 func.func @_QPsb(%arr: !fir.box<!fir.array<?xi32>> {fir.bindc_name = "arr"}) {
   %0 = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFsbEi"}
   omp.parallel   {
     %c1 = arith.constant 1 : i32
     %c50 = arith.constant 50 : i32
     omp.wsloop   for  (%indx) : i32 = (%c1) to (%c50) inclusive step (%c1) {

 // CHECK:   llvm.store %[[I1]], %[[ARR_I_REF]] : !llvm.ptr<i32>
 // CHECK: omp.yield
 // CHECK: }
 // CHECK: omp.terminator
 // CHECK: }
 // CHECK: llvm.return
 // CHECK: }
 
-// -----
-
-func.func @_QPomp_target_data() {
-  %0 = fir.alloca !fir.array<1024xi32> {bindc_name = "a", uniq_name = "_QFomp_target_dataEa"}
-  %1 = fir.alloca !fir.array<1024xi32> {bindc_name = "b", uniq_name = "_QFomp_target_dataEb"}
-  %2 = fir.alloca !fir.array<1024xi32> {bindc_name = "c", uniq_name = "_QFomp_target_dataEc"}
-  %3 = fir.alloca !fir.array<1024xi32> {bindc_name = "d", uniq_name = "_QFomp_target_dataEd"}
-  omp.target_enter_data   map((to -> %0 : !fir.ref<!fir.array<1024xi32>>), (to -> %1 : !fir.ref<!fir.array<1024xi32>>), (always, alloc -> %2 : !fir.ref<!fir.array<1024xi32>>))
-  omp.target_exit_data   map((from -> %0 : !fir.ref<!fir.array<1024xi32>>), (from -> %1 : !fir.ref<!fir.array<1024xi32>>), (release -> %2 : !fir.ref<!fir.array<1024xi32>>), (always, delete -> %3 : !fir.ref<!fir.array<1024xi32>>))
-  return
-}
-
-// CHECK-LABEL:   llvm.func @_QPomp_target_data() {
-// CHECK:           %[[VAL_0:.*]] = llvm.mlir.constant(1 : i64) : i64
-// CHECK:           %[[VAL_1:.*]] = llvm.alloca %[[VAL_0]] x !llvm.array<1024 x i32> {bindc_name = "a", in_type = !fir.array<1024xi32>, operand_segment_sizes = array<i32: 0, 0>, uniq_name = "_QFomp_target_dataEa"} : (i64) -> !llvm.ptr<array<1024 x i32>>
-// CHECK:           %[[VAL_2:.*]] = llvm.mlir.constant(1 : i64) : i64
-// CHECK:           %[[VAL_3:.*]] = llvm.alloca %[[VAL_2]] x !llvm.array<1024 x i32> {bindc_name = "b", in_type = !fir.array<1024xi32>, operand_segment_sizes = array<i32: 0, 0>, uniq_name = "_QFomp_target_dataEb"} : (i64) -> !llvm.ptr<array<1024 x i32>>
-// CHECK:           %[[VAL_4:.*]] = llvm.mlir.constant(1 : i64) : i64
-// CHECK:           %[[VAL_5:.*]] = llvm.alloca %[[VAL_4]] x !llvm.array<1024 x i32> {bindc_name = "c", in_type = !fir.array<1024xi32>, operand_segment_sizes = array<i32: 0, 0>, uniq_name = "_QFomp_target_dataEc"} : (i64) -> !llvm.ptr<array<1024 x i32>>
-// CHECK:           %[[VAL_6:.*]] = llvm.mlir.constant(1 : i64) : i64
-// CHECK:           %[[VAL_7:.*]] = llvm.alloca %[[VAL_6]] x !llvm.array<1024 x i32> {bindc_name = "d", in_type = !fir.array<1024xi32>, operand_segment_sizes = array<i32: 0, 0>, uniq_name = "_QFomp_target_dataEd"} : (i64) -> !llvm.ptr<array<1024 x i32>>
-// CHECK:           omp.target_enter_data   map((to -> %[[VAL_1]] : !llvm.ptr<array<1024 x i32>>), (to -> %[[VAL_3]] : !llvm.ptr<array<1024 x i32>>), (always, alloc -> %[[VAL_5]] : !llvm.ptr<array<1024 x i32>>))
-// CHECK:           omp.target_exit_data   map((from -> %[[VAL_1]] : !llvm.ptr<array<1024 x i32>>), (from -> %[[VAL_3]] : !llvm.ptr<array<1024 x i32>>), (release -> %[[VAL_5]] : !llvm.ptr<array<1024 x i32>>), (always, delete -> %[[VAL_7]] : !llvm.ptr<array<1024 x i32>>))
-// CHECK:           llvm.return
-// CHECK:         }
-
-// -----
-
-func.func @_QPopenmp_target_data_region() {
-  %0 = fir.alloca !fir.array<1024xi32> {bindc_name = "a", uniq_name = "_QFopenmp_target_data_regionEa"}
-  %1 = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFopenmp_target_data_regionEi"}
-  omp.target_data   map((tofrom -> %0 : !fir.ref<!fir.array<1024xi32>>)) {
-    %c1_i32 = arith.constant 1 : i32
-    %2 = fir.convert %c1_i32 : (i32) -> index
-    %c1024_i32 = arith.constant 1024 : i32
-    %3 = fir.convert %c1024_i32 : (i32) -> index
-    %c1 = arith.constant 1 : index
-    %4 = fir.convert %2 : (index) -> i32
-    %5:2 = fir.do_loop %arg0 = %2 to %3 step %c1 iter_args(%arg1 = %4) -> (index, i32) {
-      fir.store %arg1 to %1 : !fir.ref<i32>
-      %6 = fir.load %1 : !fir.ref<i32>
-      %7 = fir.load %1 : !fir.ref<i32>
-      %8 = fir.convert %7 : (i32) -> i64
-      %c1_i64 = arith.constant 1 : i64
-      %9 = arith.subi %8, %c1_i64 : i64
-      %10 = fir.coordinate_of %0, %9 : (!fir.ref<!fir.array<1024xi32>>, i64) -> !fir.ref<i32>
-      fir.store %6 to %10 : !fir.ref<i32>
-      %11 = arith.addi %arg0, %c1 : index
-      %12 = fir.convert %c1 : (index) -> i32
-      %13 = fir.load %1 : !fir.ref<i32>
-      %14 = arith.addi %13, %12 : i32
-      fir.result %11, %14 : index, i32
-    }
-    fir.store %5#1 to %1 : !fir.ref<i32>
-    omp.terminator
-  }
-  return
-}
-
-// CHECK-LABEL:   llvm.func @_QPopenmp_target_data_region() {
-// CHECK:           %[[VAL_0:.*]] = llvm.mlir.constant(1 : i64) : i64
-// CHECK:           %[[VAL_1:.*]] = llvm.alloca %[[VAL_0]] x !llvm.array<1024 x i32> {bindc_name = "a", in_type = !fir.array<1024xi32>, operand_segment_sizes = array<i32: 0, 0>, uniq_name = "_QFopenmp_target_data_regionEa"} : (i64) -> !llvm.ptr<array<1024 x i32>>
-// CHECK:           %[[VAL_2:.*]] = llvm.mlir.constant(1 : i64) : i64
-// CHECK:           %[[VAL_3:.*]] = llvm.alloca %[[VAL_2]] x i32 {bindc_name = "i", in_type = i32, operand_segment_sizes = array<i32: 0, 0>, uniq_name = "_QFopenmp_target_data_regionEi"} : (i64) -> !llvm.ptr<i32>
-// CHECK:           omp.target_data   map((tofrom -> %[[VAL_1]] : !llvm.ptr<array<1024 x i32>>)) {
-// CHECK:             %[[VAL_4:.*]] = llvm.mlir.constant(1 : i32) : i32
-// CHECK:             %[[VAL_5:.*]] = llvm.sext %[[VAL_4]] : i32 to i64
-// CHECK:             %[[VAL_6:.*]] = llvm.mlir.constant(1024 : i32) : i32
-// CHECK:             %[[VAL_7:.*]] = llvm.sext %[[VAL_6]] : i32 to i64
-// CHECK:             %[[VAL_8:.*]] = llvm.mlir.constant(1 : index) : i64
-// CHECK:             %[[VAL_9:.*]] = llvm.trunc %[[VAL_5]] : i64 to i32
-// CHECK:             %[[VAL_10:.*]] = llvm.sub %[[VAL_7]], %[[VAL_5]]  : i64
-// CHECK:             %[[VAL_11:.*]] = llvm.add %[[VAL_10]], %[[VAL_8]]  : i64
-// CHECK:             llvm.br ^bb1(%[[VAL_5]], %[[VAL_9]], %[[VAL_11]] : i64, i32, i64)
-// CHECK:           ^bb1(%[[VAL_12:.*]]: i64, %[[VAL_13:.*]]: i32, %[[VAL_14:.*]]: i64):
-// CHECK:             %[[VAL_15:.*]] = llvm.mlir.constant(0 : index) : i64
-// CHECK:             %[[VAL_16:.*]] = llvm.icmp "sgt" %[[VAL_14]], %[[VAL_15]] : i64
-// CHECK:             llvm.cond_br %[[VAL_16]], ^bb2, ^bb3
-// CHECK:           ^bb2:
-// CHECK:             llvm.store %[[VAL_13]], %[[VAL_3]] : !llvm.ptr<i32>
-// CHECK:             %[[VAL_17:.*]] = llvm.load %[[VAL_3]] : !llvm.ptr<i32>
-// CHECK:             %[[VAL_18:.*]] = llvm.load %[[VAL_3]] : !llvm.ptr<i32>
-// CHECK:             %[[VAL_19:.*]] = llvm.sext %[[VAL_18]] : i32 to i64
-// CHECK:             %[[VAL_20:.*]] = llvm.mlir.constant(1 : i64) : i64
-// CHECK:             %[[VAL_21:.*]] = llvm.sub %[[VAL_19]], %[[VAL_20]]  : i64
-// CHECK:             %[[VAL_22:.*]] = llvm.getelementptr %[[VAL_1]][0, %[[VAL_21]]] : (!llvm.ptr<array<1024 x i32>>, i64) -> !llvm.ptr<i32>
-// CHECK:             llvm.store %[[VAL_17]], %[[VAL_22]] : !llvm.ptr<i32>
-// CHECK:             %[[VAL_23:.*]] = llvm.add %[[VAL_12]], %[[VAL_8]]  : i64
-// CHECK:             %[[VAL_24:.*]] = llvm.trunc %[[VAL_8]] : i64 to i32
-// CHECK:             %[[VAL_25:.*]] = llvm.load %[[VAL_3]] : !llvm.ptr<i32>
-// CHECK:             %[[VAL_26:.*]] = llvm.add %[[VAL_25]], %[[VAL_24]]  : i32
-// CHECK:             %[[VAL_27:.*]] = llvm.add %[[VAL_12]], %[[VAL_8]]  : i64
-// CHECK:             %[[VAL_28:.*]] = llvm.mlir.constant(1 : index) : i64
-// CHECK:             %[[VAL_29:.*]] = llvm.sub %[[VAL_14]], %[[VAL_28]]  : i64
-// CHECK:             llvm.br ^bb1(%[[VAL_27]], %[[VAL_26]], %[[VAL_29]] : i64, i32, i64)
-// CHECK:           ^bb3:
-// CHECK:             llvm.store %[[VAL_13]], %[[VAL_3]] : !llvm.ptr<i32>
-// CHECK:             omp.terminator
-// CHECK:           }
-// CHECK:           llvm.return
-// CHECK:         }
-
-// -----
-
-func.func @_QPomp_target() {
-  %0 = fir.alloca !fir.array<512xi32> {bindc_name = "a", uniq_name = "_QFomp_targetEa"}
-  %c64_i32 = arith.constant 64 : i32
-  omp.target   thread_limit(%c64_i32 : i32) map((tofrom -> %0 : !fir.ref<!fir.array<512xi32>>)) {
-    %c10_i32 = arith.constant 10 : i32
-    %c1_i64 = arith.constant 1 : i64
-    %c1_i64_0 = arith.constant 1 : i64
-    %1 = arith.subi %c1_i64, %c1_i64_0 : i64
-    %2 = fir.coordinate_of %0, %1 : (!fir.ref<!fir.array<512xi32>>, i64) -> !fir.ref<i32>
-    fir.store %c10_i32 to %2 : !fir.ref<i32>
-    omp.terminator
-  }
-  return
-}
-
-// CHECK-LABEL:   llvm.func @_QPomp_target() {
-// CHECK:           %[[VAL_0:.*]] = llvm.mlir.constant(1 : i64) : i64
-// CHECK:           %[[VAL_1:.*]] = llvm.alloca %[[VAL_0]] x !llvm.array<512 x i32> {bindc_name = "a", in_type = !fir.array<512xi32>, operand_segment_sizes = array<i32: 0, 0>, uniq_name = "_QFomp_targetEa"} : (i64) -> !llvm.ptr<array<512 x i32>>
-// CHECK:           %[[VAL_2:.*]] = llvm.mlir.constant(64 : i32) : i32
-// CHECK:           omp.target   thread_limit(%[[VAL_2]] : i32) map((tofrom -> %[[VAL_1]] : !llvm.ptr<array<512 x i32>>)) {
-// CHECK:             %[[VAL_3:.*]] = llvm.mlir.constant(10 : i32) : i32
-// CHECK:             %[[VAL_4:.*]] = llvm.mlir.constant(1 : i64) : i64
-// CHECK:             %[[VAL_5:.*]] = llvm.mlir.constant(1 : i64) : i64
-// CHECK:             %[[VAL_6:.*]] = llvm.mlir.constant(0 : i64) : i64
-// CHECK:             %[[VAL_7:.*]] = llvm.getelementptr %[[VAL_1]][0, %[[VAL_6]]] : (!llvm.ptr<array<512 x i32>>, i64) -> !llvm.ptr<i32>
-// CHECK:             llvm.store %[[VAL_3]], %[[VAL_7]] : !llvm.ptr<i32>
-// CHECK:             omp.terminator
-// CHECK:           }
-// CHECK:           llvm.return
-// CHECK:         }
-
-// -----
-
 func.func @_QPsimdloop_with_nested_loop() {
   %0 = fir.alloca i32 {adapt.valuebyref}
   %1 = fir.alloca !fir.array<10xi32> {bindc_name = "a", uniq_name = "_QFsimdloop_with_nested_loopEa"}
   %2 = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFsimdloop_with_nested_loopEi"}
   %3 = fir.alloca i32 {bindc_name = "j", uniq_name = "_QFsimdloop_with_nested_loopEj"}
   %c1_i32 = arith.constant 1 : i32
   %c10_i32 = arith.constant 10 : i32
   %c1_i32_0 = arith.constant 1 : i32

 // CHECK:          }
 // CHECK:          llvm.return
 // CHECK:        }
 // CHECK:        llvm.func @_QPwork() attributes {sym_visibility = "private"}
 // CHECK:      }
 
 // -----
 
-
-func.func @_QQmain() {
-  %c0 = arith.constant 0 : index
-  %c5 = arith.constant 5 : index
-  %c1 = arith.constant 1 : index
-  %0 = fir.alloca i32
-  omp.taskgroup   {
-    %1 = fir.convert %c1 : (index) -> i32
-    cf.br ^bb1(%1, %c5 : i32, index)
-  ^bb1(%2: i32, %3: index):  // 2 preds: ^bb0, ^bb2
-    %4 = arith.cmpi sgt, %3, %c0 : index
-    cf.cond_br %4, ^bb2, ^bb3
-  ^bb2:  // pred: ^bb1
-    fir.store %2 to %0 : !fir.ref<i32>
-    omp.task   {
-      fir.call @_QFPdo_work(%0) : (!fir.ref<i32>) -> ()
-      omp.terminator
-    }
-    %5 = fir.load %0 : !fir.ref<i32>
-    %6 = arith.addi %5, %1 : i32
-    %7 = arith.subi %3, %c1 : index
-    cf.br ^bb1(%6, %7 : i32, index)
-  ^bb3:  // pred: ^bb1
-    fir.store %2 to %0 : !fir.ref<i32>
-    omp.terminator
-  }
-  return
-}
-func.func private @_QFPdo_work(!fir.ref<i32>)
-
-// CHECK-LABEL: llvm.func @_QQmain
-// CHECK:          omp.taskgroup   {
-// CHECK:            omp.task   {
-// CHECK:              llvm.call @_QFPdo_work({{.*}}) : (!llvm.ptr<i32>) -> ()
-// CHECK:              omp.terminator
-// CHECK:            }
-// CHECK:            omp.terminator
-// CHECK:          }
-// CHECK:          llvm.return
-// CHECK:        }
-// CHECK:        llvm.func @_QFPdo_work(!llvm.ptr<i32>)
-// CHECK:      }
-
-// -----
-
-func.func @_QPs() {
-  %0 = fir.address_of(@_QFsEc) : !fir.ref<i32>
-  omp.atomic.update   %0 : !fir.ref<i32> {
-  ^bb0(%arg0: i32):
-    %c1_i32 = arith.constant 1 : i32
-    %1 = arith.addi %arg0, %c1_i32 : i32
-    omp.yield(%1 : i32)
-  }
-  return
-}
-fir.global internal @_QFsEc : i32 {
-  %c10_i32 = arith.constant 10 : i32
-  fir.has_value %c10_i32 : i32
-}
-
-// CHECK-LABEL:  llvm.func @_QPs() {
-// CHECK:    %[[GLOBAL_VAR:.*]] = llvm.mlir.addressof @[[GLOBAL:.*]] : !llvm.ptr<i32>
-// CHECK:    omp.atomic.update   %[[GLOBAL_VAR]] : !llvm.ptr<i32> {
-// CHECK:    ^bb0(%[[IN_VAL:.*]]: i32):
-// CHECK:      %[[CONST_1:.*]] = llvm.mlir.constant(1 : i32) : i32
-// CHECK:      %[[OUT_VAL:.*]] = llvm.add %[[IN_VAL]], %[[CONST_1]]  : i32
-// CHECK:      omp.yield(%[[OUT_VAL]] : i32)
-// CHECK:    }
-// CHECK:    llvm.return
-// CHECK:  }
-// CHECK:  llvm.mlir.global internal @[[GLOBAL]]() {{.*}} : i32 {
-// CHECK:    %[[INIT_10:.*]] = llvm.mlir.constant(10 : i32) : i32
-// CHECK:    llvm.return %[[INIT_10]] : i32
-// CHECK:  }
-
-func.func @_QPsb() {
-  %c10 = arith.constant 10 : index
-  %c1 = arith.constant 1 : index
+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
-  %c0_i32 = arith.constant 0 : i32
-  %0 = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFsbEi"}
-  %1 = fir.alloca i32 {bindc_name = "li", uniq_name = "_QFsbEli"}
-  fir.store %c0_i32 to %1 : !fir.ref<i32>
-  omp.sections   {
-    omp.section {
-      %2 = fir.convert %c1 : (index) -> i32
-      %3:2 = fir.do_loop %arg0 = %c1 to %c10 step %c1 iter_args(%arg1 = %2) -> (index, i32) {
-        fir.store %arg1 to %0 : !fir.ref<i32>
-        %4 = fir.load %1 : !fir.ref<i32>
-        %5 = arith.addi %4, %c1_i32 : i32
-        fir.store %5 to %1 : !fir.ref<i32>
-        %6 = arith.addi %arg0, %c1 : index
-        %7 = fir.convert %c1 : (index) -> i32
-        %8 = fir.load %0 : !fir.ref<i32>
-        %9 = arith.addi %8, %7 : i32
-        fir.result %6, %9 : index, i32
-      }
-      fir.store %3#1 to %0 : !fir.ref<i32>
-      omp.terminator
-    }
+    %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:  llvm.func @_QPsb() {
-// CHECK:    %[[ONE:.*]] = llvm.mlir.constant(1 : i32) : i32
-// CHECK:    %[[SIZE:.*]] = llvm.mlir.constant(1 : i64) : i64
-// CHECK:    %[[LI_REF:.*]] = llvm.alloca %6 x i32 {bindc_name = "li", in_type = i32, operand_segment_sizes = array<i32: 0, 0>, uniq_name = "_QFsbEli"} : (i64) -> !llvm.ptr<i32>
-// CHECK:    omp.sections   {
-// CHECK:      omp.section {
-// CHECK:        llvm.br ^[[BB_ENTRY:.*]]({{.*}})
-// CHECK:      ^[[BB_ENTRY]]({{.*}}):
-// CHECK:        %[[EXIT_COND:.*]] = llvm.icmp "sgt"
-// CHECK:        llvm.cond_br %[[EXIT_COND]], ^[[BB_LOOP_BODY:.*]], ^[[BB_EXIT:.*]]
-// CHECK:      ^[[BB_LOOP_BODY]]:
-// CHECK:        %[[LI_VAL:.*]] = llvm.load %[[LI_REF]] : !llvm.ptr<i32>
-// CHECK:        %[[LI_INC:.*]] = llvm.add %[[LI_VAL]], %[[ONE]]  : i32
-// CHECK:        llvm.store %[[LI_INC]], %[[LI_REF]] : !llvm.ptr<i32>
-// CHECK:        llvm.br ^[[BB_ENTRY]]({{.*}})
-// CHECK:      ^[[BB_EXIT]]:
-// CHECK:        omp.terminator
-// CHECK:      }
+// 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:  omp.reduction.declare @[[EQV_REDUCTION:.*]] : i32 init {
-// CHECK:  ^bb0(%{{.*}}: i32):
-// CHECK:    %[[TRUE:.*]] = llvm.mlir.constant(1 : i64) : i32
-// CHECK:    omp.yield(%[[TRUE]] : i32)
-// CHECK:  } combiner {
-// CHECK:  ^bb0(%[[ARG_1:.*]]: i32, %[[ARG_2:.*]]: i32):
-// CHECK:    %[[ZERO_1:.*]] = llvm.mlir.constant(0 : i64) : i32
-// CHECK:    %[[ARGVAL_1:.*]] = llvm.icmp "ne" %[[ARG_1]], %[[ZERO_1]] : i32
-// CHECK:    %[[ZERO_2:.*]] = llvm.mlir.constant(0 : i64) : i32
-// CHECK:    %[[ARGVAL_2:.*]] = llvm.icmp "ne" %[[ARG_2]], %[[ZERO_2]] : i32
-// CHECK:    %[[RES:.*]] = llvm.icmp "eq" %[[ARGVAL_1]], %[[ARGVAL_2]] : i1
-// CHECK:    %[[RES_EXT:.*]] = llvm.zext %[[RES]] : i1 to i32
-// CHECK:    omp.yield(%[[RES_EXT]] : i32)
-// CHECK:  }
-// CHECK-LABEL:  @_QPsimple_reduction
-// CHECK-SAME: %[[ARRAY_REF:.*]]: !llvm.ptr<array<100 x i32>>
-// CHECK:    %[[RED_ACCUMULATOR:.*]] = llvm.alloca %2 x i32 {bindc_name = "x", in_type = !fir.logical<4>, operand_segment_sizes = array<i32: 0, 0>, uniq_name = "_QFsimple_reductionEx"} : (i64) -> !llvm.ptr<i32>
-// CHECK:    omp.parallel   {
-// CHECK:      omp.wsloop   reduction(@[[EQV_REDUCTION]] -> %[[RED_ACCUMULATOR]] : !llvm.ptr<i32>) for
-// CHECK:        %[[ARRAY_ELEM_REF:.*]] = llvm.getelementptr %[[ARRAY_REF]][0, %{{.*}}] : (!llvm.ptr<array<100 x i32>>, i64) -> !llvm.ptr<i32>
-// CHECK:        %[[ARRAY_ELEM:.*]] = llvm.load %[[ARRAY_ELEM_REF]] : !llvm.ptr<i32>
-// CHECK:        omp.reduction %[[ARRAY_ELEM]], %[[RED_ACCUMULATOR]] : i32, !llvm.ptr<i32>
-// CHECK:        omp.yield
-// CHECK:      omp.terminator
-// CHECK:    llvm.return
-
-omp.reduction.declare @eqv_reduction : !fir.logical<4> init {
-^bb0(%arg0: !fir.logical<4>):
-  %true = arith.constant true
-  %0 = fir.convert %true : (i1) -> !fir.logical<4>
-  omp.yield(%0 : !fir.logical<4>)
-} combiner {
-^bb0(%arg0: !fir.logical<4>, %arg1: !fir.logical<4>):
-  %0 = fir.convert %arg0 : (!fir.logical<4>) -> i1
-  %1 = fir.convert %arg1 : (!fir.logical<4>) -> i1
-  %2 = arith.cmpi eq, %0, %1 : i1
-  %3 = fir.convert %2 : (i1) -> !fir.logical<4>
-  omp.yield(%3 : !fir.logical<4>)
-}
-func.func @_QPsimple_reduction(%arg0: !fir.ref<!fir.array<100x!fir.logical<4>>> {fir.bindc_name = "y"}) {
-  %0 = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFsimple_reductionEi"}
-  %1 = fir.alloca !fir.logical<4> {bindc_name = "x", uniq_name = "_QFsimple_reductionEx"}
-  %true = arith.constant true
-  %2 = fir.convert %true : (i1) -> !fir.logical<4>
-  fir.store %2 to %1 : !fir.ref<!fir.logical<4>>
-  omp.parallel   {
-    %3 = fir.alloca i32 {adapt.valuebyref, pinned}
-    %c1_i32 = arith.constant 1 : i32
-    %c100_i32 = arith.constant 100 : i32
-    %c1_i32_0 = arith.constant 1 : i32
-    omp.wsloop   reduction(@eqv_reduction -> %1 : !fir.ref<!fir.logical<4>>) for  (%arg1) : i32 = (%c1_i32) to (%c100_i32) inclusive step (%c1_i32_0) {
-      fir.store %arg1 to %3 : !fir.ref<i32>
-      %4 = fir.load %3 : !fir.ref<i32>
-      %5 = fir.convert %4 : (i32) -> i64
-      %c1_i64 = arith.constant 1 : i64
-      %6 = arith.subi %5, %c1_i64 : i64
-      %7 = fir.coordinate_of %arg0, %6 : (!fir.ref<!fir.array<100x!fir.logical<4>>>, i64) -> !fir.ref<!fir.logical<4>>
-      %8 = fir.load %7 : !fir.ref<!fir.logical<4>>
-      omp.reduction %8, %1 : !fir.logical<4>, !fir.ref<!fir.logical<4>>
-      omp.yield
-    }
-    omp.terminator
-  }
-  return
-}
+// 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(%{{.*}}) fastmath<contract> : (!fir.ref<i32>) -> ()
+ call work(i)
+ end do
+!CHECK: omp.terminator
+ !$omp end taskloop
+
+ end subroutine