[flang-commits] [flang] [llvm] [flang][OpenMP] Decompose compound constructs, do recursive lowering (PR #90098)

[Sergio Afonso via flang-commits]

================
@@ -0,0 +1,109 @@
+//===-- Decomposer.cpp -- Compound directive decomposition ----------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Coding style: https://mlir.llvm.org/getting_started/DeveloperGuide/
+//
+//===----------------------------------------------------------------------===//
+
+#include "Decomposer.h"
+
+#include "Clauses.h"
+#include "Utils.h"
+#include "flang/Lower/PFTBuilder.h"
+#include "flang/Semantics/semantics.h"
+#include "flang/Tools/CrossToolHelpers.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Frontend/OpenMP/ClauseT.h"
+#include "llvm/Frontend/OpenMP/ConstructDecompositionT.h"
+#include "llvm/Frontend/OpenMP/OMP.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include <optional>
+#include <utility>
+#include <variant>
+
+using namespace Fortran;
+
+namespace {
+using namespace Fortran::lower::omp;
+
+struct ConstructDecomposition {
+  ConstructDecomposition(mlir::ModuleOp modOp,
+                         semantics::SemanticsContext &semaCtx,
+                         lower::pft::Evaluation &ev,
+                         llvm::omp::Directive compound,
+                         const List<Clause> &clauses)
+      : semaCtx(semaCtx), mod(modOp), eval(ev) {
+    tomp::ConstructDecompositionT decompose(getOpenMPVersionAttribute(modOp),
+                                            *this, compound,
+                                            llvm::ArrayRef(clauses));
+    output = std::move(decompose.output);
+  }
+
+  // Given an object, return its base object if one exists.
+  std::optional<Object> getBaseObject(const Object &object) {
+    return lower::omp::getBaseObject(object, semaCtx);
+  }
+
+  // Return the iteration variable of the associated loop if any.
+  std::optional<Object> getLoopIterVar() {
+    if (semantics::Symbol *symbol = getIterationVariableSymbol(eval))
+      return Object{symbol, /*designator=*/{}};
+    return std::nullopt;
+  }
+
+  semantics::SemanticsContext &semaCtx;
+  mlir::ModuleOp mod;
+  lower::pft::Evaluation &eval;
+  List<UnitConstruct> output;
+};
+} // namespace
+
+namespace Fortran::lower::omp {
+LLVM_DUMP_METHOD llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
+                                               const UnitConstruct &uc) {
+  os << llvm::omp::getOpenMPDirectiveName(uc.id);
+  for (auto [index, clause] : llvm::enumerate(uc.clauses)) {
+    os << (index == 0 ? '\t' : ' ');
+    os << llvm::omp::getOpenMPClauseName(clause.id);
+  }
+  return os;
+}
+
+ConstructQueue buildConstructQueue(
+    mlir::ModuleOp modOp, Fortran::semantics::SemanticsContext &semaCtx,
+    Fortran::lower::pft::Evaluation &eval, llvm::omp::Directive compound,
+    const List<Clause> &clauses) {
+
+  List<UnitConstruct> constructs;
+
+  ConstructDecomposition decompose(modOp, semaCtx, eval, compound, clauses);
+  assert(!decompose.output.empty());
+
+  llvm::SmallVector<llvm::omp::Directive> loweringUnits;
+  std::ignore =
+      llvm::omp::getLeafOrCompositeConstructs(compound, loweringUnits);
+
+  int leafIndex = 0;
+  for (llvm::omp::Directive dir_id : loweringUnits) {
+    constructs.push_back(UnitConstruct{dir_id});
+    UnitConstruct &uc = constructs.back();
+    llvm::ArrayRef<llvm::omp::Directive> leafsOrSelf =
+        llvm::omp::getLeafConstructsOrSelf(dir_id);
+    for (int i = 0, e = leafsOrSelf.size(); i != e; ++i) {
+      uc.clauses.append(decompose.output[leafIndex].clauses);
----------------
skatrak wrote:

My understanding of this is that `ConstructDecomposition` will split the compound directive into a list of leaf constructs, where compound directives are also split. On the other hand, `loweringUnits` will keep compound directives as a unit rather than split them to be later handled by `genOMPDispatch` and call the corresponding composite lowering function. So this piece of code will gather the clauses attached to all leaves of a composite construct into a single clause list.

The problem with it is that the lowering functions for composite constructs will then have to do themselves the work of figuring out to which leaf construct(s) each clause will apply to, after it had already been done and later undone again. I'm not sure what the best solution to that should be, but I think we need to be able to pass to these composite lowering functions a clause list per leaf construct (maybe `UnitConstruct` could hold multiple lists of constructs, or maybe we should have this function split composite constructs and then update `genOMPDispatch` to detect them by looking at multiple items and triggering a single call for them).

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