Make sure collapse is progagated correctly on the host side.

Author: jsjodin

flang/lib/Lower/OpenMP/ClauseProcessor.cpp

@@ -271,24 +271,16 @@ bool ClauseProcessor::processCancelDirectiveName(
   return true;
 }
 
-bool ClauseProcessor::processLoopNests(
-    mlir::Location currentLocation, lower::pft::Evaluation &eval,
-    mlir::omp::LoopRelatedClauseOps &result,
-    llvm::SmallVectorImpl<const semantics::Symbol *> &iv) const {
-  int64_t numCollapse = collectLoopRelatedInfo(converter, currentLocation, eval,
-                                               clauses, result, iv);
-  return numCollapse > 1;
-}
-
 bool ClauseProcessor::processCollapse(
     mlir::Location currentLocation, lower::pft::Evaluation &eval,
-    mlir::omp::LoopNestOperands &result,
+    mlir::omp::LoopRelatedClauseOps &loopResult,
+    mlir::omp::CollapseClauseOps &collapseResult,
     llvm::SmallVectorImpl<const semantics::Symbol *> &iv) const {
 
   int64_t numCollapse = collectLoopRelatedInfo(converter, currentLocation, eval,
-                                               clauses, result, iv);
+                                               clauses, loopResult, iv);
   fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
-  result.collapseNumLoops = firOpBuilder.getI64IntegerAttr(numCollapse);
+  collapseResult.collapseNumLoops = firOpBuilder.getI64IntegerAttr(numCollapse);
   return numCollapse > 1;
 }
 

flang/lib/Lower/OpenMP/ClauseProcessor.h

@@ -63,12 +63,9 @@ class ClauseProcessor {
       mlir::omp::CancelDirectiveNameClauseOps &result) const;
   bool
   processCollapse(mlir::Location currentLocation, lower::pft::Evaluation &eval,
-                  mlir::omp::LoopNestOperands &result,
+                  mlir::omp::LoopRelatedClauseOps &loopResult,
+                  mlir::omp::CollapseClauseOps &collapseResult,
                   llvm::SmallVectorImpl<const semantics::Symbol *> &iv) const;
-  bool
-  processLoopNests(mlir::Location currentLocation, lower::pft::Evaluation &eval,
-                   mlir::omp::LoopRelatedClauseOps &result,
-                   llvm::SmallVectorImpl<const semantics::Symbol *> &iv) const;
   bool processDevice(lower::StatementContext &stmtCtx,
                      mlir::omp::DeviceClauseOps &result) const;
   bool processDeviceType(mlir::omp::DeviceTypeClauseOps &result) const;

flang/lib/Lower/OpenMP/OpenMP.cpp

@@ -504,7 +504,7 @@ static void processHostEvalClauses(lower::AbstractConverter &converter,
       [[fallthrough]];
     case OMPD_distribute:
     case OMPD_distribute_simd:
-      cp.processLoopNests(loc, eval, hostInfo->ops, hostInfo->iv);
+      cp.processCollapse(loc, eval, hostInfo->ops, hostInfo->ops, hostInfo->iv);
       break;
 
     case OMPD_teams:
@@ -523,7 +523,7 @@ static void processHostEvalClauses(lower::AbstractConverter &converter,
       [[fallthrough]];
     case OMPD_target_teams_distribute:
     case OMPD_target_teams_distribute_simd:
-      cp.processLoopNests(loc, eval, hostInfo->ops, hostInfo->iv);
+      cp.processCollapse(loc, eval, hostInfo->ops, hostInfo->ops, hostInfo->iv);
       cp.processNumTeams(stmtCtx, hostInfo->ops);
       break;
 
@@ -534,7 +534,7 @@ static void processHostEvalClauses(lower::AbstractConverter &converter,
       cp.processNumTeams(stmtCtx, hostInfo->ops);
       [[fallthrough]];
     case OMPD_loop:
-      cp.processLoopNests(loc, eval, hostInfo->ops, hostInfo->iv);
+      cp.processCollapse(loc, eval, hostInfo->ops, hostInfo->ops, hostInfo->iv);
       break;
 
     case OMPD_teams_workdistribute:
@@ -1570,7 +1570,7 @@ genLoopNestClauses(lower::AbstractConverter &converter,
 
   HostEvalInfo *hostEvalInfo = getHostEvalInfoStackTop(converter);
   if (!hostEvalInfo || !hostEvalInfo->apply(clauseOps, iv))
-    cp.processCollapse(loc, eval, clauseOps, iv);
+    cp.processCollapse(loc, eval, clauseOps, clauseOps, iv);
 
   clauseOps.loopInclusive = converter.getFirOpBuilder().getUnitAttr();
   cp.processTileSizes(eval, clauseOps);

mlir/include/mlir/Dialect/OpenMP/OpenMPClauseOperands.h

@@ -40,7 +40,7 @@ struct DeviceTypeClauseOps {
 /// Clauses that correspond to operations other than omp.target, but might have
 /// to be evaluated outside of a parent target region.
 using HostEvaluatedOperands =
-    detail::Clauses<LoopRelatedClauseOps, NumTeamsClauseOps,
+    detail::Clauses<CollapseClauseOps, LoopRelatedClauseOps, NumTeamsClauseOps,
                     NumThreadsClauseOps, ThreadLimitClauseOps>;
 
 // TODO: Add `indirect` clause.