From 7e029ba1d5c3c33944fdf608b2cf965a99e7b227 Mon Sep 17 00:00:00 2001
From: Ruobing Han <hanruobing@gatech.edu>
Date: Sat, 1 Apr 2023 16:59:26 -0400
Subject: [PATCH] fix bug for AllocaInst that are wrapped in a PR

---
 .../src/x86/insert_warp_loop.cpp                 | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/compilation/KernelTranslation/src/x86/insert_warp_loop.cpp b/compilation/KernelTranslation/src/x86/insert_warp_loop.cpp
index 0819dcf..0c11c2f 100644
--- a/compilation/KernelTranslation/src/x86/insert_warp_loop.cpp
+++ b/compilation/KernelTranslation/src/x86/insert_warp_loop.cpp
@@ -319,6 +319,7 @@ void handle_local_variable_intra_warp(std::vector<ParallelRegion> PRs,
   // we should handle allocation generated by PHI
   {
     std::vector<llvm::Instruction *> instruction_to_fix;
+    std::vector<llvm::AllocaInst *> instruction_to_move;
     auto F = PRs[0].start_block->getParent();
     for (auto bb = F->begin(); bb != F->end(); bb++) {
       for (auto ii = bb->begin(); ii != bb->end(); ii++) {
@@ -343,6 +344,7 @@ void handle_local_variable_intra_warp(std::vector<ParallelRegion> PRs,
             }
           }
           if (allStoreNonDivergence) {
+            instruction_to_move.push_back(alloc);
             continue;
           }
           // Do not duplicate var used outside PRs
@@ -366,6 +368,8 @@ void handle_local_variable_intra_warp(std::vector<ParallelRegion> PRs,
           }
           if (!used_in_non_PR) {
             instruction_to_fix.push_back(alloc);
+          } else {
+            instruction_to_move.push_back(alloc);
           }
         }
       }
@@ -373,6 +377,18 @@ void handle_local_variable_intra_warp(std::vector<ParallelRegion> PRs,
     for (auto inst : instruction_to_fix) {
       AddContextSaveRestore(inst, intra_warp_loop);
     }
+    for (auto alloc : instruction_to_move) {
+      // need to move all allocInst to the entry basic block
+      IRBuilder<> builder(&*(alloc->getParent()
+                                 ->getParent()
+                                 ->getEntryBlock()
+                                 .getFirstInsertionPt()));
+      auto newAllocInst =
+          builder.CreateAlloca(alloc->getType()->getElementType(),
+                               alloc->getArraySize(), alloc->getName());
+      alloc->replaceAllUsesWith(newAllocInst);
+      alloc->removeFromParent();
+    }
   }
 
   for (auto parallel_regions : PRs) {