[llvm] a815208 - [Attributor][FIX] Ensure new BBs are registered

[Johannes Doerfert via llvm-commits]

Author: Johannes Doerfert
Date: 2023-11-01T12:12:14-07:00
New Revision: a8152086ff422d6679ce7cba96bba777eeddd60b

URL: https://github.com/llvm/llvm-project/commit/a8152086ff422d6679ce7cba96bba777eeddd60b
DIFF: https://github.com/llvm/llvm-project/commit/a8152086ff422d6679ce7cba96bba777eeddd60b.diff

LOG: [Attributor][FIX] Ensure new BBs are registered

Added: 
    openmp/libomptarget/test/offloading/generic_multiple_parallel_regions.c

Modified: 
    llvm/lib/Transforms/IPO/AttributorAttributes.cpp
    llvm/lib/Transforms/IPO/OpenMPOpt.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/IPO/AttributorAttributes.cpp b/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
index 21d83edbb18c447..c698d1babbf6354 100644
--- a/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
+++ b/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
@@ -12405,11 +12405,14 @@ struct AAIndirectCallInfoCallSite : public AAIndirectCallInfo {
       Instruction *ThenTI =
           SplitBlockAndInsertIfThen(LastCmp, IP, /* Unreachable */ false);
       BasicBlock *CBBB = CB->getParent();
+      A.registerManifestAddedBasicBlock(*ThenTI->getParent());
+      A.registerManifestAddedBasicBlock(*CBBB);
       auto *SplitTI = cast<BranchInst>(LastCmp->getNextNode());
       BasicBlock *ElseBB;
       if (IP == CB) {
         ElseBB = BasicBlock::Create(ThenTI->getContext(), "",
                                     ThenTI->getFunction(), CBBB);
+        A.registerManifestAddedBasicBlock(*ElseBB);
         IP = BranchInst::Create(CBBB, ElseBB);
         SplitTI->replaceUsesOfWith(CBBB, ElseBB);
       } else {

diff  --git a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
index 3e9705d61e8e080..74ebbcde5729247 100644
--- a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
+++ b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
@@ -4567,6 +4567,8 @@ struct AAKernelInfoFunction : AAKernelInfo {
       BasicBlock *PRNextBB =
           BasicBlock::Create(Ctx, "worker_state_machine.parallel_region.check",
                              Kernel, StateMachineEndParallelBB);
+      A.registerManifestAddedBasicBlock(*PRExecuteBB);
+      A.registerManifestAddedBasicBlock(*PRNextBB);
 
       // Check if we need to compare the pointer at all or if we can just
       // call the parallel region function.

diff  --git a/openmp/libomptarget/test/offloading/generic_multiple_parallel_regions.c b/openmp/libomptarget/test/offloading/generic_multiple_parallel_regions.c
new file mode 100644
index 000000000000000..33fe2ca127b6dc4
--- /dev/null
+++ b/openmp/libomptarget/test/offloading/generic_multiple_parallel_regions.c
@@ -0,0 +1,29 @@
+// RUN: %libomptarget-compile-run-and-check-generic
+// RUN: %libomptarget-compileopt-run-and-check-generic
+
+#include <omp.h>
+#include <stdio.h>
+
+__attribute__((optnone)) void optnone() {}
+
+int main() {
+  int i = 0;
+#pragma omp target teams num_teams(1) map(tofrom : i)
+  {
+    optnone();
+#pragma omp parallel
+    if (omp_get_thread_num() == 0)
+      ++i;
+#pragma omp parallel
+    if (omp_get_thread_num() == 0)
+      ++i;
+#pragma omp parallel
+    if (omp_get_thread_num() == 0)
+      ++i;
+#pragma omp parallel
+    if (omp_get_thread_num() == 0)
+      ++i;
+  }
+  // CHECK: 4
+  printf("%i\n", i);
+}