[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.
llvm · Sep 29, 2023 · 0b898b9 · 0b898b9
1 parent d424631
commit 0b898b9
Show file tree

Hide file tree

Showing 4 changed files with 1,169 additions and 53 deletions.
diff --git a/flang/lib/Lower/OpenMP.cpp b/flang/lib/Lower/OpenMP.cpp
@@ -497,6 +497,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;
@@ -1335,6 +1338,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;
@@ -2473,6 +2489,27 @@ genTeamsOp(Fortran::lower::AbstractConverter &converter,
                                  reductionDeclSymbols));
 }
 
+static mlir::omp::DistributeOp
+genDistributeOp(Fortran::lower::AbstractConverter &converter,
+                Fortran::lower::pft::Evaluation &eval,
+                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, 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(
@@ -2681,7 +2718,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, currentLocation, loopOpClauseList,
+                      outerCombined);
     }
     if ((llvm::omp::allParallelSet & llvm::omp::loopConstructSet)
             .test(ompDirective)) {