[Flang][Lower] Add lowering support of OpenMP distribute to MLIR

This patch adds support for lowering the OpenMP distribute directive from PFT
to MLIR. This in turn unlocks support for several related combined loop
constructs as well.

Author: skatrak

flang/lib/Lower/OpenMP.cpp

@@ -565,6 +565,9 @@ class ClauseProcessor {
   bool processDevice(Fortran::lower::StatementContext &stmtCtx,
                      mlir::Value &result) const;
   bool processDeviceType(mlir::omp::DeclareTargetDeviceType &result) const;
+  bool processDistSchedule(Fortran::lower::StatementContext &stmtCtx,
+                           mlir::UnitAttr &scheduleStatic,
+                           mlir::Value &chunkSize) const;
   bool processFinal(Fortran::lower::StatementContext &stmtCtx,
                     mlir::Value &result) const;
   bool processHint(mlir::IntegerAttr &result) const;
@@ -1490,6 +1493,19 @@ bool ClauseProcessor::processDeviceType(
   return false;
 }
 
+bool ClauseProcessor::processDistSchedule(
+    Fortran::lower::StatementContext &stmtCtx, mlir::UnitAttr &scheduleStatic,
+    mlir::Value &chunkSize) const {
+  if (auto *distScheduleClause = findUniqueClause<ClauseTy::DistSchedule>()) {
+    scheduleStatic = converter.getFirOpBuilder().getUnitAttr();
+    if (const auto *expr = Fortran::semantics::GetExpr(distScheduleClause->v)) {
+      chunkSize = fir::getBase(converter.genExprValue(*expr, stmtCtx));
+    }
+    return true;
+  }
+  return false;
+}
+
 bool ClauseProcessor::processFinal(Fortran::lower::StatementContext &stmtCtx,
                                    mlir::Value &result) const {
   const Fortran::parser::CharBlock *source = nullptr;
@@ -2997,6 +3013,27 @@ genTeamsOp(Fortran::lower::AbstractConverter &converter,
                                  reductionDeclSymbols));
 }
 
+static mlir::omp::DistributeOp
+genDistributeOp(Fortran::lower::AbstractConverter &converter,
+                Fortran::lower::pft::Evaluation &eval, bool genNested,
+                mlir::Location currentLocation,
+                const Fortran::parser::OmpClauseList &clauseList,
+                bool outerCombined = false) {
+  Fortran::lower::StatementContext stmtCtx;
+  mlir::UnitAttr scheduleStatic;
+  mlir::Value chunkSize;
+  llvm::SmallVector<mlir::Value> allocateOperands, allocatorOperands;
+
+  ClauseProcessor cp(converter, clauseList);
+  cp.processDistSchedule(stmtCtx, scheduleStatic, chunkSize);
+  cp.processAllocate(allocatorOperands, allocateOperands);
+
+  return genOpWithBody<mlir::omp::DistributeOp>(
+      converter, eval, genNested, currentLocation, outerCombined, &clauseList,
+      scheduleStatic, chunkSize, allocateOperands, allocatorOperands,
+      /*order_val=*/nullptr);
+}
+
 /// Extract the list of function and variable symbols affected by the given
 /// 'declare target' directive and return the intended device type for them.
 static mlir::omp::DeclareTargetDeviceType getDeclareTargetInfo(
@@ -3358,7 +3395,9 @@ static void genOMP(Fortran::lower::AbstractConverter &converter,
     }
     if (llvm::omp::allDistributeSet.test(ompDirective)) {
       validDirective = true;
-      TODO(currentLocation, "Distribute construct");
+      bool outerCombined = llvm::omp::topDistributeSet.test(ompDirective);
+      genDistributeOp(converter, eval, /*genNested=*/false, currentLocation,
+                      loopOpClauseList, outerCombined);
     }
     if ((llvm::omp::allParallelSet & llvm::omp::loopConstructSet)
             .test(ompDirective)) {