[Attributor][FIX] Ensure new BBs are registered

llvm · Nov 1, 2023 · a815208 · a815208
1 parent 3426d33
commit a815208
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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);
+}