[Frontend][OpenMP] Implement getLeafOrCompositeConstructs

[kparzysz]

This function will break up a construct into constituent leaf and composite constructs, e.g. if OMPD_c_d_e and OMPD_d_e are composite constructs, then OMPD_a_b_c_d_e will be broken up into the list {OMPD_a, OMPD_b, OMPD_c_d_e}.

[llvmbot]

@llvm/pr-subscribers-flang-openmp

Author: Krzysztof Parzyszek (kparzysz)

Changes

This function will break up a construct into constituent leaf and composite constructs, e.g. if OMPD_c_d_e and OMPD_d_e are composite constructs, then OMPD_a_b_c_d_e will be broken up into the list {OMPD_a, OMPD_b, OMPD_c_d_e}.

---

Full diff: https://github.com/llvm/llvm-project/pull/89104.diff

3 Files Affected:

• (modified) llvm/include/llvm/Frontend/OpenMP/OMP.h (+6)
• (modified) llvm/lib/Frontend/OpenMP/OMP.cpp (+75-13)
• (modified) llvm/unittests/Frontend/OpenMPCompositionTest.cpp (+32)

diff --git a/llvm/include/llvm/Frontend/OpenMP/OMP.h b/llvm/include/llvm/Frontend/OpenMP/OMP.h
index ec8ae68f1c2ca0..6f7a39acac1d31 100644
--- a/llvm/include/llvm/Frontend/OpenMP/OMP.h
+++ b/llvm/include/llvm/Frontend/OpenMP/OMP.h
@@ -16,9 +16,15 @@
 #include "llvm/Frontend/OpenMP/OMP.h.inc"
 
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
 
 namespace llvm::omp {
 ArrayRef<Directive> getLeafConstructs(Directive D);
+ArrayRef<Directive> getLeafConstructsOrSelf(Directive D);
+
+ArrayRef<Directive>
+getLeafOrCompositeConstructs(Directive D, SmallVectorImpl<Directive> &Output);
+
 Directive getCompoundConstruct(ArrayRef<Directive> Parts);
 
 bool isLeafConstruct(Directive D);
diff --git a/llvm/lib/Frontend/OpenMP/OMP.cpp b/llvm/lib/Frontend/OpenMP/OMP.cpp
index 4b9b7037ee4ad5..2ebadf5216a084 100644
--- a/llvm/lib/Frontend/OpenMP/OMP.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMP.cpp
@@ -25,6 +25,43 @@ using namespace llvm::omp;
 #define GEN_DIRECTIVES_IMPL
 #include "llvm/Frontend/OpenMP/OMP.inc"
 
+static iterator_range<ArrayRef<Directive>::iterator>
+getFirstCompositeRange(iterator_range<ArrayRef<Directive>::iterator> Leafs) {
+  // OpenMP Spec 5.2: [17.3, 8-9]
+  // If directive-name-A and directive-name-B both correspond to loop-
+  // associated constructs then directive-name is a composite construct
+  // otherwise directive-name is a combined construct.
+  //
+  // In the list of leaf constructs, find the first loop-associated construct,
+  // this is the beginning of the range. Then, starting from the immediately
+  // following leaf construct, find the first sequence of adjacent loop-
+  // -associated constructs. The last of those is the last one of the range.
+
+  auto firstLoopAssociated =
+      [](iterator_range<ArrayRef<Directive>::iterator> List) {
+        for (auto It = List.begin(), End = List.end(); It != End; ++It) {
+          if (getDirectiveAssociation(*It) == Association::Loop)
+            return It;
+        }
+        return List.end();
+      };
+
+  auto Begin = firstLoopAssociated(Leafs);
+  if (Begin == Leafs.end())
+    return llvm::make_range(Leafs.end(), Leafs.end());
+
+  auto End =
+      firstLoopAssociated(llvm::make_range(std::next(Begin), Leafs.end()));
+  if (End == Leafs.end())
+    return llvm::make_range(Begin, std::next(Begin));
+
+  for (; End != Leafs.end(); ++End) {
+    if (getDirectiveAssociation(*End) != Association::Loop)
+      break;
+  }
+  return llvm::make_range(Begin, End);
+}
+
 namespace llvm::omp {
 ArrayRef<Directive> getLeafConstructs(Directive D) {
   auto Idx = static_cast<std::size_t>(D);
@@ -34,6 +71,40 @@ ArrayRef<Directive> getLeafConstructs(Directive D) {
   return ArrayRef(&Row[2], &Row[2] + static_cast<int>(Row[1]));
 }
 
+ArrayRef<Directive> getLeafConstructsOrSelf(Directive D) {
+  if (auto Leafs = getLeafConstructs(D); !Leafs.empty())
+    return Leafs;
+  auto Idx = static_cast<size_t>(D);
+  assert(Idx < Directive_enumSize && "Invalid directive");
+  const auto *Row = LeafConstructTable[LeafConstructTableOrdering[Idx]];
+  // The first entry in the row is the directive itself.
+  return ArrayRef(&Row[0], &Row[0] + 1);
+}
+
+ArrayRef<Directive>
+getLeafOrCompositeConstructs(Directive D, SmallVectorImpl<Directive> &Output) {
+  using ArrayTy = ArrayRef<Directive>;
+  using IteratorTy = ArrayTy::iterator;
+  ArrayRef<Directive> Leafs = getLeafConstructsOrSelf(D);
+
+  IteratorTy Iter = Leafs.begin();
+  do {
+    auto Range = getFirstCompositeRange(llvm::make_range(Iter, Leafs.end()));
+    // All directives before the range are leaf constructs.
+    for (; Iter != Range.begin(); ++Iter)
+      Output.push_back(*Iter);
+    if (!Range.empty()) {
+      Directive Comp =
+          getCompoundConstruct(ArrayTy(Range.begin(), Range.end()));
+      assert(Comp != OMPD_unknown);
+      Output.push_back(Comp);
+    }
+    Iter = Range.end();
+  } while (Iter != Leafs.end());
+
+  return Output;
+}
+
 Directive getCompoundConstruct(ArrayRef<Directive> Parts) {
   if (Parts.empty())
     return OMPD_unknown;
@@ -88,20 +159,11 @@ Directive getCompoundConstruct(ArrayRef<Directive> Parts) {
 bool isLeafConstruct(Directive D) { return getLeafConstructs(D).empty(); }
 
 bool isCompositeConstruct(Directive D) {
-  // OpenMP Spec 5.2: [17.3, 8-9]
-  // If directive-name-A and directive-name-B both correspond to loop-
-  // associated constructs then directive-name is a composite construct
-  llvm::ArrayRef<Directive> Leafs{getLeafConstructs(D)};
-  if (Leafs.empty())
-    return false;
-  if (getDirectiveAssociation(Leafs.front()) != Association::Loop)
+  ArrayRef<Directive> Leafs = getLeafConstructsOrSelf(D);
+  if (Leafs.size() <= 1)
     return false;
-
-  size_t numLoopConstructs =
-      llvm::count_if(Leafs.drop_front(), [](Directive L) {
-        return getDirectiveAssociation(L) == Association::Loop;
-      });
-  return numLoopConstructs != 0;
+  auto Range = getFirstCompositeRange(Leafs);
+  return Range.begin() == Leafs.begin() && Range.end() == Leafs.end();
 }
 
 bool isCombinedConstruct(Directive D) {
diff --git a/llvm/unittests/Frontend/OpenMPCompositionTest.cpp b/llvm/unittests/Frontend/OpenMPCompositionTest.cpp
index 0b32e0d96dc84c..6915a0cbcaac2d 100644
--- a/llvm/unittests/Frontend/OpenMPCompositionTest.cpp
+++ b/llvm/unittests/Frontend/OpenMPCompositionTest.cpp
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/Frontend/OpenMP/OMP.h"
 #include "gtest/gtest.h"
 
@@ -38,6 +39,37 @@ TEST(Composition, GetCompoundConstruct) {
   ASSERT_EQ(C6, OMPD_parallel_for_simd);
 }
 
+TEST(Composition, GetLeafOrCompositeConstructs) {
+  SmallVector<Directive> Out1;
+  auto Ret1 = getLeafOrCompositeConstructs(
+      OMPD_target_teams_distribute_parallel_for, Out1);
+  ASSERT_EQ(Ret1, ArrayRef<Directive>(Out1));
+  ASSERT_EQ((ArrayRef<Directive>(Out1)),
+            (ArrayRef<Directive>{OMPD_target, OMPD_teams,
+                                 OMPD_distribute_parallel_for}));
+
+  SmallVector<Directive> Out2;
+  auto Ret2 =
+      getLeafOrCompositeConstructs(OMPD_parallel_masked_taskloop_simd, Out2);
+  ASSERT_EQ(Ret2, ArrayRef<Directive>(Out2));
+  ASSERT_EQ(
+      (ArrayRef<Directive>(Out2)),
+      (ArrayRef<Directive>{OMPD_parallel, OMPD_masked, OMPD_taskloop_simd}));
+
+  SmallVector<Directive> Out3;
+  auto Ret3 =
+      getLeafOrCompositeConstructs(OMPD_distribute_parallel_do_simd, Out3);
+  ASSERT_EQ(Ret3, ArrayRef<Directive>(Out3));
+  ASSERT_EQ((ArrayRef<Directive>(Out3)),
+            (ArrayRef<Directive>{OMPD_distribute_parallel_do_simd}));
+
+  SmallVector<Directive> Out4;
+  auto Ret4 = getLeafOrCompositeConstructs(OMPD_target_parallel_loop, Out4);
+  ASSERT_EQ(Ret4, ArrayRef<Directive>(Out4));
+  ASSERT_EQ((ArrayRef<Directive>(Out4)),
+            (ArrayRef<Directive>{OMPD_target, OMPD_parallel, OMPD_loop}));
+}
+
 TEST(Composition, IsLeafConstruct) {
   ASSERT_TRUE(isLeafConstruct(OMPD_loop));
   ASSERT_TRUE(isLeafConstruct(OMPD_teams));

[skatrak]

Thank you Krzysztof for working on this. I have a suggestion to simplify the implementation quite a bit, but let me know if you have issues with it.

[skatrak]

My understanding is that there can only be at most a single composite construct when decomposing any given compound construct. So there wouldn't be a "first" composite range, but rather one composite range or none. In fact, if there is a composite range, it must also be at the end of the construct.

So, something like this: <composite-construct><other-leaf> can't happen. Decomposing a single compound construct into a list of leaves while keeping composite constructs as a unit should be as simple as forwarding all block-associated leaves until finding a loop-associated leaf and then aggregating (via getCompoundConstruct()) everything from there until the end of the leaf list to append one composite or single loop construct.

[kparzysz]

This is intended to be a general utility to find the first sequence of leaf constructs that together constitute a composite construct. The input leafs aren't assumed to represent a legal construct.
It's also used in isCompositeConstruct as well, where I break up the given directive into leafs and check if the "first composite range" is equal to the sequence of leafs.

[skatrak]

Is there any reason for these functions to allow unsupported lists of directives as input? I think it's worth having a simpler implementation over allowing all sorts of invalid lists of leaf constructs as input, unless there's a good reason why we'd need that.

Not saying that the compiler should crash in that case, but they can just return empty ranges for unsupported lists of leaf constructs rather than trying to find subsets that match what we're looking for.

[kparzysz]

I created the concept that this function implements as a common element in implementing other functions. This is a utility function that has a clearly defined functionality, that only depends on what is considered a "composite" construct. It is not an API function and it doesn't have to concern itself with validating the inputs, that's left to the API functions.
There are other ways to implement the breakdown of compound constructs, yours is likely another valid approach, but I've picked this one because of its flexibility.

[skatrak]

I see what you mean. It still feels a bit overkill to me, but I don't have a big issue with it since it's a private function to this compilation unit, so it's not a blocker.

[kparzysz]

Push review updates for another PR here by mistake...

[skatrak]

I think we should check here that Range.end() == Leafs.end() and return an empty list in that case (or some other way to signal an invalid input), because otherwise we'd be allowing something like [distribute, simd, parallel] to be passed in and then return [distribute-simd, parallel]. Or is this also intended to be very lenient and callers are responsible for only dealing with valid composite constructs? Since it's checked that there's actually a valid composite construct, I think this other check should be done as well.

Do we need to return the Output argument?

[kparzysz]

This starts with an actual directive D. It is then decomposed into either leaf or composite constructs, so presumably there shouldn't be a situation where there is a composite construct followed by something else. I think we should assert on that (instead of returning an empty list), because it would indicate some kind of an internal inconsistency and this is an API function.

We don't need to return Output---I meant it to be usable in cases like

for (auto &c : getLeafOrCompositeConstructs(D, Output)) {
  ...
}

[skatrak]

I see what you mean. It still feels a bit overkill to me, but I don't have a big issue with it since it's a private function to this compilation unit, so it's not a blocker.

[skatrak]

Thanks Krzysztof for addressing my concerns. This should be good to go after adding the assert to getLeafOrCompositeConstructs, so I'm approving it already to avoid blocking it any longer. Perhaps give it until tomorrow in case anyone else has any concerns.

[kparzysz]

Windows build failure is due to fatal error C1060: compiler is out of heap space.