[Flang][OpenMP] Initial lowering of the OpenMP worksharing loop

[kiranchandramohan]

Summary

The OpenMP worksharing loop operation in the dialect is a proper loop
operation and not a container of a loop. So we have to lower the
parse-tree OpenMP loop construct and the do-loop inside the construct
to a omp.wsloop operation and there should not be a fir.do_loop inside
it. This is achieved by skipping fir.do_loop creation and calling genFIR
for the nested evaluations in the lowering of the do construct.

Note: Handling of more clauses, parallel do, storage of loop index variable etc will come in separate patches.

Part of the upstreaming effort to move LLVM Flang from fir-dev branch of
https://github.com/flang-compiler/f18-llvm-project to the LLVM Project.

Co-authored-by: Sourabh Singh Tomar <[e-mail]>
Co-authored-by: Shraiysh Vaishay <[e-mail]>

[peixin]

LGTM.

Will @Leporacanthicus upstream the collapse and schedule chunk clauses? I will upstream the ordered clause and fix that bug in fir-dev.

flang/lib/Lower/Bridge.cpp

2251

Why do you remove parser::NamelistStmt?

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

4

Since we won't add integration tests under flang/test/Lower/OpenMP, we don't need to differentiate the FIRDialect and OMPDialect, and LLVMIRDialect anymore, and can use the FileCheck without the prefix for future test cases to save our time, right?

[kiranchandramohan]

Address review comments from @peixin.

[kiranchandramohan]

flang/lib/Lower/Bridge.cpp

2251

This was an accident. Thanks.

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

4

Changed to CHECK lines.

BTW, I remembered that we have the following file for FIR&OpenMP -> LLVM dialect tests. We should add a test there for all constructs. Since there is already an entry for wsloop, not adding in this case.
https://github.com/llvm/llvm-project/blob/main/flang/test/Fir/convert-to-llvm-openmp-and-fir.fir

[peixin]

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

4

Right. I have tested the execution of some combinations of clauses for worksharing-loop in fir-dev before, even though I didn't test all combinations. The worksharing-loop with some basic support should be OK. It's no hurry to add those test cases. When all upstreaming work is done, one WIP patch adding all the constructs in convert-to-llvm-openmp-and-fir.fir is fine. Maybe let's finish the upstreaming work first.

BTW, I will add the test cases for atomic read/write including memref type conversion when updating D124610.

flang/lib/Lower/Bridge.cpp

 //===-- Bridge.cpp -- bridge to lower to MLIR -----------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
void genFIR(const Fortran::parser::OpenACCDeclarativeConstruct &accDecl) {
     builder->restoreInsertionPoint(insertPt);
   }
 
   void genFIR(const Fortran::parser::OpenMPConstruct &omp) {
     mlir::OpBuilder::InsertPoint insertPt = builder->saveInsertionPoint();
     localSymbols.pushScope();
     Fortran::lower::genOpenMPConstruct(*this, getEval(), omp);
 
-    for (Fortran::lower::pft::Evaluation &e : getEval().getNestedEvaluations())
+    // If loop is part of an OpenMP Construct then the OpenMP dialect
+    // workshare loop operation has already been created. Only the
+    // body needs to be created here and the do_loop can be skipped.
+    Fortran::lower::pft::Evaluation *curEval =
+        std::get_if<Fortran::parser::OpenMPLoopConstruct>(&omp.u)
+            ? &getEval().getFirstNestedEvaluation()
+            : &getEval();
+
+    for (Fortran::lower::pft::Evaluation &e : curEval->getNestedEvaluations())
       genFIR(e);
     localSymbols.popScope();
     builder->restoreInsertionPoint(insertPt);
   }
 
   void genFIR(const Fortran::parser::OpenMPDeclarativeConstruct &ompDecl) {
     mlir::OpBuilder::InsertPoint insertPt = builder->saveInsertionPoint();
     genOpenMPDeclarativeConstruct(*this, getEval(), ompDecl);
#endif
   void genFIR(const Fortran::parser::EndIfStmt &) {}           // nop
   void genFIR(const Fortran::parser::EndMpSubprogramStmt &) {} // nop
   void genFIR(const Fortran::parser::EndSelectStmt &) {}       // nop
   void genFIR(const Fortran::parser::EndSubroutineStmt &) {}   // nop
   void genFIR(const Fortran::parser::EntryStmt &) {}           // nop
   void genFIR(const Fortran::parser::IfStmt &) {}              // nop
   void genFIR(const Fortran::parser::IfThenStmt &) {}          // nop
   void genFIR(const Fortran::parser::NonLabelDoStmt &) {}      // nop
-
-  void genFIR(const Fortran::parser::OmpEndLoopDirective &) {
-    TODO(toLocation(), "OmpEndLoopDirective lowering");
-  }
+  void genFIR(const Fortran::parser::OmpEndLoopDirective &) {} // nop
 
   void genFIR(const Fortran::parser::NamelistStmt &) {
     TODO(toLocation(), "NamelistStmt lowering");
   }
 
   /// Generate FIR for the Evaluation `eval`.
   void genFIR(Fortran::lower::pft::Evaluation &eval,
               bool unstructuredContext = true) {
     if (unstructuredContext) {
       // When transitioning from unstructured to structured code,

flang/lib/Lower/OpenMP.cpp

 //===-- OpenMP.cpp -- Open MP directive lowering --------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
static void genObjectList(const Fortran::parser::OmpObjectList &objectList,
   }
 }
 
 template <typename Op>
 static void
 createBodyOfOp(Op &op, Fortran::lower::AbstractConverter &converter,
                mlir::Location &loc,
                const Fortran::parser::OmpClauseList *clauses = nullptr,
+               const Fortran::semantics::Symbol *arg = nullptr,
                bool outerCombined = false) {
   fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
+  // If an argument for the region is provided then create the block with that
+  // argument. Also update the symbol's address with the mlir argument value.
+  // e.g. For loops the argument is the induction variable. And all further
+  // uses of the induction variable should use this mlir value.
+  if (arg) {
+    firOpBuilder.createBlock(&op.getRegion(), {}, {converter.genType(*arg)},
+                             {loc});
+    converter.bindSymbol(*arg, op.getRegion().front().getArgument(0));
+  } else {
     firOpBuilder.createBlock(&op.getRegion());
+  }
   auto &block = op.getRegion().back();
   firOpBuilder.setInsertionPointToStart(&block);
-  // Ensure the block is well-formed.
+
+  // Insert the terminator.
+  if constexpr (std::is_same_v<Op, omp::WsLoopOp>) {
+    mlir::ValueRange results;
+    firOpBuilder.create<mlir::omp::YieldOp>(loc, results);
+  } else {
     firOpBuilder.create<mlir::omp::TerminatorOp>(loc);
+  }
+
   // Reset the insertion point to the start of the first block.
   firOpBuilder.setInsertionPointToStart(&block);
   // Handle privatization. Do not privatize if this is the outer operation.
   if (clauses && !outerCombined)
     privatizeVars(converter, *clauses);
 }
 
 static void genOMP(Fortran::lower::AbstractConverter &converter,
genOMP(Fortran::lower::AbstractConverter &converter,
 
   if (blockDirective.v == llvm::omp::OMPD_parallel) {
     // Create and insert the operation.
     auto parallelOp = firOpBuilder.create<mlir::omp::ParallelOp>(
         currentLocation, argTy, ifClauseOperand, numThreadsClauseOperand,
         allocateOperands, allocatorOperands, /*reduction_vars=*/ValueRange(),
         /*reductions=*/nullptr, procBindKindAttr);
     createBodyOfOp<omp::ParallelOp>(parallelOp, converter, currentLocation,
-                                    &opClauseList, /*isCombined=*/false);
+                                    &opClauseList);
   } else if (blockDirective.v == llvm::omp::OMPD_master) {
     auto masterOp =
         firOpBuilder.create<mlir::omp::MasterOp>(currentLocation, argTy);
     createBodyOfOp<omp::MasterOp>(masterOp, converter, currentLocation);
   } else if (blockDirective.v == llvm::omp::OMPD_single) {
     auto singleOp = firOpBuilder.create<mlir::omp::SingleOp>(
         currentLocation, allocateOperands, allocatorOperands, nowaitAttr);
     createBodyOfOp<omp::SingleOp>(singleOp, converter, currentLocation);
   } else if (blockDirective.v == llvm::omp::OMPD_ordered) {
     auto orderedOp = firOpBuilder.create<mlir::omp::OrderedRegionOp>(
         currentLocation, /*simd=*/nullptr);
     createBodyOfOp<omp::OrderedRegionOp>(orderedOp, converter, currentLocation);
   } else {
     TODO(converter.getCurrentLocation(), "Unhandled block directive");
   }
 }
 
+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;
+  mlir::Value scheduleChunkClauseOperand;
+  mlir::Attribute scheduleClauseOperand, collapseClauseOperand,
+      noWaitClauseOperand, orderedClauseOperand, orderClauseOperand;
+  const auto &wsLoopOpClauseList = std::get<Fortran::parser::OmpClauseList>(
+      std::get<Fortran::parser::OmpBeginLoopDirective>(loopConstruct.t).t);
+  if (llvm::omp::OMPD_do !=
+      std::get<Fortran::parser::OmpLoopDirective>(
+          std::get<Fortran::parser::OmpBeginLoopDirective>(loopConstruct.t).t)
+          .v) {
+    TODO(converter.getCurrentLocation(), "Combined worksharing loop construct");
+  }
+
+  Fortran::lower::pft::Evaluation *doConstructEval =
+      &eval.getFirstNestedEvaluation();
+
+  Fortran::lower::pft::Evaluation *doLoop =
+      &doConstructEval->getFirstNestedEvaluation();
+  auto *doStmt = doLoop->getIf<Fortran::parser::NonLabelDoStmt>();
+  assert(doStmt && "Expected do loop to be in the nested evaluation");
+  const auto &loopControl =
+      std::get<std::optional<Fortran::parser::LoopControl>>(doStmt->t);
+  const Fortran::parser::LoopControl::Bounds *bounds =
+      std::get_if<Fortran::parser::LoopControl::Bounds>(&loopControl->u);
+  assert(bounds && "Expected bounds for worksharing do loop");
+  Fortran::semantics::Symbol *iv = nullptr;
+  Fortran::lower::StatementContext stmtCtx;
+  lowerBound.push_back(fir::getBase(converter.genExprValue(
+      *Fortran::semantics::GetExpr(bounds->lower), stmtCtx)));
+  upperBound.push_back(fir::getBase(converter.genExprValue(
+      *Fortran::semantics::GetExpr(bounds->upper), stmtCtx)));
+  if (bounds->step) {
+    step.push_back(fir::getBase(converter.genExprValue(
+        *Fortran::semantics::GetExpr(bounds->step), stmtCtx)));
+  } else { // If `step` is not present, assume it as `1`.
+    step.push_back(firOpBuilder.createIntegerConstant(
+        currentLocation, firOpBuilder.getIntegerType(32), 1));
+  }
+  iv = bounds->name.thing.symbol;
+
+  // FIXME: Add support for following clauses:
+  // 1. linear
+  // 2. order
+  // 3. collapse
+  // 4. 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>(),
+      orderedClauseOperand.dyn_cast_or_null<IntegerAttr>(),
+      orderClauseOperand.dyn_cast_or_null<omp::ClauseOrderKindAttr>(),
+      /*inclusive=*/firOpBuilder.getUnitAttr());
+
+  // Handle attribute based clauses.
+  for (const Fortran::parser::OmpClause &clause : wsLoopOpClauseList.v) {
+    if (const auto &scheduleClause =
+            std::get_if<Fortran::parser::OmpClause::Schedule>(&clause.u)) {
+      mlir::MLIRContext *context = firOpBuilder.getContext();
+      const auto &scheduleType = scheduleClause->v;
+      const auto &scheduleKind =
+          std::get<Fortran::parser::OmpScheduleClause::ScheduleType>(
+              scheduleType.t);
+      switch (scheduleKind) {
+      case Fortran::parser::OmpScheduleClause::ScheduleType::Static:
+        wsLoopOp.schedule_valAttr(omp::ClauseScheduleKindAttr::get(
+            context, omp::ClauseScheduleKind::Static));
+        break;
+      case Fortran::parser::OmpScheduleClause::ScheduleType::Dynamic:
+        wsLoopOp.schedule_valAttr(omp::ClauseScheduleKindAttr::get(
+            context, omp::ClauseScheduleKind::Dynamic));
+        break;
+      case Fortran::parser::OmpScheduleClause::ScheduleType::Guided:
+        wsLoopOp.schedule_valAttr(omp::ClauseScheduleKindAttr::get(
+            context, omp::ClauseScheduleKind::Guided));
+        break;
+      case Fortran::parser::OmpScheduleClause::ScheduleType::Auto:
+        wsLoopOp.schedule_valAttr(omp::ClauseScheduleKindAttr::get(
+            context, omp::ClauseScheduleKind::Auto));
+        break;
+      case Fortran::parser::OmpScheduleClause::ScheduleType::Runtime:
+        wsLoopOp.schedule_valAttr(omp::ClauseScheduleKindAttr::get(
+            context, omp::ClauseScheduleKind::Runtime));
+        break;
+      }
+    }
+  }
+  // In FORTRAN `nowait` clause occur at the end of `omp do` directive.
+  // i.e
+  // !$omp do
+  // <...>
+  // !$omp end do nowait
+  if (const auto &endClauseList =
+          std::get<std::optional<Fortran::parser::OmpEndLoopDirective>>(
+              loopConstruct.t)) {
+    const auto &clauseList =
+        std::get<Fortran::parser::OmpClauseList>((*endClauseList).t);
+    for (const Fortran::parser::OmpClause &clause : clauseList.v)
+      if (std::get_if<Fortran::parser::OmpClause::Nowait>(&clause.u))
+        wsLoopOp.nowaitAttr(firOpBuilder.getUnitAttr());
+  }
+
+  createBodyOfOp<omp::WsLoopOp>(wsLoopOp, converter, currentLocation,
+                                &wsLoopOpClauseList, iv);
+}
+
 static void
 genOMP(Fortran::lower::AbstractConverter &converter,
        Fortran::lower::pft::Evaluation &eval,
        const Fortran::parser::OpenMPCriticalConstruct &criticalConstruct) {
   fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
   mlir::Location currentLocation = converter.getCurrentLocation();
   std::string name;
   const Fortran::parser::OmpCriticalDirective &cd =
std::visit(
           [&](const Fortran::parser::OpenMPSectionsConstruct
                   &sectionsConstruct) {
             genOMP(converter, eval, sectionsConstruct);
           },
           [&](const Fortran::parser::OpenMPSectionConstruct &sectionConstruct) {
             genOMP(converter, eval, sectionConstruct);
           },
           [&](const Fortran::parser::OpenMPLoopConstruct &loopConstruct) {
-            TODO(converter.getCurrentLocation(), "OpenMPLoopConstruct");
+            genOMP(converter, eval, loopConstruct);
           },
           [&](const Fortran::parser::OpenMPDeclarativeAllocate
                   &execAllocConstruct) {
             TODO(converter.getCurrentLocation(), "OpenMPDeclarativeAllocate");
           },
           [&](const Fortran::parser::OpenMPExecutableAllocate
                   &execAllocConstruct) {
             TODO(converter.getCurrentLocation(), "OpenMPExecutableAllocate");

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

+! This test checks lowering of OpenMP DO Directive (Worksharing).
+
+! RUN: bbc -fopenmp -emit-fir %s -o - | FileCheck %s
+
+!CHECK-LABEL: func @_QPsimple_loop()
+subroutine simple_loop
+  integer :: i
+  ! CHECK:  omp.parallel
+  !$OMP PARALLEL
+  ! CHECK:     %[[WS_LB:.*]] = arith.constant 1 : i32
+  ! CHECK:     %[[WS_UB:.*]] = arith.constant 9 : i32
+  ! CHECK:     %[[WS_STEP:.*]] = arith.constant 1 : i32
+  ! CHECK:     omp.wsloop for (%[[I:.*]]) : i32 = (%[[WS_LB]]) to (%[[WS_UB]]) inclusive step (%[[WS_STEP]])
+  !$OMP DO
+  do i=1, 9
+  ! CHECK:    fir.call @_FortranAioOutputInteger32({{.*}}, %[[I]]) : (!fir.ref<i8>, i32) -> i1
+    print*, i
+  end do
+  ! CHECK:       omp.yield
+  !$OMP END DO
+  ! CHECK:       omp.terminator
+  !$OMP END PARALLEL
+end subroutine
+
+!CHECK-LABEL: func @_QPsimple_loop_with_step()
+subroutine simple_loop_with_step
+  integer :: i
+  ! CHECK:  omp.parallel
+  !$OMP PARALLEL
+  ! CHECK:     %[[WS_LB:.*]] = arith.constant 1 : i32
+  ! CHECK:     %[[WS_UB:.*]] = arith.constant 9 : i32
+  ! CHECK:     %[[WS_STEP:.*]] = arith.constant 2 : i32
+  ! CHECK:     omp.wsloop for (%[[I:.*]]) : i32 = (%[[WS_LB]]) to (%[[WS_UB]]) inclusive step (%[[WS_STEP]])
+  !$OMP DO
+  do i=1, 9, 2
+  ! CHECK:    fir.call @_FortranAioOutputInteger32({{.*}}, %[[I]]) : (!fir.ref<i8>, i32) -> i1
+    print*, i
+  end do
+  ! CHECK:       omp.yield
+  !$OMP END DO
+  ! CHECK:       omp.terminator
+  !$OMP END PARALLEL
+end subroutine
+
+!CHECK-LABEL: func @_QPloop_with_schedule_nowait()
+subroutine loop_with_schedule_nowait
+  integer :: i
+  ! CHECK:  omp.parallel
+  !$OMP PARALLEL
+  ! CHECK:     %[[WS_LB:.*]] = arith.constant 1 : i32
+  ! CHECK:     %[[WS_UB:.*]] = arith.constant 9 : i32
+  ! CHECK:     %[[WS_STEP:.*]] = arith.constant 1 : i32
+  ! CHECK:     omp.wsloop schedule(runtime) nowait for (%[[I:.*]]) : i32 = (%[[WS_LB]]) to (%[[WS_UB]]) inclusive step (%[[WS_STEP]])
+  !$OMP DO SCHEDULE(runtime)
+  do i=1, 9
+  ! CHECK:    fir.call @_FortranAioOutputInteger32({{.*}}, %[[I]]) : (!fir.ref<i8>, i32) -> i1
+    print*, i
+  end do
+  ! CHECK:       omp.yield
+  !$OMP END DO NOWAIT
+  ! CHECK:       omp.terminator
+  !$OMP END PARALLEL
+end subroutine