refactor code to parse in OpenACC.cpp

Author: SusanTan

flang/include/flang/Lower/OpenACC.h

@@ -122,6 +122,9 @@ void genOpenACCTerminator(fir::FirOpBuilder &, mlir::Operation *,
 /// clause.
 uint64_t getLoopCountForCollapseAndTile(const Fortran::parser::AccClauseList &);
 
+/// Returns only the collapse(N) depth (defaults to 1 when absent).
+uint64_t getLoopCountForCollapse(const Fortran::parser::AccClauseList &);
+
 /// Checks whether the current insertion point is inside OpenACC loop.
 bool isInOpenACCLoop(fir::FirOpBuilder &);
 

flang/lib/Lower/Bridge.cpp

@@ -3222,15 +3222,17 @@ class FirConverter : public Fortran::lower::AbstractConverter {
         const Fortran::parser::AccClauseList &clauseList =
             std::get<Fortran::parser::AccClauseList>(beginLoopDir.t);
         loopCount = Fortran::lower::getLoopCountForCollapseAndTile(clauseList);
-        std::tie(collapseForce, collapseDepth) = parseCollapse(clauseList);
+        collapseDepth = Fortran::lower::getLoopCountForCollapse(clauseList);
+        std::tie(collapseForce, std::ignore) = parseCollapse(clauseList);
       } else if (accCombined) {
         const Fortran::parser::AccBeginCombinedDirective &beginCombinedDir =
             std::get<Fortran::parser::AccBeginCombinedDirective>(
                 accCombined->t);
         const Fortran::parser::AccClauseList &clauseList =
             std::get<Fortran::parser::AccClauseList>(beginCombinedDir.t);
         loopCount = Fortran::lower::getLoopCountForCollapseAndTile(clauseList);
-        std::tie(collapseForce, collapseDepth) = parseCollapse(clauseList);
+        collapseDepth = Fortran::lower::getLoopCountForCollapse(clauseList);
+        std::tie(collapseForce, std::ignore) = parseCollapse(clauseList);
       }
 
       if (curEval->lowerAsStructured()) {

flang/lib/Lower/OpenACC.cpp

@@ -4889,6 +4889,20 @@ uint64_t Fortran::lower::getLoopCountForCollapseAndTile(
   return collapseLoopCount;
 }
 
+uint64_t Fortran::lower::getLoopCountForCollapse(
+    const Fortran::parser::AccClauseList &clauseList) {
+  for (const Fortran::parser::AccClause &clause : clauseList.v) {
+    if (const auto *collapseClause =
+            std::get_if<Fortran::parser::AccClause::Collapse>(&clause.u)) {
+      const Fortran::parser::AccCollapseArg &arg = collapseClause->v;
+      const auto &collapseValue =
+          std::get<Fortran::parser::ScalarIntConstantExpr>(arg.t);
+      return *Fortran::semantics::GetIntValue(collapseValue);
+    }
+  }
+  return 1;
+}
+
 /// Create an ACC loop operation for a DO construct when inside ACC compute
 /// constructs This serves as a bridge between regular DO construct handling and
 /// ACC loop creation