[llvm] [ObjCARC] Don't sink objc_retain past releasing atomics (PR #184113)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Marina Taylor (citymarina)

<details>
<summary>Changes</summary>

The releasing atomic may cause another thread to objc_release the pointer, so consider these instructions as potentially decrementing refcounts.

This patch is a synthesis of several AI models' outputs, as well as work by David Kilzer.

Co-authored-by: David Kilzer <ddkilzer@<!-- -->apple.com>
Assisted-by: claude

rdar://152185192

---
Full diff: https://github.com/llvm/llvm-project/pull/184113.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/ObjCARC/DependencyAnalysis.cpp (+10-1) 
- (added) llvm/test/Transforms/ObjCARC/sink-past-atomic.ll (+97) 


``````````diff
diff --git a/llvm/lib/Transforms/ObjCARC/DependencyAnalysis.cpp b/llvm/lib/Transforms/ObjCARC/DependencyAnalysis.cpp
index b4cc00033e720..9706093102dae 100644
--- a/llvm/lib/Transforms/ObjCARC/DependencyAnalysis.cpp
+++ b/llvm/lib/Transforms/ObjCARC/DependencyAnalysis.cpp
@@ -24,6 +24,7 @@
 #include "ProvenanceAnalysis.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/IR/CFG.h"
+#include "llvm/Support/AtomicOrdering.h"
 
 using namespace llvm;
 using namespace llvm::objcarc;
@@ -67,7 +68,15 @@ bool llvm::objcarc::CanDecrementRefCount(const Instruction *Inst,
                                          const Value *Ptr,
                                          ProvenanceAnalysis &PA,
                                          ARCInstKind Class) {
-  // First perform a quick check if Class can not touch ref counts.
+  // Atomic RMW and CmpXchg instructions with release or stronger ordering
+  // publish memory to other threads, which may then read the stored pointer and
+  // release it. Treat these as potentially decrementing refcounts.
+  if (const auto *RMW = dyn_cast<AtomicRMWInst>(Inst))
+    return isReleaseOrStronger(RMW->getOrdering());
+  if (const auto *CmpXchg = dyn_cast<AtomicCmpXchgInst>(Inst))
+    return isReleaseOrStronger(CmpXchg->getSuccessOrdering());
+
+  // Perform a quick check if Class can not touch ref counts.
   if (!CanDecrementRefCount(Class))
     return false;
 
diff --git a/llvm/test/Transforms/ObjCARC/sink-past-atomic.ll b/llvm/test/Transforms/ObjCARC/sink-past-atomic.ll
new file mode 100644
index 0000000000000..796909cb41c76
--- /dev/null
+++ b/llvm/test/Transforms/ObjCARC/sink-past-atomic.ll
@@ -0,0 +1,97 @@
+; RUN: opt -passes=objc-arc -S < %s | FileCheck %s
+
+target datalayout = "e-m:o-i64:64-i128:128-n32:64-S128"
+
+declare ptr @llvm.objc.retain(ptr)
+declare void @llvm.objc.release(ptr)
+
+; Retain must not sink past an atomicrmw with release or stronger ordering.
+
+define void @test_atomicrmw_release(ptr %obj, ptr %atomic_slot) {
+; CHECK-LABEL: @test_atomicrmw_release
+; CHECK: call ptr @llvm.objc.retain
+; CHECK: atomicrmw
+entry:
+  %obj_as_int = ptrtoint ptr %obj to i64
+  %retained = call ptr @llvm.objc.retain(ptr %obj)
+  %old_value = atomicrmw xchg ptr %atomic_slot, i64 %obj_as_int release, align 8
+  %old_obj = inttoptr i64 %old_value to ptr
+  call void @llvm.objc.release(ptr %old_obj)
+  call void @llvm.objc.release(ptr %obj)
+  ret void
+}
+
+define void @test_atomicrmw_acqrel(ptr %obj, ptr %atomic_slot) {
+; CHECK-LABEL: @test_atomicrmw_acqrel
+; CHECK: call ptr @llvm.objc.retain
+; CHECK: atomicrmw
+entry:
+  %obj_as_int = ptrtoint ptr %obj to i64
+  %retained = call ptr @llvm.objc.retain(ptr %obj)
+  %old_value = atomicrmw xchg ptr %atomic_slot, i64 %obj_as_int acq_rel, align 8
+  %old_obj = inttoptr i64 %old_value to ptr
+  call void @llvm.objc.release(ptr %old_obj)
+  call void @llvm.objc.release(ptr %obj)
+  ret void
+}
+
+define void @test_atomicrmw_seqcst(ptr %obj, ptr %atomic_slot) {
+; CHECK-LABEL: @test_atomicrmw_seqcst
+; CHECK: call ptr @llvm.objc.retain
+; CHECK: atomicrmw
+entry:
+  %obj_as_int = ptrtoint ptr %obj to i64
+  %retained = call ptr @llvm.objc.retain(ptr %obj)
+  %old_value = atomicrmw xchg ptr %atomic_slot, i64 %obj_as_int seq_cst, align 8
+  %old_obj = inttoptr i64 %old_value to ptr
+  call void @llvm.objc.release(ptr %old_obj)
+  call void @llvm.objc.release(ptr %obj)
+  ret void
+}
+
+; Retain must not sink past a cmpxchg with release or stronger success ordering.
+
+define void @test_cmpxchg_release(ptr %obj, ptr %atomic_slot, i64 %expected) {
+; CHECK-LABEL: @test_cmpxchg_release
+; CHECK: call ptr @llvm.objc.retain
+; CHECK: cmpxchg
+entry:
+  %obj_as_int = ptrtoint ptr %obj to i64
+  %retained = call ptr @llvm.objc.retain(ptr %obj)
+  %res = cmpxchg ptr %atomic_slot, i64 %expected, i64 %obj_as_int release seq_cst, align 8
+  %old_int = extractvalue { i64, i1 } %res, 0
+  %old_obj = inttoptr i64 %old_int to ptr
+  call void @llvm.objc.release(ptr %old_obj)
+  call void @llvm.objc.release(ptr %obj)
+  ret void
+}
+
+define void @test_cmpxchg_acqrel(ptr %obj, ptr %atomic_slot, i64 %expected) {
+; CHECK-LABEL: @test_cmpxchg_acqrel
+; CHECK: call ptr @llvm.objc.retain
+; CHECK: cmpxchg
+entry:
+  %obj_as_int = ptrtoint ptr %obj to i64
+  %retained = call ptr @llvm.objc.retain(ptr %obj)
+  %res = cmpxchg ptr %atomic_slot, i64 %expected, i64 %obj_as_int acq_rel seq_cst, align 8
+  %old_int = extractvalue { i64, i1 } %res, 0
+  %old_obj = inttoptr i64 %old_int to ptr
+  call void @llvm.objc.release(ptr %old_obj)
+  call void @llvm.objc.release(ptr %obj)
+  ret void
+}
+
+define void @test_cmpxchg_seqcst(ptr %obj, ptr %atomic_slot, i64 %expected) {
+; CHECK-LABEL: @test_cmpxchg_seqcst
+; CHECK: call ptr @llvm.objc.retain
+; CHECK: cmpxchg
+entry:
+  %obj_as_int = ptrtoint ptr %obj to i64
+  %retained = call ptr @llvm.objc.retain(ptr %obj)
+  %res = cmpxchg ptr %atomic_slot, i64 %expected, i64 %obj_as_int seq_cst seq_cst, align 8
+  %old_int = extractvalue { i64, i1 } %res, 0
+  %old_obj = inttoptr i64 %old_int to ptr
+  call void @llvm.objc.release(ptr %old_obj)
+  call void @llvm.objc.release(ptr %obj)
+  ret void
+}

``````````

</details>


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