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

Atomic RMW Xchg wasn't handled before when searching for known base
pointers in the IR.

Author: Il-Capitano

llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

@@ -577,8 +577,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,

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()