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

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.

Reviewed By: peixin, shraiysh

Differential Revision: https://reviews.llvm.org/D126383

flang/lib/Lower/OpenMP.cpp

@@ -577,6 +577,21 @@ static void genOMP(Fortran::lower::AbstractConverter &converter,
         &*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);
@@ -592,7 +607,6 @@ static void genOMP(Fortran::lower::AbstractConverter &converter,
   // 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,

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

@@ -0,0 +1,69 @@
+! 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"}
+
+!$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:           {{.*}} = 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_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_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