[flang][OpenMP] Enable tiling

flang/include/flang/Lower/OpenMP.h

@@ -79,7 +79,6 @@ void genOpenMPDeclarativeConstruct(AbstractConverter &,
 void genOpenMPSymbolProperties(AbstractConverter &converter,
                                const pft::Variable &var);
 
-int64_t getCollapseValue(const Fortran::parser::OmpClauseList &clauseList);
 void genThreadprivateOp(AbstractConverter &, const pft::Variable &);
 void genDeclareTargetIntGlobal(AbstractConverter &, const pft::Variable &);
 bool isOpenMPTargetConstruct(const parser::OpenMPConstruct &);

flang/include/flang/Parser/parse-tree.h

@@ -5024,8 +5024,10 @@ struct OpenMPBlockConstruct {
 struct OpenMPLoopConstruct {
   TUPLE_CLASS_BOILERPLATE(OpenMPLoopConstruct);
   OpenMPLoopConstruct(OmpBeginLoopDirective &&a)
-      : t({std::move(a), std::nullopt, std::nullopt}) {}
+      : t({std::move(a), std::nullopt, std::nullopt, std::nullopt}) {}
   std::tuple<OmpBeginLoopDirective, std::optional<DoConstruct>,
+      // Inner loop construct used to handle tiling for now.
+      std::optional<common::Indirection<OpenMPLoopConstruct>>,
       std::optional<OmpEndLoopDirective>>
       t;
 };

flang/lib/Lower/OpenMP/OpenMP.cpp

@@ -45,6 +45,7 @@
 
 using namespace Fortran::lower::omp;
 using namespace Fortran::common::openmp;
+using namespace Fortran::semantics;
 
 static llvm::cl::opt<bool> DumpAtomicAnalysis("fdebug-dump-atomic-analysis");
 
@@ -456,6 +457,7 @@ static void processHostEvalClauses(lower::AbstractConverter &converter,
       return;
 
     const parser::OmpClauseList *beginClauseList = nullptr;
+    const parser::OmpClauseList *middleClauseList = nullptr;
     const parser::OmpClauseList *endClauseList = nullptr;
     common::visit(
         common::visitors{
@@ -473,6 +475,23 @@ static void processHostEvalClauses(lower::AbstractConverter &converter,
               beginClauseList =
                   &std::get<parser::OmpClauseList>(beginDirective.t);
 
+              // FIXME(JAN): For now we check if there is an inner
+              // OpenMPLoopConstruct, and extract the size clause from there
+              const auto &innerOptional = std::get<std::optional<
+                  common::Indirection<parser::OpenMPLoopConstruct>>>(
+                  ompConstruct.t);
+              if (innerOptional.has_value()) {
+                const auto &innerLoopDirective = innerOptional.value().value();
+                const auto &innerBegin =
+                    std::get<parser::OmpBeginLoopDirective>(
+                        innerLoopDirective.t);
+                const auto &innerDirective =
+                    std::get<parser::OmpLoopDirective>(innerBegin.t);
+                if (innerDirective.v == llvm::omp::Directive::OMPD_tile) {
+                  middleClauseList =
+                      &std::get<parser::OmpClauseList>(innerBegin.t);
+                }
+              }
               if (auto &endDirective =
                       std::get<std::optional<parser::OmpEndLoopDirective>>(
                           ompConstruct.t))
@@ -485,6 +504,9 @@ static void processHostEvalClauses(lower::AbstractConverter &converter,
     assert(beginClauseList && "expected begin directive");
     clauses.append(makeClauses(*beginClauseList, semaCtx));
 
+    if (middleClauseList)
+      clauses.append(makeClauses(*middleClauseList, semaCtx));
+
     if (endClauseList)
       clauses.append(makeClauses(*endClauseList, semaCtx));
   };
@@ -960,6 +982,7 @@ static void genLoopVars(
     storeOp =
         createAndSetPrivatizedLoopVar(converter, loc, indexVal, argSymbol);
   }
+
   firOpBuilder.setInsertionPointAfter(storeOp);
 }
 
@@ -1712,6 +1735,30 @@ genLoopNestClauses(lower::AbstractConverter &converter,
     cp.processCollapse(loc, eval, clauseOps, iv);
 
   clauseOps.loopInclusive = converter.getFirOpBuilder().getUnitAttr();
+
+  fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
+  for (auto &clause : clauses) {
+    if (clause.id == llvm::omp::Clause::OMPC_collapse) {
+      const auto &collapse = std::get<clause::Collapse>(clause.u);
+      int64_t collapseValue = evaluate::ToInt64(collapse.v).value();
+      clauseOps.numCollapse = firOpBuilder.getI64IntegerAttr(collapseValue);
+    } else if (clause.id == llvm::omp::Clause::OMPC_sizes) {
+      // This case handles the stand-alone tiling construct
+      const auto &sizes = std::get<clause::Sizes>(clause.u);
+      llvm::SmallVector<int64_t> sizeValues;
+      for (auto &size : sizes.v) {
+        int64_t sizeValue = evaluate::ToInt64(size).value();
+        sizeValues.push_back(sizeValue);
+      }
+      clauseOps.tileSizes = sizeValues;
+    }
+  }
+
+  llvm::SmallVector<int64_t> sizeValues;
+  auto *ompCons{eval.getIf<parser::OpenMPConstruct>()};
+  collectTileSizesFromOpenMPConstruct(ompCons, sizeValues, semaCtx);
+  if (sizeValues.size() > 0)
+    clauseOps.tileSizes = sizeValues;
 }
 
 static void genLoopClauses(
@@ -2085,9 +2132,9 @@ static mlir::omp::LoopNestOp genLoopNestOp(
     return llvm::SmallVector<const semantics::Symbol *>(iv);
   };
 
-  auto *nestedEval =
-      getCollapsedLoopEval(eval, getCollapseValue(item->clauses));
-
+  uint64_t nestValue = getCollapseValue(item->clauses);
+  nestValue = nestValue < iv.size() ? iv.size() : nestValue;
+  auto *nestedEval = getCollapsedLoopEval(eval, nestValue);
   return genOpWithBody<mlir::omp::LoopNestOp>(
       OpWithBodyGenInfo(converter, symTable, semaCtx, loc, *nestedEval,
                         directive)
@@ -3610,6 +3657,8 @@ static void genOMPDispatch(lower::AbstractConverter &converter,
                        item);
     break;
   case llvm::omp::Directive::OMPD_tile:
+    newOp = genLoopOp(converter, symTable, semaCtx, eval, loc, queue, item);
+    break;
   case llvm::omp::Directive::OMPD_unroll: {
     unsigned version = semaCtx.langOptions().OpenMPVersion;
     TODO(loc, "Unhandled loop directive (" +
@@ -4186,6 +4235,7 @@ static void genOMP(lower::AbstractConverter &converter, lower::SymMap &symTable,
       std::get<parser::OmpBeginLoopDirective>(loopConstruct.t);
   List<Clause> clauses = makeClauses(
       std::get<parser::OmpClauseList>(beginLoopDirective.t), semaCtx);
+
   if (auto &endLoopDirective =
           std::get<std::optional<parser::OmpEndLoopDirective>>(
               loopConstruct.t)) {
@@ -4292,18 +4342,6 @@ void Fortran::lower::genOpenMPSymbolProperties(
     lower::genDeclareTargetIntGlobal(converter, var);
 }
 
-int64_t
-Fortran::lower::getCollapseValue(const parser::OmpClauseList &clauseList) {
-  for (const parser::OmpClause &clause : clauseList.v) {
-    if (const auto &collapseClause =
-            std::get_if<parser::OmpClause::Collapse>(&clause.u)) {
-      const auto *expr = semantics::GetExpr(collapseClause->v);
-      return evaluate::ToInt64(*expr).value();
-    }
-  }
-  return 1;
-}
-
 void Fortran::lower::genThreadprivateOp(lower::AbstractConverter &converter,
                                         const lower::pft::Variable &var) {
   fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();

flang/lib/Lower/OpenMP/Utils.cpp

@@ -15,6 +15,7 @@
 #include "Clauses.h"
 
 #include "ClauseFinder.h"
+#include "flang/Evaluate/fold.h"
 #include <flang/Lower/AbstractConverter.h>
 #include <flang/Lower/ConvertType.h>
 #include <flang/Lower/DirectivesCommon.h>
@@ -24,10 +25,30 @@
 #include <flang/Parser/parse-tree.h>
 #include <flang/Parser/tools.h>
 #include <flang/Semantics/tools.h>
+#include <flang/Semantics/type.h>
 #include <llvm/Support/CommandLine.h>
 
 #include <iterator>
 
+using namespace Fortran::semantics;
+
+template <typename T>
+MaybeIntExpr EvaluateIntExpr(SemanticsContext &context, const T &expr) {
+  if (MaybeExpr maybeExpr{
+          Fold(context.foldingContext(), AnalyzeExpr(context, expr))}) {
+    if (auto *intExpr{Fortran::evaluate::UnwrapExpr<SomeIntExpr>(*maybeExpr)}) {
+      return std::move(*intExpr);
+    }
+  }
+  return std::nullopt;
+}
+
+template <typename T>
+std::optional<std::int64_t> EvaluateInt64(SemanticsContext &context,
+                                          const T &expr) {
+  return Fortran::evaluate::ToInt64(EvaluateIntExpr(context, expr));
+}
+
 llvm::cl::opt<bool> treatIndexAsSection(
     "openmp-treat-index-as-section",
     llvm::cl::desc("In the OpenMP data clauses treat `a(N)` as `a(N:N)`."),
@@ -38,14 +59,21 @@ namespace lower {
 namespace omp {
 
 int64_t getCollapseValue(const List<Clause> &clauses) {
-  auto iter = llvm::find_if(clauses, [](const Clause &clause) {
-    return clause.id == llvm::omp::Clause::OMPC_collapse;
-  });
-  if (iter != clauses.end()) {
-    const auto &collapse = std::get<clause::Collapse>(iter->u);
-    return evaluate::ToInt64(collapse.v).value();
+  int64_t collapseValue = 1;
+  int64_t numTileSizes = 0;
+  for (auto &clause : clauses) {
+    if (clause.id == llvm::omp::Clause::OMPC_collapse) {
+      const auto &collapse = std::get<clause::Collapse>(clause.u);
+      collapseValue = evaluate::ToInt64(collapse.v).value();
+    } else if (clause.id == llvm::omp::Clause::OMPC_sizes) {
+      const auto &sizes = std::get<clause::Sizes>(clause.u);
+      numTileSizes = sizes.v.size();
+    }
   }
-  return 1;
+
+  collapseValue = collapseValue - numTileSizes;
+  int64_t result = collapseValue > numTileSizes ? collapseValue : numTileSizes;
+  return result;
 }
 
 void genObjectList(const ObjectList &objects,
@@ -606,11 +634,48 @@ static void convertLoopBounds(lower::AbstractConverter &converter,
   }
 }
 
+// Populates the sizes vector with values if the given OpenMPConstruct
+// Contains a loop construct with an inner tiling construct.
+void collectTileSizesFromOpenMPConstruct(
+    const parser::OpenMPConstruct *ompCons,
+    llvm::SmallVectorImpl<int64_t> &tileSizes, SemanticsContext &semaCtx) {
+  if (!ompCons)
+    return;
+
+  if (auto *ompLoop{std::get_if<parser::OpenMPLoopConstruct>(&ompCons->u)}) {
+    const auto &innerOptional = std::get<
+        std::optional<common::Indirection<parser::OpenMPLoopConstruct>>>(
+        ompLoop->t);
+    if (innerOptional.has_value()) {
+      const auto &innerLoopDirective = innerOptional.value().value();
+      const auto &innerBegin =
+          std::get<parser::OmpBeginLoopDirective>(innerLoopDirective.t);
+      const auto &innerDirective =
+          std::get<parser::OmpLoopDirective>(innerBegin.t).v;
+
+      if (innerDirective == llvm::omp::Directive::OMPD_tile) {
+        // Get the size values from parse tree and convert to a vector
+        const auto &innerClauseList{
+            std::get<parser::OmpClauseList>(innerBegin.t)};
+        for (const auto &clause : innerClauseList.v)
+          if (const auto tclause{
+                  std::get_if<parser::OmpClause::Sizes>(&clause.u)}) {
+            for (auto &tval : tclause->v) {
+              if (const auto v{EvaluateInt64(semaCtx, tval)})
+                tileSizes.push_back(*v);
+            }
+          }
+      }
+    }
+  }
+}
+
 bool collectLoopRelatedInfo(
     lower::AbstractConverter &converter, mlir::Location currentLocation,
     lower::pft::Evaluation &eval, const omp::List<omp::Clause> &clauses,
     mlir::omp::LoopRelatedClauseOps &result,
     llvm::SmallVectorImpl<const semantics::Symbol *> &iv) {
+
   bool found = false;
   fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
 
@@ -627,6 +692,38 @@ bool collectLoopRelatedInfo(
     found = true;
   }
 
+  // Collect sizes from tile directive if present
+  std::int64_t sizesLengthValue = 0l;
+  if (auto *ompCons{eval.getIf<parser::OpenMPConstruct>()}) {
+    if (auto *ompLoop{std::get_if<parser::OpenMPLoopConstruct>(&ompCons->u)}) {
+      const auto &innerOptional = std::get<
+          std::optional<common::Indirection<parser::OpenMPLoopConstruct>>>(
+          ompLoop->t);
+      if (innerOptional.has_value()) {
+        const auto &innerLoopDirective = innerOptional.value().value();
+        const auto &innerBegin =
+            std::get<parser::OmpBeginLoopDirective>(innerLoopDirective.t);
+        const auto &innerDirective =
+            std::get<parser::OmpLoopDirective>(innerBegin.t).v;
+
+        if (innerDirective == llvm::omp::Directive::OMPD_tile) {
+          // Get the size values from parse tree and convert to a vector
+          const auto &innerClauseList{
+              std::get<parser::OmpClauseList>(innerBegin.t)};
+          for (const auto &clause : innerClauseList.v)
+            if (const auto tclause{
+                    std::get_if<parser::OmpClause::Sizes>(&clause.u)}) {
+              sizesLengthValue = tclause->v.size();
+              found = true;
+            }
+        }
+      }
+    }
+  }
+
+  collapseValue = collapseValue - sizesLengthValue;
+  collapseValue =
+      collapseValue < sizesLengthValue ? sizesLengthValue : collapseValue;
   std::size_t loopVarTypeSize = 0;
   do {
     lower::pft::Evaluation *doLoop =
@@ -659,7 +756,6 @@ bool collectLoopRelatedInfo(
   } while (collapseValue > 0);
 
   convertLoopBounds(converter, currentLocation, result, loopVarTypeSize);
-
   return found;
 }
 } // namespace omp

flang/lib/Lower/OpenMP/Utils.h

@@ -166,6 +166,12 @@ bool collectLoopRelatedInfo(
     lower::pft::Evaluation &eval, const omp::List<omp::Clause> &clauses,
     mlir::omp::LoopRelatedClauseOps &result,
     llvm::SmallVectorImpl<const semantics::Symbol *> &iv);
+
+void collectTileSizesFromOpenMPConstruct(
+    const parser::OpenMPConstruct *ompCons,
+    llvm::SmallVectorImpl<int64_t> &tileSizes,
+    Fortran::semantics::SemanticsContext &semaCtx);
+
 } // namespace omp
 } // namespace lower
 } // namespace Fortran

flang/lib/Parser/unparse.cpp

@@ -2926,6 +2926,8 @@ class UnparseVisitor {
     Walk(std::get<OmpBeginLoopDirective>(x.t));
     Put("\n");
     EndOpenMP();
+    Walk(
+        std::get<std::optional<common::Indirection<OpenMPLoopConstruct>>>(x.t));
     Walk(std::get<std::optional<DoConstruct>>(x.t));
     Walk(std::get<std::optional<OmpEndLoopDirective>>(x.t));
   }