[flang] [llvm] [mlir][llvm][OpenMP] Hoist __atomic_load alloca (PR #132888)

[via llvm-commits]

https://github.com/NimishMishra updated https://github.com/llvm/llvm-project/pull/132888

>From 978e715c39ca0fc06f25a65cc2fde241d0f17b63 Mon Sep 17 00:00:00 2001
From: Nimish Mishra <neelam.nimish at gmail.com>
Date: Thu, 27 Mar 2025 17:55:35 +0530
Subject: [PATCH] Emit stack save/restore for alloca in __atomic_load

---
 flang/test/Integration/OpenMP/atomic-capture-complex.f90 | 2 ++
 llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp                | 3 +++
 2 files changed, 5 insertions(+)

diff --git a/flang/test/Integration/OpenMP/atomic-capture-complex.f90 b/flang/test/Integration/OpenMP/atomic-capture-complex.f90
index 4ffd18097d79e..55035e9012309 100644
--- a/flang/test/Integration/OpenMP/atomic-capture-complex.f90
+++ b/flang/test/Integration/OpenMP/atomic-capture-complex.f90
@@ -16,6 +16,7 @@
 !CHECK: br label %entry
 
 !CHECK: entry:
+!CHECK: %[[STACK_SAVE_PTR:.*]] = call ptr @llvm.stacksave.p0()
 !CHECK: %[[ATOMIC_TEMP_LOAD:.*]] = alloca { float, float }, align 8
 !CHECK: call void @__atomic_load(i64 8, ptr %[[ORIG_VAL]], ptr %[[ATOMIC_TEMP_LOAD]], i32 0)
 !CHECK: %[[PHI_NODE_ENTRY_1:.*]] = load { float, float }, ptr %[[ATOMIC_TEMP_LOAD]], align 8
@@ -33,6 +34,7 @@
 !CHECK: %[[VAL_11:.*]] = call i1 @__atomic_compare_exchange(i64 8, ptr %[[ORIG_VAL]], ptr %[[ATOMIC_TEMP_LOAD]], ptr %[[X_NEW_VAL]],
 !i32 2, i32 2)
 !CHECK: %[[VAL_12:.*]] = load { float, float }, ptr %[[ATOMIC_TEMP_LOAD]], align 4
+!CHECK: call void @llvm.stackrestore.p0(ptr %[[STACK_SAVE_PTR]])
 !CHECK: br i1 %[[VAL_11]], label %.atomic.exit, label %.atomic.cont
 
 !CHECK:   .atomic.exit
diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 2e5ce5308eea5..1d05cbc4b0ae5 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -8749,6 +8749,7 @@ Expected<std::pair<Value *, Value *>> OpenMPIRBuilder::emitAtomicUpdate(
       Res.second = emitRMWOpAsInstruction(Res.first, Expr, RMWOp);
   } else if (RMWOp == llvm::AtomicRMWInst::BinOp::BAD_BINOP &&
              XElemTy->isStructTy()) {
+    auto *StackSave = Builder.CreateStackSave();
     LoadInst *OldVal =
         Builder.CreateLoad(XElemTy, X, X->getName() + ".atomic.load");
     OldVal->setAtomic(AO);
@@ -8760,6 +8761,7 @@ Expected<std::pair<Value *, Value *>> OpenMPIRBuilder::emitAtomicUpdate(
         &Builder, XElemTy, LoadSize * 8, LoadSize * 8, OldVal->getAlign(),
         OldVal->getAlign(), true /* UseLibcall */, X);
     auto AtomicLoadRes = atomicInfo.EmitAtomicLoadLibcall(AO);
+    
     BasicBlock *CurBB = Builder.GetInsertBlock();
     Instruction *CurBBTI = CurBB->getTerminator();
     CurBBTI = CurBBTI ? CurBBTI : Builder.CreateUnreachable();
@@ -8785,6 +8787,7 @@ Expected<std::pair<Value *, Value *>> OpenMPIRBuilder::emitAtomicUpdate(
     auto Result = atomicInfo.EmitAtomicCompareExchangeLibcall(
         AtomicLoadRes.second, NewAtomicAddr, AO, Failure);
     LoadInst *PHILoad = Builder.CreateLoad(XElemTy, Result.first);
+    Builder.CreateStackRestore(StackSave);
     PHI->addIncoming(PHILoad, Builder.GetInsertBlock());
     Builder.CreateCondBr(Result.second, ExitBB, ContBB);
     OldVal->eraseFromParent();