diff --git a/flang/lib/Lower/OpenMP.cpp b/flang/lib/Lower/OpenMP.cpp --- a/flang/lib/Lower/OpenMP.cpp +++ b/flang/lib/Lower/OpenMP.cpp @@ -577,6 +577,21 @@ &*std::next(doConstructEval->getNestedEvaluations().begin()); } while (collapseValue > 0); + for (const auto &clause : wsLoopOpClauseList.v) { + if (const auto &scheduleClause = + std::get_if(&clause.u)) { + if (const auto &chunkExpr = + std::get>( + 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 @@ // FIXME: Add support for following clauses: // 1. linear // 2. order - // 3. schedule (with chunk) auto wsLoopOp = firOpBuilder.create( currentLocation, lowerBound, upperBound, step, linearVars, linearStepVars, reductionVars, /*reductions=*/nullptr, diff --git a/flang/test/Lower/OpenMP/omp-wsloop-chunks.f90 b/flang/test/Lower/OpenMP/omp-wsloop-chunks.f90 new file mode 100644 --- /dev/null +++ b/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, 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, 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 +! 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 +! 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, i32) -> i1 +! CHECK: omp.yield +! CHECK: } +! CHECK: return +! CHECK: } + +end