Ke Wen
7e51592129
Add support for bfloat16. Add ncclAvg reduction operation. Improve performance for aggregated operations. Improve performance for tree. Improve network error reporting. Add NCCL_NET parameter to force a specific network. Add NCCL_IB_QPS_PER_CONNECTION parameter to split IB traffic onto multiple queue pairs. Fix topology detection error in WSL2. Fix proxy memory elements affinity (improve alltoall performance). Fix graph search on cubemesh topologies. Fix hang in cubemesh during NVB connections.
69 lines
2.2 KiB
C++
69 lines
2.2 KiB
C++
/*************************************************************************
|
|
* Copyright (c) 2016-2019, NVIDIA CORPORATION. All rights reserved.
|
|
*
|
|
* See LICENSE.txt for license information
|
|
************************************************************************/
|
|
|
|
#ifndef OP128_H_
|
|
#define OP128_H_
|
|
|
|
inline __device__ void load128(const uint64_t* ptr, uint64_t &v0, uint64_t &v1) {
|
|
asm volatile("ld.volatile.global.v2.u64 {%0,%1}, [%2];"
|
|
: "=l"(v0), "=l"(v1) : "l"(ptr));
|
|
}
|
|
|
|
inline __device__ void store128(uint64_t* ptr, uint64_t v0, uint64_t v1) {
|
|
asm volatile("st.volatile.global.v2.u64 [%2], {%0,%1};"
|
|
:: "l"(v0), "l"(v1), "l"(ptr));
|
|
}
|
|
|
|
inline __device__ uint64_t* shmemCvtPtr(volatile uint64_t* shmemGenericPtr) {
|
|
uint64_t* shmemAsmPtr;
|
|
asm volatile("cvta.to.shared.u64 %0, %1;" : "=l"(shmemAsmPtr) : "l"(shmemGenericPtr));
|
|
return shmemAsmPtr;
|
|
}
|
|
|
|
inline __device__ void loadShmem128(uint64_t* shmemAsmPtr, uint64_t &v0, uint64_t &v1) {
|
|
asm volatile("ld.volatile.shared.v2.u64 {%0,%1}, [%2];"
|
|
: "=l"(v0), "=l"(v1) : "l"(shmemAsmPtr));
|
|
}
|
|
|
|
inline __device__ void storeShmem128(uint64_t* shmemAsmPtr, uint64_t v0, uint64_t v1) {
|
|
asm volatile("st.volatile.shared.v2.u64 [%2], {%0,%1};"
|
|
:: "l"(v0), "l"(v1), "l"(shmemAsmPtr));
|
|
}
|
|
|
|
template<typename T>
|
|
inline __device__ void loadShmemMisaligned128(T *ptr, uint64_t &v0, uint64_t &v1) {
|
|
union {
|
|
uint32_t tmp4[4];
|
|
uint64_t tmp8[2];
|
|
};
|
|
if(sizeof(T) < 4) {
|
|
uint32_t *ptr4 = reinterpret_cast<uint32_t*>(reinterpret_cast<uintptr_t>(ptr) & -uintptr_t(4));
|
|
#pragma unroll
|
|
for(int e=0; e < 4; e++) {
|
|
// Produce 4 bytes of sub-register type by reading 2 4-byte
|
|
// aligned values and shifting.
|
|
uint32_t lo, hi;
|
|
asm("ld.shared.b32 %0,[%1];" : "=r"(lo) : "l"(ptr4+e+0));
|
|
asm("ld.shared.b32 %0,[%1];" : "=r"(hi) : "l"(ptr4+e+1));
|
|
tmp4[e] = __funnelshift_r(lo, hi, 8*(int(reinterpret_cast<uintptr_t>(ptr))%4));
|
|
}
|
|
}
|
|
else if(sizeof(T) == 4) {
|
|
#pragma unroll
|
|
for(int e=0; e < 4; e++)
|
|
asm("ld.shared.b32 %0,[%1];" : "=r"(tmp4[e]) : "l"(ptr+e));
|
|
}
|
|
else /*sizeof(T)==8*/ {
|
|
#pragma unroll
|
|
for(int e=0; e < 2; e++)
|
|
asm("ld.shared.b64 %0,[%1];" : "=l"(tmp8[e]) : "l"(ptr+e));
|
|
}
|
|
v0 = tmp8[0];
|
|
v1 = tmp8[1];
|
|
}
|
|
|
|
#endif
|