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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Csanád Hajdú (Il-Capitano)

<details>
<summary>Changes</summary>

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

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


2 Files Affected:

- (modified) llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp (+9-2) 
- (added) llvm/test/Transforms/RewriteStatepointsForGC/base-atomicrmw.ll (+14) 


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

``````````

</details>


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