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

[via llvm-commits]

Author: Csanád Hajdú
Date: 2024-11-08T10:23:27-08:00
New Revision: e53c46a9084caac115d7f694e5f16f904b0d7124

URL: https://github.com/llvm/llvm-project/commit/e53c46a9084caac115d7f694e5f16f904b0d7124
DIFF: https://github.com/llvm/llvm-project/commit/e53c46a9084caac115d7f694e5f16f904b0d7124.diff

LOG: [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.

Added: 
    llvm/test/Transforms/RewriteStatepointsForGC/base-atomicrmw.ll

Modified: 
    llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp b/llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp
index ef7ef8ef7911e0..4a39a484e41c29 100644
--- a/llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp
+++ b/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,

diff  --git a/llvm/test/Transforms/RewriteStatepointsForGC/base-atomicrmw.ll b/llvm/test/Transforms/RewriteStatepointsForGC/base-atomicrmw.ll
new file mode 100644
index 00000000000000..9d33a5760ce619
--- /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()