Enable LL128 protocol support (#605)
* Enable LL128 protocol support * Use shared memory object directly when possible
This commit is contained in:
Wenkai Du
GitHub
@@ -574,7 +574,6 @@ template<typename T, typename RedOp>
|
||||
struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE> {
|
||||
__device__ __attribute__((noinline)) void run(ncclWorkElem *args) {
|
||||
using Proto = ProtoSimple<ALLREDUCE_CHUNKSTEPS/ALLREDUCE_SLICESTEPS, ALLREDUCE_SLICESTEPS>;
|
||||
if (threadIdx.x == 0) __insert_timestamp(__LINE__);
|
||||
runRing<T, RedOp, Proto>(args);
|
||||
}
|
||||
};
|
||||
@@ -582,7 +581,6 @@ struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SI
|
||||
template<typename T, typename RedOp>
|
||||
struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_TREE, NCCL_PROTO_SIMPLE> {
|
||||
__device__ __attribute__((noinline)) void run(ncclWorkElem *args) {
|
||||
if (threadIdx.x == 0) __insert_timestamp(__LINE__);
|
||||
runTreeUpDown<T, RedOp, ProtoSimple<1, 1>>(args);
|
||||
}
|
||||
};
|
||||
@@ -688,7 +686,6 @@ struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_COLLNET, NCCL_PROTO
|
||||
template<typename T, typename RedOp>
|
||||
struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_LL> {
|
||||
__device__ __attribute__((noinline)) void run(ncclWorkElem *args) {
|
||||
if (threadIdx.x == 0) __insert_timestamp(__LINE__);
|
||||
runRing<T, RedOp, ProtoLL>(args);
|
||||
}
|
||||
};
|
||||
@@ -696,7 +693,6 @@ struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_LL
|
||||
template<typename T, typename RedOp>
|
||||
struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_TREE, NCCL_PROTO_LL> {
|
||||
__device__ __attribute__((noinline)) void run(ncclWorkElem *args) {
|
||||
if (threadIdx.x == 0) __insert_timestamp(__LINE__);
|
||||
if (args->pad_0 == 0) runTreeUpDown<T, RedOp, ProtoLL>(args);
|
||||
else runTreeSplit<T, RedOp, ProtoLL>(args);
|
||||
}
|
||||
|
||||
+207
-111
@@ -102,26 +102,112 @@ static const __device__ constexpr ncclKernelFunc_t ncclFuncs[]{
|
||||
#endif
|
||||
};
|
||||
|
||||
template<unsigned short f, unsigned short l>
|
||||
#define NCCL_FUNC5_LL128(func, algo, devredop, type, nullify) \
|
||||
MACRO_IF(nullify, nullptr, NCCL_FUNC_NAME(func, algo, LL, devredop, type)), \
|
||||
MACRO_IF(nullify, nullptr, NCCL_FUNC_NAME(func, algo, LL128, devredop, type)), \
|
||||
MACRO_IF(nullify, nullptr, NCCL_FUNC_NAME(func, algo, SIMPLE, devredop, type))
|
||||
|
||||
#define NCCL_FUNC4_LL128(func, devredop, type, nullify) \
|
||||
NCCL_FUNC5_LL128(func, TREE, devredop, type, nullify), \
|
||||
NCCL_FUNC5_LL128(func, RING, devredop, type, nullify), \
|
||||
NCCL_FUNC5_LL128(func, COLLNET, devredop, type, nullify)
|
||||
|
||||
// Must be consistent with ncclDataType_t
|
||||
#define NCCL_FUNCS3A_LL128(func, devredop, nullForFloat) \
|
||||
NCCL_FUNC4_LL128(func, devredop, int8_t, 0), \
|
||||
NCCL_FUNC4_LL128(func, devredop, uint8_t, 0), \
|
||||
NCCL_FUNC4_LL128(func, devredop, int32_t, 0), \
|
||||
NCCL_FUNC4_LL128(func, devredop, uint32_t, 0), \
|
||||
NCCL_FUNC4_LL128(func, devredop, int64_t, 0), \
|
||||
NCCL_FUNC4_LL128(func, devredop, uint64_t, 0), \
|
||||
NCCL_FUNC4_LL128(func, devredop, half, nullForFloat), \
|
||||
NCCL_FUNC4_LL128(func, devredop, float, nullForFloat), \
|
||||
NCCL_FUNC4_LL128(func, devredop, double, nullForFloat), \
|
||||
NCCL_FUNC4_LL128(func, devredop, rccl_bfloat16, nullForFloat)
|
||||
#define NCCL_FUNCS3B_LL128(func, devredop) \
|
||||
NCCL_FUNC4_LL128(func, devredop, int8_t, 0), \
|
||||
NCCL_FUNC4_LL128(func, devredop, int8_t, 0), \
|
||||
NCCL_FUNC4_LL128(func, devredop, int8_t, 0), \
|
||||
NCCL_FUNC4_LL128(func, devredop, int8_t, 0), \
|
||||
NCCL_FUNC4_LL128(func, devredop, int8_t, 0), \
|
||||
NCCL_FUNC4_LL128(func, devredop, int8_t, 0), \
|
||||
NCCL_FUNC4_LL128(func, devredop, int8_t, 0), \
|
||||
NCCL_FUNC4_LL128(func, devredop, int8_t, 0), \
|
||||
NCCL_FUNC4_LL128(func, devredop, int8_t, 0), \
|
||||
NCCL_FUNC4_LL128(func, devredop, int8_t, 0)
|
||||
|
||||
// Must be consistent with ncclRedOp_t
|
||||
#define NCCL_FUNCS2A_LL128(func) \
|
||||
NCCL_FUNCS3A_LL128(func, Sum, /*nullForFloat=*/0), \
|
||||
NCCL_FUNCS3A_LL128(func, Prod, /*nullForFloat=*/0), \
|
||||
NCCL_FUNCS3A_LL128(func, Max, /*nullForFloat=*/0), \
|
||||
NCCL_FUNCS3A_LL128(func, Min, /*nullForFloat=*/0), \
|
||||
NCCL_FUNCS3A_LL128(func, PreMulSum, /*nullForFloat=*/0), \
|
||||
NCCL_FUNCS3A_LL128(func, SumPostDiv, /*nullForFloat=*/1)
|
||||
|
||||
#define NCCL_FUNCS2B_LL128(func) \
|
||||
NCCL_FUNCS3B_LL128(func, Sum), \
|
||||
NCCL_FUNCS3B_LL128(func, Sum), \
|
||||
NCCL_FUNCS3B_LL128(func, Sum), \
|
||||
NCCL_FUNCS3B_LL128(func, Sum), \
|
||||
NCCL_FUNCS3B_LL128(func, Sum), \
|
||||
NCCL_FUNCS3B_LL128(func, Sum)
|
||||
|
||||
// Must be consistent with the ncclFuncSet enum
|
||||
using ncclKernelFunc_t = void (*)();
|
||||
|
||||
static const __device__ constexpr ncclKernelFunc_t ncclFuncs_ll128[]{
|
||||
// Don't try to initialize the host shadow copy of this device-side global
|
||||
// variable. There is no host pointer to a device-side function, which
|
||||
// confuses clang. This will be fixed in the next clang release.
|
||||
#if defined(__HIP_DEVICE_COMPILE__)
|
||||
#if defined(BUILD_ALLREDUCE_ONLY)
|
||||
NCCL_FUNC4_LL128(AllReduce, Sum, float, 0),
|
||||
#else
|
||||
NCCL_FUNCS2B_LL128(Broadcast),
|
||||
NCCL_FUNCS2A_LL128(Reduce),
|
||||
NCCL_FUNCS2B_LL128(AllGather),
|
||||
NCCL_FUNCS2A_LL128(ReduceScatter),
|
||||
NCCL_FUNCS2A_LL128(AllReduce),
|
||||
NCCL_ONERANK_REDUCE_NAME(PreMulSum, int8_t),
|
||||
NCCL_ONERANK_REDUCE_NAME(PreMulSum, uint8_t),
|
||||
NCCL_ONERANK_REDUCE_NAME(PreMulSum, int32_t),
|
||||
NCCL_ONERANK_REDUCE_NAME(PreMulSum, uint32_t),
|
||||
NCCL_ONERANK_REDUCE_NAME(PreMulSum, int64_t),
|
||||
NCCL_ONERANK_REDUCE_NAME(PreMulSum, uint64_t),
|
||||
NCCL_ONERANK_REDUCE_NAME(PreMulSum, half),
|
||||
NCCL_ONERANK_REDUCE_NAME(PreMulSum, float),
|
||||
NCCL_ONERANK_REDUCE_NAME(PreMulSum, double),
|
||||
#if defined(RCCL_BFLOAT16)
|
||||
NCCL_ONERANK_REDUCE_NAME(PreMulSum, rccl_bfloat16),
|
||||
#endif
|
||||
NCCL_FUNC_NAME(SendRecv, RING, SIMPLE, Sum, int8_t),
|
||||
NCCL_FUNC_NAME(AllToAllPivot, RING, SIMPLE, Sum, int8_t),
|
||||
#endif
|
||||
#endif
|
||||
};
|
||||
|
||||
template<unsigned short f, unsigned short l, bool u>
|
||||
struct Caller {
|
||||
static __device__ __host__
|
||||
void call(unsigned short funcIndex) noexcept
|
||||
{
|
||||
constexpr unsigned short m = f + (l - f) / 2;
|
||||
|
||||
return (funcIndex < m) ? Caller<f, m>::call(funcIndex) : Caller<m, l>::call(funcIndex);
|
||||
return (funcIndex < m) ? Caller<f, m, u>::call(funcIndex) : Caller<m, l, u>::call(funcIndex);
|
||||
}
|
||||
};
|
||||
|
||||
template<unsigned short f>
|
||||
struct Caller<f, f + 1>{
|
||||
template<unsigned short f, bool u>
|
||||
struct Caller<f, f + 1, u>{
|
||||
static __device__ __host__
|
||||
void call(unsigned short funcIndex) noexcept { ncclFuncs[f](); }
|
||||
void call(unsigned short funcIndex) noexcept { if (u) ncclFuncs_ll128[f](); else ncclFuncs[f](); }
|
||||
};
|
||||
|
||||
static_assert(FUNC_INDEX_P2P == 2710, "Wrong P2P function index");
|
||||
static_assert(FUNC_INDEX_ALLTOALL_PIVOT == 2711, "Wrong AllToAllPivot function index");
|
||||
|
||||
template<bool USING_LL128>
|
||||
inline
|
||||
__device__
|
||||
void NCCL_CALL_FUNCTIONS(unsigned short funcIndex) noexcept {
|
||||
@@ -130,43 +216,59 @@ void NCCL_CALL_FUNCTIONS(unsigned short funcIndex) noexcept {
|
||||
ncclFunction_AllReduce_RING_SIMPLE_Sum_float();
|
||||
else if (funcIndex == FUNC_INDEX(ncclFuncAllReduce, ncclSum, ncclFloat32, NCCL_ALGO_RING, NCCL_PROTO_LL))
|
||||
ncclFunction_AllReduce_RING_LL_Sum_float();
|
||||
else if (funcIndex == FUNC_INDEX(ncclFuncAllReduce, ncclSum, ncclFloat32, NCCL_ALGO_RING, NCCL_PROTO_LL128))
|
||||
else if (USING_LL128 && funcIndex == FUNC_INDEX(ncclFuncAllReduce, ncclSum, ncclFloat32, NCCL_ALGO_RING, NCCL_PROTO_LL128))
|
||||
ncclFunction_AllReduce_RING_LL128_Sum_float();
|
||||
else if (!USING_LL128 && funcIndex == FUNC_INDEX(ncclFuncAllReduce, ncclSum, ncclFloat32, NCCL_ALGO_RING, NCCL_PROTO_LL128))
|
||||
ncclFunction_AllReduce_RING_LL_Sum_float();
|
||||
else if (funcIndex == FUNC_INDEX(ncclFuncAllReduce, ncclSum, ncclFloat32, NCCL_ALGO_TREE, NCCL_PROTO_SIMPLE))
|
||||
ncclFunction_AllReduce_TREE_SIMPLE_Sum_float();
|
||||
else if (funcIndex == FUNC_INDEX(ncclFuncAllReduce, ncclSum, ncclFloat32, NCCL_ALGO_TREE, NCCL_PROTO_LL))
|
||||
ncclFunction_AllReduce_TREE_LL_Sum_float();
|
||||
else if (USING_LL128 && funcIndex == FUNC_INDEX(ncclFuncAllReduce, ncclSum, ncclFloat32, NCCL_ALGO_TREE, NCCL_PROTO_LL128))
|
||||
ncclFunction_AllReduce_TREE_LL128_Sum_float();
|
||||
else if (!USING_LL128 && funcIndex == FUNC_INDEX(ncclFuncAllReduce, ncclSum, ncclFloat32, NCCL_ALGO_TREE, NCCL_PROTO_LL128))
|
||||
ncclFunction_AllReduce_TREE_LL_Sum_float();
|
||||
else if (funcIndex == FUNC_INDEX(ncclFuncAllReduce, ncclSum, ncclFloat32, NCCL_ALGO_COLLNET, NCCL_PROTO_SIMPLE))
|
||||
ncclFunction_AllReduce_COLLNET_SIMPLE_Sum_float();
|
||||
else if (funcIndex == FUNC_INDEX(ncclFuncAllReduce, ncclSum, ncclFloat32, NCCL_ALGO_COLLNET, NCCL_PROTO_LL))
|
||||
ncclFunction_AllReduce_COLLNET_LL_Sum_float();
|
||||
else if (USING_LL128 && funcIndex == FUNC_INDEX(ncclFuncAllReduce, ncclSum, ncclFloat32, NCCL_ALGO_COLLNET, NCCL_PROTO_LL128))
|
||||
ncclFunction_AllReduce_COLLNET_LL128_Sum_float();
|
||||
else if (!USING_LL128 && funcIndex == FUNC_INDEX(ncclFuncAllReduce, ncclSum, ncclFloat32, NCCL_ALGO_COLLNET, NCCL_PROTO_LL128))
|
||||
ncclFunction_AllReduce_COLLNET_LL_Sum_float();
|
||||
else
|
||||
assert("Unsupported function index");
|
||||
#else
|
||||
if (funcIndex < 540) {
|
||||
if (funcIndex % 9 == 0) ncclFunction_Broadcast_TREE_LL_Sum_int8_t();
|
||||
else if (funcIndex % 9 == 1) ncclFunction_Broadcast_TREE_LL_Sum_int8_t();
|
||||
else if (USING_LL128 && funcIndex % 9 == 1) ncclFunction_Broadcast_TREE_LL128_Sum_int8_t();
|
||||
else if (!USING_LL128 && funcIndex % 9 == 1) ncclFunction_Broadcast_TREE_LL_Sum_int8_t();
|
||||
else if (funcIndex % 9 == 2) ncclFunction_Broadcast_TREE_SIMPLE_Sum_int8_t();
|
||||
else if (funcIndex % 9 == 3) ncclFunction_Broadcast_RING_LL_Sum_int8_t();
|
||||
else if (funcIndex % 9 == 4) ncclFunction_Broadcast_RING_LL_Sum_int8_t();
|
||||
else if (USING_LL128 && funcIndex % 9 == 4) ncclFunction_Broadcast_RING_LL128_Sum_int8_t();
|
||||
else if (!USING_LL128 && funcIndex % 9 == 4) ncclFunction_Broadcast_RING_LL_Sum_int8_t();
|
||||
else if (funcIndex % 9 == 5) ncclFunction_Broadcast_RING_SIMPLE_Sum_int8_t();
|
||||
else if (funcIndex % 9 == 6) ncclFunction_Broadcast_COLLNET_LL_Sum_int8_t();
|
||||
else if (funcIndex % 9 == 7) ncclFunction_Broadcast_COLLNET_LL_Sum_int8_t();
|
||||
else if (USING_LL128 && funcIndex % 9 == 7) ncclFunction_Broadcast_COLLNET_LL128_Sum_int8_t();
|
||||
else if (!USING_LL128 && funcIndex % 9 == 7) ncclFunction_Broadcast_COLLNET_LL_Sum_int8_t();
|
||||
else ncclFunction_Broadcast_COLLNET_SIMPLE_Sum_int8_t();
|
||||
}
|
||||
else if (funcIndex < 1080) Caller<540, 1080>::call(funcIndex);
|
||||
else if (funcIndex < 1080) Caller<540, 1080, USING_LL128>::call(funcIndex);
|
||||
else if (funcIndex < 1620) {
|
||||
if (funcIndex % 9 == 0) ncclFunction_AllGather_TREE_LL_Sum_int8_t();
|
||||
else if (funcIndex % 9 == 1) ncclFunction_AllGather_TREE_LL_Sum_int8_t();
|
||||
else if (USING_LL128 && funcIndex % 9 == 1) ncclFunction_AllGather_TREE_LL128_Sum_int8_t();
|
||||
else if (!USING_LL128 && funcIndex % 9 == 1) ncclFunction_AllGather_TREE_LL_Sum_int8_t();
|
||||
else if (funcIndex % 9 == 2) ncclFunction_AllGather_TREE_SIMPLE_Sum_int8_t();
|
||||
else if (funcIndex % 9 == 3) ncclFunction_AllGather_RING_LL_Sum_int8_t();
|
||||
else if (funcIndex % 9 == 4) ncclFunction_AllGather_RING_LL_Sum_int8_t();
|
||||
else if (USING_LL128 && funcIndex % 9 == 4) ncclFunction_AllGather_RING_LL128_Sum_int8_t();
|
||||
else if (!USING_LL128 && funcIndex % 9 == 4) ncclFunction_AllGather_RING_LL_Sum_int8_t();
|
||||
else if (funcIndex % 9 == 5) ncclFunction_AllGather_RING_SIMPLE_Sum_int8_t();
|
||||
else if (funcIndex % 9 == 6) ncclFunction_AllGather_COLLNET_LL_Sum_int8_t();
|
||||
else if (funcIndex % 9 == 7) ncclFunction_AllGather_COLLNET_LL_Sum_int8_t();
|
||||
else if (USING_LL128 && funcIndex % 9 == 7) ncclFunction_AllGather_COLLNET_LL128_Sum_int8_t();
|
||||
else if (!USING_LL128 && funcIndex % 9 == 7) ncclFunction_AllGather_COLLNET_LL_Sum_int8_t();
|
||||
else ncclFunction_AllGather_COLLNET_SIMPLE_Sum_int8_t();
|
||||
}
|
||||
else if (funcIndex < 2700) Caller<1620, 2700>::call(funcIndex);
|
||||
else if (funcIndex < 2700) Caller<1620, 2700, USING_LL128>::call(funcIndex);
|
||||
else {
|
||||
switch (funcIndex - 2700) {
|
||||
case 0:
|
||||
@@ -219,52 +321,52 @@ class ncclFunction {
|
||||
|
||||
#ifdef ENABLE_COLLTRACE
|
||||
#define traceColl(elem,launch_type) \
|
||||
uint32_t pos = __atomic_fetch_add(ncclShmem->comm.collTraceTail, 1, __ATOMIC_SEQ_CST)%COLLTRACE_NUM_ITEMS; \
|
||||
ncclShmem->comm.collTrace[pos].timeStamp = __builtin_amdgcn_s_memrealtime(); \
|
||||
ncclShmem->comm.collTrace[pos].bid = blockIdx.x; \
|
||||
ncclShmem->comm.collTrace[pos].funcIndex = ncclShmem->work.header.funcIndex; \
|
||||
asm volatile ("s_getreg_b32 %0, hwreg(HW_REG_HW_ID)" : "=s" (ncclShmem->comm.collTrace[pos].data_0)); \
|
||||
uint32_t pos = __atomic_fetch_add(shmem.comm.collTraceTail, 1, __ATOMIC_SEQ_CST)%COLLTRACE_NUM_ITEMS; \
|
||||
shmem.comm.collTrace[pos].timeStamp = __builtin_amdgcn_s_memrealtime(); \
|
||||
shmem.comm.collTrace[pos].bid = blockIdx.x; \
|
||||
shmem.comm.collTrace[pos].funcIndex = shmem.work.header.funcIndex; \
|
||||
asm volatile ("s_getreg_b32 %0, hwreg(HW_REG_HW_ID)" : "=s" (shmem.comm.collTrace[pos].data_0)); \
|
||||
if (elem.header.type == ncclWorkTypeP2p) { \
|
||||
struct ncclWorkElemP2p *p2pElems = (struct ncclWorkElemP2p *)&elem; \
|
||||
ncclShmem->comm.collTrace[pos].p2p[0].connIndex = p2pElems[0].connIndex; \
|
||||
ncclShmem->comm.collTrace[pos].p2pOpCount[0] = p2pElems[0].opCount; \
|
||||
ncclShmem->comm.collTrace[pos].p2p[0].ngroups = p2pElems[0].ngroups; \
|
||||
ncclShmem->comm.collTrace[pos].p2p[0].nWarps = p2pElems[0].nWarps; \
|
||||
ncclShmem->comm.collTrace[pos].p2p[0].warpStart = p2pElems[0].warpStart; \
|
||||
ncclShmem->comm.collTrace[pos].p2p[0].peer = (uint16_t)(p2pElems[0].peer); \
|
||||
ncclShmem->comm.collTrace[pos].p2p[1].connIndex = p2pElems[1].connIndex; \
|
||||
ncclShmem->comm.collTrace[pos].p2pOpCount[1] = p2pElems[1].opCount; \
|
||||
ncclShmem->comm.collTrace[pos].p2p[1].ngroups = p2pElems[1].ngroups; \
|
||||
ncclShmem->comm.collTrace[pos].p2p[1].nWarps = p2pElems[1].nWarps; \
|
||||
ncclShmem->comm.collTrace[pos].p2p[1].warpStart = p2pElems[1].warpStart; \
|
||||
ncclShmem->comm.collTrace[pos].p2p[1].peer = (uint16_t)(p2pElems[1].peer); \
|
||||
ncclShmem->comm.collTrace[pos].type = (ncclCollTraceP2pElemType|launch_type); \
|
||||
shmem.comm.collTrace[pos].p2p[0].connIndex = p2pElems[0].connIndex; \
|
||||
shmem.comm.collTrace[pos].p2pOpCount[0] = p2pElems[0].opCount; \
|
||||
shmem.comm.collTrace[pos].p2p[0].ngroups = p2pElems[0].ngroups; \
|
||||
shmem.comm.collTrace[pos].p2p[0].nWarps = p2pElems[0].nWarps; \
|
||||
shmem.comm.collTrace[pos].p2p[0].warpStart = p2pElems[0].warpStart; \
|
||||
shmem.comm.collTrace[pos].p2p[0].peer = (uint16_t)(p2pElems[0].peer); \
|
||||
shmem.comm.collTrace[pos].p2p[1].connIndex = p2pElems[1].connIndex; \
|
||||
shmem.comm.collTrace[pos].p2pOpCount[1] = p2pElems[1].opCount; \
|
||||
shmem.comm.collTrace[pos].p2p[1].ngroups = p2pElems[1].ngroups; \
|
||||
shmem.comm.collTrace[pos].p2p[1].nWarps = p2pElems[1].nWarps; \
|
||||
shmem.comm.collTrace[pos].p2p[1].warpStart = p2pElems[1].warpStart; \
|
||||
shmem.comm.collTrace[pos].p2p[1].peer = (uint16_t)(p2pElems[1].peer); \
|
||||
shmem.comm.collTrace[pos].type = (ncclCollTraceP2pElemType|launch_type); \
|
||||
} else { \
|
||||
ncclShmem->comm.collTrace[pos].opCount = elem.opCount; \
|
||||
ncclShmem->comm.collTrace[pos].coll.nWarps = elem.header.nWarps; \
|
||||
ncclShmem->comm.collTrace[pos].coll.bid = elem.bid; \
|
||||
ncclShmem->comm.collTrace[pos].coll.nChannels = elem.nChannels; \
|
||||
ncclShmem->comm.collTrace[pos].type = (ncclCollTraceCollElemType|launch_type); \
|
||||
shmem.comm.collTrace[pos].opCount = elem.opCount; \
|
||||
shmem.comm.collTrace[pos].coll.nWarps = elem.header.nWarps; \
|
||||
shmem.comm.collTrace[pos].coll.bid = elem.bid; \
|
||||
shmem.comm.collTrace[pos].coll.nChannels = elem.nChannels; \
|
||||
shmem.comm.collTrace[pos].type = (ncclCollTraceCollElemType|launch_type); \
|
||||
}
|
||||
|
||||
#define traceKernelLaunch(elem,firstLaunch) { \
|
||||
traceColl(elem,(firstLaunch?ncclCollTraceKernelLaunchType:ncclCollTraceCollLaunchType)); \
|
||||
}
|
||||
#define traceKernelEnd() { \
|
||||
uint32_t pos = __atomic_fetch_add(ncclShmem->comm.collTraceTail, 1, __ATOMIC_SEQ_CST)%COLLTRACE_NUM_ITEMS; \
|
||||
ncclShmem->comm.collTrace[pos].timeStamp = __builtin_amdgcn_s_memrealtime(); \
|
||||
ncclShmem->comm.collTrace[pos].bid = bid; \
|
||||
ncclShmem->comm.collTrace[pos].type = ncclCollTraceKernelEndType; \
|
||||
uint32_t pos = __atomic_fetch_add(shmem.comm.collTraceTail, 1, __ATOMIC_SEQ_CST)%COLLTRACE_NUM_ITEMS; \
|
||||
shmem.comm.collTrace[pos].timeStamp = __builtin_amdgcn_s_memrealtime(); \
|
||||
shmem.comm.collTrace[pos].bid = bid; \
|
||||
shmem.comm.collTrace[pos].type = ncclCollTraceKernelEndType; \
|
||||
}
|
||||
#define traceAbort() { \
|
||||
uint32_t pos = __atomic_fetch_add(ncclShmem->comm.collTraceTail, 1, __ATOMIC_SEQ_CST)%COLLTRACE_NUM_ITEMS; \
|
||||
ncclShmem->comm.collTrace[pos].timeStamp = __builtin_amdgcn_s_memrealtime(); \
|
||||
ncclShmem->comm.collTrace[pos].bid = bid; \
|
||||
ncclShmem->comm.collTrace[pos].type = ncclCollTraceAbortType; \
|
||||
uint32_t pos = __atomic_fetch_add(shmem.comm.collTraceTail, 1, __ATOMIC_SEQ_CST)%COLLTRACE_NUM_ITEMS; \
|
||||
shmem.comm.collTrace[pos].timeStamp = __builtin_amdgcn_s_memrealtime(); \
|
||||
shmem.comm.collTrace[pos].bid = bid; \
|
||||
shmem.comm.collTrace[pos].type = ncclCollTraceAbortType; \
|
||||
}
|
||||
// traceData(int16_t data2, uint32_t data4, uint64_t data8_0, uint64_t data8_1)
|
||||
#define traceData(data2, data4, data8_0, data8_1) { \
|
||||
uint32_t pos = __atomic_fetch_add(ncclShmem->comm.collTraceTail, 1, __ATOMIC_SEQ_CST)%COLLTRACE_NUM_ITEMS; \
|
||||
uint32_t pos = atomicAdd(ncclShmem->comm.collTraceTail, 1)%COLLTRACE_NUM_ITEMS; \
|
||||
ncclShmem->comm.collTrace[pos].bid = blockIdx.x; \
|
||||
ncclShmem->comm.collTrace[pos].timeStamp = __builtin_amdgcn_s_memrealtime(); \
|
||||
ncclShmem->comm.collTrace[pos].funcIndex = data2; \
|
||||
@@ -281,22 +383,16 @@ class ncclFunction {
|
||||
|
||||
#ifdef ENABLE_PROFILING
|
||||
#define __insert_timestamp(line_num) do { \
|
||||
if (ncclShmem->prof.count < PROFILE_NUM_ITEMS) { \
|
||||
ncclShmem->prof.elem[ncclShmem->prof.count].line = line_num; \
|
||||
ncclShmem->prof.elem[ncclShmem->prof.count].timeStamp = __builtin_amdgcn_s_memrealtime(); \
|
||||
ncclShmem->prof.count++; \
|
||||
if (shmem.prof.count < PROFILE_NUM_ITEMS) { \
|
||||
shmem.prof.elem[shmem.prof.count].line = line_num; \
|
||||
shmem.prof.elem[shmem.prof.count].timeStamp = __builtin_amdgcn_s_memrealtime(); \
|
||||
shmem.prof.count++; \
|
||||
} \
|
||||
} while(0);
|
||||
#else
|
||||
#define __insert_timestamp(line_num)
|
||||
#endif
|
||||
|
||||
__device__ inline bool barrierReduceAny(int bit, uint32_t* abortCount) {
|
||||
if (bit) atomicAdd(abortCount, 1);
|
||||
__syncthreads();
|
||||
return atomicAdd(abortCount, 0) != 0;
|
||||
}
|
||||
|
||||
// Copy src to dst and fill extra size with zeroes
|
||||
template<typename Tdst, typename Tsrc>
|
||||
__device__ void copyToShmem(Tdst *dst, Tsrc const *src, int tid, int nthreads) {
|
||||
@@ -385,19 +481,18 @@ struct ncclShmemGroup {
|
||||
|
||||
struct ncclShmemData {
|
||||
union {
|
||||
uint64_t ll128warp[NCCL_MAX_GROUPS][NCCL_MAX_GROUPS];
|
||||
struct ncclShmemGroup groups[NCCL_MAX_GROUPS];
|
||||
};
|
||||
uint32_t sync[NCCL_MAX_GROUPS];
|
||||
uint64_t redOpArgs[NCCL_MAX_DIRECT_ARITY+1];
|
||||
struct ncclDevComm comm;
|
||||
struct ncclChannel channel;
|
||||
uint64_t pad;
|
||||
uint64_t pad[2];
|
||||
struct ncclWork work;
|
||||
#ifdef ENABLE_PROFILING
|
||||
struct ncclProf prof;
|
||||
#endif
|
||||
};
|
||||
static_assert(offsetof(struct ncclShmemData, work)%16 == 0, "shmem.work needs to be 16B aligned");
|
||||
|
||||
static __device__ void ncclRedopPtrDeref(struct ncclWorkElem* we) {
|
||||
if (we->header.type != ncclWorkTypeUnused && we->redOpArgIsPtr) {
|
||||
@@ -422,16 +517,14 @@ static __device__ void ncclRedopPtrDeref(struct ncclWorkElem* we) {
|
||||
|
||||
extern __device__ struct ncclShmemData *ncclShmem;
|
||||
|
||||
template<ncclFunc_t Fn, typename T, typename RedOp, int Algo, int Proto, int FnIndex, bool COLLTRACE>
|
||||
__device__ void ncclKernel(struct ncclDevComm* comm, ncclWorkElem first) {
|
||||
template<ncclFunc_t Fn, typename T, typename RedOp, int Algo, int Proto, int FnIndex, bool COLLTRACE, bool USING_LL128>
|
||||
__device__ void ncclKernel(struct ncclDevComm* comm) {
|
||||
int tid = threadIdx.x;
|
||||
int nthreads = blockDim.x;
|
||||
int bid = blockIdx.x;
|
||||
__shared__ struct ncclShmemData shmem;
|
||||
ncclShmem = &shmem;
|
||||
__shared__ uint32_t abortCount;
|
||||
if (tid == 0) {
|
||||
abortCount = 0;
|
||||
for (auto i = 0; i < NCCL_MAX_GROUPS; i++) {
|
||||
shmem.groups[i].barrier = 0;
|
||||
for (auto j = 0; j < NCCL_MAX_GROUPS; j++) shmem.groups[i].barrier_next[j] = 0;
|
||||
@@ -439,48 +532,38 @@ __device__ void ncclKernel(struct ncclDevComm* comm, ncclWorkElem first) {
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
int turn = copyToShmem(&ncclShmem->comm, comm);
|
||||
int turn = copyToShmem(&shmem.comm, comm);
|
||||
#ifdef ENABLE_PROFILING
|
||||
if (tid == 0) {
|
||||
ncclShmem->prof.count = 0;
|
||||
ncclShmem->prof.seq = ncclShmem->comm.devProf[bid].seq;
|
||||
shmem.prof.count = 0;
|
||||
shmem.prof.seq = shmem.comm.devProf[bid].seq;
|
||||
}
|
||||
#endif
|
||||
if (tid == 0) __insert_timestamp(__LINE__);
|
||||
// get address of channel without incurring indirect load from ncclDevCom::channels
|
||||
ncclChannel *channel = &((ncclDevCommAndChannels*)comm)->channels[bid];
|
||||
turn = copyToShmem(&ncclShmem->channel, channel, turn);
|
||||
turn = copyToShmem(&shmem.channel, channel, turn);
|
||||
|
||||
// To optimize for latency, (only) the first operation is passed as argument.
|
||||
if (bid == 0 && first.header.type != ncclWorkTypeUnused) {
|
||||
// Copy first elem to work and zero out the rest
|
||||
copyToShmem(&ncclShmem->work, &first, tid, nthreads);
|
||||
}
|
||||
__syncthreads(); // publish ncclShmem
|
||||
if (tid == 0) __insert_timestamp(__LINE__);
|
||||
if (tid == 0) __insert_timestamp(__LINE__);
|
||||
|
||||
ncclWork *workFifoHost = ncclShmem->channel.workFifo;
|
||||
ncclWork *workFifoDev = ncclShmem->channel.workFifoDev;
|
||||
int workFifoIx = ncclShmem->channel.index;
|
||||
|
||||
bool skipLoadWork = false, firstLaunch = true;
|
||||
if (bid == 0 && first.header.type != ncclWorkTypeUnused)
|
||||
skipLoadWork = true;
|
||||
ncclWork *workFifoHost = shmem.channel.workFifo;
|
||||
ncclWork *workFifoDev = shmem.channel.workFifoDev;
|
||||
int workFifoIx = shmem.channel.index;
|
||||
bool firstLaunch = true;
|
||||
|
||||
while (true) {
|
||||
if (!skipLoadWork) {
|
||||
copyToShmem(&ncclShmem->work, &workFifoDev[workFifoIx], tid, nthreads);
|
||||
if (tid == 0) __insert_timestamp(__LINE__);
|
||||
{ // Check whether the last operation was aborted and make sure all threads exit
|
||||
int aborted = tid == 0 ? *comm->abortFlag : 0;
|
||||
if (barrierReduceAny(aborted, &abortCount)) { // publish ncclShmem->work
|
||||
if (COLLTRACE && tid == 0) traceAbort();
|
||||
break;
|
||||
}
|
||||
if (tid == 0)
|
||||
workFifoHost[workFifoIx].header.type = ncclWorkTypeUnused;
|
||||
copyToShmem(&shmem.work, &workFifoDev[workFifoIx], tid, nthreads);
|
||||
if (tid == 0) __insert_timestamp(__LINE__);
|
||||
{ // Check whether the last operation was aborted and make sure all threads exit
|
||||
int aborted = tid == 0 ? *comm->abortFlag : 0;
|
||||
if (__any(aborted)) { // publish shmem.work
|
||||
if (COLLTRACE && tid == 0) traceAbort();
|
||||
break;
|
||||
}
|
||||
if (tid == 0)
|
||||
workFifoHost[workFifoIx].header.type = ncclWorkTypeUnused;
|
||||
}
|
||||
if (tid == 0) __insert_timestamp(__LINE__);
|
||||
|
||||
@@ -489,47 +572,59 @@ __device__ void ncclKernel(struct ncclDevComm* comm, ncclWorkElem first) {
|
||||
channel->index = workFifoIx; // write back to real channel, not shmem shadow
|
||||
|
||||
__syncwarp();
|
||||
if (ncclShmem->work.header.type == ncclWorkTypeColl) {
|
||||
if (tid < NCCL_MAX_WORK_ELEMENTS) ncclRedopPtrDeref(&ncclShmem->work.elems[tid]);
|
||||
} else if (ncclShmem->work.header.type == ncclWorkTypeRegColl) {
|
||||
if (tid < NCCL_MAX_WORK_ELEMENTS_REG) ncclRedopPtrDeref(&ncclShmem->work.regElems[tid].elem);
|
||||
if (shmem.work.header.type == ncclWorkTypeColl) {
|
||||
if (tid < NCCL_MAX_WORK_ELEMENTS) ncclRedopPtrDeref(&shmem.work.elems[tid]);
|
||||
} else if (shmem.work.header.type == ncclWorkTypeRegColl) {
|
||||
if (tid < NCCL_MAX_WORK_ELEMENTS_REG) ncclRedopPtrDeref(&shmem.work.regElems[tid].elem);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (COLLTRACE && tid == 0) {
|
||||
traceKernelLaunch(ncclShmem->work.elems[0],firstLaunch);
|
||||
traceKernelLaunch(shmem.work.elems[0],firstLaunch);
|
||||
firstLaunch = false;
|
||||
#pragma unroll 1
|
||||
for(int e=1; e < NCCL_MAX_WORK_ELEMENTS && ncclShmem->work.elems[e].header.type != ncclWorkTypeUnused; e ++) {
|
||||
traceColl(ncclShmem->work.elems[e], 0);
|
||||
for(int e=1; e < NCCL_MAX_WORK_ELEMENTS && shmem.work.elems[e].header.type != ncclWorkTypeUnused; e ++) {
|
||||
traceColl(shmem.work.elems[e], 0);
|
||||
}
|
||||
}
|
||||
if (tid == 0) __insert_timestamp(__LINE__);
|
||||
if (ncclShmem->work.header.funcIndex == FnIndex)
|
||||
RunWork<Fn, T, RedOp, Algo, Proto>().run(&ncclShmem->work);
|
||||
if (shmem.work.header.funcIndex == FnIndex)
|
||||
RunWork<Fn, T, RedOp, Algo, Proto>().run(&shmem.work);
|
||||
else
|
||||
NCCL_CALL_FUNCTIONS(ncclShmem->work.header.funcIndex);
|
||||
NCCL_CALL_FUNCTIONS<USING_LL128>(shmem.work.header.funcIndex);
|
||||
|
||||
if (ncclShmem->work.header.isLast) break;
|
||||
if (shmem.work.header.isLast) break;
|
||||
__syncthreads();
|
||||
skipLoadWork = false;
|
||||
}
|
||||
if (COLLTRACE && tid == 0) traceKernelEnd()
|
||||
#ifdef ENABLE_PROFILING
|
||||
if (ncclShmem->comm.devProf->seq < PROFILE_NUM_LAUNCHES) {
|
||||
copyToShmem(ncclShmem->comm.devProf+MAXCHANNELS*ncclShmem->prof.seq+blockIdx.x, &ncclShmem->prof);
|
||||
if (tid == 0) ncclShmem->comm.devProf[bid].seq++;
|
||||
if (shmem.comm.devProf->seq < PROFILE_NUM_LAUNCHES) {
|
||||
__syncthreads();
|
||||
copyToShmem(shmem.comm.devProf+MAXCHANNELS*shmem.prof.seq+blockIdx.x, &shmem.prof);
|
||||
if (tid == 0) shmem.comm.devProf[bid].seq++;
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
#define IMPL_COLL_KERN(func, algo, proto, devredop, type, fIndex) \
|
||||
__launch_bounds__(NCCL_MAX_NTHREADS, 1) \
|
||||
__global__ void NCCL_KERN_NAME(func, algo, proto, devredop, type)(struct ncclDevComm* comm, ncclWorkElem first) { \
|
||||
if (comm->collTraceThread) \
|
||||
ncclKernel<ncclFunc##func, type, Func##devredop<type>, NCCL_ALGO_##algo, NCCL_PROTO_##proto, fIndex, true>(comm, first); \
|
||||
else \
|
||||
ncclKernel<ncclFunc##func, type, Func##devredop<type>, NCCL_ALGO_##algo, NCCL_PROTO_##proto, fIndex, false>(comm, first); \
|
||||
__global__ void NCCL_KERN_NAME(func, algo, proto, devredop, type)(struct ncclDevComm* comm) { \
|
||||
ncclKernel<ncclFunc##func, type, Func##devredop<type>, NCCL_ALGO_##algo, NCCL_PROTO_##proto, fIndex, false, false>(comm); \
|
||||
} \
|
||||
\
|
||||
__launch_bounds__(NCCL_MAX_NTHREADS, 1) \
|
||||
__global__ void NCCL_KERN_NAME_DEBUG(func, algo, proto, devredop, type)(struct ncclDevComm* comm) { \
|
||||
ncclKernel<ncclFunc##func, type, Func##devredop<type>, NCCL_ALGO_##algo, NCCL_PROTO_##proto, fIndex, true, false>(comm); \
|
||||
} \
|
||||
\
|
||||
__launch_bounds__(NCCL_MAX_NTHREADS, 1) \
|
||||
__global__ void NCCL_KERN_NAME_LL128(func, algo, proto, devredop, type)(struct ncclDevComm* comm) { \
|
||||
ncclKernel<ncclFunc##func, type, Func##devredop<type>, NCCL_ALGO_##algo, NCCL_PROTO_##proto, fIndex, false, true>(comm); \
|
||||
} \
|
||||
\
|
||||
__launch_bounds__(NCCL_MAX_NTHREADS, 1) \
|
||||
__global__ void NCCL_KERN_NAME_LL128_DEBUG(func, algo, proto, devredop, type)(struct ncclDevComm* comm) { \
|
||||
ncclKernel<ncclFunc##func, type, Func##devredop<type>, NCCL_ALGO_##algo, NCCL_PROTO_##proto, fIndex, true, true>(comm); \
|
||||
}
|
||||
|
||||
// Examples : AllReduce, RING, LL, Sum, uint8
|
||||
@@ -542,7 +637,8 @@ __device__ __attribute__((noinline)) void NCCL_FUNC_NAME(func, algo, proto, dev
|
||||
// Only generate inline kernels for LL
|
||||
#define IMPL_COLL4(func, algo, devredop, type, ncclType) \
|
||||
IMPL_COLL_FUNC(func, algo, LL, devredop, type) \
|
||||
IMPL_COLL_FUNC(func, algo, SIMPLE, devredop, type) \
|
||||
IMPL_COLL_FUNC(func, algo, LL128, devredop, type) \
|
||||
IMPL_COLL_FUNC(func, algo, SIMPLE, devredop, type)
|
||||
|
||||
#define IMPL_COLL3(func, devredop, type, ncclType) \
|
||||
IMPL_COLL4(func, TREE, devredop, type, ncclType) \
|
||||
|
||||
@@ -9,25 +9,6 @@
|
||||
#define OP128_H_
|
||||
|
||||
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
|
||||
inline __device__ uint64_t* shmemCvtPtr(volatile uint64_t* shmemGenericPtr) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
inline __device__ void load128(const uint64_t* ptr, uint64_t &v0, uint64_t &v1) {
|
||||
}
|
||||
|
||||
inline __device__ void store128(uint64_t* ptr, uint64_t v0, uint64_t v1) {
|
||||
}
|
||||
|
||||
inline __device__ void loadShmem128(uint64_t* shmemAsmPtr, uint64_t &v0, uint64_t &v1) {
|
||||
}
|
||||
|
||||
inline __device__ void storeShmem128(uint64_t* shmemAsmPtr, uint64_t v0, uint64_t v1) {
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
inline __device__ void loadShmemMisaligned128(T *ptr, uint64_t &v0, uint64_t &v1) {
|
||||
}
|
||||
#else
|
||||
inline __device__ void load128(const uint64_t* ptr, uint64_t &v0, uint64_t &v1) {
|
||||
asm volatile("ld.volatile.global.v2.u64 {%0,%1}, [%2];"
|
||||
|
||||
@@ -9,7 +9,6 @@
|
||||
|
||||
#define NCCL_LL128_FLAGTHREAD (NCCL_LL128_LINEELEMS-1)
|
||||
|
||||
#define __any_sync(WARP_MASK, needReload) (true)
|
||||
|
||||
template<typename T, typename RedOp, typename Fan, int Direct, int P2p>
|
||||
class Primitives<T, RedOp, Fan, Direct, ProtoLL128, P2p>:
|
||||
@@ -51,20 +50,25 @@ class Primitives<T, RedOp, Fan, Direct, ProtoLL128, P2p>:
|
||||
inline __device__ uint64_t recvFlag(int i) { return recvStep[i]+1; }
|
||||
inline __device__ uint64_t sendFlag(int i) { return sendStep[i]+1; }
|
||||
|
||||
uint64_t* barriers;
|
||||
uint64_t* barrier_next;
|
||||
|
||||
inline __device__ void barrier() {
|
||||
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
|
||||
__syncthreads();
|
||||
if (nthreads != WARP_SIZE)
|
||||
barrier_by_group();
|
||||
#else
|
||||
asm volatile ("bar.sync %1, %0;" :: "r"(nthreads), "r"(15-group));
|
||||
#endif
|
||||
}
|
||||
|
||||
uint32_t abort = 0;
|
||||
uint32_t* sync;
|
||||
|
||||
inline __device__ int checkAbort(int &spins, int i, int send) {
|
||||
spins++;
|
||||
if (abort == 0 && spins == NCCL_SPINS_BEFORE_CHECK_ABORT) {
|
||||
abort = *ncclShmem->comm.abortFlag;
|
||||
abort = __atomic_load_n(ncclShmem->comm.abortFlag, __ATOMIC_SEQ_CST);
|
||||
spins = 0;
|
||||
}
|
||||
return abort;
|
||||
@@ -74,67 +78,69 @@ class Primitives<T, RedOp, Fan, Direct, ProtoLL128, P2p>:
|
||||
if (sendConnHeadPtr) {
|
||||
int spins = 0;
|
||||
while (sendConnHeadCache + NCCL_STEPS < sendConnHead + 1) {
|
||||
sendConnHeadCache = *sendConnHeadPtr;
|
||||
__builtin_amdgcn_s_sleep(8);
|
||||
sendConnHeadCache = atomicAdd_system((unsigned long long *)sendConnHeadPtr, 0);
|
||||
if (checkAbort(spins, wid, 1)) break;
|
||||
}
|
||||
__asm__ __volatile__("s_wakeup");
|
||||
if (sendConnFifoPtr) {
|
||||
sendConnFifoPtr[sendStep[wid]%NCCL_STEPS] = nbytes;
|
||||
__atomic_store_n(sendConnFifoPtr+sendStep[wid]%NCCL_STEPS, nbytes, __ATOMIC_SEQ_CST);
|
||||
}
|
||||
sendConnHead += 1;
|
||||
}
|
||||
}
|
||||
|
||||
inline __device__ void postRecv() {
|
||||
if (recvConnHeadPtr) *recvConnHeadPtr = recvConnHead += 1;
|
||||
if (recvConnHeadPtr) atomicExch_system((unsigned long long *)recvConnHeadPtr, recvConnHead += 1);
|
||||
}
|
||||
inline __device__ void postSend() {
|
||||
if (sendConnTailPtr) { __threadfence(); *sendConnTailPtr = sendConnTail += 1; }
|
||||
if (sendConnTailPtr) { __threadfence(); atomicExch_system((unsigned long long *)sendConnTailPtr, sendConnTail += 1); }
|
||||
}
|
||||
|
||||
template<int WordPerThread>
|
||||
__device__ __forceinline__ void loadRegsBegin(uint64_t(®s)[WordPerThread], T const *src, int eltN) {
|
||||
constexpr int EltPer16B = 16/sizeof(T);
|
||||
if(reinterpret_cast<uintptr_t>(src)%16 == 0) {
|
||||
/* We are aligned to 16 bytes, so load directly to registers no shmem.
|
||||
* Flag threads load half as much data which gets shuffled to the even
|
||||
* registers during Finish. The point of splitting into two phases is to
|
||||
* defer that shuffle, which incurs a dependency stall, until after other
|
||||
* memops are launched by the caller.
|
||||
*/
|
||||
#pragma unroll
|
||||
for(int g=0; g < WordPerThread/2; g++) {
|
||||
int ix = g*WARP_SIZE - 4*(g/2) + wid - (g%2)*(wid/8);
|
||||
if(!flagThread || g%2==0) {
|
||||
if(ix*EltPer16B < eltN)
|
||||
load128((uint64_t*)(src + ix*EltPer16B), regs[2*g+0], regs[2*g+1]);
|
||||
}
|
||||
}
|
||||
}
|
||||
else {
|
||||
// Not aligned. Stage the smallest 16 byte aligned region subsuming the
|
||||
// buffer into shmem.
|
||||
int misalignment = reinterpret_cast<uintptr_t>(src) % 16;
|
||||
uint64_t *src8 = reinterpret_cast<uint64_t*>(reinterpret_cast<uintptr_t>(src) & -uintptr_t(16));
|
||||
uint64_t *shm8 = shmemCvtPtr(ncclShmem->ll128warp[warp]);
|
||||
#pragma unroll
|
||||
for(int g=0; g < WordPerThread/2; g++)
|
||||
if((g*WARP_SIZE + wid)*16 < misalignment + eltN*sizeof(T))
|
||||
load128(src8 + 2*(g*WARP_SIZE + wid), regs[2*g+0], regs[2*g+1]);
|
||||
#pragma unroll
|
||||
for(int g=0; g < WordPerThread/2; g++)
|
||||
storeShmem128(shm8 + 2*(g*WARP_SIZE + wid), regs[2*g+0], regs[2*g+1]);
|
||||
|
||||
__syncwarp();
|
||||
|
||||
// Now load from shmem stage to regs. Preserve the same pre-shuffled layout
|
||||
// as the aligned case since Finish() will be applied regardless.
|
||||
T *shm = (T*)shm8 + misalignment/sizeof(T);
|
||||
#pragma unroll
|
||||
for(int g=0; g < WordPerThread/2; g++) {
|
||||
int ix = g*WARP_SIZE - 4*(g/2) + wid - (g%2)*(wid/8);
|
||||
if(!flagThread || g%2==0) {
|
||||
if(ix*EltPer16B < eltN)
|
||||
loadShmemMisaligned128(shm + ix*EltPer16B, regs[2*g+0], regs[2*g+1]);
|
||||
/* We are aligned to 16 bytes, so load directly to registers no shmem.
|
||||
* Flag threads load half as much data which gets shuffled to the even
|
||||
* registers during Finish. The point of splitting into two phases is to
|
||||
* defer that shuffle, which incurs a dependency stall, until after other
|
||||
* memops are launched by the caller.
|
||||
*/
|
||||
#pragma unroll
|
||||
for(int g=0; g < WordPerThread/2; g++) {
|
||||
int ix = g*WARP_SIZE - 16*(g/2) + wid - (g%2)*(wid/4);
|
||||
if(!flagThread || g%2==0) {
|
||||
if(ix*EltPer16B < eltN) {
|
||||
if(reinterpret_cast<uintptr_t>(src)%4 == 0) {
|
||||
regs[2*g+0] = __builtin_nontemporal_load((uint64_t*)(src + ix*EltPer16B));
|
||||
regs[2*g+1] = __builtin_nontemporal_load((uint64_t*)(src + ix*EltPer16B)+1);
|
||||
} else {
|
||||
union {
|
||||
uint64_t regs64[WordPerThread];
|
||||
uint32_t regs32[WordPerThread*2];
|
||||
uint16_t regs16[WordPerThread*4];
|
||||
uint8_t regs8[WordPerThread*8];
|
||||
};
|
||||
if (sizeof(T) == 8) {
|
||||
uint64_t *src64 = (uint64_t*)(src+ix*EltPer16B);
|
||||
for (int i=0; i < 2; i++)
|
||||
regs64[2*g+i] = __builtin_nontemporal_load(src64+i);
|
||||
} else if (sizeof(T) == 4) {
|
||||
uint32_t *src32 = (uint32_t*)(src+ix*EltPer16B);
|
||||
for (int i=0; i < 2*sizeof(uint64_t)/sizeof(T); i++)
|
||||
regs32[2*g+i] = __builtin_nontemporal_load(src32+i);
|
||||
} else if (sizeof(T) == 2) {
|
||||
uint16_t *src16 = (uint16_t*)(src+ix*EltPer16B);
|
||||
for (int i=0; i < 2*sizeof(uint64_t)/sizeof(T); i++)
|
||||
regs16[2*g+i] = __builtin_nontemporal_load(src16+i);
|
||||
} else if (sizeof(T) == 1) {
|
||||
uint8_t *src8 = (uint8_t*)(src+ix*EltPer16B);
|
||||
for (int i=0; i < 2*sizeof(uint64_t)/sizeof(T); i++)
|
||||
regs8[2*g+i] = __builtin_nontemporal_load(src8+i);
|
||||
}
|
||||
regs[2*g+0] = regs64[2*g+0];
|
||||
regs[2*g+1] = regs64[2*g+1];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -157,26 +163,45 @@ class Primitives<T, RedOp, Fan, Direct, ProtoLL128, P2p>:
|
||||
for (int g=1; g < WordPerThread/2; g+=2) {
|
||||
if (flagThread) regs[2*g-1] = regs[2*g];
|
||||
}
|
||||
// Write to dst if 16-byte aligned, shmem otherwise.
|
||||
int misalignment = reinterpret_cast<uintptr_t>(dst)%16;
|
||||
uint64_t *shm8 = shmemCvtPtr(ncclShmem->ll128warp[warp]);
|
||||
// Write to dst if 4-byte aligned, shmem otherwise.
|
||||
int misalignment = reinterpret_cast<uintptr_t>(dst)%4;
|
||||
#pragma unroll
|
||||
for(int g=0; g < WordPerThread/2; g++) {
|
||||
int ix = g*WARP_SIZE - 4*(g/2) + wid - (g%2)*(wid/8);
|
||||
int ix = g*WARP_SIZE - 16*(g/2) + wid - (g%2)*(wid/4);
|
||||
if (!flagThread || g%2==0) {
|
||||
if(misalignment == 0 && (ix+1)*EltPer16B <= eltN)
|
||||
store128((uint64_t*)(dst + ix*EltPer16B), regs[2*g+0], regs[2*g+1]);
|
||||
else
|
||||
storeShmem128(shm8+2*ix, regs[2*g+0], regs[2*g+1]);
|
||||
if(misalignment == 0 && (ix+1)*EltPer16B <= eltN) {
|
||||
__builtin_nontemporal_store(regs[2*g+0], (uint64_t*)(dst + ix*EltPer16B));
|
||||
__builtin_nontemporal_store(regs[2*g+1], (uint64_t*)(dst + ix*EltPer16B)+1);
|
||||
} else {
|
||||
union {
|
||||
uint64_t regs64[WordPerThread];
|
||||
uint32_t regs32[WordPerThread*2];
|
||||
uint16_t regs16[WordPerThread*4];
|
||||
uint8_t regs8[WordPerThread*8];
|
||||
};
|
||||
regs64[2*g+0] = regs[2*g+0];
|
||||
regs64[2*g+1] = regs[2*g+1];
|
||||
int remaining = eltN - ix*EltPer16B;
|
||||
if (sizeof(T) == 8) {
|
||||
uint64_t *dst64 = (uint64_t*)(dst+ix*EltPer16B);
|
||||
for (int i=0; i < 2 && i < remaining; i++)
|
||||
__builtin_nontemporal_store(regs64[2*g+i], dst64+i);
|
||||
} else if (sizeof(T) == 4) {
|
||||
uint32_t *dst32 = (uint32_t*)(dst+ix*EltPer16B);
|
||||
for (int i=0; i < 2*sizeof(uint64_t)/sizeof(T) && i < remaining; i++)
|
||||
__builtin_nontemporal_store(regs32[2*g+i], dst32+i);
|
||||
} else if (sizeof(T) == 2) {
|
||||
uint16_t *dst16 = (uint16_t*)(dst+ix*EltPer16B);
|
||||
for (int i=0; i < 2*sizeof(uint64_t)/sizeof(T) && i < remaining; i++)
|
||||
__builtin_nontemporal_store(regs16[2*g+i], dst16+i);
|
||||
} else if (sizeof(T) == 1) {
|
||||
uint8_t *dst8 = (uint8_t*)(dst+ix*EltPer16B);
|
||||
for (int i=0; i < 2*sizeof(uint64_t)/sizeof(T) && i < remaining; i++)
|
||||
__builtin_nontemporal_store(regs8[2*g+i], dst8+i);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncwarp();
|
||||
// Write rest from shmem to dst. No need to coalesce stores to 16-bytes,
|
||||
// the hardware keeps up fine.
|
||||
T *shm = (T*)ncclShmem->ll128warp[warp];
|
||||
int skip = misalignment == 0 ? eltN & -EltPer16B : 0;
|
||||
for(int i=skip+wid; i < eltN; i += WARP_SIZE)
|
||||
dst[i] = shm[i];
|
||||
}
|
||||
|
||||
#define WARP_MASK 0xffffffff
|
||||
@@ -197,10 +222,11 @@ class Primitives<T, RedOp, Fan, Direct, ProtoLL128, P2p>:
|
||||
needReload = false;
|
||||
#pragma unroll
|
||||
for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
|
||||
load128(ptr+u*WARP_SIZE, vr[u], vr[u+1]);
|
||||
vr[u] = __builtin_nontemporal_load(ptr+u*WARP_SIZE);
|
||||
vr[u+1] = __builtin_nontemporal_load(ptr+u*WARP_SIZE+1);
|
||||
needReload |= flagThread && (vr[u+1] != flag);
|
||||
}
|
||||
} while (__any_sync(WARP_MASK, needReload) && checkAbort(spins, 0, 0) == 0);
|
||||
} while (__any(needReload) && checkAbort(spins, 0, 0) == 0);
|
||||
}
|
||||
|
||||
/************* Finish register load **************/
|
||||
@@ -238,10 +264,11 @@ class Primitives<T, RedOp, Fan, Direct, ProtoLL128, P2p>:
|
||||
needReload = false;
|
||||
#pragma unroll
|
||||
for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
|
||||
load128(ptr+u*WARP_SIZE, vr[u], vr[u+1]);
|
||||
vr[u] = __builtin_nontemporal_load(ptr+u*WARP_SIZE);
|
||||
vr[u+1] = __builtin_nontemporal_load(ptr+u*WARP_SIZE+1);
|
||||
needReload |= flagThread && (vr[u+1] != flag);
|
||||
}
|
||||
} while (__any_sync(WARP_MASK, needReload) && checkAbort(spins, i, 0) == 0);
|
||||
} while (__any(needReload) && checkAbort(spins, i, 0) == 0);
|
||||
|
||||
#pragma unroll
|
||||
for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
|
||||
@@ -260,6 +287,9 @@ class Primitives<T, RedOp, Fan, Direct, ProtoLL128, P2p>:
|
||||
}
|
||||
}
|
||||
|
||||
#if !defined(__gfx1030__)
|
||||
if (tid == 0) __asm__ __volatile__("buffer_wbinvl1_vol");
|
||||
#endif
|
||||
/************************ Send **************************/
|
||||
if (SEND) {
|
||||
for (int i=1; i<MaxSend && i<fan.nsend(); i++) {
|
||||
@@ -267,14 +297,16 @@ class Primitives<T, RedOp, Fan, Direct, ProtoLL128, P2p>:
|
||||
uint64_t* ptr = sendPtr(i)+ll128Offset;
|
||||
#pragma unroll
|
||||
for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
|
||||
store128(ptr+u*WARP_SIZE, v[u], flagThread ? flag : v[u+1]);
|
||||
__builtin_nontemporal_store(v[u], ptr+u*WARP_SIZE);
|
||||
__builtin_nontemporal_store(flagThread ? flag : v[u+1], ptr+u*WARP_SIZE+1);
|
||||
}
|
||||
}
|
||||
uint64_t flag = sendFlag(0);
|
||||
uint64_t* ptr = sendPtr(0)+ll128Offset;
|
||||
#pragma unroll
|
||||
for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
|
||||
store128(ptr+u*WARP_SIZE, v[u], flagThread ? flag : v[u+1]);
|
||||
__builtin_nontemporal_store(v[u], ptr+u*WARP_SIZE);
|
||||
__builtin_nontemporal_store(flagThread ? flag : v[u+1], ptr+u*WARP_SIZE+1);
|
||||
}
|
||||
}
|
||||
/********************** End Send ************************/
|
||||
@@ -361,12 +393,14 @@ class Primitives<T, RedOp, Fan, Direct, ProtoLL128, P2p>:
|
||||
public:
|
||||
__device__ Primitives(
|
||||
const int tid, const int nthreads, int const *recvPeers, int const *sendPeers,
|
||||
void const *inputBuf, void *outputBuf, uint64_t redOpArg, int group=0, int connIndex=0
|
||||
void const *inputBuf, void *outputBuf, uint64_t redOpArg, int group=0
|
||||
):
|
||||
redOp(redOpArg),
|
||||
tid(tid), nthreads(nthreads), wid(tid%WARP_SIZE), warp(tid/WARP_SIZE),
|
||||
flagThread((tid%8)==7), group(group),
|
||||
flagThread((tid%4)==3), group(group&(uint16_t)0xFFFF),
|
||||
stepSize(ncclShmem->comm.buffSizes[NCCL_PROTO_LL128]/NCCL_STEPS/sizeof(uint64_t)) {
|
||||
barriers = &ncclShmem->groups[this->group].barrier;
|
||||
barrier_next = ncclShmem->groups[this->group].barrier_next;
|
||||
|
||||
auto *channel = &ncclShmem->channel;
|
||||
int nrecv=0, nsend=0;
|
||||
|
||||
Reference in New Issue
Block a user