Skip to content

fix(rocshmem): use scoped atomics in Anvil SDMA device code #11591

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
dlamd1dai wants to merge 1 commit into from
Closed

fix(rocshmem): use scoped atomics in Anvil SDMA device code #11591

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 32 additions & 21 deletions projects/rocshmem/src/sdma/anvil_device.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -161,11 +161,26 @@ __device__ __forceinline__ SDMA_PKT_FENCE_64B_MI4 CreateFence64BPacketMI4(uint64

#endif // SDMA_IS_OSS7

// RCCL device builds poison __hip_atomic_*; use explicit scoped atomics instead.
__device__ __forceinline__ int anvilScopedMemoryScope(int hipScope) {
return hipScope == __HIP_MEMORY_SCOPE_SYSTEM ? __MEMORY_SCOPE_SYSTEM : __MEMORY_SCOPE_DEVICE;
}

template <typename T>
__device__ __forceinline__ T anvilScopedLoad(T* ptr, int order, int hipScope) {
return __scoped_atomic_load_n(ptr, order, anvilScopedMemoryScope(hipScope));
}

template <typename T>
__device__ __forceinline__ void anvilScopedStore(T* ptr, T val, int order, int hipScope) {
__scoped_atomic_store_n(ptr, val, order, anvilScopedMemoryScope(hipScope));
}

// Spin-poll until *addr >= expected (agent-scope relaxed load)
template <int64_t MAX_SPIN_COUNT = -1>
__device__ __forceinline__ void poll_until_ge(uint64_t* addr, uint64_t expected) {
[[maybe_unused]] int64_t spin_count = 0;
while (__hip_atomic_load(addr, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT) < expected) {
while (anvilScopedLoad(addr, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT) < expected) {
spin_count++;
assert(MAX_SPIN_COUNT < 0 || spin_count != MAX_SPIN_COUNT);
}
Expand All @@ -186,7 +201,7 @@ struct SdmaQueueDeviceHandle {
return true;
}
// Only read hardware register if the queue is full based on cached index
cachedHwReadIndex = __hip_atomic_load(rptr, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
cachedHwReadIndex = anvilScopedLoad(rptr, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
__atomic_signal_fence(__ATOMIC_SEQ_CST);
return (uptoIndex - cachedHwReadIndex) < queue_size_in_bytes;
}
Expand All @@ -199,7 +214,7 @@ struct SdmaQueueDeviceHandle {
int retries = 0;

while (true) {
cur_index = __hip_atomic_load(cachedWptr, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
cur_index = anvilScopedLoad(cachedWptr, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
offset = 0;

// Wraparound and Pad NOPs on remaining bytes
Expand All @@ -209,9 +224,8 @@ struct SdmaQueueDeviceHandle {
uint64_t new_index = cur_index + size_in_bytes + offset;

if (CanWriteUpto(new_index)) {
if (__hip_atomic_compare_exchange_strong(cachedWptr, &cur_index, new_index,
__ATOMIC_RELAXED, __ATOMIC_RELAXED,
__HIP_MEMORY_SCOPE_AGENT)) {
if (__scoped_atomic_compare_exchange_n(cachedWptr, &cur_index, new_index, false, __ATOMIC_RELAXED,
__ATOMIC_RELAXED, __MEMORY_SCOPE_DEVICE)) {
break;
}
}
Expand Down Expand Up @@ -240,23 +254,22 @@ struct SdmaQueueDeviceHandle {
// First DWORD encodes the NOP count per SDMA spec; remaining DWORDs are zero
for (uint32_t i = 0; i < numOffsetDwords; i++) {
uint32_t val = (i == 0) ? (((numOffsetDwords - 1) & 0xFFFF) << 16) : 0;
__hip_atomic_store(queueBuf + base_index_in_dwords + i, val, __ATOMIC_RELAXED,
__HIP_MEMORY_SCOPE_AGENT);
anvilScopedStore(queueBuf + base_index_in_dwords + i, val, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
}
pendingWptr += offset;
base_index_in_dwords = WrapIntoRing(pendingWptr) / sizeof(uint32_t);

for (uint32_t i = 0; i < numDwords; i++) {
__hip_atomic_store(queueBuf + base_index_in_dwords + i, packetPtr[i], __ATOMIC_RELAXED,
__HIP_MEMORY_SCOPE_AGENT);
anvilScopedStore(queueBuf + base_index_in_dwords + i, packetPtr[i], __ATOMIC_RELAXED,
__HIP_MEMORY_SCOPE_AGENT);
}
pendingWptr += sizeof(PacketType);
}

__device__ __forceinline__ void submitPacket(uint64_t base, uint64_t pendingWptr) {
int retries = 0;
while (true) {
uint64_t val = __hip_atomic_load(committedWptr, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
uint64_t val = anvilScopedLoad(committedWptr, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
__atomic_signal_fence(__ATOMIC_SEQ_CST);
if (val == base) {
// All stores inside the loop to avoid SIMD reconvergence deadlock:
Expand All @@ -270,7 +283,7 @@ struct SdmaQueueDeviceHandle {
__builtin_amdgcn_wave_barrier();
__atomic_signal_fence(__ATOMIC_SEQ_CST);

__hip_atomic_store(wptr, pendingWptr, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
anvilScopedStore(wptr, pendingWptr, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);

#if defined(__GFX12__)
asm volatile("s_wait_loadcnt 0x0\n s_wait_storecnt 0x0" ::: "memory");
Expand All @@ -280,7 +293,7 @@ struct SdmaQueueDeviceHandle {
__builtin_amdgcn_wave_barrier();
__atomic_signal_fence(__ATOMIC_SEQ_CST);

__hip_atomic_store(doorbell, pendingWptr, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_SYSTEM);
anvilScopedStore(doorbell, pendingWptr, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_SYSTEM);

#if defined(__GFX12__)
asm volatile("s_wait_loadcnt 0x0\n s_wait_storecnt 0x0" ::: "memory");
Expand All @@ -290,7 +303,7 @@ struct SdmaQueueDeviceHandle {
__builtin_amdgcn_wave_barrier();
__atomic_signal_fence(__ATOMIC_SEQ_CST);

__hip_atomic_store(committedWptr, pendingWptr, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
anvilScopedStore(committedWptr, pendingWptr, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);

#if defined(__GFX12__)
asm volatile("s_wait_loadcnt 0x0\n s_wait_storecnt 0x0" ::: "memory");
Expand All @@ -302,8 +315,7 @@ struct SdmaQueueDeviceHandle {

// Relaxed agent-scope atomic store writes directly to GL2, making
// maxWritePtr visible to quiet callers on other CUs.
__hip_atomic_store(&maxWritePtr, pendingWptr, __ATOMIC_RELAXED,
__HIP_MEMORY_SCOPE_AGENT);
anvilScopedStore(&maxWritePtr, pendingWptr, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
break;
}
__builtin_amdgcn_s_sleep(1);
Expand All @@ -321,19 +333,18 @@ struct SdmaQueueDeviceHandle {
__device__ __forceinline__ void flushTo(uint64_t upToIndex) {
uint64_t hw_read_index;
do {
hw_read_index = __hip_atomic_load(rptr, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
hw_read_index = anvilScopedLoad(rptr, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
} while (hw_read_index < upToIndex);
}

// Spin-poll rptr until HW has consumed all submitted packets
__device__ __forceinline__ void quietAll() {
// One agent-scope load to read maxWritePtr set by a potentially different CU.
// Held in a register for the loop — it does not need updated during quietAll.
uint64_t target = __hip_atomic_load(&maxWritePtr, __ATOMIC_RELAXED,
__HIP_MEMORY_SCOPE_AGENT);
uint64_t target = anvilScopedLoad(&maxWritePtr, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
uint64_t hw_read_index;
do {
hw_read_index = __hip_atomic_load(rptr, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
hw_read_index = anvilScopedLoad(rptr, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
} while (hw_read_index < target);
}

Expand Down Expand Up @@ -609,7 +620,7 @@ __device__ __forceinline__ void quiet(SdmaQueueSingleProducerDeviceHandle& handl
__device__ __forceinline__ bool waitForSignal(HSAuint64* addr, uint64_t expected) {
int retries = 0;
while (true) {
uint64_t value = __hip_atomic_load(addr, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
uint64_t value = anvilScopedLoad(addr, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
if (value >= expected) { // >= not == to avoid infinite spin if signal overshoots
return true;
}
Expand Down
Loading