[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 2df1197d00ad47507560f928072399fc0e77e92b Mon Sep 17 00:00:00 2001
From: Nimish Mishra <neelam.nimish at gmail.com>
Date: Thu, 27 Mar 2025 17:57:26 +0530
Subject: [PATCH 1/2] 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..89a87823ac9cf 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();

>From 63d3d76295587b68a6859c5b283da64887aab572 Mon Sep 17 00:00:00 2001
From: Nimish Mishra <neelam.nimish at gmail.com>
Date: Thu, 27 Mar 2025 17:58:29 +0530
Subject: [PATCH 2/2] Delete stray newline

---
 llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp | 1 -
 1 file changed, 1 deletion(-)

diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 89a87823ac9cf..c0cd8fcdd058a 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -8761,7 +8761,6 @@ 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();