Use shared helper function to find the sizes clause.

Author: jsjodin

flang/lib/Lower/OpenMP/Utils.cpp

@@ -599,15 +599,13 @@ 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(
+// Helper function that finds the sizes clause in a inner OMPD_tile directive
+// and passes the sizes clause to the callback function if found.
+static void processTileSizesFromOpenMPConstruct(
     const parser::OpenMPConstruct *ompCons,
-    llvm::SmallVectorImpl<int64_t> &tileSizes,
-    Fortran::semantics::SemanticsContext &semaCtx) {
+    std::function<void(const parser::OmpClause::Sizes *)> processFun) {
   if (!ompCons)
     return;
-
   if (auto *ompLoop{std::get_if<parser::OpenMPLoopConstruct>(&ompCons->u)}) {
     const auto &nestedOptional =
         std::get<std::optional<parser::NestedConstruct>>(ompLoop->t);
@@ -624,16 +622,13 @@ void collectTileSizesFromOpenMPConstruct(
           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
+        // 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);
-            }
+            processFun(tclause);
             break;
           }
         }
@@ -642,6 +637,20 @@ void collectTileSizesFromOpenMPConstruct(
   }
 }
 
+/// 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,
+    Fortran::semantics::SemanticsContext &semaCtx) {
+  processTileSizesFromOpenMPConstruct(
+      ompCons, [&](const parser::OmpClause::Sizes *tclause) {
+        for (auto &tval : tclause->v)
+          if (const auto v{EvaluateInt64(semaCtx, tval)})
+            tileSizes.push_back(*v);
+      });
+}
+
 int64_t collectLoopRelatedInfo(
     lower::AbstractConverter &converter, mlir::Location currentLocation,
     lower::pft::Evaluation &eval, const omp::List<omp::Clause> &clauses,
@@ -663,37 +672,13 @@ int64_t collectLoopRelatedInfo(
     numCollapse = collapseValue;
   }
 
-  // Collect sizes from tile directive if present
+  // 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 &nestedOptional =
-          std::get<std::optional<parser::NestedConstruct>>(ompLoop->t);
-      assert(nestedOptional.has_value() &&
-             "Expected a DoConstruct or OpenMPLoopConstruct");
-      const auto *innerConstruct =
-          std::get_if<common::Indirection<parser::OpenMPLoopConstruct>>(
-              &(nestedOptional.value()));
-      if (innerConstruct) {
-        const auto &innerLoopDirective = innerConstruct->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();
-              break;
-            }
-        }
-      }
-    }
+    processTileSizesFromOpenMPConstruct(
+        ompCons, [&](const parser::OmpClause::Sizes *tclause) {
+          sizesLengthValue = tclause->v.size();
+        });
   }
 
   collapseValue = std::max(collapseValue, sizesLengthValue);