[flang-commits] [flang] [flang][OpenMP] Map `teams loop` to `teams distribute` when required. (PR #127489)

[Sergio Afonso via flang-commits]

================
@@ -118,6 +119,56 @@ class GenericLoopConversionPattern
     return result;
   }
 
+  /// Checks whether a `teams loop` construct can be rewriten to `teams
+  /// distribute parallel do` or it has to be converted to `teams distribute`.
+  ///
+  /// This checks similar constrains to what is checked by `TeamsLoopChecker` in
+  /// SemaOpenMP.cpp in clang.
+  static bool teamsLoopCanBeParallelFor(mlir::omp::LoopOp loopOp) {
+    bool canBeParallelFor = true;
+    loopOp.walk([&](mlir::omp::LoopOp nestedLoopOp) {
+      if (nestedLoopOp == loopOp)
+        mlir::WalkResult::advance();
+
+      GenericLoopCombinedInfo combinedInfo =
+          findGenericLoopCombineInfo(nestedLoopOp);
+
+      // Worksharing loops cannot be nested inside each other. Therefore, if the
+      // current `loop` directive nests another `loop` whose `bind` modifier is
+      // `parallel`, this `loop` directive cannot be mapped to `distribute
+      // parallel for` but rather only to `distribute`.
+      if (combinedInfo == GenericLoopCombinedInfo::Standalone &&
+          nestedLoopOp.getBindKind() &&
+          *nestedLoopOp.getBindKind() == mlir::omp::ClauseBindKind::Parallel)
+        canBeParallelFor = false;
+
+      // TODO check for combined `parallel loop` when we support it.
+
+      return canBeParallelFor ? mlir::WalkResult::advance()
+                              : mlir::WalkResult::interrupt();
+    });
----------------
skatrak wrote:

Nit: I think this is simpler, but feel free to ignore if you disagree.
```suggestion
    bool canBeParallelFor = !loopOp.walk<WalkOrder::PreOrder>([&](mlir::omp::LoopOp nestedLoopOp) {
      if (nestedLoopOp == loopOp)
        return mlir::WalkResult::advance();

      GenericLoopCombinedInfo combinedInfo =
          findGenericLoopCombineInfo(nestedLoopOp);

      // Worksharing loops cannot be nested inside each other. Therefore, if the
      // current `loop` directive nests another `loop` whose `bind` modifier is
      // `parallel`, this `loop` directive cannot be mapped to `distribute
      // parallel for` but rather only to `distribute`.
      if (combinedInfo == GenericLoopCombinedInfo::Standalone &&
          nestedLoopOp.getBindKind() &&
          *nestedLoopOp.getBindKind() == mlir::omp::ClauseBindKind::Parallel)
        return mlir::WalkResult::interrupt();

      // TODO check for combined `parallel loop` when we support it.

      return mlir::WalkResult::advance();
    }).wasInterrupted();
```

https://github.com/llvm/llvm-project/pull/127489