Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,7 @@ bool ncclAllGatherDdaFabricEligible(ncclComm* comm, const void* sendbuff, void*
if (comm->nRanks < 2 || comm->nRanks > dda::common::kDdaMaxNranks) {
return false;
}
if (datatype != ncclFloat32 && datatype != ncclFloat16 && datatype != ncclBfloat16) {
if (!rcclDdaDatatypeSupported(comm, datatype)) {
return false;
}

Expand All @@ -135,7 +135,7 @@ uint32_t ncclAllGatherDdaFabricBlocks(ncclComm* comm, size_t sendcount, ncclData

ncclResult_t ncclAllGatherDdaFabric(const void* sendbuff, void* recvbuff, size_t sendcount, ncclDataType_t datatype,
ncclComm* comm, cudaStream_t stream) {
if (datatype != ncclFloat32 && datatype != ncclFloat16 && datatype != ncclBfloat16) {
if (!rcclDdaDatatypeSupported(comm, datatype)) {
return ncclInvalidArgument;
}
int typeSize = ncclTypeSize(datatype);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -125,7 +125,7 @@ bool ncclAllGatherDdaFabricLLEligible(ncclComm* comm, const void* sendbuff, void
if (comm->nRanks < 2 || comm->nRanks > dda::common::kDdaMaxNranks) {
return false;
}
if (datatype != ncclFloat32 && datatype != ncclFloat16 && datatype != ncclBfloat16) {
if (!rcclDdaDatatypeSupported(comm, datatype)) {
return false;
}

Expand Down Expand Up @@ -154,7 +154,7 @@ uint32_t ncclAllGatherDdaFabricLLBlocks(ncclComm* comm, size_t sendcount, ncclDa

ncclResult_t ncclAllGatherDdaFabricLL(const void* sendbuff, void* recvbuff, size_t sendcount, ncclDataType_t datatype,
ncclComm* comm, cudaStream_t stream) {
if (datatype != ncclFloat32 && datatype != ncclFloat16 && datatype != ncclBfloat16) {
if (!rcclDdaDatatypeSupported(comm, datatype)) {
return ncclInvalidArgument;
}
// AllGather is a pure copy, so the payload moves as raw bytes: instantiate the
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -152,7 +152,7 @@ bool ncclAllGatherDdaFabricLL128Eligible(ncclComm* comm, const void* sendbuff, v
if (comm->nRanks < 2 || comm->nRanks > dda::common::kDdaMaxNranks) {
return false;
}
if (datatype != ncclFloat32 && datatype != ncclFloat16 && datatype != ncclBfloat16) {
if (!rcclDdaDatatypeSupported(comm, datatype)) {
return false;
}
const size_t perRankBytes = sendcount * ncclTypeSize(datatype);
Expand Down Expand Up @@ -180,7 +180,7 @@ uint32_t ncclAllGatherDdaFabricLL128Blocks(ncclComm* comm, size_t sendcount, ncc

ncclResult_t ncclAllGatherDdaFabricLL128(const void* sendbuff, void* recvbuff, size_t sendcount,
ncclDataType_t datatype, ncclComm* comm, cudaStream_t stream) {
if (datatype != ncclFloat32 && datatype != ncclFloat16 && datatype != ncclBfloat16) {
if (!rcclDdaDatatypeSupported(comm, datatype)) {
return ncclInvalidArgument;
}
// AllGather is a pure copy, so the payload moves as raw bytes: instantiate the
Expand Down
34 changes: 22 additions & 12 deletions projects/rccl/src/algorithms/dda/all_reduce/all_reduce_dda_ll.h
Original file line number Diff line number Diff line change
Expand Up @@ -92,20 +92,30 @@ __launch_bounds__(512)
}
}

// Phase 2: poll my slots for the other ranks, reduce with my own data.
// Phase 2: poll my slots for the other ranks and fold every rank's payload in
// rank order, so each rank rounds the same partial sums and all get the same
// result. My own payload comes straight from sendbuff.
LLPacket16* myBase = reinterpret_cast<LLPacket16*>(peerScratch[selfRank]) + bankOffsetPkts;
for (size_t pk = gtid; pk < nPk; pk += stride) {
uint32_t acc0 = in[2 * pk];
uint32_t acc1 = in[2 * pk + 1];
for (int r = 1; r < nRanks; ++r) {
const int peer = (selfRank + r) % nRanks;
volatile LLPacket16* src = myBase + (size_t)peer * slot;
uint32_t d0, f0, d1, f1;
do {
ddaLLLoadLineB128(reinterpret_cast<const uint32_t*>(const_cast<LLPacket16*>(&src[pk])), d0, f0, d1, f1);
} while (f0 != flag || f1 != flag);
acc0 = vecElementAdd<T>(acc0, d0);
acc1 = vecElementAdd<T>(acc1, d1);
uint32_t acc0 = 0, acc1 = 0;
for (int peer = 0; peer < nRanks; ++peer) {
uint32_t d0, d1;
if (peer == selfRank) {
d0 = in[2 * pk];
d1 = in[2 * pk + 1];
} else {
volatile LLPacket16* src = myBase + (size_t)peer * slot;
uint32_t f0, f1;
do {
ddaLLLoadLineB128(reinterpret_cast<const uint32_t*>(const_cast<LLPacket16*>(&src[pk])), d0, f0, d1, f1);
} while (f0 != flag || f1 != flag);
}
if (peer == 0) {
acc0 = d0;
acc1 = d1;
} else {
ddaLLAddPayload<T>(acc0, acc1, d0, d1);
}
}
out[2 * pk] = acc0;
out[2 * pk + 1] = acc1;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -39,13 +39,13 @@ constexpr size_t ddaLL128ArSlotWords(size_t bankSize, int nRanks) {
//
// Phase 1 (publish): pack this rank's slice into registers once, then push it
// onto the wire in every peer's scratch at slot selfRank, flag word embedded.
// Phase 2 (reduce): for each peer, poll that peer's slot in our own scratch and
// fold the arriving words into an accumulator seeded with our own payload, then
// unpack the accumulator into recvbuff. Flag polling supplies the cross-rank
// ordering, so no GPU barrier is used.
// Phase 2 (reduce): for each rank in rank order, poll that rank's slot in our own
// scratch and fold the arriving words into an accumulator, then unpack the
// accumulator into recvbuff. Flag polling supplies the cross-rank ordering, so
// no GPU barrier is used.
//
// Self does not round-trip through scratch; its contribution seeds the
// accumulator straight from sendbuff. Scratch is double buffered: bank = flag & 1.
// Self does not round-trip through scratch; its contribution is packed straight
// from sendbuff. Scratch is double buffered: bank = flag & 1.
template <typename T, int NRANKS_CT>
#if defined(USE_ROCM)
__launch_bounds__(512)
Expand Down Expand Up @@ -112,22 +112,25 @@ __launch_bounds__(512)
const int eltInSlice =
rem < (size_t)kDdaLL128DataBytesPerSlice ? (int)rem : kDdaLL128DataBytesPerSlice;

// Seed with our own payload; peers are folded on top.
// Fold every rank's payload in rank order, so each rank rounds the same
// partial sums and all get the same result. Our own payload is packed from
// sendbuff into the same register layout the wire delivers.
uint64_t acc[kDdaLL128WordsPerThread] = {};
ddaLL128LoadRegs<int8_t>(acc, srcBytes + dataByte, eltInSlice, lane, flagLane);

for (int r = 1; r < nRanks; ++r) {
const int peer = (selfRank + r) % nRanks;
const uint64_t* gatherSlot = gatherBase + (uint64_t)peer * (uint64_t)slotWords;
uint64_t vr[kDdaLL128WordsPerThread];
ddaLL128PollWire(gatherSlot + s * (size_t)kDdaLL128WireWordsPerSlice + 2 * lane, vr, flag, lane);
for (int peer = 0; peer < nRanks; ++peer) {
uint64_t vr[kDdaLL128WordsPerThread] = {};
if (peer == selfRank) {
ddaLL128LoadRegs<int8_t>(vr, srcBytes + dataByte, eltInSlice, lane, flagLane);
} else {
const uint64_t* gatherSlot = gatherBase + (uint64_t)peer * (uint64_t)slotWords;
ddaLL128PollWire(gatherSlot + s * (size_t)kDdaLL128WireWordsPerSlice + 2 * lane, vr, flag, lane);
}
// On a flag lane the odd words hold flags rather than payload, so the sums
// landing there are meaningless -- ddaLL128StoreRegs re-derives those slots
// from the even ones and never writes them out, so folding blind is cheaper
// than predicating the loop.
#pragma unroll
for (int u = 0; u < kDdaLL128WordsPerThread; ++u) {
acc[u] = ddaLL128AddWord<T>(acc[u], vr[u]);
acc[u] = peer == 0 ? vr[u] : ddaLL128AddWord<T>(acc[u], vr[u]);
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -141,8 +141,7 @@ __launch_bounds__(1024)
do {
ddaLLLoadLineB128(reinterpret_cast<const uint32_t*>(const_cast<LLPacket16*>(&src[pk])), d0, f0, d1, f1);
} while (f0 != flag || f1 != flag);
acc0 = vecElementAdd<T>(acc0, d0);
acc1 = vecElementAdd<T>(acc1, d1);
ddaLLAddPayload<T>(acc0, acc1, d0, d1);
}
out[2 * pk] = acc0;
out[2 * pk + 1] = acc1;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -158,7 +158,7 @@ bool ncclAllReduceDdaFabricEligible(ncclComm* comm, const void* sendbuff, void*
if (op != ncclSum) {
return false;
}
if (datatype != ncclFloat32 && datatype != ncclFloat16 && datatype != ncclBfloat16) {
if (!rcclDdaDatatypeSupported(comm, datatype)) {
return false;
}
const size_t bytes = count * ncclTypeSize(datatype);
Expand Down Expand Up @@ -194,6 +194,10 @@ ncclResult_t ncclAllReduceDdaFabric(const void* sendbuff, void* recvbuff, size_t
return ncclAllReduceDdaFabricTyped<half>(sendbuff, recvbuff, count, comm, stream);
case ncclBfloat16:
return ncclAllReduceDdaFabricTyped<bf16>(sendbuff, recvbuff, count, comm, stream);
case ncclFloat8e4m3:
return ncclAllReduceDdaFabricTyped<dda::common::fp8e4m3>(sendbuff, recvbuff, count, comm, stream);
case ncclFloat8e5m2:
return ncclAllReduceDdaFabricTyped<dda::common::fp8e5m2>(sendbuff, recvbuff, count, comm, stream);
default:
return ncclInvalidArgument;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -337,7 +337,7 @@ bool ddaLLArOneShotEligible(ncclComm* comm, const void* sendbuff, void* recvbuff
if (op != ncclSum) {
return false;
}
if (datatype != ncclFloat32 && datatype != ncclFloat16 && datatype != ncclBfloat16) {
if (!rcclDdaDatatypeSupported(comm, datatype)) {
return false;
}

Expand Down Expand Up @@ -385,7 +385,7 @@ bool ddaLLArTwoShotEligible(ncclComm* comm, const void* sendbuff, void* recvbuff
if (op != ncclSum) {
return false;
}
if (datatype != ncclFloat32 && datatype != ncclFloat16 && datatype != ncclBfloat16) {
if (!rcclDdaDatatypeSupported(comm, datatype)) {
return false;
}

Expand Down Expand Up @@ -439,7 +439,7 @@ bool ddaLL128ArOneShotEligible(ncclComm* comm, const void* sendbuff, void* recvb
if (op != ncclSum) {
return false;
}
if (datatype != ncclFloat32 && datatype != ncclFloat16 && datatype != ncclBfloat16) {
if (!rcclDdaDatatypeSupported(comm, datatype)) {
return false;
}

Expand Down Expand Up @@ -492,7 +492,7 @@ bool ddaLL128ArTwoShotEligible(ncclComm* comm, const void* sendbuff, void* recvb
if (op != ncclSum) {
return false;
}
if (datatype != ncclFloat32 && datatype != ncclFloat16 && datatype != ncclBfloat16) {
if (!rcclDdaDatatypeSupported(comm, datatype)) {
return false;
}

Expand Down Expand Up @@ -555,6 +555,10 @@ ncclResult_t ncclAllReduceDdaFabricLL(const void* sendbuff, void* recvbuff, size
return ncclAllReduceDdaFabricLLTyped<half>(sendbuff, recvbuff, count, comm, stream);
case ncclBfloat16:
return ncclAllReduceDdaFabricLLTyped<bf16>(sendbuff, recvbuff, count, comm, stream);
case ncclFloat8e4m3:
return ncclAllReduceDdaFabricLLTyped<dda::common::fp8e4m3>(sendbuff, recvbuff, count, comm, stream);
case ncclFloat8e5m2:
return ncclAllReduceDdaFabricLLTyped<dda::common::fp8e5m2>(sendbuff, recvbuff, count, comm, stream);
Comment thread
alex-breslow-amd marked this conversation as resolved.
Resolved
default:
return ncclInvalidArgument;
}
Expand All @@ -570,6 +574,10 @@ ncclResult_t ncclAllReduceDdaFabricLL(const void* sendbuff, void* recvbuff, size
return ncclAllReduceDdaFabricLLTwoShotTyped<half>(sendbuff, recvbuff, count, comm, stream);
case ncclBfloat16:
return ncclAllReduceDdaFabricLLTwoShotTyped<bf16>(sendbuff, recvbuff, count, comm, stream);
case ncclFloat8e4m3:
return ncclAllReduceDdaFabricLLTwoShotTyped<dda::common::fp8e4m3>(sendbuff, recvbuff, count, comm, stream);
case ncclFloat8e5m2:
return ncclAllReduceDdaFabricLLTwoShotTyped<dda::common::fp8e5m2>(sendbuff, recvbuff, count, comm, stream);
default:
return ncclInvalidArgument;
}
Expand Down Expand Up @@ -597,6 +605,10 @@ ncclResult_t ncclAllReduceDdaFabricLL128(const void* sendbuff, void* recvbuff, s
return ncclAllReduceDdaFabricLL128OneShotTyped<half>(sendbuff, recvbuff, count, comm, stream);
case ncclBfloat16:
return ncclAllReduceDdaFabricLL128OneShotTyped<bf16>(sendbuff, recvbuff, count, comm, stream);
case ncclFloat8e4m3:
return ncclAllReduceDdaFabricLL128OneShotTyped<dda::common::fp8e4m3>(sendbuff, recvbuff, count, comm, stream);
case ncclFloat8e5m2:
return ncclAllReduceDdaFabricLL128OneShotTyped<dda::common::fp8e5m2>(sendbuff, recvbuff, count, comm, stream);
default:
return ncclInvalidArgument;
}
Expand All @@ -613,6 +625,10 @@ ncclResult_t ncclAllReduceDdaFabricLL128(const void* sendbuff, void* recvbuff, s
return ncclAllReduceDdaFabricLL128TwoShotTyped<half>(sendbuff, recvbuff, count, comm, stream);
case ncclBfloat16:
return ncclAllReduceDdaFabricLL128TwoShotTyped<bf16>(sendbuff, recvbuff, count, comm, stream);
case ncclFloat8e4m3:
return ncclAllReduceDdaFabricLL128TwoShotTyped<dda::common::fp8e4m3>(sendbuff, recvbuff, count, comm, stream);
case ncclFloat8e5m2:
return ncclAllReduceDdaFabricLL128TwoShotTyped<dda::common::fp8e5m2>(sendbuff, recvbuff, count, comm, stream);
default:
return ncclInvalidArgument;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -114,7 +114,7 @@ bool ncclAllToAllDdaFabricEligible(ncclComm* comm, const void* sendbuff, void* r
if (comm->nRanks < 2 || comm->nRanks > dda::common::kDdaMaxNranks) {
return false;
}
if (datatype != ncclFloat32 && datatype != ncclFloat16 && datatype != ncclBfloat16) {
if (!rcclDdaDatatypeSupported(comm, datatype)) {
return false;
}

Expand All @@ -139,7 +139,7 @@ uint32_t ncclAllToAllDdaFabricBlocks(ncclComm* comm, size_t count, ncclDataType_

ncclResult_t ncclAllToAllDdaFabric(const void* sendbuff, void* recvbuff, size_t count, ncclDataType_t datatype,
ncclComm* comm, cudaStream_t stream) {
if (datatype != ncclFloat32 && datatype != ncclFloat16 && datatype != ncclBfloat16) {
if (!rcclDdaDatatypeSupported(comm, datatype)) {
return ncclInvalidArgument;
}
int typeSize = ncclTypeSize(datatype);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -122,7 +122,7 @@ bool ncclAllToAllDdaFabricLLEligible(ncclComm* comm, const void* sendbuff, void*
if (comm->nRanks < 2 || comm->nRanks > dda::common::kDdaMaxNranks) {
return false;
}
if (datatype != ncclFloat32 && datatype != ncclFloat16 && datatype != ncclBfloat16) {
if (!rcclDdaDatatypeSupported(comm, datatype)) {
return false;
}

Expand All @@ -149,7 +149,7 @@ uint32_t ncclAllToAllDdaFabricLLBlocks(ncclComm* comm, size_t count, ncclDataTyp

ncclResult_t ncclAllToAllDdaFabricLL(const void* sendbuff, void* recvbuff, size_t count, ncclDataType_t datatype,
ncclComm* comm, cudaStream_t stream) {
if (datatype != ncclFloat32 && datatype != ncclFloat16 && datatype != ncclBfloat16) {
if (!rcclDdaDatatypeSupported(comm, datatype)) {
return ncclInvalidArgument;
}
// All-to-all moves raw bytes, so instantiate once for int8_t and scale the
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -165,7 +165,7 @@ bool ncclAllToAllDdaFabricLL128Eligible(ncclComm* comm, const void* sendbuff, vo
if (comm->nRanks < 2 || comm->nRanks > dda::common::kDdaMaxNranks) {
return false;
}
if (datatype != ncclFloat32 && datatype != ncclFloat16 && datatype != ncclBfloat16) {
if (!rcclDdaDatatypeSupported(comm, datatype)) {
return false;
}

Expand All @@ -190,7 +190,7 @@ uint32_t ncclAllToAllDdaFabricLL128Blocks(ncclComm* comm, size_t count, ncclData

ncclResult_t ncclAllToAllDdaFabricLL128(const void* sendbuff, void* recvbuff, size_t count, ncclDataType_t datatype,
ncclComm* comm, cudaStream_t stream) {
if (datatype != ncclFloat32 && datatype != ncclFloat16 && datatype != ncclBfloat16) {
if (!rcclDdaDatatypeSupported(comm, datatype)) {
return ncclInvalidArgument;
}
const int typeSize = ncclTypeSize(datatype);
Expand Down
Loading
Loading