[llvm] Be more aggressive with ObjCARCContract.cpp (PR #130630)

[via llvm-commits]

https://github.com/AZero13 created https://github.com/llvm/llvm-project/pull/130630

If Inst is a retain, we can be more aggressive by checking it

>From cafa666c4f61f99ff454f95f3d1f742fd6b8a173 Mon Sep 17 00:00:00 2001
From: Rose <gfunni234 at gmail.com>
Date: Mon, 10 Mar 2025 12:17:11 -0400
Subject: [PATCH] Be more aggressive with ObjCARCContract.cpp

If Inst is a retain, we can be more aggressive by checking it
---
 .../Transforms/ObjCARC/ObjCARCContract.cpp    | 22 +++++++++----
 llvm/test/Transforms/ObjCARC/contract.ll      | 33 ++++++++++++++++++-
 2 files changed, 47 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/Transforms/ObjCARC/ObjCARCContract.cpp b/llvm/lib/Transforms/ObjCARC/ObjCARCContract.cpp
index 311d3b1cfc0a0..2d8ad1ac66825 100644
--- a/llvm/lib/Transforms/ObjCARC/ObjCARCContract.cpp
+++ b/llvm/lib/Transforms/ObjCARC/ObjCARCContract.cpp
@@ -243,13 +243,21 @@ static StoreInst *findSafeStoreForStoreStrongContraction(LoadInst *Load,
 
     // Ok, now we know we have not seen a store yet.
 
-    // If Inst is a retain, we don't care about it as it doesn't prevent moving
-    // the load to the store.
-    //
-    // TODO: This is one area where the optimization could be made more
-    // aggressive.
-    if (IsRetain(Class))
-      continue;
+    // If Inst is a retain, we can be more aggressive by checking if:
+    // 1. The retain is on an unrelated object (different RC identity root)
+    // 2. The retain's argument doesn't alias with our load location
+    // In these cases, we can safely move past the retain.
+    if (IsRetain(Class)) {
+      Value *RetainArg = GetArgRCIdentityRoot(Inst);
+      Value *LoadRoot = GetRCIdentityRoot(Load);
+      if (RetainArg != LoadRoot && !AA->alias(MemoryLocation::get(Load), 
+          MemoryLocation::getBeforeOrAfter(RetainArg))) {
+        continue;
+      }
+      // If the retain is on the same object or we can't prove no aliasing,
+      // be conservative and bail out
+      return nullptr;
+    }
 
     // See if Inst can write to our load location, if it can not, just ignore
     // the instruction.
diff --git a/llvm/test/Transforms/ObjCARC/contract.ll b/llvm/test/Transforms/ObjCARC/contract.ll
index 70bd57a0c719a..869d6cca7b59f 100644
--- a/llvm/test/Transforms/ObjCARC/contract.ll
+++ b/llvm/test/Transforms/ObjCARC/contract.ll
@@ -32,7 +32,7 @@ entry:
   ret void
 }
 
-; Merge objc_retain and objc_autorelease into objc_retainAutorelease.
+; Merge objc_retain and objc_autorelease into objc_retainAutoreleasedReturnValue.
 
 ; CHECK-LABEL: define void @test2(
 ; CHECK: tail call ptr @llvm.objc.retainAutorelease(ptr %x) [[NUW:#[0-9]+]]
@@ -238,3 +238,34 @@ declare void @llvm.objc.clang.arc.use(...) nounwind
 declare void @llvm.objc.clang.arc.noop.use(...) nounwind
 
 ; CHECK: attributes [[NUW]] = { nounwind }
+
+; CHECK-LABEL: define void @test_aggressive_retain_opt(
+; CHECK: tail call ptr @llvm.objc.retainAutoreleasedReturnValue(ptr %p)
+; CHECK: }
+define void @test_aggressive_retain_opt() {
+  %p = call ptr @returner()
+  %unrelated = call ptr @returner()
+  tail call ptr @llvm.objc.retain(ptr %unrelated) nounwind
+  tail call ptr @llvm.objc.retain(ptr %p) nounwind
+  ret void
+}
+
+; CHECK-LABEL: define void @test_aggressive_retain_opt_alias(
+; CHECK: tail call ptr @llvm.objc.retain(ptr %p)
+; CHECK: }
+define void @test_aggressive_retain_opt_alias(ptr %p, ptr %q) {
+  store ptr %p, ptr %q
+  tail call ptr @llvm.objc.retain(ptr %p) nounwind
+  ret void
+}
+
+; CHECK-LABEL: define void @test_aggressive_retain_opt_no_alias(
+; CHECK: tail call ptr @llvm.objc.retainAutoreleasedReturnValue(ptr %p)
+; CHECK: }
+define void @test_aggressive_retain_opt_no_alias() {
+  %p = call ptr @returner()
+  %q = alloca ptr
+  store ptr null, ptr %q
+  tail call ptr @llvm.objc.retain(ptr %p) nounwind
+  ret void
+}