[llvm-branch-commits] [clang] [clang][OpenMP] Rewrite `getOpenMPCaptureRegions` in term of leafs (PR #97110)

[Krzysztof Parzyszek via llvm-branch-commits]

https://github.com/kparzysz created https://github.com/llvm/llvm-project/pull/97110

Replace the switch in `getOpenMPCaptureRegions` with a loop collecting capture regions based on the constituent directives.

>From 2d25e0d32672ecae3dc3ad42c50446e651eceb06 Mon Sep 17 00:00:00 2001
From: Krzysztof Parzyszek <Krzysztof.Parzyszek at amd.com>
Date: Fri, 28 Jun 2024 15:27:42 -0500
Subject: [PATCH] [clang][OpenMP] Rewrite `getOpenMPCaptureRegions` in term of
 leafs

Replace the switch in `getOpenMPCaptureRegions` with a loop collecting
capture regions based on the constituent directives.
---
 clang/lib/Basic/OpenMPKinds.cpp | 170 ++++++++++++++------------------
 1 file changed, 72 insertions(+), 98 deletions(-)

diff --git a/clang/lib/Basic/OpenMPKinds.cpp b/clang/lib/Basic/OpenMPKinds.cpp
index 30c34c207ae23..152891dfa27dc 100644
--- a/clang/lib/Basic/OpenMPKinds.cpp
+++ b/clang/lib/Basic/OpenMPKinds.cpp
@@ -747,105 +747,79 @@ void clang::getOpenMPCaptureRegions(
   assert(unsigned(DKind) < llvm::omp::Directive_enumSize);
   assert(isOpenMPCapturingDirective(DKind));
 
-  switch (DKind) {
-  case OMPD_metadirective:
-    CaptureRegions.push_back(OMPD_metadirective);
-    break;
-  case OMPD_parallel:
-  case OMPD_parallel_for:
-  case OMPD_parallel_for_simd:
-  case OMPD_parallel_master:
-  case OMPD_parallel_masked:
-  case OMPD_parallel_sections:
-  case OMPD_distribute_parallel_for:
-  case OMPD_distribute_parallel_for_simd:
-  case OMPD_parallel_loop:
-    CaptureRegions.push_back(OMPD_parallel);
-    break;
-  case OMPD_target_teams:
-  case OMPD_target_teams_distribute:
-  case OMPD_target_teams_distribute_simd:
-    CaptureRegions.push_back(OMPD_task);
-    CaptureRegions.push_back(OMPD_target);
-    CaptureRegions.push_back(OMPD_teams);
-    break;
-  case OMPD_teams:
-  case OMPD_teams_distribute:
-  case OMPD_teams_distribute_simd:
-    CaptureRegions.push_back(OMPD_teams);
-    break;
-  case OMPD_target:
-  case OMPD_target_simd:
-    CaptureRegions.push_back(OMPD_task);
-    CaptureRegions.push_back(OMPD_target);
-    break;
-  case OMPD_teams_loop:
-  case OMPD_teams_distribute_parallel_for:
-  case OMPD_teams_distribute_parallel_for_simd:
-    CaptureRegions.push_back(OMPD_teams);
-    CaptureRegions.push_back(OMPD_parallel);
-    break;
-  case OMPD_target_parallel:
-  case OMPD_target_parallel_for:
-  case OMPD_target_parallel_for_simd:
-  case OMPD_target_parallel_loop:
-    CaptureRegions.push_back(OMPD_task);
-    CaptureRegions.push_back(OMPD_target);
-    CaptureRegions.push_back(OMPD_parallel);
-    break;
-  case OMPD_task:
-  case OMPD_target_enter_data:
-  case OMPD_target_exit_data:
-  case OMPD_target_update:
-    CaptureRegions.push_back(OMPD_task);
-    break;
-  case OMPD_taskloop:
-  case OMPD_taskloop_simd:
-  case OMPD_master_taskloop:
-  case OMPD_master_taskloop_simd:
-  case OMPD_masked_taskloop:
-  case OMPD_masked_taskloop_simd:
-    CaptureRegions.push_back(OMPD_taskloop);
-    break;
-  case OMPD_parallel_masked_taskloop:
-  case OMPD_parallel_masked_taskloop_simd:
-  case OMPD_parallel_master_taskloop:
-  case OMPD_parallel_master_taskloop_simd:
-    CaptureRegions.push_back(OMPD_parallel);
-    CaptureRegions.push_back(OMPD_taskloop);
-    break;
-  case OMPD_target_teams_loop:
-  case OMPD_target_teams_distribute_parallel_for:
-  case OMPD_target_teams_distribute_parallel_for_simd:
-    CaptureRegions.push_back(OMPD_task);
-    CaptureRegions.push_back(OMPD_target);
-    CaptureRegions.push_back(OMPD_teams);
-    CaptureRegions.push_back(OMPD_parallel);
-    break;
-  case OMPD_nothing:
-    CaptureRegions.push_back(OMPD_nothing);
-    break;
-  case OMPD_loop:
-    // TODO: 'loop' may require different capture regions depending on the bind
-    // clause or the parent directive when there is no bind clause. Use
-    // OMPD_unknown for now.
-  case OMPD_simd:
-  case OMPD_for:
-  case OMPD_for_simd:
-  case OMPD_sections:
-  case OMPD_single:
-  case OMPD_taskgroup:
-  case OMPD_distribute:
-  case OMPD_ordered:
-  case OMPD_target_data:
-  case OMPD_distribute_simd:
-  case OMPD_scope:
-  case OMPD_dispatch:
+  auto getRegionsForLeaf = [&](OpenMPDirectiveKind LKind) {
+    assert(isLeafConstruct(LKind) && "Epecting leaf directive");
+    switch (LKind) {
+    case OMPD_metadirective:
+      CaptureRegions.push_back(OMPD_metadirective);
+      break;
+    case OMPD_nothing:
+      CaptureRegions.push_back(OMPD_nothing);
+      break;
+    case OMPD_parallel:
+      CaptureRegions.push_back(OMPD_parallel);
+      break;
+    case OMPD_target:
+      CaptureRegions.push_back(OMPD_task);
+      CaptureRegions.push_back(OMPD_target);
+      break;
+    case OMPD_task:
+    case OMPD_target_enter_data:
+    case OMPD_target_exit_data:
+    case OMPD_target_update:
+      CaptureRegions.push_back(OMPD_task);
+      break;
+    case OMPD_teams:
+      CaptureRegions.push_back(OMPD_teams);
+      break;
+    case OMPD_taskloop:
+      CaptureRegions.push_back(OMPD_taskloop);
+      break;
+    case OMPD_loop:
+      // TODO: 'loop' may require different capture regions depending on the
+      // bind clause or the parent directive when there is no bind clause.
+      // If any of the directives that push regions here are parents of 'loop',
+      // assume 'parallel'. Otherwise do nothing.
+      if (!CaptureRegions.empty() &&
+          !llvm::is_contained(CaptureRegions, OMPD_parallel))
+        CaptureRegions.push_back(OMPD_parallel);
+      break;
+    case OMPD_dispatch:
+    case OMPD_distribute:
+    case OMPD_for:
+    case OMPD_masked:
+    case OMPD_master:
+    case OMPD_ordered:
+    case OMPD_scope:
+    case OMPD_sections:
+    case OMPD_simd:
+    case OMPD_single:
+    case OMPD_target_data:
+    case OMPD_taskgroup:
+      // These directives (when standalone) use OMPD_unknown as the region,
+      // but when they're constituents of a compound directive, and other
+      // leafs from that directive have specific regions, then these directives
+      // add no additional regions.
+      break;
+    default:
+      llvm::errs() << getOpenMPDirectiveName(LKind) << '\n';
+      llvm_unreachable("Unexpected directive");
+    }
+  };
+
+  for (OpenMPDirectiveKind L : getLeafConstructsOrSelf(DKind))
+    getRegionsForLeaf(L);
+
+  // If no regions were added, then only the leafs using "unknown" were
+  // present. Push a single OMPD_unknown as the capture region.
+  if (CaptureRegions.empty())
     CaptureRegions.push_back(OMPD_unknown);
-    break;
-  default:
-    llvm_unreachable("Unhandled OpenMP directive");
-  }
+
+  // OMPD_unknown is only expected as the only region. If other regions
+  // are present OMPD_unknown should not be present.
+  assert((CaptureRegions[0] == OMPD_unknown ||
+          !llvm::is_contained(CaptureRegions, OMPD_unknown)) &&
+         "Misplaced OMPD_unknown");
 }
 
 bool clang::checkFailClauseParameter(OpenMPClauseKind FailClauseParameter) {