Cherry-pick PR llvm#102613

ROCm · Aug 16, 2024 · bce6b27 · bce6b27
1 parent d2b3307
commit bce6b27
Show file tree

Hide file tree

Showing 7 changed files with 133 additions and 509 deletions.
diff --git a/flang/lib/Lower/OpenMP/Decomposer.cpp b/flang/lib/Lower/OpenMP/Decomposer.cpp
@@ -22,7 +22,6 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Frontend/OpenMP/ClauseT.h"
-#include "llvm/Frontend/OpenMP/ConstructCompositionT.h"
 #include "llvm/Frontend/OpenMP/ConstructDecompositionT.h"
 #include "llvm/Frontend/OpenMP/OMP.h"
 #include "llvm/Support/raw_ostream.h"
@@ -68,12 +67,6 @@ struct ConstructDecomposition {
 };
 } // namespace
 
-static UnitConstruct mergeConstructs(uint32_t version,
-                                     llvm::ArrayRef<UnitConstruct> units) {
-  tomp::ConstructCompositionT compose(version, units);
-  return compose.merged;
-}
-
 namespace Fortran::lower::omp {
 LLVM_DUMP_METHOD llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
                                                const UnitConstruct &uc) {
@@ -90,38 +83,37 @@ ConstructQueue buildConstructQueue(
     Fortran::lower::pft::Evaluation &eval, const parser::CharBlock &source,
     llvm::omp::Directive compound, const List<Clause> &clauses) {
 
-  List<UnitConstruct> constructs;
-
   ConstructDecomposition decompose(modOp, semaCtx, eval, compound, clauses);
   assert(!decompose.output.empty() && "Construct decomposition failed");
 
-  llvm::SmallVector<llvm::omp::Directive> loweringUnits;
-  std::ignore =
-      llvm::omp::getLeafOrCompositeConstructs(compound, loweringUnits);
-  uint32_t version = getOpenMPVersionAttribute(modOp);
-
-  int leafIndex = 0;
-  for (llvm::omp::Directive dir_id : loweringUnits) {
-    llvm::ArrayRef<llvm::omp::Directive> leafsOrSelf =
-        llvm::omp::getLeafConstructsOrSelf(dir_id);
-    size_t numLeafs = leafsOrSelf.size();
-
-    llvm::ArrayRef<UnitConstruct> toMerge{&decompose.output[leafIndex],
-                                          numLeafs};
-    auto &uc = constructs.emplace_back(mergeConstructs(version, toMerge));
-
-    if (!transferLocations(clauses, uc.clauses)) {
-      // If some clauses are left without source information, use the
-      // directive's source.
-      for (auto &clause : uc.clauses) {
-        if (clause.source.empty())
-          clause.source = source;
-      }
-    }
-    leafIndex += numLeafs;
+  for (UnitConstruct &uc : decompose.output) {
+    assert(getLeafConstructs(uc.id).empty() && "unexpected compound directive");
+    //  If some clauses are left without source information, use the directive's
+    //  source.
+    for (auto &clause : uc.clauses)
+      if (clause.source.empty())
+        clause.source = source;
+  }
+
+  return decompose.output;
+}
+
+bool matchLeafSequence(ConstructQueue::const_iterator item,
+                       const ConstructQueue &queue,
+                       llvm::omp::Directive directive) {
+  llvm::ArrayRef<llvm::omp::Directive> leafDirs =
+      llvm::omp::getLeafConstructsOrSelf(directive);
+
+  for (auto [dir, leaf] :
+       llvm::zip_longest(leafDirs, llvm::make_range(item, queue.end()))) {
+    if (!dir.has_value() || !leaf.has_value())
+      return false;
+
+    if (*dir != leaf->id)
+      return false;
   }
 
-  return constructs;
+  return true;
 }
 
 bool isLastItemInQueue(ConstructQueue::iterator item,

diff --git a/flang/lib/Lower/OpenMP/Decomposer.h b/flang/lib/Lower/OpenMP/Decomposer.h
@@ -10,7 +10,6 @@
 
 #include "flang/Lower/OpenMP/Clauses.h"
 #include "mlir/IR/BuiltinOps.h"
-#include "llvm/Frontend/OpenMP/ConstructCompositionT.h"
 #include "llvm/Frontend/OpenMP/ConstructDecompositionT.h"
 #include "llvm/Frontend/OpenMP/OMP.h"
 #include "llvm/Support/Compiler.h"
@@ -49,6 +48,15 @@ ConstructQueue buildConstructQueue(mlir::ModuleOp modOp,
 
 bool isLastItemInQueue(ConstructQueue::iterator item,
                        const ConstructQueue &queue);
+
+/// Try to match the leaf constructs conforming the given \c directive to the
+/// range of leaf constructs starting from \c item to the end of the \c queue.
+/// If \c directive doesn't represent a compound directive, check that \c item
+/// matches that directive and is the only element before the end of the
+/// \c queue.
+bool matchLeafSequence(ConstructQueue::const_iterator item,
+                       const ConstructQueue &queue,
+                       llvm::omp::Directive directive);
 } // namespace Fortran::lower::omp
 
 #endif // FORTRAN_LOWER_OPENMP_DECOMPOSER_H