[OpenMP] Pass chunk-size to MLIR while lowering from parse-tree

[Leporacanthicus]

Summary

Test that chunk size is passed to the static init function.
Using three different variations:

1. Single constant.
2. Expression with constants.
3. Variable value.

[shraiysh]

The code LGTM. The test should be from F90 to MLIR (FIR + OpenMP Dialect) only. Lowering to LLVM IR from FIR and OpenMP Dialect is tested separately.

[Leporacanthicus]

Update test to MLIR

[peixin]

LGTM

[shraiysh]

Yup this looks okay to me now. Thank you!

[Leporacanthicus]

Removed some unnecessary test-lines [testing for variables used nowhere else
and constants used to set up the I/O, where the call itself is not tested]

flang/lib/Lower/OpenMP.cpp

do {
     loopVarTypeSize = std::max(loopVarTypeSize,
                                bounds->name.thing.symbol->GetUltimate().size());
 
     collapseValue--;
     doConstructEval =
         &*std::next(doConstructEval->getNestedEvaluations().begin());
   } while (collapseValue > 0);
 
+  for (const auto &clause : wsLoopOpClauseList.v) {
+    if (const auto &scheduleClause =
+            std::get_if<Fortran::parser::OmpClause::Schedule>(&clause.u)) {
+      if (const auto &chunkExpr =
+              std::get<std::optional<Fortran::parser::ScalarIntExpr>>(
+                  scheduleClause->v.t)) {
+        if (const auto *expr = Fortran::semantics::GetExpr(*chunkExpr)) {
+          Fortran::lower::StatementContext stmtCtx;
+          scheduleChunkClauseOperand =
+              fir::getBase(converter.genExprValue(*expr, 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]);
     upperBound[it] = firOpBuilder.createConvert(currentLocation, loopVarType,
                                                 upperBound[it]);
     step[it] =
         firOpBuilder.createConvert(currentLocation, loopVarType, step[it]);
   }
 
   // FIXME: Add support for following clauses:
   // 1. linear
   // 2. order
-  // 3. schedule (with chunk)
   auto wsLoopOp = firOpBuilder.create<mlir::omp::WsLoopOp>(
       currentLocation, lowerBound, upperBound, step, linearVars, linearStepVars,
       reductionVars, /*reductions=*/nullptr,
       scheduleClauseOperand.dyn_cast_or_null<omp::ClauseScheduleKindAttr>(),
       scheduleChunkClauseOperand, /*schedule_modifiers=*/nullptr,
       /*simd_modifier=*/nullptr,
       collapseClauseOperand.dyn_cast_or_null<IntegerAttr>(),
       noWaitClauseOperand.dyn_cast_or_null<UnitAttr>(),

flang/test/Lower/OpenMP/omp-wsloop-chunks.f90

+! This test checks that chunk size is passed correctly when lowering of
+! OpenMP DO Directive(Worksharing) with chunk size
+
+! RUN: bbc -fopenmp -emit-fir %s -o - | FileCheck %s
+
+program wsloop
+        integer :: i
+        integer :: chunk
+
+! CHECK-LABEL: func.func @_QQmain() {
+! CHECK:         %[[VAL_0:.*]] = fir.alloca i32 {bindc_name = "chunk", uniq_name = "_QFEchunk"}
+! CHECK:         %[[VAL_1:.*]] = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFEi"}
+
+!$OMP DO SCHEDULE(static, 4)
+
+do i=1, 9
+  print*, i
+
+! CHECK:         %[[VAL_2:.*]] = arith.constant 1 : i32
+! CHECK:         %[[VAL_3:.*]] = arith.constant 9 : i32
+! CHECK:         %[[VAL_4:.*]] = arith.constant 1 : i32
+! CHECK:         %[[VAL_5:.*]] = arith.constant 4 : i32
+! CHECK:         omp.wsloop   schedule(static = %[[VAL_5]] : i32) nowait for  (%[[VAL_6:.*]]) : i32 = (%[[VAL_2]]) to (%[[VAL_3]]) inclusive step (%[[VAL_4]]) {
+! CHECK:           %[[VAL_7:.*]] = arith.constant -1 : i32
+! CHECK:           {{.*}} = fir.call @_FortranAioOutputInteger32({{.*}}, %[[VAL_6]]) : (!fir.ref<i8>, i32) -> i1
+! CHECK:           omp.yield
+! CHECK:         }
+
+end do
+!$OMP END DO NOWAIT
+!$OMP DO SCHEDULE(static, 2+2)
+
+do i=1, 9
+  print*, i*2
+
+! CHECK:         %[[VAL_14:.*]] = arith.constant 1 : i32
+! CHECK:         %[[VAL_15:.*]] = arith.constant 9 : i32
+! CHECK:         %[[VAL_16:.*]] = arith.constant 1 : i32
+! CHECK:         %[[VAL_17:.*]] = arith.constant 4 : i32
+! CHECK:         omp.wsloop   schedule(static = %[[VAL_17]] : i32) nowait for  (%[[VAL_18:.*]]) : i32 = (%[[VAL_14]]) to (%[[VAL_15]]) inclusive step (%[[VAL_16]]) {
+! CHECK:           %[[VAL_19:.*]] = arith.constant -1 : i32
+! CHECK:           %[[VAL_24:.*]] = arith.constant 2 : i32
+! CHECK:           %[[VAL_25:.*]] = arith.muli %[[VAL_24]], %[[VAL_18]] : i32
+! CHECK:           {{.*}} = fir.call @_FortranAioOutputInteger32({{.*}}, %[[VAL_25]]) : (!fir.ref<i8>, i32) -> i1
+! CHECK:           omp.yield
+! CHECK:         }
+  
+end do
+!$OMP END DO NOWAIT
+chunk = 6
+!$OMP DO SCHEDULE(static, chunk)
+
+do i=1, 9
+   print*, i*3
+end do
+!$OMP END DO NOWAIT
+! CHECK:         %[[VAL_28:.*]] = arith.constant 6 : i32
+! CHECK:         fir.store %[[VAL_28]] to %[[VAL_0]] : !fir.ref<i32>
+! CHECK:         %[[VAL_29:.*]] = arith.constant 1 : i32
+! CHECK:         %[[VAL_30:.*]] = arith.constant 9 : i32
+! CHECK:         %[[VAL_31:.*]] = arith.constant 1 : i32
+! CHECK:         %[[VAL_32:.*]] = fir.load %[[VAL_0]] : !fir.ref<i32>
+! CHECK:         omp.wsloop   schedule(static = %[[VAL_32]] : i32) nowait for  (%[[VAL_33:.*]]) : i32 = (%[[VAL_29]]) to (%[[VAL_30]]) inclusive step (%[[VAL_31]]) {
+! CHECK:           %[[VAL_34:.*]] = arith.constant -1 : i32
+! CHECK:           %[[VAL_39:.*]] = arith.constant 3 : i32
+! CHECK:           %[[VAL_40:.*]] = arith.muli %[[VAL_39]], %[[VAL_33]] : i32
+! CHECK:           {{.*}} = fir.call @_FortranAioOutputInteger32({{.*}}, %[[VAL_40]]) : (!fir.ref<i8>, i32) -> i1
+! CHECK:           omp.yield
+! CHECK:         }
+! CHECK:         return
+! CHECK:       }
+
+end