Merge pull request #969 from Leporacanthicus/collapse

[flang][OpenMP] Add support for collapse clause

Author: schweitzpgi

flang/include/flang/Lower/OpenMP.h

@@ -13,10 +13,13 @@
 #ifndef FORTRAN_LOWER_OPENMP_H
 #define FORTRAN_LOWER_OPENMP_H
 
+#include <cinttypes>
+
 namespace Fortran {
 namespace parser {
 struct OpenMPConstruct;
 struct OmpEndLoopDirective;
+struct OmpClauseList;
 } // namespace parser
 
 namespace lower {
@@ -29,6 +32,9 @@ struct Evaluation;
 
 void genOpenMPConstruct(AbstractConverter &, pft::Evaluation &,
                         const parser::OpenMPConstruct &);
+
+int64_t getCollapseValue(const Fortran::parser::OmpClauseList &clauseList);
+
 } // namespace lower
 } // namespace Fortran
 

flang/lib/Lower/Bridge.cpp

@@ -1347,13 +1347,28 @@ class FirConverter : public Fortran::lower::AbstractConverter {
     auto insertPt = builder->saveInsertionPoint();
     localSymbols.pushScope();
     genOpenMPConstruct(*this, getEval(), omp);
+
+    auto ompLoop = std::get_if<Fortran::parser::OpenMPLoopConstruct>(&omp.u);
+
     // 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();
+    // Skip the number of collapsed loops, which is 1 when there is a
+    // no collapse requested.
+
+    Fortran::lower::pft::Evaluation *curEval = &getEval();
+    if (ompLoop) {
+      const auto &wsLoopOpClauseList = std::get<Fortran::parser::OmpClauseList>(
+          std::get<Fortran::parser::OmpBeginLoopDirective>(ompLoop->t).t);
+      int64_t collapseValue =
+          Fortran::lower::getCollapseValue(wsLoopOpClauseList);
+
+      curEval = &curEval->getFirstNestedEvaluation();
+      for (auto i = 1; i < collapseValue; i++) {
+        curEval = &*std::next(curEval->getNestedEvaluations().begin());
+      }
+    }
+
     for (auto &e : curEval->getNestedEvaluations())
       genFIR(e);
     localSymbols.popScope();

flang/lib/Lower/OpenMP.cpp

@@ -90,23 +90,29 @@ 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) {
+static void createBodyOfOp(
+    Op &op, Fortran::lower::AbstractConverter &converter, mlir::Location &loc,
+    const Fortran::parser::OmpClauseList *clauses = nullptr,
+    const SmallVector<const Fortran::semantics::Symbol *> &args = {}) {
   auto &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)});
-    fir::ExtendedValue exval = op.getRegion().front().getArgument(0);
-    [[maybe_unused]] bool success = converter.bindSymbol(*arg, exval);
-    assert(
-        success &&
-        "Existing binding prevents setting MLIR value for the index variable");
+  if (args.size()) {
+    SmallVector<Type> tiv;
+    int argIndex = 0;
+    for (auto &arg : args) {
+      tiv.push_back(converter.genType(*arg));
+    }
+    firOpBuilder.createBlock(&op.getRegion(), {}, tiv);
+    for (auto &arg : args) {
+      fir::ExtendedValue exval = op.getRegion().front().getArgument(argIndex);
+      [[maybe_unused]] bool success = converter.bindSymbol(*arg, exval);
+      assert(success && "Existing binding prevents setting MLIR value for the "
+                        "index variable");
+      argIndex++;
+    }
   } else {
     firOpBuilder.createBlock(&op.getRegion());
   }
@@ -397,6 +403,18 @@ getSIMDModifier(const Fortran::parser::OmpScheduleClause &x) {
   return mlir::omp::ScheduleModifier::none;
 }
 
+int64_t Fortran::lower::getCollapseValue(
+    const Fortran::parser::OmpClauseList &clauseList) {
+  for (const auto &clause : clauseList.v) {
+    if (const auto &collapseClause =
+            std::get_if<Fortran::parser::OmpClause::Collapse>(&clause.u)) {
+      const auto *expr = Fortran::semantics::GetExpr(collapseClause->v);
+      return Fortran::evaluate::ToInt64(*expr).value();
+    }
+  }
+  return 1;
+}
+
 static void genOMP(Fortran::lower::AbstractConverter &converter,
                    Fortran::lower::pft::Evaluation &eval,
                    const Fortran::parser::OpenMPLoopConstruct &loopConstruct) {
@@ -437,34 +455,8 @@ static void genOMP(Fortran::lower::AbstractConverter &converter,
       genObjectList(ompObjectList, converter, lastPrivateClauseOperands);
     }
   }
-  // FIXME: Can be done in a better way ?
-  auto &doConstructEval =
-      eval.getFirstNestedEvaluation().getFirstNestedEvaluation();
-  auto *doStmt = doConstructEval.getIf<Fortran::parser::NonLabelDoStmt>();
-
-  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);
-  Fortran::semantics::Symbol *iv = nullptr;
-  if (bounds) {
-    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)));
-    }
-    // If `step` is not present, assume it as `1`.
-    else {
-      step.push_back(firOpBuilder.createIntegerConstant(
-          currentLocation, firOpBuilder.getIntegerType(32), 1));
-    }
-    iv = bounds->name.thing.symbol;
-  }
 
+  int64_t collapseValue = Fortran::lower::getCollapseValue(wsLoopOpClauseList);
   for (const auto &clause : wsLoopOpClauseList.v) {
     if (const auto &scheduleClause =
             std::get_if<Fortran::parser::OmpClause::Schedule>(&clause.u)) {
@@ -480,6 +472,41 @@ static void genOMP(Fortran::lower::AbstractConverter &converter,
     }
   }
 
+  // Collect the loops to collapse.
+  auto *doConstructEval = &eval.getFirstNestedEvaluation();
+
+  SmallVector<const Fortran::semantics::Symbol *> iv;
+  do {
+    auto *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);
+    if (bounds) {
+      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)));
+      }
+      // If `step` is not present, assume it as `1`.
+      else {
+        step.push_back(firOpBuilder.createIntegerConstant(
+            currentLocation, firOpBuilder.getIntegerType(32), 1));
+      }
+      iv.push_back(bounds->name.thing.symbol);
+    }
+
+    collapseValue--;
+    doConstructEval =
+        &*std::next(doConstructEval->getNestedEvaluations().begin());
+  } while (collapseValue > 0);
+
   // FIXME: Add support for following clauses:
   // 1. linear
   // 2. order