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,90 @@ +! 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 - | tco | FileCheck %s + +program wsloop + integer :: i + integer :: chunk + +!CHECK: define void @_QQmain() +!CHECK: %[[CHUNK_ADDR:.*]] = alloca i32, i64 1, align 4, !dbg !11 + +!$OMP DO SCHEDULE(static, 4) + +!CHECK: omp_loop.preheader: ; preds = %entry +!CHECK: call i32 @__kmpc_global_thread_num{{.*}} +!CHECK: call void @__kmpc_for_static_init_4u(%{{.*}}, i32 %{{.*}}, i32 33, i32* %{{.*}}, i32* %{{.*}}, i32* %{{.*}}, i32 1, i32 4) +!CHECK: br label %omp_dispatch.header +!CHECK: omp_dispatch.header: ; preds = %omp_dispatch.inc, %omp_dispatch.preheader +!CHECK: %omp_dispatch.iv = phi i32 [ 0, %omp_dispatch.preheader ], [ %omp_dispatch.next, %omp_dispatch.inc ] + +do i=1, 9 + print*, i + + +!CHECK: br label %omp_dispatch.cond +!CHECK: omp_dispatch.cond: ; preds = %omp_dispatch.header +!CHECK: %omp_dispatch.cmp = icmp ult i32 %{{.*}}, %{{.*}} +!CHECK: br i1 %omp_dispatch.cmp, label %omp_dispatch.body, label %omp_dispatch.exit +!CHECK: omp_dispatch.exit: ; preds = %omp_dispatch.cond +!CHECK: call void @__kmpc_for_static_fini(%struct.ident_t* @{{.*}}, i32 %omp_global_thread_num) + +!CHECK: omp_loop.body: ; preds = %omp_loop.cond +!CHECK: %{{.*}} = add i32 %{{.*}}, %{{.*}} +!CHECK: %{{.*}} = mul i32 %{{.*}}, 1 +!CHECK: %{{.*}} = add i32 %{{.*}}, 1 +!CHECK: br label %omp.wsloop.region +!CHECK: omp.wsloop.region: ; preds = %omp_loop.body + + +end do +!$OMP END DO NOWAIT +!$OMP DO SCHEDULE(static, 2+2) + +!CHECK: omp_loop.preheader{{.*}}: ; preds = %omp_loop.after +!CHECK: call i32 @__kmpc_global_thread_num +!CHECK: call void @__kmpc_for_static_init_4u(%{{.*}}, i32 %{{.*}}, i32 33, i32* %{{.*}}, i32* %{{.*}}, i32* %{{.*}}, i32 1, i32 4) +!CHECK: br label %omp_dispatch.header{{.*}} +!CHECK: omp_dispatch.header{{.*}}: ; preds = %omp_dispatch.inc{{.*}}, %omp_dispatch.preheader{{.*}} +!CHECK: %omp_dispatch.iv{{.*}} = phi i32 [ 0, %omp_dispatch.preheader{{.*}} ], [ %omp_dispatch.next{{.*}}, %omp_dispatch.inc{{.*}} ] + +do i=1, 9 + print*, i*2 + +!CHECK: br label %omp_dispatch.cond{{.*}} +!CHECK: omp_dispatch.cond{{.*}}: ; preds = %omp_dispatch.header{{.*}} +!CHECK: %omp_dispatch.cmp{{.*}} = icmp ult i32 %{{.*}}, %{{.*}} +!CHECK: br i1 %omp_dispatch.cmp{{.*}}, label %omp_dispatch.body{{.*}}, label %omp_dispatch.exit{{.*}} +!CHECK: omp_dispatch.exit{{.*}}: ; preds = %omp_dispatch.cond{{.*}} +!CHECK: call void @__kmpc_for_static_fini(%struct.ident_t* @{{.*}}, i32 %omp_global_thread_num{{.*}}) + +end do +!$OMP END DO NOWAIT +chunk = 6 +!$OMP DO SCHEDULE(static, chunk) +!CHECK: omp_loop.after{{.*}}: ; preds = %omp_dispatch.after{{.*}} +!CHECK: store i32 6, i32* %[[CHUNK_ADDR]], align 4 +!CHECK: %[[CHUNK_VAL:.*]] = load i32, i32* %[[CHUNK_ADDR]], align 4 + +!CHECK: omp_loop.preheader{{.*}}: ; preds = %omp_loop.after +!CHECK: call i32 @__kmpc_global_thread_num +!CHECK: call void @__kmpc_for_static_init_4u(%{{.*}}, i32 %{{.*}}, i32 33, i32* %{{.*}}, i32* %{{.*}}, i32* %{{.*}}, i32 1, i32 %[[CHUNK_VAL]]) +!CHECK: br label %omp_loop.header{{.*}} +!CHECK: omp_loop.header{{.*}}: ; preds = %omp_loop.inc{{.*}}, %omp_loop.preheader{{.*}} +!CHECK: %omp_loop.iv{{.*}} = phi i32 [ 0, %omp_loop.preheader{{.*}} ], [ %omp_loop.next{{.*}}, %omp_loop.inc{{.*}} ] + +do i=1, 9 + print*, i*3 +end do +!$OMP END DO NOWAIT + +!CHECK: omp_loop.body{{.*}}: ; preds = %omp_loop.cond +!CHECK: %{{.*}} = add i32 %{{.*}}, %{{.*}} +!CHECK: %{{.*}} = mul i32 %{{.*}}, 1 +!CHECK: %{{.*}} = add i32 %{{.*}}, 1 +!CHECK: br label %omp.wsloop.region{{.*}} +!CHECK: omp.wsloop.region{{.*}}: ; preds = %omp_loop.body{{.*}} + + +end