AI/rocm-systems
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Revert "Use relaxed atomics for LL on GFX11 (#859)" (#1148)

Browse Source 
This reverts commit 6a0a6a37d9.

Use inline asm for 128b load on GFX11 for better peformance.
This commit is contained in:
Wenkai Du
2024-04-26 07:49:55 -07:00
committed by GitHub
parent 0ec41f1386
commit f330b82985
1 changed files with 47 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/device/prims_ll.h
+47 -28
View File
@@ -10,18 +10,6 @@
#include "npkit/npkit.h"
#endif
#ifdef __GFX11__
#define LL_STORE(SRC, DST) \
__atomic_store_n((DST), (SRC), __ATOMIC_RELAXED)
#define LL_LOAD(SRC) \
__atomic_load_n(SRC, __ATOMIC_RELAXED)
#else
#define LL_STORE(SRC, DST) \
__builtin_nontemporal_store((SRC), (DST))
#define LL_LOAD(SRC) \
__builtin_nontemporal_load(SRC)
#endif
template<typename T, typename RedOp, typename Fan, int Direct, int P2p>
class Primitives<T, RedOp, Fan, Direct, ProtoLL, P2p>:
public PrimitivesWithoutDirect<Primitives<T, RedOp, Fan, Direct, ProtoLL, P2p>> {
@@ -162,8 +150,13 @@ private:
#if defined(__HIP_PLATFORM_AMD__) || defined(__HCC__) || defined(__HIPCC__)
union ncclLLFifoLine i4;
do {
i4.v[0] = LL_LOAD(src->v);
i4.v[1] = LL_LOAD(src->v+1);
#ifdef __GFX11__
asm volatile ("global_load_b128 %0, %1, off glc slc dlc\n"
"s_waitcnt vmcnt(0)\n" : "=v"(i4.i4) : "v"(&src->i4));
#else
i4.v[0] = __builtin_nontemporal_load(src->v);
i4.v[1] = __builtin_nontemporal_load(src->v+1);
#endif
#if defined(ENABLE_NPKIT) && (defined(ENABLE_NPKIT_EVENT_PRIM_LL_DATA_PROCESS_ENTRY) && defined(ENABLE_NPKIT_EVENT_PRIM_LL_DATA_PROCESS_EXIT) || defined(ENABLE_NPKIT_PRIM_COLLECT_DATA_PROCESS_TIME))
npkitWaitRecvSpins++;
#endif
@@ -198,8 +191,13 @@ private:
if (i < fan.nrecv()) {
union ncclLLFifoLine* src = recvPtr(i) + offset;
#if defined(__HIP_PLATFORM_AMD__) || defined(__HCC__) || defined(__HIPCC__)
line[i].v[0] = LL_LOAD(src->v);
line[i].v[1] = LL_LOAD(src->v+1);
#ifdef __GFX11__
asm volatile ("global_load_b128 %0, %1, off glc slc dlc\n"
"s_waitcnt vmcnt(0)\n" : "=v"(line[i].i4) : "v"(&src->i4));
#else
line[i].v[0] = __builtin_nontemporal_load(src->v);
line[i].v[1] = __builtin_nontemporal_load(src->v+1);
#endif
#else
asm("ld.volatile.global.v4.u32 {%0,%1,%2,%3}, [%4];" : "=r"(line[i].data1), "=r"(line[i].flag1), "=r"(line[i].data2), "=r"(line[i].flag2) : "l"(&src->i4));
#endif
@@ -220,8 +218,13 @@ private:
do {
#if defined(__HIP_PLATFORM_AMD__) || defined(__HCC__) || defined(__HIPCC__)
line[i].v[0] = LL_LOAD(src->v);
line[i].v[1] = LL_LOAD(src->v+1);
#ifdef __GFX11__
asm volatile ("global_load_b128 %0, %1, off glc slc dlc\n"
"s_waitcnt vmcnt(0)\n" : "=v"(line[i].i4) : "v"(&src->i4));
#else
line[i].v[0] = __builtin_nontemporal_load(src->v);
line[i].v[1] = __builtin_nontemporal_load(src->v+1);
#endif
#else
asm("ld.volatile.global.v4.u32 {%0,%1,%2,%3}, [%4];" : "=r"(line[i].data1), "=r"(line[i].flag1), "=r"(line[i].data2), "=r"(line[i].flag2) : "l"(&src->i4));
#endif
@@ -249,8 +252,8 @@ private:
i4.flag1 = flag;
i4.data2 = (val >> 32);
i4.flag2 = flag;
LL_STORE(i4.v[0], dst->v);
LL_STORE(i4.v[1], dst->v+1);
__builtin_nontemporal_store(i4.v[0], dst->v);
__builtin_nontemporal_store(i4.v[1], dst->v+1);
#else
asm volatile("st.volatile.global.v4.u32 [%0], {%1,%2,%3,%4};" :: "l"(&dst->i4), "r"((uint32_t)val), "r"(flag), "r"((uint32_t)(val >> 32)), "r"(flag));
#endif
@@ -269,13 +272,29 @@ private:
};
#if defined(__HIP_PLATFORM_AMD__) || defined(__HCC__) || defined(__HIPCC__)
if(sizeof(U) == 1)
u1 = LL_LOAD((uint8_t*)src);
#ifdef __GFX11__
u1 = __atomic_load_n((uint8_t*)src, __ATOMIC_RELAXED);
#else
u1 = __builtin_nontemporal_load((uint8_t*)src);
#endif
else if(sizeof(U) == 2)
u2 = LL_LOAD((uint16_t*)src);
#ifdef __GFX11__
u2 = __atomic_load_n((uint16_t*)src, __ATOMIC_RELAXED);
#else
u2 = __builtin_nontemporal_load((uint16_t*)src);
#endif
else if(sizeof(U) == 4)
u4 = LL_LOAD((uint32_t*)src);
#ifdef __GFX11__
u4 = __atomic_load_n((uint32_t*)src, __ATOMIC_RELAXED);
#else
u4 = __builtin_nontemporal_load((uint32_t*)src);
#endif
else
u8 = LL_LOAD((uint64_t*)src);
#ifdef __GFX11__
u8 = __atomic_load_n((uint64_t*)src, __ATOMIC_RELAXED);
#else
u8 = __builtin_nontemporal_load((uint64_t*)src);
#endif
#else
if(sizeof(U) == 1)
asm("ld.volatile.global.b8 %0,[%1];" : "=r"(u4) : "l"(src));
@@ -301,13 +320,13 @@ private:
elt = val;
#if defined(__HIP_PLATFORM_AMD__) || defined(__HCC__) || defined(__HIPCC__)
if(sizeof(U) == 1)
LL_STORE(u1, (uint8_t*)dst);
__builtin_nontemporal_store(u1, (uint8_t*)dst);
else if(sizeof(U) == 2)
LL_STORE(u2, (uint16_t*)dst);
__builtin_nontemporal_store(u2, (uint16_t*)dst);
else if(sizeof(U) == 4)
LL_STORE(u4, (uint32_t*)dst);
__builtin_nontemporal_store(u4, (uint32_t*)dst);
else
LL_STORE(u8, (uint64_t*)dst);
__builtin_nontemporal_store(u8, (uint64_t*)dst);
#else
if(sizeof(U) == 1)
asm("st.volatile.global.b8 [%0],%1;" :: "l"(dst), "r"(u4));