[llvm] [Statepoint] Treat result of atomicrmw xchg as a base pointer (PR #97280)

[Csanád Hajdú via llvm-commits]

https://github.com/Il-Capitano updated https://github.com/llvm/llvm-project/pull/97280

>From 62b479da96470430967f00c9e3aced46a2d17794 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Csan=C3=A1d=20Hajd=C3=BA?= <csanad.hajdu at arm.com>
Date: Thu, 20 Jun 2024 16:09:17 +0200
Subject: [PATCH] [Statepoint] Treat result of atomicrmw xchg as a base pointer

Atomic RMW Xchg wasn't handled before when searching for known base
pointers in the IR.
---
 .../Transforms/Scalar/RewriteStatepointsForGC.cpp  | 11 +++++++++--
 .../RewriteStatepointsForGC/base-atomicrmw.ll      | 14 ++++++++++++++
 2 files changed, 23 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/Transforms/RewriteStatepointsForGC/base-atomicrmw.ll

diff --git a/llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp b/llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp
index 2b99e28acb4e9..bd823d33989ff 100644
--- a/llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp
+++ b/llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp
@@ -578,8 +578,15 @@ static Value *findBaseDefiningValue(Value *I, DefiningValueMapTy &Cache,
     return I;
   }
 
-  assert(!isa<AtomicRMWInst>(I) && "Xchg handled above, all others are "
-                                   "binary ops which don't apply to pointers");
+  if (auto *RMWI = dyn_cast<AtomicRMWInst>(I)) {
+    assert(RMWI->getOperation() == AtomicRMWInst::Xchg &&
+           "Only Xchg is allowed for pointer values");
+    // A RMW Xchg is a combined atomic load and store, so we can treat the
+    // loaded value as a base pointer.
+    Cache[I] = I;
+    setKnownBase(I, /* IsKnownBase */ true, KnownBases);
+    return I;
+  }
 
   // The aggregate ops.  Aggregates can either be in the heap or on the
   // stack, but in either case, this is simply a field load.  As a result,
diff --git a/llvm/test/Transforms/RewriteStatepointsForGC/base-atomicrmw.ll b/llvm/test/Transforms/RewriteStatepointsForGC/base-atomicrmw.ll
new file mode 100644
index 0000000000000..9d33a5760ce61
--- /dev/null
+++ b/llvm/test/Transforms/RewriteStatepointsForGC/base-atomicrmw.ll
@@ -0,0 +1,14 @@
+; RUN: opt < %s -passes=rewrite-statepoints-for-gc -S 2>&1 | FileCheck %s
+
+define ptr addrspace(1) @test(ptr %a, ptr addrspace(1) %b) gc "statepoint-example" {
+; CHECK-LABEL: @test
+; CHECK-NEXT:    [[RES:%.*]] = atomicrmw xchg ptr %a, ptr addrspace(1) %b seq_cst
+; CHECK-NEXT:    [[STATEPOINT_TOKEN:%.*]] = call token (i64, i32, ptr, i32, i32, ...) @llvm.experimental.gc.statepoint.p0(i64 2882400000, i32 0, ptr elementtype(void ()) @foo, i32 0, i32 0, i32 0, i32 0) [ "gc-live"(ptr addrspace(1) [[RES]]) ]
+; CHECK-NEXT:    [[RES_RELOCATED:%.*]] = call coldcc ptr addrspace(1) @llvm.experimental.gc.relocate.p1(token [[STATEPOINT_TOKEN]], i32 0, i32 0)
+; CHECK-NEXT:    ret ptr addrspace(1) [[RES_RELOCATED]]
+  %res = atomicrmw xchg ptr %a, ptr addrspace(1) %b seq_cst
+  call void @foo()
+  ret ptr addrspace(1) %res
+}
+
+declare void @foo()