[llvm] [mlir] [Utils][mlir] Fix interaction between CodeExtractor and OpenMPIRBuilder (PR #145051)

[Kajetan Puchalski via llvm-commits]

================
@@ -1868,8 +1866,14 @@ CallInst *CodeExtractor::emitReplacerCall(
   // This takes place of the original loop
   BasicBlock *codeReplacer =
       BasicBlock::Create(Context, "codeRepl", oldFunction, ReplIP);
+  // In cases with multiple levels of outlining, e.g. with OpenMP,
+  // AllocationBlock may end up in a function different than oldFunction. We
+  // need to make sure we do not use it in those cases, otherwise the alloca
+  // will end up in a different function from its users and break the module.
   BasicBlock *AllocaBlock =
-      AllocationBlock ? AllocationBlock : &oldFunction->getEntryBlock();
+      (AllocationBlock && oldFunction == AllocationBlock->getParent())
+          ? AllocationBlock
+          : &oldFunction->getEntryBlock();
----------------
mrkajetanp wrote:

If we go up the chain, the problem originates over here:
https://github.com/llvm/llvm-project/blob/fccab5d757778204666d70e2f1592952fc8b336d/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp#L501
`findAllocaInsertPoint` does a stackWalk to find a preexisting alloca insert point, and it does find it - but it's in a different function. We could put a check here to make sure that we only use it if it's in the same function as `builder.getInsertBlock()`? 
This would fix the issue, but fixing it inside CodeExtractor would make it it more robust for other cases when a wrong block might be passed on accident. I think I'd lean towards keeping it as-is on that account, but if other people have strong preferences or arguments to the contrary then I am happy to change it.

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