[Mlir-commits] [mlir] [OpenMP][MLIR] Hoist static `alloca`s emitted by private `init` regions to the allocation IP of the construct (PR #171597)

[Kareem Ergawy]

================
@@ -1596,10 +1596,64 @@ static llvm::Expected<llvm::Value *> initPrivateVar(
   return phis[0];
 }
 
+/// Beginning with \p startBlock, this function visits all reachable successor
+/// blocks. For each such block, static alloca instructions (i.e. non-array
+/// allocas) are collected. Then, these collected alloca instructions are moved
+/// to the \p allocaIP insertion point.
+///
+/// This is useful in cases where, for example, more than one allocatable or
+/// array are privatized. In such cases, we allocate a number of temporary
+/// descriptors to handle the initialization logic. Additonally, for each
+/// private value, there is branching logic based on the value of the origianl
+/// private variable's allocation state. Therefore, we end up with descriptor
+/// alloca instructions preceded by conditional branches which casues runtime
+/// issues at least on the GPU.
+static void hoistStaticAllocasToAllocaIP(
+    llvm::BasicBlock *startBlock,
+    const llvm::OpenMPIRBuilder::InsertPointTy &allocaIP) {
+  llvm::SmallVector<llvm::BasicBlock *> inlinedBlocks{startBlock};
+  llvm::SmallPtrSet<llvm::BasicBlock *, 4> seenBlocks;
+  llvm::SmallVector<llvm::Instruction *> staticAllocas;
+
+  while (!inlinedBlocks.empty()) {
+    llvm::BasicBlock *curBlock = inlinedBlocks.front();
+    inlinedBlocks.erase(inlinedBlocks.begin());
+    llvm::Instruction *terminator = curBlock->getTerminator();
+
+    for (llvm::Instruction &inst : *curBlock) {
+      if (auto *allocaInst = mlir::dyn_cast<llvm::AllocaInst>(&inst)) {
+        if (!allocaInst->isArrayAllocation()) {
----------------
ergawy wrote:

Can you elaborate a bit further? The problem is not specific to a single `init` region but with inlining multiple regions one after the other.

Maybe the solution I am proposing here is not general enough but I think it does cover the Fortran case which is related to descriptors. Maybe a more general solution would be to check dominance info and move the alloca only if (1) its blocks post-dominates the entry block and (2) any SSA arguments to the alloca are within that same block?

https://github.com/llvm/llvm-project/pull/171597