AI/rocm-systems
2
0
Форкнуть 0
Код Задачи Запросы на слияние Действия Пакеты Проекты Релизы Вики Активность

Merge remote-tracking branch 'nccl/master' into develop

Просмотр исходного кода
Этот коммит содержится в:
BertanDogancay
2025-08-28 15:45:42 -05:00
родитель a0ec15bafe 72d2432094
Коммит 08a7be231b
108 изменённых файлов: 7754 добавлений и 2129 удалений
Скачать Patch файл Скачать Diff файл Показать все файлы Свернуть все файлы
src/device/Makefile
+127
Просмотреть файл
@@ -0,0 +1,127 @@
#
# Copyright (c) 2015-2021, NVIDIA CORPORATION. All rights reserved.
#
# See LICENSE.txt for license information
#
SHELL := /usr/bin/env bash
MAKEFLAGS += -r
.SUFFIXES:
.SECONDARY:
NCCLDIR := ../..
include $(NCCLDIR)/makefiles/common.mk
include $(NCCLDIR)/makefiles/version.mk
BUILDDIR ?= $(abspath ../../build)
OBJDIR := $(BUILDDIR)/obj/device
MANIFEST := $(OBJDIR)/manifest
DEVGLUE_OBJ := $(OBJDIR)/device_glue.o
INCFLAGS = -I. -I.. -I$(BUILDDIR)/include -I../include
NVCUFLAGS += $(INCFLAGS) --compiler-options "-fPIC -fvisibility=hidden"
CXXFLAGS += $(INCFLAGS)
NVCUFLAGS_SYM := -ccbin $(CXX) $(CXXSTD) --expt-extended-lambda -Xptxas -maxrregcount=128 -Xfatbin -compress-all
NVCUFLAGS_SYM += $(INCFLAGS) --compiler-options "-fPIC -fvisibility=hidden"
SAY = @bash -c 'path="$$2"; [[ "$$(realpath "$$2")" =~ ^$(subst .,\.,$(abspath $(NCCLDIR)))/(.*)$$ ]] && path="$${BASH_REMATCH[1]}"; printf "%-15s %s\n" "$$1" "$$path"' SAY
COMPILE.cu = $(NVCC) $(NVCUFLAGS) -dc $2 -o $1
COMPILE.cc = $(CXX) $(CXXFLAGS) -c $2 -o $1
define COMPILE
@$(SAY) "Compiling" $2;\
mkdir -p $(dir $1);\
$(call COMPILE$(or $3,$(suffix $2)),$1,$2)
endef
ifeq ($(shell echo "$$((1000*$(CUDA_MAJOR) + 10*$(CUDA_MINOR) >= 12080))"),1)
NVCC_GENCODE_LDMC_FP8 = -gencode=arch=compute_100a,code=sm_100a \
-gencode=arch=compute_120a,code=sm_120a
else ifeq ($(shell echo "$$((1000*$(CUDA_MAJOR) + 10*$(CUDA_MINOR) >= 12070))"),1)
NVCC_GENCODE_LDMC_FP8 = -gencode=arch=compute_100a,code=sm_100a
else
NVCC_GENCODE_LDMC_FP8 =
endif
define COMPILE_SYM
@$(SAY) "Compiling" $2;\
mkdir -p $(dir $1);\
$(NVCC) $(NVCUFLAGS_SYM) $3 -dw $2 -o $1
endef
DEPENDS.cu = $(NVCC) $(NVCUFLAGS) -M -dc $1
DEPENDS.cc = $(CXX) $(CXXFLAGS) -M -c $1
define DEPENDS
@$(SAY) "Dependencies" $2;\
mkdir -p $(dir $1);\
mk=$$($(call DEPENDS$(suffix $2),$2));\
[[ $$mk =~ ^[^:]*:(.*)$$ ]];\
files=$${BASH_REMATCH[1]};\
files=$$(for x in $$files; do case "$$x" in '\'|$$'\t') ;; *) echo "$$x"; esac; done);\
files=$$(for x in $$files; do [[ "$$(realpath "$$x")" == "$$(realpath "$(NCCLDIR)")"* ]] && echo "$$x"; done);\
echo "$(patsubst %.d,%.o,$1) $1: " $$files > $1
endef
all: $(MANIFEST)
$(OBJDIR)/gensrc: generate.py
@mkdir -p $@
(which python3 >/dev/null || \
(bar='!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!'; \
printf "\n$${bar}\nERROR: Building NCCL requires a Python 3 installation invokable as 'python3'.\n$${bar}\n\n" 1>&2; \
exit 1)) \
&& ./generate.py $@ "$(ONLY_FUNCS)"
$(OBJDIR)/gensrc/symmetric: $(OBJDIR)/gensrc symmetric/generate.py
@mkdir -p $@
./symmetric/generate.py $@
# The trailing ";" is necessary to make this an "empty recipe":
# https://www.gnu.org/software/make/manual/html_node/Empty-Recipes.html
$(OBJDIR)/gensrc/rules.mk: $(OBJDIR)/gensrc ;
$(OBJDIR)/gensrc/symmetric/rules.mk: $(OBJDIR)/gensrc/symmetric ;
-include $(OBJDIR)/gensrc/rules.mk
# "gensrc/rules.mk" populates $(LIB_OBJS_GEN)
-include $(OBJDIR)/gensrc/symmetric/rules.mk
# "gensrc/symmetric/rules.mk" populates $(LIB_OBJS_SYM_GEN)
SRCS = common.cu onerank.cu
LIB_OBJS = $(patsubst %, $(OBJDIR)/%.o, $(SRCS)) $(LIB_OBJS_GEN) $(LIB_OBJS_SYM_GEN)
$(OBJDIR)/%.o: % $(OBJDIR)/%.d
$(call COMPILE,$@,$<)
$(OBJDIR)/genobj/%.o: $(OBJDIR)/gensrc $(OBJDIR)/genobj/%.d
$(call COMPILE,$@,$(OBJDIR)/gensrc/$*)
$(OBJDIR)/genobj/symmetric/%.o: $(OBJDIR)/gensrc/symmetric $(OBJDIR)/genobj/symmetric/%.d
$(call COMPILE,$@,$(OBJDIR)/gensrc/symmetric/$*)
$(OBJDIR)/%.d: %
$(call DEPENDS,$@,$<)
$(OBJDIR)/genobj/%.d: $(OBJDIR)/gensrc/%
$(call DEPENDS,$@,$<)
$(OBJDIR)/genobj/symmetric/%.d: $(OBJDIR)/gensrc/symmetric/%
$(call DEPENDS,$@,$<)
$(DEVGLUE_OBJ): $(LIB_OBJS)
$(NVCC) $(NVCUFLAGS) -dlink $^ -o $@
$(MANIFEST): $(LIB_OBJS) $(DEVGLUE_OBJ)
@echo $^ > $@
-include $(wildcard $(OBJDIR)/*.d)
-include $(wildcard $(OBJDIR)/genobj/*.d)
-include $(wildcard $(OBJDIR)/genobj/symmetric/*.d)
.PHONY: clean
clean:
rm -rf $(OBJDIR)
src/device/all_gather.h
+206 -58
Просмотреть файл
@@ -244,7 +244,7 @@ struct RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_PAT, NCCL_PROTO_SIMPLE
while (1) {
struct ncclPatStep* ps = shmem->patSteps+(step%NCCL_SHMEM_PAT_STEPS);
int* poll = &ps->flags;
while (__hip_atomic_load(poll, __ATOMIC_ACQUIRE, __HIP_MEMORY_SCOPE_WORKGROUP) != 0) {
while (__hip_atomic_load(poll, __ATOMIC_ACQUIRE, __HIP_MEMORY_SCOPE_WORKGROUP) != 0) {
pollCount++ ;// Wait for workers to be done with step 'step-NCCL_SHMEM_PAT_STEPS'
}
patAlgo.getNextOp(ps);
@@ -272,7 +272,7 @@ struct RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_PAT, NCCL_PROTO_SIMPLE
struct ncclPatStep* ps = shmem->patSteps+(step%NCCL_SHMEM_PAT_STEPS);
int* poll = &ps->flags;
while (__hip_atomic_load(poll, __ATOMIC_ACQUIRE, __HIP_MEMORY_SCOPE_WORKGROUP) == 0){
pollCount++; // Wait for compute thread
pollCount++; // Wait for compute thread
}
int last = ps->last;
prims.patCopy(ps, shmem);
@@ -286,73 +286,221 @@ struct RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_PAT, NCCL_PROTO_SIMPLE
template<typename T, typename RedOp>
struct RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_NVLS, NCCL_PROTO_SIMPLE> {
template<bool BcastSendNotRecv>
struct Scatterer {
struct ncclDevWorkColl* work;
ssize_t chunkSize;
ssize_t railGridOffset;
template<int SlicePerChunk, int MinSrcs, int MaxSrcs, int MinDsts, int MaxDsts, int MultimemSrcs, int MultimemDsts>
__device__ __forceinline__ void operator()(
int tid, int tn, int slice, int maxSliceSize,
int nSrcs, void** srcPtrs, int nDsts, void** dstPtrs, int32_t* dstSizes, uint32_t sendDirectFlag, uint32_t recvDirectFlag
) {
static_assert(SlicePerChunk==1, "require: SlicePerChunk==1");
static_assert(MaxDsts<=1 || MaxSrcs<=1, "require: MaxDsts<=1 || MaxSrcs<=1");
struct ncclNvls* nvls = &ncclShmem.channel.nvls;
int nNodes = ncclShmem.comm.nNodes;
int nRails = nvls->nHeads;
int part = ncclShmem.channelId - work->channelLo;
char* inbuf = (char*)work->sendbuff;
char* outbuf = (char*)work->recvbuff;
ssize_t countPerRank = work->collnet.count;
bool inPlace = (inbuf == outbuf + ncclShmem.comm.rank * countPerRank);
ssize_t railAllBeg = min(railGridOffset + part * chunkSize, nNodes * countPerRank);
ssize_t railAllEnd = min(railAllBeg + chunkSize, nNodes * countPerRank);
int railAllSize = railAllEnd - railAllBeg;
int rail = 0;
int src = 0;
if (BcastSendNotRecv) {
rail = nvls->headRank;
} else {
if (work->regUsed) return;
rail = 0;
}
if (tid < nDsts) dstSizes[tid] = railAllSize;
do {
int node = railAllBeg / countPerRank;
int railAllOffset = 0;
while (railAllOffset < railAllSize) {
ssize_t railOneBeg = node * countPerRank;
ssize_t railOneEnd = railOneBeg + countPerRank;
ssize_t railOneOffset = (railAllBeg + railAllOffset) - railOneBeg;
int delta = min(railAllEnd, railOneEnd) - (railAllBeg + railAllOffset);
int rank = ncclShmem.comm.collNetDenseToUserRank[node * nRails + rail];
ssize_t userOneBeg = rank * countPerRank + railOneOffset;
int outIsDst = (inPlace && rank == ncclShmem.comm.rank) || BcastSendNotRecv || work->regUsed ? 0 : 1;
if (nSrcs != 0 && outIsDst + nDsts != 0) {
reduceCopy<ncclCollUnroll(), USE_ACC, RedOp, T,
/*MultimemSrcs,MinSrcs,MaxSrcs=*/MultimemSrcs, 1, 1,
/*MultimemDsts=*/MultimemDsts, 0 + MultimemDsts + MinDsts, 1 + MaxDsts,
/*PreOpSrcs=*/0>
(tid, tn, 0, nullptr, false,
/*nSrcs=*/1, [=]__device__(int s/*==0*/) -> void* {
return (char*)srcPtrs[src] + railAllOffset;
},
/*nDsts=*/outIsDst + nDsts, [=]__device__(int d) -> void* {
return d < outIsDst ? outbuf + userOneBeg
: work->regUsed ? (char*)dstPtrs[d - outIsDst] + userOneBeg
: (char*)dstPtrs[d - outIsDst] + railAllOffset;
}, delta);
}
railAllOffset += delta;
node += 1;
}
rail += 1;
src += 1;
} while (!BcastSendNotRecv && src < nRails);
}
};
__device__ __forceinline__ void run(int tid, int/*nthreads*/, struct ncclDevWorkColl* work) {
struct ncclNvls* nvls = &ncclShmem.channel.nvls;
const ssize_t rank = ncclShmem.comm.rank;
size_t count, gridOffset, channelCount;
size_t chunkCount;
ncclCollCbdPart(work, ncclShmem.channelId, NCCL_PROTO_SIMPLE, sizeof(T), &count, &gridOffset, &channelCount, &chunkCount);
size_t offset;
int nelem;
const int nThreadsBcast = work->regUsed ? (NCCL_MAX_NTHREADS - WARP_SIZE) : 4 * WARP_SIZE;
const int nThreadsGather = work->regUsed ? WARP_SIZE : NCCL_MAX_NTHREADS - nThreadsBcast;
const int tidEndGather = nThreadsGather;
const int tidEndBcast = tidEndGather + nThreadsBcast;
const int nThreadsNetSend = work->oneNode ? 0 : (work->netRegUsed ? WARP_SIZE : 6 * WARP_SIZE);
const int nThreadsGather = work->regUsed ? roundUp(nvls->nHeads << 2, WARP_SIZE) : 8 * WARP_SIZE;
const int nThreadsBcast = NCCL_MAX_NTHREADS - nThreadsNetSend - nThreadsGather;
if (!work->regUsed) {
if (tid < tidEndGather) {
// Gather
using Proto = ProtoSimple<1, 1, USE_ACC, COLL_UNROLL>;
Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_NVLS_ARITY, 0>, /*Direct=*/0, Proto, 0>
prims(tid, nThreadsGather, nvls->up, NULL, NULL, work->recvbuff,
work->redOpArg, 0 * Proto::MaxGroupWidth, 1, 1);
for (size_t elemOffset = 0; elemOffset < channelCount; elemOffset += chunkCount) {
offset = gridOffset + elemOffset;
nelem = min(chunkCount, channelCount - elemOffset);
prims.gather(offset, nvls->nHeads * count, nelem, count, -1, 0);
const int tidEndGather = nThreadsGather;
const int tidEndNetSend = tidEndGather + nThreadsNetSend;
const int tidEndBcast = tidEndNetSend + nThreadsBcast;
if (work->oneNode) {
const ssize_t rank = ncclShmem.comm.rank;
size_t count, gridOffset, channelCount, offset, chunkCount;
ncclCollCbdPart(work, ncclShmem.channelId, NCCL_PROTO_SIMPLE, sizeof(T), &count, &gridOffset, &channelCount, &chunkCount);
if (!work->regUsed) {
if (tid < tidEndGather) {
// Gather
using Proto = ProtoSimple<1, 1, USE_ACC, COLL_UNROLL>;
Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_NVLS_ARITY, 0>, /*Direct=*/0, Proto, 0>
prims(tid, nThreadsGather, nvls->up, NULL, NULL, work->recvbuff,
work->redOpArg, 0 * Proto::MaxGroupWidth, 1, 1);
for (size_t elemOffset = 0; elemOffset < channelCount; elemOffset += chunkCount) {
offset = gridOffset + elemOffset;
nelem = min(chunkCount, channelCount - elemOffset);
prims.gather(offset, nvls->nHeads * count, nelem, count, -1, 0);
}
// coverity[overrun-call] => Coverity think prims.index can be greater than 1
} else if (tid < tidEndBcast) {
// Bcast through NVLS
using Proto = ProtoSimple<1, 1, USE_ACC, COLL_UNROLL, 0, 1>;
Primitives<T, RedOp, FanAsymmetric<0, 1>, /*Direct=*/0, Proto, 0>
prims(tid - tidEndGather, nThreadsBcast, NULL, &nvls->down, work->sendbuff, NULL,
work->redOpArg, 3 * Proto::MaxGroupWidth, 0, 0);
for (size_t elemOffset = 0; elemOffset < channelCount; elemOffset += chunkCount) {
offset = gridOffset + elemOffset;
nelem = min(chunkCount, channelCount - elemOffset);
prims.send(offset, nelem);
}
// coverity[overrun-call] => Coverity think prims.index can be greater than 1
}
// coverity[overrun-call] => Coverity think prims.index can be greater than 1
} else if (tid < tidEndBcast) {
// Bcast through NVLS
using Proto = ProtoSimple<1, 1, COLL_UNROLL, 0, 1>;
Primitives<T, RedOp, FanAsymmetric<0, 1>, /*Direct=*/0, Proto, 0>
prims(tid - tidEndGather, nThreadsBcast, NULL, &nvls->down, work->sendbuff, NULL,
work->redOpArg, 3 * Proto::MaxGroupWidth, 0, 0);
for (size_t elemOffset = 0; elemOffset < channelCount; elemOffset += chunkCount) {
offset = gridOffset + elemOffset;
nelem = min(chunkCount, channelCount - elemOffset);
prims.send(offset, nelem);
} else {
if (tid < tidEndGather) {
using Proto = ProtoSimple<1, 1, USE_ACC, COLL_UNROLL>;
Primitives<T, RedOp, FanSymmetric<NCCL_MAX_NVLS_ARITY>, /*Direct=*/0, Proto, 0>
prims(tid, nThreadsGather, nvls->up, nvls->up, NULL, NULL,
work->redOpArg, 0 * Proto::MaxGroupWidth, 1, 1);
/* used as sync */
prims.scatter(0, 0, 0, 0, -1, 0);
for (size_t elemOffset = 0; elemOffset < channelCount; elemOffset += chunkCount) {
prims.gather(0, 0, 0, 0, -1, 0);
}
} else if (tid < tidEndBcast) {
using Proto = ProtoSimple<1, 1, USE_ACC, COLL_UNROLL, 0, 1>;
Primitives<T, RedOp, FanSymmetric<1>, /*Direct=*/1, Proto, 0>
prims(tid - tidEndGather, nThreadsBcast, &nvls->down, &nvls->down, work->sendbuff, NULL,
work->redOpArg, 1 * Proto::MaxGroupWidth, 0, 0, work);
/* used as sync */
prims.recv(0, 0);
for (size_t elemOffset = 0; elemOffset < channelCount; elemOffset += chunkCount) {
ssize_t inpOffset = gridOffset + elemOffset;
ssize_t outOffset = inpOffset + rank * count;
nelem = min(chunkCount, channelCount - elemOffset);
prims.directSend(inpOffset, outOffset, nelem);
}
}
// coverity[overrun-call] => Coverity think prims.index can be greater than 1
}
} else {
/* direct allgather */
// NVLS + IB SHARP
int nNodes = ncclShmem.comm.nNodes;
int part = ncclShmem.channelId - work->channelLo;
ssize_t countPerRank = work->collnet.count;
const int nChannels = work->channelHi - work->channelLo + 1;
ssize_t chunkCount = work->collnet.chunkCount;
if (tid < tidEndGather) {
using Proto = ProtoSimple<1, 1, USE_ACC, COLL_UNROLL>;
Primitives<T, RedOp, FanSymmetric<NCCL_MAX_NVLS_ARITY>, /*Direct=*/0, Proto, 0>
prims(tid, nThreadsGather, nvls->up, nvls->up, NULL, NULL,
work->redOpArg, 0 * Proto::MaxGroupWidth, 1, 1);
/* used as sync */
prims.scatter(0, 0, 0, 0, -1, 0);
for (size_t elemOffset = 0; elemOffset < channelCount; elemOffset += chunkCount) {
prims.gather(0, 0, 0, 0, -1, 0);
Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_NVLS_ARITY, 0>, /*Direct=*/1, Proto, 0>
prims(tid, nThreadsGather, nvls->up, nullptr, nullptr, work->recvbuff,
/*redOpArg=*/0, 1 * Proto::MaxGroupWidth, 1, 1, work);
for (ssize_t railGridOffset = 0; railGridOffset < nNodes * countPerRank; railGridOffset += nChannels * chunkCount) {
Scatterer</*BcastSendNotRecv=*/false> scat;
scat.work = work;
scat.chunkSize = chunkCount;
scat.railGridOffset = railGridOffset;
prims.template process</*Recv=*/1, /*Send=*/0>(scat);
}
} else if (tid < tidEndBcast) {
using Proto = ProtoSimple<1, 1, COLL_UNROLL, 0, 1>;
Primitives<T, RedOp, FanSymmetric<1>, /*Direct=*/1, Proto, 0>
prims(tid - tidEndGather, nThreadsBcast, &nvls->down, &nvls->down, work->sendbuff, NULL,
work->redOpArg, 1 * Proto::MaxGroupWidth, 0, 0, work);
/* used as sync */
prims.recv(0, 0);
for (size_t elemOffset = 0; elemOffset < channelCount; elemOffset += chunkCount) {
ssize_t inpOffset = gridOffset + elemOffset;
ssize_t outOffset = inpOffset + rank * count;
nelem = min(chunkCount, channelCount - elemOffset);
prims.directSend(inpOffset, outOffset, nelem);
} else {
if (work->netRegUsed) {
using ProtoSend = ProtoSimple<1, 1, USE_ACC, COLL_UNROLL>;
using ProtoBcast = ProtoSimple<1, 1, USE_ACC, COLL_UNROLL, 0, 1>;
int maxSteps = (int)divUp(nNodes * countPerRank, nChannels * chunkCount);
int curSteps = -1;
int postThread = tid - tidEndGather == 0 ? 1 : 0;
// for UB, we need to control the send speed to avoid net congestion.
// first unroll 2 steps, then unroll the rest steps when the data is received.
if (postThread) {
curSteps = min(2, maxSteps);
Primitives<T, RedOp, FanAsymmetric<0, 1>, /*Direct=*/1, ProtoSend, 0>::sendPeerNotify(nvls->out, 1, curSteps);
}
Primitives<T, RedOp, FanSymmetric<1>, /*Direct=*/1, ProtoBcast, 0>
prims(tid - tidEndGather, nThreadsNetSend + nThreadsBcast, &nvls->out, &nvls->down, nullptr, nullptr,
/*redOpArg=*/0, 2 * ProtoBcast::MaxGroupWidth, 0, 0, work);
for (ssize_t railGridOffset = 0; railGridOffset < nNodes * countPerRank; railGridOffset += nChannels * chunkCount) {
Scatterer</*BcastSendNotRecv=*/true> scat;
scat.work = work;
scat.chunkSize = chunkCount;
scat.railGridOffset = railGridOffset;
prims.template process</*Recv=*/1, /*Send=*/1>(scat);
if (postThread && curSteps < maxSteps) {
curSteps++;
Primitives<T, RedOp, FanAsymmetric<0, 1>, /*Direct=*/1, ProtoSend, 0>::sendPeerNotify(nvls->out, 1, 1);
}
}
} else {
if (tid < tidEndNetSend) {
using Proto = ProtoSimple<1, 1, USE_ACC, COLL_UNROLL>;
Primitives<T, RedOp, FanAsymmetric<0, 1>, /*Direct=*/0, Proto, 0>
prims(tid - tidEndGather, nThreadsNetSend, nullptr, &nvls->out, work->sendbuff, nullptr,
/*redOpArg=*/0, 0 * Proto::MaxGroupWidth, 1, 1);
for (ssize_t railGridOffset = 0; railGridOffset < nNodes * countPerRank; railGridOffset += nChannels * chunkCount) {
ssize_t railAllBeg = railGridOffset + part * chunkCount;
ssize_t railAllEnd = min(railAllBeg + chunkCount, nNodes * countPerRank);
ssize_t railOneBeg = ncclShmem.comm.node * countPerRank;
ssize_t railOneEnd = railOneBeg + countPerRank;
ssize_t beg = max(railAllBeg, railOneBeg);
ssize_t end = min(railAllEnd, railOneEnd);
prims.send(beg - railOneBeg, max(ssize_t(0), end - beg));
}
} else {
using Proto = ProtoSimple<1, 1, USE_ACC, COLL_UNROLL, 0, 1>;
Primitives<T, RedOp, FanSymmetric<1>, /*Direct=*/0, Proto, 0>
prims(tid - tidEndNetSend, nThreadsBcast, &nvls->out, &nvls->down, nullptr, nullptr,
/*redOpArg=*/0, 2 * Proto::MaxGroupWidth, 0, 0);
for (ssize_t railGridOffset = 0; railGridOffset < nNodes * countPerRank; railGridOffset += nChannels * chunkCount) {
Scatterer</*BcastSendNotRecv=*/true> scat;
scat.work = work;
scat.chunkSize = chunkCount;
scat.railGridOffset = railGridOffset;
prims.template process</*Recv=*/1, /*Send=*/1>(scat);
}
}
}
}
}
@@ -367,7 +515,7 @@ struct RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NCCL_P
ssize_t chunkSize;
ssize_t railGridOffset;
template<int SlicePerChunk, int MinSrcs, int MaxSrcs, int MinDsts, int MaxDsts>
template<int SlicePerChunk, int MinSrcs, int MaxSrcs, int MinDsts, int MaxDsts, int MultimemSrcs, int MultimemDsts>
__device__ __forceinline__ void operator()(
int tid, int tn, int slice, int maxSliceSize,
int nSrcs, void** srcPtrs, int nDsts, void** dstPtrs, int32_t* dstSizes, uint32_t sendDirectFlag, uint32_t recvDirectFlag
src/device/all_reduce.h
+1 -1
Просмотреть файл
@@ -791,7 +791,7 @@ struct RunWorkColl<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_NVLS, NCCL_PROTO_SIMPL
}
} else if (tid < tidEndReduce && nvls->headRank != -1) {
// Reduce, broadcast through NVLS
using Proto = ProtoSimple<1, 1, COLL_UNROLL, 1, 1>;
using Proto = ProtoSimple<1, 1, USE_ACC, COLL_UNROLL, 1, 1>;
Primitives<T, RedOp, FanSymmetric<1>, /*Direct=*/1, Proto, 0>
prims(tid - tidEndGather, nThreadsReduce, &nvls->down, &nvls->down, NULL, NULL,
work->redOpArg, 2 * Proto::MaxGroupWidth, 0, 0, work);
src/device/common.h
+30 -42
Просмотреть файл
@@ -145,7 +145,6 @@ struct ncclShmemData {
uint16_t funcId;
int nWorks;
int workSize;
uint32_t workConsumed;
uint64_t workCounter;
bool profilerEnabled;
struct ncclShmemGroup groups[NCCL_MAX_GROUPS];
@@ -331,7 +330,6 @@ __device__ __forceinline__ void loadWorkBatchToShmem(
}
if (tid == 0) {
ncclShmem.workSize = workSize;
ncclShmem.workConsumed = batch.offsetBase + (64-__clzll(batch.offsetBitset))*workSize;
}
// We deliberately replicate these div and mod calculations into the case
// blocks above so that they get constant divisor optimizations by the compiler.
@@ -392,6 +390,16 @@ __device__ __forceinline__ void loadWorkBatchToShmem(
}
}
__device__ __forceinline__ unsigned long long int globaltimer() {
#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
return wall_clock64();
#else
unsigned long long int timer;
asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(timer));
return timer;
#endif
}
template<ncclFunc_t Fn, typename T, typename RedOp, int Algo, int Proto, int USE_ACC, int COLL_UNROLL, int Pipeline>
struct RunWorkColl {
__device__ void run(int tid, int tn, struct ncclDevWorkColl* work) {
@@ -446,40 +454,30 @@ struct RunWorkBatch {
#define STOP 1
#define FINI 2
__device__ __forceinline__ bool profilerEnabled(void) {
// Check if any of the workItems in the batch is profiled. If so, there is an equivalent
// profiler ProxyOp waiting for the counter update in the host thread. If this check was
// done only for the first workItem the profiler counter for other workItems in the batch
// could never be updated, leaving the host thread spinning forever for the counter update
// and causing a hang.
bool enabled = false;
for (int i = 0; i < ncclShmem.nWorks && !enabled; i++) {
if (ncclShmem.workType == ncclDevWorkTypeP2p)
enabled = ((struct ncclDevWorkP2p*)ncclShmem.workStorage)[i].profilerEnabled;
else
enabled = ((struct ncclDevWorkColl*)ncclShmem.workStorage)[i].profilerEnabled;
}
return enabled;
__device__ __forceinline__ bool profilerEnabled(int workItemIdx) {
return (ncclShmem.workType == ncclDevWorkTypeP2p) ?
((struct ncclDevWorkP2p*)ncclShmem.workStorage)[workItemIdx].profilerEnabled :
((struct ncclDevWorkColl*)ncclShmem.workStorage)[workItemIdx].profilerEnabled;
}
__device__ __forceinline__ void profiler(int action) {
if (action == START) {
if (threadIdx.x == 0) {
// increment workCounter regardless of the profiler being active or not
if (threadIdx.x == 0) {
int idx = 0;
uint64_t wc = ncclShmem.channel.workCounter + 1;
if (action == START) {
for (; wc <= ncclShmem.channel.workCounter + ncclShmem.nWorks; wc++) {
if (!profilerEnabled(idx++)) continue;
ncclShmem.comm.workStarted[ncclShmem.channelId].data[wc%MAX_PROFILER_EVENTS_PER_CHANNEL].timestamp = globaltimer();
ncclShmem.comm.workStarted[ncclShmem.channelId].data[wc%MAX_PROFILER_EVENTS_PER_CHANNEL].counter = wc;
}
} else {
for (; wc <= ncclShmem.channel.workCounter + ncclShmem.nWorks; wc++) {
if (!profilerEnabled(idx++)) continue;
ncclShmem.comm.workCompleted[ncclShmem.channelId].data[wc%MAX_PROFILER_EVENTS_PER_CHANNEL].timestamp = globaltimer();
ncclShmem.comm.workCompleted[ncclShmem.channelId].data[wc%MAX_PROFILER_EVENTS_PER_CHANNEL].counter = wc;
}
ncclShmem.channel.workCounter += ncclShmem.nWorks;
if(!profilerEnabled()) return;
ncclShmem.comm.workStarted[ncclShmem.channelId] = ncclShmem.channel.workCounter;
}
} else if (action == STOP) {
if (threadIdx.x == 0 && profilerEnabled()) {
ncclShmem.comm.workCompleted[ncclShmem.channelId] = ncclShmem.channel.workCounter;
}
} else { // FINI
if (threadIdx.x == 0) {
// store the workCounter back to vidmem regardless of the profiler being active or not
((ncclDevCommAndChannels*)ncclShmem.args.comm)->channels[ncclShmem.channelId].workCounter = ncclShmem.channel.workCounter;
if (!profilerEnabled()) return;
ncclShmem.comm.workCompleted[ncclShmem.channelId] = ncclShmem.channel.workCounter;
if (action == FINI) ((ncclDevCommAndChannels*)ncclShmem.args.comm)->channels[ncclShmem.channelId].workCounter = ncclShmem.channel.workCounter;
}
}
}
@@ -597,11 +595,6 @@ __device__ __forceinline__ void ncclKernelMain(struct ncclDevKernelArgs const* a
if (tid == 0) __insert_timestamp(__LINE__);
if (COLLTRACE && tid%WARP_SIZE == 0) traceKernelLaunch(ncclCollTraceKernelLaunchType, 0);
if (tid == 0 && ncclShmem.args.workStorageType == ncclDevWorkStorageTypeFifo) {
// ncclShmem.workConsumed written by loadWorkBatchToShmem before __syncthreads()
ncclShmem.comm.workConsumed[ncclShmem.channelId] = ncclShmem.workConsumed;
}
while (ncclShmem.aborted == 0) {
if (tid == 0) __insert_timestamp(__LINE__);
profiler(START);
@@ -641,11 +634,6 @@ __device__ __forceinline__ void ncclKernelMain(struct ncclDevKernelArgs const* a
profiler(STOP);
loadWorkBatchToShmem(tid%WARP_SIZE, tn, args, batchIx);
__syncthreads();
if (tid == 0 && ncclShmem.args.workStorageType == ncclDevWorkStorageTypeFifo) {
// ncclShmem.workConsumed written by loadWorkBatchToShmem before __syncthreads()
ncclShmem.comm.workConsumed[ncclShmem.channelId] = ncclShmem.workConsumed;
}
if (COLLTRACE && tid%WARP_SIZE == 0) traceKernelLaunch(ncclCollTraceCollLaunchType, batchIx);
}
if (COLLTRACE && tid%WARP_SIZE == 0) traceKernelEnd(ncclCollTraceKernelEndType);
src/device/op128.h
+114 -13
Просмотреть файл
@@ -93,37 +93,65 @@ template<>
union BytePack<0> {};
template<>
union BytePack<1> {
uint8_t u8, native;
uint8_t u8[1], native;
};
template<>
union BytePack<2> {
BytePack<1> half[2];
BytePack<1> b1[2];
uint8_t u8[2];
uint16_t u16, native;
uint16_t u16[1], native;
};
template<>
union BytePack<4> {
BytePack<2> half[2];
BytePack<1> b1[4];
BytePack<2> b2[2];
uint8_t u8[4];
uint16_t u16[2];
uint32_t u32, native;
uint32_t u32[1], native;
inline __device__ BytePack<4>() = default;
inline __device__ BytePack<4>(const BytePack<4>& other) {
*this = other;
}
inline __device__ BytePack<4>& operator=(const BytePack<4>& other) {
u32[0] = other.u32[0];
return *this;
}
};
template<>
union BytePack<8> {
BytePack<4> half[2];
BytePack<1> b1[8];
BytePack<2> b2[4];
BytePack<4> b4[2];
uint8_t u8[8];
uint16_t u16[4];
uint32_t u32[2];
uint64_t u64, native;
uint64_t u64[1], native;
inline __device__ BytePack<8>() = default;
inline __device__ BytePack<8>(const BytePack<8>& other) {
*this = other;
}
inline __device__ BytePack<8>& operator=(const BytePack<8>& other) {
u64[0] = other.u64[0];
return *this;
}
};
template<>
union alignas(16) BytePack<16> {
BytePack<8> half[2];
BytePack<1> b1[16];
BytePack<2> b2[8];
BytePack<4> b4[4];
BytePack<8> b8[2];
uint8_t u8[16];
uint16_t u16[8];
uint32_t u32[4];
uint64_t u64[2];
ulong2 ul2, native;
ulong2 ul2[1], native;
#if !defined(USE_INDIRECT_FUNCTION_CALL) || defined(__gfx942__) || defined(__gfx950__)
inline __device__ BytePack<16>() = default;
inline __device__ BytePack<16>(const BytePack<16>& other) {
@@ -136,6 +164,30 @@ union alignas(16) BytePack<16> {
}
#endif
};
template<int Size>
union BytePack {
BytePack<Size/2> half[2];
BytePack<1> b1[Size];
BytePack<2> b2[Size/2];
BytePack<4> b4[Size/4];
BytePack<8> b8[Size/8];
BytePack<16> b16[Size/16];
uint8_t u8[Size];
uint16_t u16[Size/2];
uint32_t u32[Size/4];
uint64_t u64[Size/8];
inline __device__ BytePack<Size>() = default;
inline __device__ BytePack<Size>(const BytePack<Size>& other) {
*this = other;
}
inline __device__ BytePack<Size>& operator=(const BytePack<Size>& other) {
for (int i = 0; i < Size/8; i++) {
u64[i] = other.u64[i];
}
return *this;
}
};
template<typename T>
struct BytePackOf {
@@ -343,19 +395,19 @@ __device__ __forceinline__ void multimem_st_global<0>(uintptr_t addr, BytePack<0
}
template<>
__device__ __forceinline__ void multimem_st_global<1>(uintptr_t addr, BytePack<1> val) {
asm volatile("st.global.b8 [%0], %1;" :: "l"(addr), "r"((uint32_t)val.u8) : "memory");
asm volatile("st.global.b8 [%0], %1;" :: "l"(addr), "r"((uint32_t)val.native) : "memory");
}
template<>
__device__ __forceinline__ void multimem_st_global<2>(uintptr_t addr, BytePack<2> val) {
asm volatile("st.global.b16 [%0], %1;" :: "l"(addr), "h"(val.u16) : "memory");
asm volatile("st.global.b16 [%0], %1;" :: "l"(addr), "h"(val.native) : "memory");
}
template<>
__device__ __forceinline__ void multimem_st_global<4>(uintptr_t addr, BytePack<4> val) {
asm volatile("multimem.st.global.b32 [%0], %1;" :: "l"(addr), "r"(val.u32) : "memory");
asm volatile("multimem.st.global.b32 [%0], %1;" :: "l"(addr), "r"(val.native) : "memory");
}
template<>
__device__ __forceinline__ void multimem_st_global<8>(uintptr_t addr, BytePack<8> val) {
asm volatile("multimem.st.global.b64 [%0], %1;" :: "l"(addr), "l"(val.u64) : "memory");
asm volatile("multimem.st.global.b64 [%0], %1;" :: "l"(addr), "l"(val.native) : "memory");
}
template<>
__device__ __forceinline__ void multimem_st_global<16>(uintptr_t addr, BytePack<16> val) {
@@ -370,6 +422,55 @@ __device__ __forceinline__ void multimem_st_global(uintptr_t addr, BytePack<Size
}
#endif
// Load pack starting at index in array. Ignore elements past end (length of array).
template<typename Pack, typename T>
__device__ __forceinline__ Pack loadPack(T* ptr, int ix, int end) {
constexpr int Size = sizeof(Pack);
ptr += ix;
int n = end - ix;
if (alignof(T) == Size && sizeof(T) == Size) {
return *(Pack*)ptr;
} else if ((Size+3)/4 + 1 < Size/sizeof(T)) {
union { Pack ans; uint32_t part[Size/4]; };
int misalign = reinterpret_cast<uintptr_t>(ptr) % 4;
uint32_t* down = reinterpret_cast<uint32_t*>(reinterpret_cast<uintptr_t>(ptr) & -uintptr_t(4));
int i;
#pragma unroll
for (i=0; i < Size/4; i++) {
if (i*4/sizeof(T) < 1 || i*4/sizeof(T) < n) part[i] = down[i];
}
uint32_t extra;
if (misalign) extra = down[i];
#pragma unroll
for (i=0; i < Size/4; i++) {
part[i] = __funnelshift_r(part[i], part[i+1], 8*misalign);
}
if (misalign) part[i] = __funnelshift_r(part[i], extra, 8*misalign);
return ans;
} else {
union { Pack ans; BytePack<sizeof(T)> part[Size/sizeof(T)]; };
#pragma unroll
for (int i=0; i < Size/sizeof(T); i++) {
if (i < 1 || i < n) part[i] = ((BytePack<sizeof(T)>*)ptr)[i];
}
return ans;
}
}
// Store pack starting at index in array. Ignore elements past end (length of array).
template<typename Pack, typename T>
__device__ __forceinline__ void storePack(T* ptr, int ix, int end, Pack val) {
constexpr int Size = sizeof(Pack);
union { Pack tmp; BytePack<sizeof(T)> part[Size/sizeof(T)]; };
tmp = val;
ptr += ix;
int n = end - ix;
#pragma unroll
for (int i=0; i < Size/sizeof(T); i++) {
if (i < 1 || i < n) ((BytePack<sizeof(T)>*)ptr)[i] = part[i];
}
}
#if __CUDA_ARCH__ >= 900 && CUDART_VERSION >= 12010
// Warp-uniform memory copy from shared address (not generic) to global memory.
// The number of bytes copied is `min(MaxBytes, nBytesAhead)`, a negative value
@@ -413,10 +514,10 @@ __device__ __forceinline__ void copyGlobalShared_WarpUnrolled(
b4[3] = ld_shared<4>(srcAddr + 3*4);
if (srcMisalign != 0) {
BytePack<4> b4_4 = ld_shared<4>(srcAddr + 4*4);
b4[0].u32 = __funnelshift_r(b4[0].u32, b4[1].u32, srcMisalign*8);
b4[1].u32 = __funnelshift_r(b4[1].u32, b4[2].u32, srcMisalign*8);
b4[2].u32 = __funnelshift_r(b4[2].u32, b4[3].u32, srcMisalign*8);
b4[3].u32 = __funnelshift_r(b4[3].u32, b4_4.u32, srcMisalign*8);
b4[0].native = __funnelshift_r(b4[0].native, b4[1].native, srcMisalign*8);
b4[1].native = __funnelshift_r(b4[1].native, b4[2].native, srcMisalign*8);
b4[2].native = __funnelshift_r(b4[2].native, b4[3].native, srcMisalign*8);
b4[3].native = __funnelshift_r(b4[3].native, b4_4.native, srcMisalign*8);
}
if (Multimem) multimem_st_global<16>(dstAddr, b16);
else st_global<16>(dstAddr, b16);
src/device/prims_simple.h
+12 -6
Просмотреть файл
@@ -155,7 +155,7 @@ private:
void **ptrs = isSendNotRecv ? (ncclShmem.groups[group].dsts + Dst)
: (ncclShmem.groups[group].srcs + Src);
if (flags & NetRegMode) {
if ((flags & NetRegMode) && ((!isSendNotRecv && DirectRecv) || (isSendNotRecv && DirectSend))) {
if (P2p) {
ptrs[index] = NULL;
} else {
@@ -506,7 +506,7 @@ public:
}
template<int Recv, int Send, typename Fn>
__device__ __forceinline__ void process(Fn &&fn, uint32_t sendDirectFlag, uint32_t recvDirectFlag) {
__device__ __forceinline__ void process(Fn &&fn, uint32_t sendDirectFlag = 0, uint32_t recvDirectFlag = 0) {
#pragma unroll 1
for (int slice=0; slice < SlicePerChunk; slice++) {
if (tid < nworkers) {
@@ -530,7 +530,7 @@ public:
} else if (flags & DirectRead) { // empty send
ptrs[index] = nullptr;
} else {
ptrs[index] = connEltsFifo + (step%NCCL_STEPS)*stepSize;
ptrs[index] = connEltsFifo + (step%NCCL_STEPS)*connStepSize;
}
} else {
if (flags & DirectRead) {
@@ -541,11 +541,11 @@ public:
else
ptrs[index] = nullptr;
} else {
ptrs[index] = connEltsFifo + (step%NCCL_STEPS)*stepSize;
ptrs[index] = connEltsFifo + (step%NCCL_STEPS)*connStepSize;
}
}
} else {
ptrs[index] = connEltsFifo + (step%NCCL_STEPS)*stepSize;
ptrs[index] = connEltsFifo + (step%NCCL_STEPS)*connStepSize;
}
}
subBarrier();
@@ -560,7 +560,7 @@ public:
} else {
nsend = fan.nsend();
}
fn.template operator() < SlicePerChunk, 0, Recv*MaxRecv, 0, Send*MaxSend >
fn.template operator()<SlicePerChunk, 0, Recv*MaxRecv, 0, Send*MaxSend, MultimemSrcs, MultimemDsts>
(tid, nworkers, slice, stepSize * StepPerSlice,
nrecv, ncclShmem.groups[group].srcs,
nsend, ncclShmem.groups[group].dsts, ncclShmem.groups[group].dstSizes, sendDirectFlag, recvDirectFlag);
@@ -1083,6 +1083,12 @@ public:
__device__ __forceinline__ void directRecvDirectSend(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
genericOp<1, 1, 1, 1, -1, -1>(inpIx, outIx, eltN, postOp);
}
__device__ __forceinline__ void recvDirectSend(intptr_t outIx, int eltN, bool postOp=false) {
genericOp<0, 1, 1, 1, -1, -1>(-1, outIx, eltN, postOp);
}
__device__ __forceinline__ void directRecvSend(intptr_t outIx, int eltN, bool postOp=false) {
genericOp<1, 0, 1, 1, -1, -1>(outIx, outIx, eltN, postOp);
}
__device__ __forceinline__ void recvCopyDirectSend(intptr_t outIx, int eltN, bool postOp=false) {
genericOp<0, 1, 1, 1, -1, Output>(-1, outIx, eltN, postOp);
}
src/device/reduce_kernel.h
+318 -128
Просмотреть файл
@@ -42,18 +42,18 @@ struct IsFloatingPoint<double>: std::true_type {};
// 3. Have constructor taking `uint64_t opArg`.
template<typename T>
struct FuncCopy { using EltType = T; __device__ FuncCopy(uint64_t opArg=0) {}; };
struct FuncCopy { using EltType = T; __device__ __forceinline__ FuncCopy(uint64_t opArg=0) {}; };
template<typename T>
struct FuncSum { using EltType = T; __device__ FuncSum(uint64_t opArg=0) {}; };
struct FuncSum { using EltType = T; __device__ __forceinline__ FuncSum(uint64_t opArg=0) {}; };
template<typename T>
struct FuncProd { using EltType = T; __device__ FuncProd(uint64_t opArg=0) {}; };
struct FuncProd { using EltType = T; __device__ __forceinline__ FuncProd(uint64_t opArg=0) {}; };
template<typename T>
struct FuncMinMax {
using EltType = T;
BytePack<sizeof(T)> xormask; // only used by integers
bool isMinNotMax; // only used by floats
__device__ FuncMinMax(uint64_t opArg=0) {
__device__ __forceinline__ FuncMinMax(uint64_t opArg=0) {
xormask.native = opArg;
isMinNotMax = (opArg&1)==0;
}
@@ -68,13 +68,13 @@ template<typename T> struct FuncSumPostDiv;
template<typename Fn>
struct RedOpArg { // default case: no argument
static constexpr bool ArgUsed = false;
__device__ static uint64_t loadArg(void *ptr) { return 0; }
__device__ __forceinline__ static uint64_t loadArg(void *ptr) { return 0; }
};
template<typename T>
struct RedOpArg<FuncMinMax<T>> {
static constexpr bool ArgUsed = true;
__device__ static uint64_t loadArg(void *ptr) {
__device__ __forceinline__ static uint64_t loadArg(void *ptr) {
union { uint64_t u64; T val; };
u64 = 0;
val = *(T*)ptr;
@@ -88,6 +88,11 @@ struct RedOpArg<FuncMinMax<T>> {
// of elements. These classes are intended to be specialized for specific
// combinations of reduction function and pack size.
template<typename A, typename B, int EltPerPackA>
struct Apply_Cast/*{
static BytePack<EltPerPackA*sizeof(B)/sizeof(A)> cast(BytePack<EltPerPackA*sizeof(A)> a);
}*/;
template<typename Fn, int EltPerPack>
struct Apply_Reduce /*{
static BytePack<EltPerPack*sizeof(T)> reduce(
@@ -115,16 +120,60 @@ struct Apply_LoadMultimem/*{
static BytePack<BytePerPack> load(Fn fn, uintptr_t addr);
}*/;
// Helpers for dealing with BytePack<0>'s
template<typename A, typename B, int EltPerPack>
struct Apply_Cast_MaybeEmpty: Apply_Cast<A, B, EltPerPack> {};
template<typename A, typename B>
struct Apply_Cast_MaybeEmpty<A, B, /*EltPerPack=*/0> {
__device__ constexpr static BytePack<0> cast(BytePack<0> a) { return {}; }
};
template<typename Fn, int EltPerPack>
struct Apply_Reduce_MaybeEmpty: Apply_Reduce<Fn, EltPerPack> {};
template<typename Fn>
struct Apply_Reduce_MaybeEmpty<Fn, 0> {
__device__ constexpr static BytePack<0> reduce(Fn fn, BytePack<0> a, BytePack<0> b) { return {}; }
};
template<typename Fn, int EltPerPack>
struct Apply_PreOp_MaybeEmpty: Apply_PreOp<Fn, EltPerPack> {};
template<typename Fn>
struct Apply_PreOp_MaybeEmpty<Fn, 0> {
static constexpr bool IsIdentity = true;
__device__ constexpr static BytePack<0> preOp(Fn fn, BytePack<0> a) { return {}; }
};
template<typename Fn, int EltPerPack>
struct Apply_PostOp_MaybeEmpty: Apply_PostOp<Fn, EltPerPack> {};
template<typename Fn>
struct Apply_PostOp_MaybeEmpty<Fn, 0> {
static constexpr bool IsIdentity = true;
__device__ constexpr static BytePack<0> postOp(Fn fn, BytePack<0> a) { return {}; }
};
template<typename Fn, int BytePerPack>
struct Apply_LoadMultimem_MaybeEmpty: Apply_LoadMultimem<Fn, BytePerPack> {};
template<typename Fn>
struct Apply_LoadMultimem_MaybeEmpty<Fn, 0> {
__device__ constexpr static BytePack<0> load(Fn fn, uintptr_t addr) { return {}; }
};
////////////////////////////////////////////////////////////////////////////////
// Public API for calling the trait classes. These take the data elements as a
// pack of any type, which could be a BytePack<?> or any integral type (uint64_t,
// uint32_t, etc.), and will return a new pack where each element has been
// transformed appropriately.
template<typename A, typename B, typename PackA>
__device__ __forceinline__ BytePack<BytePackOf<PackA>::Size*sizeof(B)/sizeof(A)> applyCast(PackA a) {
return Apply_Cast_MaybeEmpty<A, B, BytePackOf<PackA>::Size/sizeof(A)>::cast(toPack(a));
}
template<typename Fn, typename Pack>
__device__ __forceinline__ Pack applyReduce(Fn fn, Pack a, Pack b) {
return fromPack<Pack>(
Apply_Reduce<Fn, BytePackOf<Pack>::Size/sizeof(typename Fn::EltType)>
Apply_Reduce_MaybeEmpty<Fn, BytePackOf<Pack>::Size/sizeof(typename Fn::EltType)>
::reduce(fn, toPack(a), toPack(b))
);
}
@@ -132,7 +181,7 @@ __device__ __forceinline__ Pack applyReduce(Fn fn, Pack a, Pack b) {
template<typename Fn, typename Pack>
__device__ __forceinline__ Pack applyPreOp(Fn fn, Pack a) {
return fromPack<Pack>(
Apply_PreOp<Fn, BytePackOf<Pack>::Size/sizeof(typename Fn::EltType)>
Apply_PreOp_MaybeEmpty<Fn, BytePackOf<Pack>::Size/sizeof(typename Fn::EltType)>
::preOp(fn, toPack(a))
);
}
@@ -140,23 +189,107 @@ __device__ __forceinline__ Pack applyPreOp(Fn fn, Pack a) {
template<typename Fn, typename Pack>
__device__ __forceinline__ Pack applyPostOp(Fn fn, Pack a) {
return fromPack<Pack>(
Apply_PostOp<Fn, BytePackOf<Pack>::Size/sizeof(typename Fn::EltType)>
Apply_PostOp_MaybeEmpty<Fn, BytePackOf<Pack>::Size/sizeof(typename Fn::EltType)>
::postOp(fn, toPack(a))
);
}
template<typename Fn, int BytePerPack>
__device__ __forceinline__ BytePack<BytePerPack> applyLoadMultimem(Fn fn, uintptr_t addr) {
return Apply_LoadMultimem<Fn, BytePerPack>::load(fn, addr);
return Apply_LoadMultimem_MaybeEmpty<Fn, BytePerPack>::load(fn, addr);
}
////////////////////////////////////////////////////////////////////////////////
// Apply_Cast
template<typename A, typename B, int EltPerPack>
struct Apply_Cast {
__device__ __forceinline__ static BytePack<EltPerPack*sizeof(B)> cast(BytePack<EltPerPack*sizeof(A)> a) {
BytePack<EltPerPack*sizeof(B)> b;
b.half[0] = Apply_Cast<A, B, EltPerPack/2>::cast(a.half[0]);
b.half[1] = Apply_Cast<A, B, EltPerPack/2>::cast(a.half[1]);
return b;
}
};
template<typename A, typename B>
struct Apply_Cast<A, B, /*EltPerPack=*/1> {
__device__ __forceinline__ static BytePack<sizeof(B)> cast(BytePack<sizeof(A)> a) {
return toPack(B(fromPack<A>(a)));
}
};
template<>
struct Apply_Cast<__half, float, /*EltPerPack=*/1> {
__device__ __forceinline__ static BytePack<sizeof(float)> cast(BytePack<sizeof(__half)> a) {
return toPack(__half2float(fromPack<__half>(a)));
}
};
template<>
struct Apply_Cast<float, __half, /*EltPerPack=*/1> {
__device__ __forceinline__ static BytePack<sizeof(__half)> cast(BytePack<sizeof(float)> a) {
return toPack(__float2half_rn(fromPack<float>(a)));
}
};
template<>
struct Apply_Cast<__half, float, /*EltPerPack=*/2> {
__device__ __forceinline__ static BytePack<4*2> cast(BytePack<2*2> a) {
return toPack(__half22float2(fromPack<__half2>(a)));
}
};
template<>
struct Apply_Cast<float, __half, /*EltPerPack=*/2> {
__device__ __forceinline__ static BytePack<2*2> cast(BytePack<4*2> a) {
return toPack(__float22half2_rn(fromPack<float2>(a)));
}
};
#if defined(__CUDA_BF16_TYPES_EXIST__) && (CUDART_RUNTIME >= 12000 || __CUDA_ARCH__ >= 800)
template<>
struct Apply_Cast<__nv_bfloat16, float, /*EltPerPack=*/2> {
__device__ __forceinline__ static BytePack<4*2> cast(BytePack<2*2> a) {
return toPack(__bfloat1622float2(fromPack<__nv_bfloat162>(a)));
}
};
template<>
struct Apply_Cast<float ,__nv_bfloat16, /*EltPerPack=*/2> {
__device__ __forceinline__ static BytePack<2*2> cast(BytePack<4*2> a) {
return toPack(__float22bfloat162_rn(fromPack<float2>(a)));
}
};
#endif
#define EASY_CAST(A, B, EltPerPack, VecA, VecB) \
template<> \
struct Apply_Cast<A, B, EltPerPack> { \
__device__ __forceinline__ static BytePack<sizeof(B)*EltPerPack> cast(BytePack<sizeof(A)*EltPerPack> a) { \
return toPack(VecB(fromPack<VecA>(a))); \
} \
}; \
template<> \
struct Apply_Cast<B, A, EltPerPack> { \
__device__ __forceinline__ static BytePack<sizeof(A)*EltPerPack> cast(BytePack<sizeof(B)*EltPerPack> b) { \
return toPack(VecA(fromPack<VecB>(b))); \
} \
};
#if defined(__CUDA_FP8_TYPES_EXIST__)
EASY_CAST(__nv_fp8_e5m2, float, 2, __nv_fp8x2_e5m2, float2)
EASY_CAST(__nv_fp8_e5m2, float, 4, __nv_fp8x4_e5m2, float4)
EASY_CAST(__nv_fp8_e4m3, float, 2, __nv_fp8x2_e4m3, float2)
EASY_CAST(__nv_fp8_e4m3, float, 4, __nv_fp8x4_e4m3, float4)
#endif
#undef EASY_CAST
////////////////////////////////////////////////////////////////////////////////
// Apply_Reduce
// Nonsensical base case
template<typename Fn>
struct Apply_Reduce<Fn, /*EltPerPack=*/0> {
__device__ static BytePack<0> reduce(Fn fn, BytePack<0> a, BytePack<0> b) {
__device__ __forceinline__ static BytePack<0> reduce(Fn fn, BytePack<0> a, BytePack<0> b) {
return {};
}
};
@@ -168,7 +301,7 @@ struct Apply_Reduce<Fn, /*EltPerPack=*/0> {
template<typename Fn, int EltPerPack>
struct Apply_Reduce {
template<int Size>
__device__ static BytePack<Size> reduce(Fn fn, BytePack<Size> a, BytePack<Size> b) {
__device__ __forceinline__ static BytePack<Size> reduce(Fn fn, BytePack<Size> a, BytePack<Size> b) {
a.half[0] = Apply_Reduce<Fn, EltPerPack/2>::reduce(fn, a.half[0], b.half[0]);
a.half[1] = Apply_Reduce<Fn, EltPerPack/2>::reduce(fn, a.half[1], b.half[1]);
return a;
@@ -178,25 +311,25 @@ struct Apply_Reduce {
// Base case definitions (EltPerPack == 1)
template<typename T>
struct Apply_Reduce<FuncCopy<T>, /*EltPerPack=*/1> {
__device__ static BytePack<sizeof(T)> reduce(FuncCopy<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
__device__ __forceinline__ static BytePack<sizeof(T)> reduce(FuncCopy<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
return a;
}
};
template<typename T>
struct Apply_Reduce<FuncSum<T>, /*EltPerPack=*/1> {
__device__ static BytePack<sizeof(T)> reduce(FuncSum<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
__device__ __forceinline__ static BytePack<sizeof(T)> reduce(FuncSum<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
return toPack<T>(fromPack<T>(a) + fromPack<T>(b));
}
};
template<typename T>
struct Apply_Reduce<FuncProd<T>, /*EltPerPack=*/1> {
__device__ static BytePack<sizeof(T)> reduce(FuncProd<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
__device__ __forceinline__ static BytePack<sizeof(T)> reduce(FuncProd<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
return toPack<T>(fromPack<T>(a) * fromPack<T>(b));
}
};
template<typename T>
struct Apply_Reduce<FuncMinMax<T>, /*EltPerPack=*/1> {
__device__ static BytePack<sizeof(T)> reduce(FuncMinMax<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
__device__ __forceinline__ static BytePack<sizeof(T)> reduce(FuncMinMax<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
return (a.native ^ fn.xormask.native) < (b.native ^ fn.xormask.native) ? a : b;
}
};
@@ -204,7 +337,7 @@ struct Apply_Reduce<FuncMinMax<T>, /*EltPerPack=*/1> {
// Optimizations for specfic types and element count combinations:
template<>
struct Apply_Reduce<FuncSum<uint8_t>, /*EltPerPack=*/4> {
__device__ static BytePack<4> reduce(FuncSum<uint8_t> fn, BytePack<4> a, BytePack<4> b) {
__device__ __forceinline__ static BytePack<4> reduce(FuncSum<uint8_t> fn, BytePack<4> a, BytePack<4> b) {
constexpr uint32_t even = 0x00ff00ffu;
uint32_t x = (a.native & even) + (b.native & even);
uint32_t y = (a.native & ~even) + (b.native & ~even);
@@ -240,7 +373,7 @@ struct Apply_Reduce<FuncMinMax<uint8_t>, /*EltPerPack=*/4> {
// template<>
// struct Apply_Reduce<FuncProd<uint8_t>, /*EltPerPack=*/4> {
// __device__ static BytePack<4> reduce(FuncProd<uint8_t> fn, BytePack<4> apack, BytePack<4> bpack) {
// __device__ __forceinline__ static BytePack<4> reduce(FuncProd<uint8_t> fn, BytePack<4> apack, BytePack<4> bpack) {
// uint32_t a = apack.native;
// uint32_t b = bpack.native;
// uint32_t ab0 = (a*b) & 0xffu;
@@ -326,7 +459,7 @@ template<typename Fn, int EltPerPack>
struct Apply_PreOp {
static constexpr bool IsIdentity = Apply_PreOp<Fn, EltPerPack/2>::IsIdentity;
template<int Size>
__device__ static BytePack<Size> preOp(Fn fn, BytePack<Size> a) {
__device__ __forceinline__ static BytePack<Size> preOp(Fn fn, BytePack<Size> a) {
#if __cpp_if_constexpr
if constexpr(!IsIdentity) {
#else
@@ -346,7 +479,7 @@ template<typename Fn>
struct Apply_PreOp<Fn, /*EltPerPack=*/1> {
static constexpr bool IsIdentity = true;
template<int Size>
__device__ static BytePack<Size> preOp(Fn fn, BytePack<Size> a) {
__device__ __forceinline__ static BytePack<Size> preOp(Fn fn, BytePack<Size> a) {
return a;
}
};
@@ -354,7 +487,7 @@ struct Apply_PreOp<Fn, /*EltPerPack=*/1> {
template<typename Fn>
struct Apply_PreOp<Fn, /*EltPerPack=*/0> {
static constexpr bool IsIdentity = true;
__device__ static BytePack<0> preOp(Fn fn, BytePack<0> a) {
__device__ __forceinline__ static BytePack<0> preOp(Fn fn, BytePack<0> a) {
return {};
}
};
@@ -367,7 +500,7 @@ template<typename Fn, int EltPerPack>
struct Apply_PostOp {
static constexpr bool IsIdentity = Apply_PostOp<Fn, EltPerPack/2>::IsIdentity;
template<int Size>
__device__ static BytePack<Size> postOp(Fn fn, BytePack<Size> a) {
__device__ __forceinline__ static BytePack<Size> postOp(Fn fn, BytePack<Size> a) {
#if __cpp_if_constexpr
if constexpr(!IsIdentity) {
#else
@@ -387,7 +520,7 @@ template<typename Fn>
struct Apply_PostOp<Fn, /*EltPerPack=*/1> {
static constexpr bool IsIdentity = true;
template<int Size>
__device__ static BytePack<Size> postOp(Fn fn, BytePack<Size> a) {
__device__ __forceinline__ static BytePack<Size> postOp(Fn fn, BytePack<Size> a) {
return a;
}
};
@@ -395,7 +528,7 @@ struct Apply_PostOp<Fn, /*EltPerPack=*/1> {
template<typename Fn>
struct Apply_PostOp<Fn, /*EltPerPack=*/0> {
static constexpr bool IsIdentity = true;
__device__ static BytePack<0> postOp(Fn fn, BytePack<0> a) {
__device__ __forceinline__ static BytePack<0> postOp(Fn fn, BytePack<0> a) {
return {};
}
};
@@ -407,7 +540,7 @@ struct Apply_PostOp<Fn, /*EltPerPack=*/0> {
template<typename T>
struct RedOpArg<FuncPreMulSum<T>> {
static constexpr bool ArgUsed = true;
__device__ static uint64_t loadArg(void *ptr) {
__device__ __forceinline__ static uint64_t loadArg(void *ptr) {
union { uint64_t u64; T val; };
u64 = 0;
val = *(T*)ptr;
@@ -420,7 +553,7 @@ template<typename T>
struct FuncPreMulSum {
using EltType = T;
T scalar;
__device__ FuncPreMulSum(uint64_t opArg=0) {
__device__ __forceinline__ FuncPreMulSum(uint64_t opArg=0) {
union { uint64_t u64; T val; };
u64 = opArg;
scalar = val;
@@ -434,7 +567,7 @@ template<>
struct FuncPreMulSum<half> {
using EltType = half;
half2 scalar;
__device__ FuncPreMulSum(uint64_t opArg=0) {
__device__ __forceinline__ FuncPreMulSum(uint64_t opArg=0) {
union { uint64_t u64; __half val; };
u64 = opArg;
scalar.x = val;
@@ -451,7 +584,7 @@ struct FuncPreMulSum<half> {
using EltType = hip_bfloat16;
#if __CUDA_ARCH__ >= 800
__nv_bfloat162 scalar;
__device__ FuncPreMulSum(uint64_t opArg=0) {
__device__ __forceinline__ FuncPreMulSum(uint64_t opArg=0) {
union { uint64_t u64; __nv_bfloat16 val; };
u64 = opArg;
scalar.x = val;
@@ -459,7 +592,7 @@ struct FuncPreMulSum<half> {
}
#else
float scalar;
__device__ FuncPreMulSum(uint64_t opArg=0) {
__device__ __forceinline__ FuncPreMulSum(uint64_t opArg=0) {
union { uint64_t u64; hip_bfloat16 val; };
u64 = opArg;
scalar = (float)(val);
@@ -474,7 +607,7 @@ struct FuncPreMulSum<half> {
struct FuncPreMulSum<__nv_fp8_e4m3> {
using EltType = __nv_fp8_e4m3;
__half2 scalar2;
__device__ FuncPreMulSum(uint64_t opArg) {
__device__ __forceinline__ FuncPreMulSum(uint64_t opArg) {
union { uint64_t u64; __nv_fp8_storage_t val; };
u64 = opArg;
scalar2.x = __half(__nv_cvt_fp8_to_halfraw(val, __NV_E4M3));
@@ -486,7 +619,7 @@ struct FuncPreMulSum<half> {
struct FuncPreMulSum<__nv_fp8_e5m2> {
using EltType = __nv_fp8_e5m2;
__half2 scalar2;
__device__ FuncPreMulSum(uint64_t opArg) {
__device__ __forceinline__ FuncPreMulSum(uint64_t opArg) {
union { uint64_t u64; __nv_fp8_storage_t val; };
u64 = opArg;
scalar2.x = __half(__nv_cvt_fp8_to_halfraw(val, __NV_E5M2));
@@ -528,7 +661,7 @@ struct FuncPreMulSum<half> {
template<typename T, int EltPerPack>
struct Apply_Reduce<FuncPreMulSum<T>, EltPerPack> {
__device__ static BytePack<EltPerPack*sizeof(T)> reduce(FuncPreMulSum<T> fn, BytePack<EltPerPack*sizeof(T)> a, BytePack<EltPerPack*sizeof(T)> b) {
__device__ __forceinline__ static BytePack<EltPerPack*sizeof(T)> reduce(FuncPreMulSum<T> fn, BytePack<EltPerPack*sizeof(T)> a, BytePack<EltPerPack*sizeof(T)> b) {
// FuncPreMulSum reduce dispatches to FuncSum.
return Apply_Reduce<FuncSum<T>, EltPerPack>::reduce(FuncSum<T>(), a, b);
}
@@ -538,7 +671,7 @@ struct Apply_Reduce<FuncPreMulSum<T>, EltPerPack> {
template<typename T>
struct Apply_PreOp<FuncPreMulSum<T>, /*EltPerPack=*/1> {
static constexpr bool IsIdentity = false;
__device__ static BytePack<sizeof(T)> preOp(FuncPreMulSum<T> fn, BytePack<sizeof(T)> a) {
__device__ __forceinline__ static BytePack<sizeof(T)> preOp(FuncPreMulSum<T> fn, BytePack<sizeof(T)> a) {
return toPack<T>(fromPack<T>(a) * fn.scalar);
}
};
@@ -549,7 +682,7 @@ struct Apply_PreOp<FuncPreMulSum<T>, /*EltPerPack=*/1> {
template<>
struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
static constexpr bool IsIdentity = false;
__device__ static BytePack<sizeof(half)> preOp(FuncPreMulSum<half> fn, BytePack<sizeof(half)> a) {
__device__ __forceinline__ static BytePack<sizeof(half)> preOp(FuncPreMulSum<half> fn, BytePack<sizeof(half)> a) {
return toPack<half>(__hmul(fromPack<half>(a), fn.scalar.x));
}
};
@@ -557,7 +690,7 @@ struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
template<>
struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/2> {
static constexpr bool IsIdentity = false;
__device__ static BytePack<sizeof(half2)> preOp(FuncPreMulSum<half> fn, BytePack<sizeof(half2)> a) {
__device__ __forceinline__ static BytePack<sizeof(half2)> preOp(FuncPreMulSum<half> fn, BytePack<sizeof(half2)> a) {
return toPack<half2>(__hmul2(fromPack<half2>(a), fn.scalar));
}
};
@@ -570,7 +703,7 @@ struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
template<>
struct Apply_PreOp<FuncPreMulSum<hip_bfloat16>, /*EltPerPack=*/1> {
static constexpr bool IsIdentity = false;
__device__ static BytePack<sizeof(hip_bfloat16)> preOp(
__device__ __forceinline__ static BytePack<sizeof(hip_bfloat16)> preOp(
FuncPreMulSum<hip_bfloat16> fn, BytePack<sizeof(hip_bfloat16)> a
) {
#if __CUDA_ARCH__ >= 800
@@ -584,7 +717,7 @@ struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
template<>
struct Apply_PreOp<FuncPreMulSum<hip_bfloat16>, /*EltPerPack=*/2> {
static constexpr bool IsIdentity = false;
__device__ static BytePack<sizeof(__nv_bfloat162)> preOp(
__device__ __forceinline__ static BytePack<sizeof(__nv_bfloat162)> preOp(
FuncPreMulSum<__nv_bfloat16> fn, BytePack<sizeof(__nv_bfloat162)> a
) {
return toPack<__nv_bfloat162>(__hmul2(fromPack<__nv_bfloat162>(a), fn.scalar));
@@ -601,7 +734,7 @@ struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
template<>
struct Apply_PreOp<FuncPreMulSum<__nv_fp8_e4m3>, /*EltPerPack=*/1> {
static constexpr bool IsIdentity = false;
__device__ static BytePack<sizeof(__nv_fp8_e4m3)> preOp(
__device__ __forceinline__ static BytePack<sizeof(__nv_fp8_e4m3)> preOp(
FuncPreMulSum<__nv_fp8_e4m3> fn, BytePack<sizeof(__nv_fp8_e4m3)> a
) {
return toPack<__nv_fp8_e4m3>(__nv_fp8_e4m3(__hmul(__half(fromPack<__nv_fp8_e4m3>(a)), fn.scalar2.x)));
@@ -610,7 +743,7 @@ struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
template<>
struct Apply_PreOp<FuncPreMulSum<__nv_fp8_e4m3>, /*EltPerPack=*/2> {
static constexpr bool IsIdentity = false;
__device__ static BytePack<sizeof(__nv_fp8x2_e4m3)> preOp(
__device__ __forceinline__ static BytePack<sizeof(__nv_fp8x2_e4m3)> preOp(
FuncPreMulSum<__nv_fp8_e4m3> fn, BytePack<sizeof(__nv_fp8x2_e4m3)> a
) {
return toPack<__nv_fp8x2_e4m3>(__nv_fp8x2_e4m3(__hmul2(__half2(fromPack<__nv_fp8x2_e4m3>(a)), fn.scalar2)));
@@ -620,7 +753,7 @@ struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
template<>
struct Apply_PreOp<FuncPreMulSum<__nv_fp8_e5m2>, /*EltPerPack=*/1> {
static constexpr bool IsIdentity = false;
__device__ static BytePack<sizeof(__nv_fp8_e5m2)> preOp(
__device__ __forceinline__ static BytePack<sizeof(__nv_fp8_e5m2)> preOp(
FuncPreMulSum<__nv_fp8_e5m2> fn, BytePack<sizeof(__nv_fp8_e5m2)> a
) {
return toPack<__nv_fp8_e5m2>(__nv_fp8_e5m2(__hmul(__half(fromPack<__nv_fp8_e5m2>(a)), fn.scalar2.x)));
@@ -629,7 +762,7 @@ struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
template<>
struct Apply_PreOp<FuncPreMulSum<__nv_fp8_e5m2>, /*EltPerPack=*/2> {
static constexpr bool IsIdentity = false;
__device__ static BytePack<sizeof(__nv_fp8x2_e5m2)> preOp(
__device__ __forceinline__ static BytePack<sizeof(__nv_fp8x2_e5m2)> preOp(
FuncPreMulSum<__nv_fp8_e5m2> fn, BytePack<sizeof(__nv_fp8x2_e5m2)> a
) {
return toPack<__nv_fp8x2_e5m2>(__nv_fp8x2_e5m2(__hmul2(__half2(fromPack<__nv_fp8x2_e5m2>(a)), fn.scalar2)));
@@ -666,7 +799,7 @@ struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
template<typename T>
struct RedOpArg<FuncSumPostDiv<T>> {
static constexpr bool ArgUsed = true;
__device__ static uint64_t loadArg(void *ptr) {
__device__ __forceinline__ static uint64_t loadArg(void *ptr) {
return *(uint64_t*)ptr;
}
};
@@ -709,12 +842,12 @@ struct FuncSumPostDiv {
uint32_t divisor:31, isSigned:1;
UintType recip;
__device__ FuncSumPostDiv(uint64_t opArg=0) {
__device__ __forceinline__ FuncSumPostDiv(uint64_t opArg=0) {
isSigned = opArg & 1;
divisor = opArg >> 1;
recip = Divider<UintType>::divide(UintType(-1), divisor);
}
__device__ T divide(T x) {
__device__ __forceinline__ T divide(T x) {
// x is negative iff we are in signed mode and the top bit is set
bool xneg = isSigned && (x & ~(T(-1)>>1));
// Compute abs(x):
@@ -736,7 +869,7 @@ struct FuncSumPostDiv {
template<typename T, int EltPerPack>
struct Apply_Reduce<FuncSumPostDiv<T>, EltPerPack>:
Apply_Reduce<FuncSum<T>, EltPerPack> {
__device__ static BytePack<EltPerPack*sizeof(T)> reduce(FuncSumPostDiv<T> fn, BytePack<EltPerPack*sizeof(T)> a, BytePack<EltPerPack*sizeof(T)> b) {
__device__ __forceinline__ static BytePack<EltPerPack*sizeof(T)> reduce(FuncSumPostDiv<T> fn, BytePack<EltPerPack*sizeof(T)> a, BytePack<EltPerPack*sizeof(T)> b) {
// FuncSumPostDiv reduce dispatches to FuncSum.
return Apply_Reduce<FuncSum<T>, EltPerPack>::reduce(FuncSum<T>(), a, b);
}
@@ -745,7 +878,7 @@ struct Apply_Reduce<FuncSumPostDiv<T>, EltPerPack>:
template<typename T>
struct Apply_PostOp<FuncSumPostDiv<T>, /*EltPerPack=*/1> {
static constexpr bool IsIdentity = false;
__device__ static BytePack<sizeof(T)> postOp(FuncSumPostDiv<T> fn, BytePack<sizeof(T)> a) {
__device__ __forceinline__ static BytePack<sizeof(T)> postOp(FuncSumPostDiv<T> fn, BytePack<sizeof(T)> a) {
return toPack<T>(fn.divide(fromPack<T>(a)));
}
};
@@ -753,120 +886,145 @@ struct Apply_PostOp<FuncSumPostDiv<T>, /*EltPerPack=*/1> {
////////////////////////////////////////////////////////////////////////////////
// Apply_LoadMultimem
#define SIZEOF_BytePack_field_u16 2
#define PTX_REG_BytePack_field_u16 "h"
#define RegCode_for_size_1 "r"
#define RegCode_for_size_2 "h"
#define RegCode_for_size_4 "r"
#define RegCode_for_size_8 "l"
#define SIZEOF_BytePack_field_u32 4
#define PTX_REG_BytePack_field_u32 "r"
#define RegSize_for_size_1 4
#define RegSize_for_size_2 2
#define RegSize_for_size_4 4
#define RegSize_for_size_8 8
#define SIZEOF_BytePack_field_u64 8
#define PTX_REG_BytePack_field_u64 "l"
#define PtxAcc_for_u32
#define PtxAcc_for_s32
#define PtxAcc_for_s64
#define PtxAcc_for_u64
#define PtxAcc_for_f32
#define PtxAcc_for_f64
#if CUDART_VERSION >= 12020
#define PtxAcc_for_f16 ".acc::f32"
#define PtxAcc_for_bf16 ".acc::f32"
#define PtxAcc_for_f16x2 ".acc::f32"
#define PtxAcc_for_bf16x2 ".acc::f32"
#else
#define PtxAcc_for_f16
#define PtxAcc_for_bf16
#define PtxAcc_for_f16x2
#define PtxAcc_for_bf16x2
#endif
#define PtxAcc_for_e4m3 ".acc::f16"
#define PtxAcc_for_e5m2 ".acc::f16"
#define PtxAcc_for_e4m3x4 ".acc::f16"
#define PtxAcc_for_e5m2x4 ".acc::f16"
#define DEFINE_Apply_LoadMultimem_sum(T, ptx_ty, pack_field) \
#define DEFINE_Apply_LoadMultimem_sum(T, ptx_ty, PackSize) \
template<> \
struct Apply_LoadMultimem<FuncSum<T>, SIZEOF_BytePack_field_##pack_field> { \
static constexpr int PackSize = SIZEOF_BytePack_field_##pack_field; \
__device__ static BytePack<PackSize> load(FuncSum<T> fn, uintptr_t addr) { \
BytePack<PackSize> ans; \
asm volatile("multimem.ld_reduce.relaxed.sys.global.add." #ptx_ty " %0, [%1];" \
: "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field) \
struct Apply_LoadMultimem<FuncSum<T>, PackSize> { \
__device__ __forceinline__ static BytePack<PackSize> load(FuncSum<T> fn, uintptr_t addr) { \
BytePack<RegSize_for_size_##PackSize> reg; \
asm volatile("multimem.ld_reduce.relaxed.sys.global.add" PtxAcc_for_##ptx_ty "." #ptx_ty " %0, [%1];" \
: "=" RegCode_for_size_##PackSize(reg.native) \
: "l"(addr) : "memory"); \
BytePack<PackSize> ans; \
ans.native = reg.native; \
return ans; \
} \
};
#define DEFINE_Apply_LoadMultimem_minmax(T, ptx_ty, pack_field) \
#define DEFINE_Apply_LoadMultimem_minmax(T, ptx_ty, PackSize) \
template<> \
struct Apply_LoadMultimem<FuncMinMax<T>, SIZEOF_BytePack_field_##pack_field> { \
static constexpr int PackSize = SIZEOF_BytePack_field_##pack_field; \
__device__ static BytePack<PackSize> load(FuncMinMax<T> fn, uintptr_t addr) { \
BytePack<PackSize> ans; \
struct Apply_LoadMultimem<FuncMinMax<T>, PackSize> { \
__device__ __forceinline__ static BytePack<PackSize> load(FuncMinMax<T> fn, uintptr_t addr) { \
BytePack<RegSize_for_size_##PackSize> reg; \
if (fn.isMinNotMax) { \
asm volatile("multimem.ld_reduce.relaxed.sys.global.min." #ptx_ty " %0, [%1];" \
: "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field) \
: "=" RegCode_for_size_##PackSize(reg.native) \
: "l"(addr) : "memory"); \
} else { \
asm volatile("multimem.ld_reduce.relaxed.sys.global.max." #ptx_ty " %0, [%1];" \
: "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field) \
: "=" RegCode_for_size_##PackSize(reg.native) \
: "l"(addr) : "memory"); \
} \
BytePack<PackSize> ans; \
ans.native = reg.native; \
return ans; \
} \
};
#define DEFINE_Apply_LoadMultimem_sum_v4(T, ptx_ty, pack_field) \
#define DEFINE_Apply_LoadMultimem_sum_v4(T, ptx_ty, VecEltSize) \
template<> \
struct Apply_LoadMultimem<FuncSum<T>, 4*(SIZEOF_BytePack_field_##pack_field)> { \
static constexpr int PackSize = 4*(SIZEOF_BytePack_field_##pack_field); \
__device__ static BytePack<PackSize> load(FuncSum<T> fn, uintptr_t addr) { \
BytePack<PackSize> ans; \
asm volatile("multimem.ld_reduce.relaxed.sys.global.add.v4." #ptx_ty " {%0,%1,%2,%3}, [%4];" \
: "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[0]), \
"=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[1]), \
"=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[2]), \
"=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[3]) \
struct Apply_LoadMultimem<FuncSum<T>, 4*(VecEltSize)> { \
static constexpr int PackSize = 4*(VecEltSize); \
__device__ __forceinline__ static BytePack<PackSize> load(FuncSum<T> fn, uintptr_t addr) { \
union { BytePack<PackSize> ans; BytePack<VecEltSize> elts[4]; }; \
asm volatile("multimem.ld_reduce.relaxed.sys.global.add" PtxAcc_for_##ptx_ty ".v4." #ptx_ty " {%0,%1,%2,%3}, [%4];" \
: "=" RegCode_for_size_##VecEltSize(elts[0].native), \
"=" RegCode_for_size_##VecEltSize(elts[1].native), \
"=" RegCode_for_size_##VecEltSize(elts[2].native), \
"=" RegCode_for_size_##VecEltSize(elts[3].native) \
: "l"(addr) : "memory"); \
return ans; \
} \
};
#define DEFINE_Apply_LoadMultimem_minmax_v4(T, ptx_ty, pack_field) \
#define DEFINE_Apply_LoadMultimem_minmax_v4(T, ptx_ty, VecEltSize) \
template<> \
struct Apply_LoadMultimem<FuncMinMax<T>, 4*(SIZEOF_BytePack_field_##pack_field)> { \
static constexpr int PackSize = 4*(SIZEOF_BytePack_field_##pack_field); \
__device__ static BytePack<PackSize> load(FuncMinMax<T> fn, uintptr_t addr) { \
BytePack<PackSize> ans; \
struct Apply_LoadMultimem<FuncMinMax<T>, 4*(VecEltSize)> { \
static constexpr int PackSize = 4*(VecEltSize); \
__device__ __forceinline__ static BytePack<PackSize> load(FuncMinMax<T> fn, uintptr_t addr) { \
union { BytePack<PackSize> ans; BytePack<VecEltSize> elts[4]; }; \
if (fn.isMinNotMax) { \
asm volatile("multimem.ld_reduce.relaxed.sys.global.min.v4." #ptx_ty " {%0,%1,%2,%3}, [%4];" \
: "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[0]), \
"=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[1]), \
"=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[2]), \
"=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[3]) \
: "=" RegCode_for_size_##VecEltSize(elts[0].native), \
"=" RegCode_for_size_##VecEltSize(elts[1].native), \
"=" RegCode_for_size_##VecEltSize(elts[2].native), \
"=" RegCode_for_size_##VecEltSize(elts[3].native) \
: "l"(addr) : "memory"); \
} else { \
asm volatile("multimem.ld_reduce.relaxed.sys.global.max.v4." #ptx_ty " {%0,%1,%2,%3}, [%4];" \
: "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[0]), \
"=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[1]), \
"=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[2]), \
"=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[3]) \
: "=" RegCode_for_size_##VecEltSize(elts[0].native), \
"=" RegCode_for_size_##VecEltSize(elts[1].native), \
"=" RegCode_for_size_##VecEltSize(elts[2].native), \
"=" RegCode_for_size_##VecEltSize(elts[3].native) \
: "l"(addr) : "memory"); \
} \
return ans; \
} \
};
#define DEFINE_Apply_LoadMultimem_sum_v4x2_and_subhalf(T, ptx_ty, pack_field) \
DEFINE_Apply_LoadMultimem_sum_v4(T, ptx_ty, pack_field) \
#define DEFINE_Apply_LoadMultimem_sum_v4_and_xparts(T, ptx_ty, VecEltSize) \
DEFINE_Apply_LoadMultimem_sum_v4(T, ptx_ty, VecEltSize) \
template<> \
struct Apply_LoadMultimem<FuncSum<T>, sizeof(T)> { \
__device__ static BytePack<sizeof(T)> load(FuncSum<T> fn, uintptr_t addr) { \
BytePack<2*sizeof(T)> tmp; \
asm volatile("multimem.ld_reduce.relaxed.sys.global.add." #ptx_ty " %0, [%1];" \
: "=" PTX_REG_BytePack_field_##pack_field(tmp.pack_field) \
: "l"(addr & -uintptr_t(2*sizeof(T))) : "memory"); \
return tmp.half[(addr/sizeof(T))%2]; \
__device__ __forceinline__ static BytePack<sizeof(T)> load(FuncSum<T> fn, uintptr_t addr) { \
union { BytePack<VecEltSize> tmp; BytePack<sizeof(T)> elts[(VecEltSize)/sizeof(T)]; }; \
asm volatile("multimem.ld_reduce.relaxed.sys.global.add" PtxAcc_for_##ptx_ty "." #ptx_ty " %0, [%1];" \
: "=" RegCode_for_size_##VecEltSize(tmp.native) \
: "l"(addr & -uintptr_t(VecEltSize)) : "memory"); \
return elts[(addr/sizeof(T))%((VecEltSize)/sizeof(T))]; \
} \
};
#define DEFINE_Apply_LoadMultimem_minmax_v4x2_and_subhalf(T, ptx_ty, pack_field) \
DEFINE_Apply_LoadMultimem_minmax_v4(T, ptx_ty, pack_field) \
#define DEFINE_Apply_LoadMultimem_minmax_v4_and_xparts(T, ptx_ty, VecEltSize) \
DEFINE_Apply_LoadMultimem_minmax_v4(T, ptx_ty, VecEltSize) \
template<> \
struct Apply_LoadMultimem<FuncMinMax<T>, sizeof(T)> { \
__device__ static BytePack<sizeof(T)> load(FuncMinMax<T> fn, uintptr_t addr) { \
BytePack<2*sizeof(T)> tmp; \
__device__ __forceinline__ static BytePack<sizeof(T)> load(FuncMinMax<T> fn, uintptr_t addr) { \
union { BytePack<VecEltSize> tmp; BytePack<sizeof(T)> elts[(VecEltSize)/sizeof(T)]; }; \
if (fn.isMinNotMax) { \
asm volatile("multimem.ld_reduce.relaxed.sys.global.min." #ptx_ty " %0, [%1];" \
: "=" PTX_REG_BytePack_field_##pack_field(tmp.pack_field) \
: "l"(addr & -uintptr_t(2*sizeof(T))) : "memory"); \
: "=" RegCode_for_size_##VecEltSize(tmp.native) \
: "l"(addr & -uintptr_t(VecEltSize)) : "memory"); \
} else { \
asm volatile("multimem.ld_reduce.relaxed.sys.global.max." #ptx_ty " %0, [%1];" \
: "=" PTX_REG_BytePack_field_##pack_field(tmp.pack_field) \
: "l"(addr & -uintptr_t(2*sizeof(T))) : "memory"); \
: "=" RegCode_for_size_##VecEltSize(tmp.native) \
: "l"(addr & -uintptr_t(VecEltSize)) : "memory"); \
} \
return tmp.half[(addr/sizeof(T))%2]; \
return elts[(addr/sizeof(T))%((VecEltSize)/sizeof(T))]; \
} \
};
template<typename Fn, int BytePerPack>
struct Apply_LoadMultimem {
__device__ static BytePack<BytePerPack> load(Fn fn, uintptr_t addr) {
__device__ __forceinline__ static BytePack<BytePerPack> load(Fn fn, uintptr_t addr) {
//__trap();
return {};
}
@@ -889,29 +1047,43 @@ struct Apply_LoadMultimem {
/*multimem.ld_reduce not supported:*/ 0;
};
DEFINE_Apply_LoadMultimem_sum(uint32_t, u32, u32)
DEFINE_Apply_LoadMultimem_minmax(uint32_t, u32, u32)
DEFINE_Apply_LoadMultimem_sum(uint32_t, u32, 4)
DEFINE_Apply_LoadMultimem_minmax(uint32_t, u32, 4)
DEFINE_Apply_LoadMultimem_sum(int32_t, s32, u32)
DEFINE_Apply_LoadMultimem_minmax(int32_t, s32, u32)
DEFINE_Apply_LoadMultimem_sum(int32_t, s32, 4)
DEFINE_Apply_LoadMultimem_minmax(int32_t, s32, 4)
DEFINE_Apply_LoadMultimem_sum(uint64_t, u64, u64)
DEFINE_Apply_LoadMultimem_minmax(uint64_t, u64, u64)
DEFINE_Apply_LoadMultimem_sum(uint64_t, u64, 8)
DEFINE_Apply_LoadMultimem_minmax(uint64_t, u64, 8)
DEFINE_Apply_LoadMultimem_sum(int64_t, u64, u64)
DEFINE_Apply_LoadMultimem_minmax(int64_t, s64, u64)
DEFINE_Apply_LoadMultimem_sum(int64_t, u64, 8)
DEFINE_Apply_LoadMultimem_minmax(int64_t, s64, 8)
DEFINE_Apply_LoadMultimem_sum(float, f32, u32)
DEFINE_Apply_LoadMultimem_sum_v4(float, f32, u32)
DEFINE_Apply_LoadMultimem_sum(float, f32, 4)
DEFINE_Apply_LoadMultimem_sum_v4(float, f32, 4)
DEFINE_Apply_LoadMultimem_sum(double, f64, u64)
DEFINE_Apply_LoadMultimem_sum(double, f64, 8)
DEFINE_Apply_LoadMultimem_sum_v4x2_and_subhalf(half, f16x2, u32)
DEFINE_Apply_LoadMultimem_minmax_v4x2_and_subhalf(half, f16x2, u32)
DEFINE_Apply_LoadMultimem_sum_v4_and_xparts(half, f16x2, 4)
DEFINE_Apply_LoadMultimem_minmax_v4_and_xparts(half, f16x2, 4)
#if defined(RCCL_BFLOAT16)
DEFINE_Apply_LoadMultimem_sum_v4x2_and_subhalf(hip_bfloat16, bf16x2, u32)
DEFINE_Apply_LoadMultimem_minmax_v4x2_and_subhalf(hip_bfloat16, bf16x2, u32)
DEFINE_Apply_LoadMultimem_sum_v4_and_xparts(hip_bfloat16, bf16x2, 4)
DEFINE_Apply_LoadMultimem_minmax_v4_and_xparts(hip_bfloat16, bf16x2, 4)
#endif
#if defined(RCCL_BFLOAT16)
#if NCCL_CUDA_ARCH_FAMILY_SPECIFIC == 1000 || NCCL_CUDA_ARCH_FAMILY_SPECIFIC == 1010 || NCCL_CUDA_ARCH_SPECIFIC == 1200 || NCCL_CUDA_ARCH_SPECIFIC == 1210
DEFINE_Apply_LoadMultimem_sum_v4_and_xparts(__nv_fp8_e4m3, e4m3x4, 4)
DEFINE_Apply_LoadMultimem_minmax_v4_and_xparts(__nv_fp8_e4m3, e4m3x4, 4)
DEFINE_Apply_LoadMultimem_sum_v4_and_xparts(__nv_fp8_e5m2, e5m2x4, 4)
DEFINE_Apply_LoadMultimem_minmax_v4_and_xparts(__nv_fp8_e5m2, e5m2x4, 4)
#else
DEFINE_Apply_LoadMultimem_sum_v4_and_xparts(rccl_float8, e4m3x4, 4)
DEFINE_Apply_LoadMultimem_minmax_v4_and_xparts(rccl_float8, e4m3x4, 4)
DEFINE_Apply_LoadMultimem_sum_v4_and_xparts(rccl_bfloat8, e5m2x4, 4)
DEFINE_Apply_LoadMultimem_minmax_v4_and_xparts(rccl_bfloat8, e5m2x4, 4)
#endif
#endif
#else
template<typename Fn>
@@ -923,11 +1095,29 @@ struct Apply_LoadMultimem {
#undef DEFINE_Apply_LoadMultimem
#undef DEFINE_Apply_LoadMultimem_v4
#undef DEFINE_Apply_LoadMultimem_v4x2_and_subhalf
#undef SIZEOF_BytePack_field_u64
#undef PTX_REG_BytePack_field_u64
#undef SIZEOF_BytePack_field_u32
#undef PTX_REG_BytePack_field_u32
#undef SIZEOF_BytePack_field_u16
#undef PTX_REG_BytePack_field_u16
#undef RegCode_for_size_2
#undef RegCode_for_size_4
#undef RegCode_for_size_8
#undef RegSize_for_size_1
#undef RegSize_for_size_2
#undef RegSize_for_size_4
#undef RegSize_for_size_8
#undef PtxAcc_for_u32
#undef PtxAcc_for_s32
#undef PtxAcc_for_s64
#undef PtxAcc_for_u64
#undef PtxAcc_for_f32
#undef PtxAcc_for_f64
#undef PtxAcc_for_f16
#undef PtxAcc_for_bf16
#undef PtxAcc_for_f16x2
#undef PtxAcc_for_bf16x2
#undef PtxAcc_for_e4m3
#undef PtxAcc_for_e5m2
#undef PtxAcc_for_e4m3x4
#undef PtxAcc_for_e5m2x4
#endif // REDUCE_KERNEL_H_
src/device/reduce_scatter.h
+186 -62
Просмотреть файл
@@ -235,82 +235,206 @@ struct RunWorkColl<ncclFuncReduceScatter, T, RedOp, NCCL_ALGO_PAT, NCCL_PROTO_SI
template<typename T, typename RedOp>
struct RunWorkColl<ncclFuncReduceScatter, T, RedOp, NCCL_ALGO_NVLS, NCCL_PROTO_SIMPLE> {
template<bool ReduceSendNotRecv>
struct Scatterer {
struct ncclDevWorkColl* work;
int chunkCount;
ssize_t railGridOffset;
template<int SlicePerChunk, int MinSrcs, int MaxSrcs, int MinDsts, int MaxDsts, int MultimemSrcs, int MultimemDsts>
__device__ __forceinline__ void operator()(
int tid, int tn, int slice, int maxSliceSize,
int nSrcs, void** srcPtrs, int nDsts, void** dstPtrs, int32_t* dstSizes, uint32_t sendDirectFlag, uint32_t recvDirectFlag
) {
static_assert(SlicePerChunk == 1, "require: SlicePerChunk==1");
static_assert(MaxDsts <= 1 || MaxSrcs <= 1, "require: MaxDsts<=1 || MaxSrcs<=1");
struct ncclNvls* nvls = &ncclShmem.channel.nvls;
int nNodes = ncclShmem.comm.nNodes;
int nRails = nvls->nHeads;
int part = ncclShmem.channelId - work->channelLo;
void* inbuf = (void*)work->sendbuff;
ssize_t countPerRank = work->collnet.count;
ssize_t railAllBeg = min(railGridOffset + part * chunkCount, nNodes * countPerRank);
ssize_t railAllEnd = min(railAllBeg + chunkCount, nNodes * countPerRank);
int railAllSize = railAllEnd - railAllBeg;
int rail = nvls->headRank;
int dst = 0;
if (ReduceSendNotRecv) {
if (work->regUsed) return;
rail = 0;
nSrcs = 1;
} else {
rail = nvls->headRank;
}
if (tid < nDsts) dstSizes[tid] = railAllSize;
do {
int node = railAllBeg / countPerRank;
int railAllOffset = 0;
while (railAllOffset < railAllSize) {
ssize_t railOneBeg = node * countPerRank;
ssize_t railOneEnd = railOneBeg + countPerRank;
ssize_t railOneOffset = (railAllBeg + railAllOffset) - railOneBeg;
int delta = min(railAllEnd, railOneEnd) - (railAllBeg + railAllOffset);
int rank = ncclShmem.comm.collNetDenseToUserRank[node * nRails + rail];
ssize_t userOneBeg = rank * countPerRank + railOneOffset;
if (nDsts != 0) {
reduceCopy<ncclCollUnroll(), USE_ACC, RedOp, T,
/*MultimemSrcs=*/MultimemSrcs, 1, 1 + MaxSrcs,
/*MultimemDsts,MinDsts,MaxDsts=*/MultimemDsts, 1, 1,
/*PreOpSrcs=*/1>
(tid, tn, work->redOpArg, &work->redOpArg, false,
/*nSrcs=*/nSrcs, [=]__device__(int s) {
return work->regUsed ? (T*)srcPtrs[s] + userOneBeg :
!ReduceSendNotRecv ? (T*)srcPtrs[s] + railAllOffset:
(T*)inbuf + userOneBeg;
},
/*nDsts=*/1, [=]__device__(int d/*==0*/) {
return (T*)dstPtrs[dst] + railAllOffset;
}, delta);
}
railAllOffset += delta;
node += 1;
}
dst += 1;
rail += 1;
} while (ReduceSendNotRecv && dst < nRails);
}
};
__device__ __forceinline__ void run(int tid, int/*nthreads*/, struct ncclDevWorkColl* work) {
struct ncclNvls* nvls = &ncclShmem.channel.nvls;
size_t count;
size_t gridOffset;
size_t channelCount;
size_t chunkCount;
ncclCollCbdPart(work, ncclShmem.channelId, NCCL_PROTO_SIMPLE, sizeof(T), &count, &gridOffset, &channelCount, &chunkCount);
const int rank = ncclShmem.comm.rank;
const int nranks = ncclShmem.comm.nRanks;
size_t offset;
int nelem;
/* if we are direct NVLS, we only need to allocate 1 warp to scatter for sync;
* if not, based on #ranks, we allocate 7 or 5 warps to reduce to saturate bandwidth
* and the rest are allocated to scatter. */
const int nThreadsReduce = work->regUsed ? (NCCL_MAX_NTHREADS - WARP_SIZE) : (nranks <= 6 ? 7 * WARP_SIZE : 5 * WARP_SIZE);
const int nThreadsScatter = work->regUsed ? WARP_SIZE : (NCCL_MAX_NTHREADS - nThreadsReduce);
const int tidEndScatter = nThreadsScatter;
const int nThreadsNetRecv = work->oneNode ? 0 : (work->netRegUsed ? WARP_SIZE : 6 * WARP_SIZE);
const int nThreadsScatter = work->regUsed ? roundUp(nvls->nHeads << 2, WARP_SIZE) : 8 * WARP_SIZE;
const int nThreadsReduce = NCCL_MAX_NTHREADS - nThreadsNetRecv - nThreadsScatter;
const int tidEndNetRecv = nThreadsNetRecv;
const int tidEndScatter = tidEndNetRecv + nThreadsScatter;
const int tidEndReduce = tidEndScatter + nThreadsReduce;
if (!work->regUsed) {
if (tid < tidEndScatter) {
// Scatter
using Proto = ProtoSimple<1, 1, USE_ACC, COLL_UNROLL>;
Primitives<T, RedOp, FanAsymmetric<0, NCCL_MAX_NVLS_ARITY>, /*Direct=*/0, Proto, 0>
prims(tid, nThreadsScatter, NULL, nvls->up, work->sendbuff, NULL,
work->redOpArg, 0 * Proto::MaxGroupWidth, 1, 1);
for (size_t elemOffset = 0; elemOffset < channelCount; elemOffset += chunkCount) {
offset = gridOffset + elemOffset;
nelem = min(chunkCount, channelCount - elemOffset);
prims.scatter(offset, nvls->nHeads * count, nelem, count, -1, 0);
if (work->oneNode) {
const int rank = ncclShmem.comm.rank;
size_t offset;
size_t count, gridOffset, channelCount, chunkCount;
ncclCollCbdPart(work, ncclShmem.channelId, NCCL_PROTO_SIMPLE, sizeof(T), &count, &gridOffset, &channelCount, &chunkCount);
if (!work->regUsed) {
if (tid < tidEndScatter) {
// Scatter
using Proto = ProtoSimple<1, 1, USE_ACC, COLL_UNROLL>;
Primitives<T, RedOp, FanAsymmetric<0, NCCL_MAX_NVLS_ARITY>, /*Direct=*/0, Proto, 0>
prims(tid, nThreadsScatter, NULL, nvls->up, work->sendbuff, NULL,
work->redOpArg, 0 * Proto::MaxGroupWidth, 1, 1);
for (size_t elemOffset = 0; elemOffset < channelCount; elemOffset += chunkCount) {
offset = gridOffset + elemOffset;
nelem = min(chunkCount, channelCount - elemOffset);
prims.scatter(offset, nvls->nHeads * count, nelem, count, -1, 0);
}
// coverity[overrun-call] => Coverity think prims.index can be greater than 1
} else if (tid < tidEndReduce) {
// Reduce through NVLS
using Proto = ProtoSimple<1, 1, USE_ACC, COLL_UNROLL, 1, 0>;
Primitives<T, RedOp, FanAsymmetric<1, 0>, /*Direct=*/0, Proto, 0>
prims(tid - tidEndScatter, nThreadsReduce, &nvls->down, NULL, NULL, work->recvbuff,
work->redOpArg, 3 * Proto::MaxGroupWidth, 0, 0);
for (size_t elemOffset = 0; elemOffset < channelCount; elemOffset += chunkCount) {
offset = gridOffset + elemOffset;
nelem = min(chunkCount, channelCount - elemOffset);
prims.recv(offset, nelem);
}
}
// coverity[overrun-call] => Coverity think prims.index can be greater than 1
} else if (tid < tidEndReduce) {
// Reduce through NVLS
using Proto = ProtoSimple<1, 1, USE_ACC, COLL_UNROLL, 1, 0>;
Primitives<T, RedOp, FanAsymmetric<1, 0>, /*Direct=*/0, Proto, 0>
prims(tid - tidEndScatter, nThreadsReduce, &nvls->down, NULL, NULL, work->recvbuff,
work->redOpArg, 3 * Proto::MaxGroupWidth, 0, 0);
for (size_t elemOffset = 0; elemOffset < channelCount; elemOffset += chunkCount) {
offset = gridOffset + elemOffset;
nelem = min(chunkCount, channelCount - elemOffset);
prims.recv(offset, nelem);
} else {
if (tid < tidEndScatter) {
// Scatter
using Proto = ProtoSimple<1, 1, USE_ACC, COLL_UNROLL>;
Primitives<T, RedOp, FanSymmetric<NCCL_MAX_NVLS_ARITY>, /*Direct=*/0, Proto, 0>
prims(tid, nThreadsScatter, nvls->up, nvls->up, NULL, NULL,
work->redOpArg, 0 * Proto::MaxGroupWidth, 1, 1);
for (size_t elemOffset = 0; elemOffset < channelCount; elemOffset += chunkCount) {
prims.scatter(0, 0, 0, 0, -1, 0);
}
/* gather used as sync */
prims.gather(0, 0, 0, 0, -1, 0);
} else if (tid < tidEndReduce) {
// Reduce through NVLS
using Proto = ProtoSimple<1, 1, USE_ACC, COLL_UNROLL, 1, 0>;
Primitives<T, RedOp, FanSymmetric<1>, /*Direct=*/1, Proto, 0>
prims(tid - tidEndScatter, nThreadsReduce, &nvls->down, &nvls->down, NULL, work->recvbuff,
work->redOpArg, 3 * Proto::MaxGroupWidth, 0, 0, work);
for (size_t elemOffset = 0; elemOffset < channelCount; elemOffset += chunkCount) {
size_t outOffset = gridOffset + elemOffset;
size_t inpOffset = outOffset + rank * count;
nelem = min(chunkCount, channelCount - elemOffset);
// Coverity complains about a possible overrun inside the method invoked below, but that's actually
// a false positive.
// coverity[overrun-call:FALSE]
prims.directRecvCopy(inpOffset, outOffset, nelem);
}
/* send for sync */
prims.send(0, 0);
}
}
} else {
if (tid < tidEndScatter) {
// Scatter
// multi-node
int nNodes = ncclShmem.comm.nNodes;
int part = ncclShmem.channelId - work->channelLo;
ssize_t countPerRank = work->collnet.count;
const int nChannels = work->channelHi - work->channelLo + 1;
ssize_t chunkCount = work->collnet.chunkCount;
if (tid < tidEndNetRecv) {
using Proto = ProtoSimple<1, 1, USE_ACC, COLL_UNROLL>;
Primitives<T, RedOp, FanSymmetric<NCCL_MAX_NVLS_ARITY>, /*Direct=*/0, Proto, 0>
prims(tid, nThreadsScatter, nvls->up, nvls->up, NULL, NULL,
work->redOpArg, 0 * Proto::MaxGroupWidth, 1, 1);
for (size_t elemOffset = 0; elemOffset < channelCount; elemOffset += chunkCount) {
prims.scatter(0, 0, 0, 0, -1, 0);
if (work->netRegUsed) {
if (tid == 0) {
int steps = (int)divUp(nNodes * countPerRank, nChannels * chunkCount);
Primitives<T, RedOp, FanAsymmetric<1, 0>, /*Direct=*/0, Proto, 0>::recvPeerNotify(nvls->out, 0, steps);
}
__syncwarp();
} else {
Primitives<T, RedOp, FanAsymmetric<1, 0>, /*Direct=*/0, Proto, 0>
prims(tid, nThreadsNetRecv, &nvls->out, nullptr, nullptr, work->recvbuff,
work->redOpArg, 0 * Proto::MaxGroupWidth, 0, 0);
for (ssize_t railGridOffset = 0; railGridOffset < nNodes * countPerRank; railGridOffset += nChannels * chunkCount) {
ssize_t railAllBeg = railGridOffset + part * chunkCount;
ssize_t railAllEnd = min(railAllBeg + chunkCount, nNodes * countPerRank);
ssize_t railOneBeg = ncclShmem.comm.node * countPerRank;
ssize_t railOneEnd = railOneBeg + countPerRank;
ssize_t beg = max(railAllBeg, railOneBeg);
ssize_t end = min(railAllEnd, railOneEnd);
prims.recv(beg - railOneBeg, max(ssize_t(0), end - beg), /*postOp=*/true);
}
}
/* gather used as sync */
prims.gather(0, 0, 0, 0, -1, 0);
} else if (tid < tidEndReduce) {
// Reduce through NVLS
using Proto = ProtoSimple<1, 1, USE_ACC, COLL_UNROLL, 1, 0>;
Primitives<T, RedOp, FanSymmetric<1>, /*Direct=*/1, Proto, 0>
prims(tid - tidEndScatter, nThreadsReduce, &nvls->down, &nvls->down, NULL, work->recvbuff,
work->redOpArg, 3 * Proto::MaxGroupWidth, 0, 0, work);
for (size_t elemOffset = 0; elemOffset < channelCount; elemOffset += chunkCount) {
size_t outOffset = gridOffset + elemOffset;
size_t inpOffset = outOffset + rank * count;
nelem = min(chunkCount, channelCount - elemOffset);
// Coverity complains about a possible overrun inside the method invoked below, but that's actually
// a false positive.
// coverity[overrun-call:FALSE]
prims.directRecvCopy(inpOffset, outOffset, nelem);
} else {
if (tid < tidEndScatter) {
using Proto = ProtoSimple<1, 1, USE_ACC, COLL_UNROLL>;
Primitives<T, RedOp, FanAsymmetric<0, NCCL_MAX_NVLS_ARITY>, /*Direct=*/1, Proto, 0>
prims(tid - tidEndNetRecv, nThreadsScatter, nullptr, nvls->up, work->sendbuff, nullptr,
work->redOpArg, 1 * Proto::MaxGroupWidth, 1, 1, work);
for (ssize_t railGridOffset = 0; railGridOffset < nNodes * countPerRank; railGridOffset += nChannels * chunkCount) {
Scatterer</*ReduceSendNotRecv=*/true> scat;
scat.work = work;
scat.chunkCount = chunkCount;
scat.railGridOffset = railGridOffset;
prims.template process</*Recv=*/0, /*Send=*/1>(scat);
}
} else if (tid < tidEndReduce) {
using Proto = ProtoSimple<1, 1, USE_ACC, COLL_UNROLL, 1, 0>;
Primitives<T, RedOp, FanSymmetric<1>, /*Direct=*/1, Proto, 0>
prims(tid - tidEndScatter, nThreadsReduce, &nvls->down, &nvls->out, nullptr, nullptr,
work->redOpArg, 2 * Proto::MaxGroupWidth, 0, 1, work);
for (ssize_t railGridOffset = 0; railGridOffset < nNodes * countPerRank; railGridOffset += nChannels * chunkCount) {
Scatterer</*ReduceSendNotRecv=*/false> scat;
scat.work = work;
scat.chunkCount = chunkCount;
scat.railGridOffset = railGridOffset;
prims.template process</*Recv=*/1, /*Send=*/1>(scat);
}
}
/* send for sync */
prims.send(0, 0);
}
}
}
@@ -324,7 +448,7 @@ struct RunWorkColl<ncclFuncReduceScatter, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NC
int chunkSize;
ssize_t railGridOffset;
template<int SlicePerChunk, int MinSrcs, int MaxSrcs, int MinDsts, int MaxDsts>
template<int SlicePerChunk, int MinSrcs, int MaxSrcs, int MinDsts, int MaxDsts, int MultimemSrcs, int MultimemDsts>
__device__ __forceinline__ void operator()(
int tid, int tn, int slice, int maxSliceSize,
int nSrcs, void** srcPtrs, int nDsts, void** dstPtrs, int32_t* dstSizes, uint32_t sendDirectFlag, uint32_t recvDirectFlag
@@ -363,7 +487,7 @@ struct RunWorkColl<ncclFuncReduceScatter, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NC
int rank = ncclShmem.comm.collNetDenseToUserRank[node*nRails + rail];
ssize_t userOneBeg = rank*countPerRank + railOneOffset;
if (nDsts != 0) {
reduceCopy<ncclCollUnroll(), RedOp, T,
reduceCopy<ncclCollUnroll(), USE_ACC, RedOp, T,
/*MultimemSrcs=*/0, 1+MinSrcs, 1+MaxSrcs,
/*MultimemDsts,MinDsts,MaxDsts=*/0,1,1,
/*PreOpSrcs=*/1>
src/device/symmetric/all_gather.cuh
+367
Просмотреть файл
@@ -0,0 +1,367 @@
#include "symmetric.h"
#include "symmetric/kernel.h"
#include "symmetric/primitives.h"
template<int BytePerPack, int UnrollPacks, int UnrollPeers>
static __device__ void bcastDeep(
ncclSymPrims& prim, int tn, int t, bool waitNeeded,
char* inputHere, char* outputRank0, bool inPlace, int nIters
) {
using Pack = BytePack<BytePerPack>;
int wn = tn/WARP_SIZE;
int w = t/WARP_SIZE;
int lane = t%WARP_SIZE;
int const& rank = prim.rank;
int const& nRanks = prim.nRanks;
uint32_t const& stride4G = prim.stride4G;
Pack* inpHere = (Pack*)inputHere + intptr_t(w)*UnrollPacks*WARP_SIZE + lane;
Pack* outRank0 = (Pack*)outputRank0 + intptr_t(w)*UnrollPacks*WARP_SIZE + lane;
Pack tmp[UnrollPacks];
nIters -= w;
if (0 < nIters) {
#pragma unroll
for (int u=0; u < UnrollPacks; u++) {
tmp[u] = inpHere[u*WARP_SIZE];
}
}
if (waitNeeded) prim.barrierWait(ncclCoopCta(), /*acquire=*/false);
if (0 < nIters) {
while (true) {
int dr = inPlace ? 1 : 0;
int r = rank + dr;
if (r == nRanks) r = 0;
#pragma unroll 2
for (int partial=0; partial <= 1; partial++) {
#pragma unroll 1
for (int i = 0;
partial ? i < 1 : (dr + UnrollPeers <= nRanks);
partial ? i++ : (dr += UnrollPeers)) {
#pragma unroll
for (int ur=0; ur < UnrollPeers-partial; ur++) {
if (partial && dr == nRanks) break;
#pragma unroll UnrollPacks
for (int u=0; u < UnrollPacks; u++) {
add4G(outRank0, r*stride4G)[u*WARP_SIZE] = tmp[u];
}
if (++r == nRanks) r = 0;
}
}
}
inpHere += intptr_t(wn)*UnrollPacks*WARP_SIZE;
outRank0 += intptr_t(wn)*UnrollPacks*WARP_SIZE;
nIters -= wn;
if (nIters <= 0) break;
// Load data for next iteration.
#pragma unroll
for (int u=0; u < UnrollPacks; u++) {
tmp[u] = inpHere[u*WARP_SIZE];
}
}
}
}
template<int UnrollPeers, typename T>
static __device__ void bcastEnds(
ncclSymPrims& prim, int tn, int t,
T* inputHere, T* outputRank0, bool inPlace, size_t nElts, uint32_t nPreElts, size_t nSufElts
) {
int const& rank = prim.rank;
int const& nRanks = prim.nRanks;
uint32_t const& stride4G = prim.stride4G;
BytePack<sizeof(T)>* inpHere = (BytePack<sizeof(T)>*)inputHere;
BytePack<sizeof(T)>* outRank0 = (BytePack<sizeof(T)>*)outputRank0;
#pragma unroll 1
for (size_t i = t; i < nPreElts+nSufElts; i += tn) {
size_t elt = i < nPreElts ? i : nElts-nPreElts-nSufElts+i;
BytePack<sizeof(T)> tmp = inpHere[elt];
int dr = inPlace ? 1 : 0;
int r = rank + dr;
if (r == nRanks) r = 0;
#pragma unroll 1
for (; dr + UnrollPeers <= nRanks; dr += UnrollPeers) {
#pragma unroll UnrollPeers
for (int u=0; u < UnrollPeers; u++) {
*add4G(outRank0+elt, r*stride4G) = tmp;
if (++r == nRanks) r = 0;
}
}
#pragma unroll UnrollPeers
for (int u=0; u < UnrollPeers; u++) {
if (dr+u == nRanks) break;
*add4G(outRank0+elt, r*stride4G) = tmp;
if (++r == nRanks) r = 0;
}
}
}
template<typename T>
static __device__ void bcast(
ncclSymPrims& prim, int tn, int t, bool waitNeeded, T* input, T* output, size_t nElts
) {
bool inPlace = (input == output);
// Mpve to rank=0
output = prim.peerPtr(0, output);
uintptr_t inputUptr = reinterpret_cast<uintptr_t>(input);
uintptr_t outputUptr = reinterpret_cast<uintptr_t>(output);
size_t nBytes = nElts*sizeof(T);
uint32_t nPreBytes = (128u - inputUptr)%128u;
nPreBytes = min((size_t)nPreBytes, nBytes);
uintptr_t cursor = nPreBytes;
constexpr int MinWarpPerBlock = 4;
if ((inputUptr-outputUptr)%16 == 0) {
constexpr int BytePerPack = 16, UnrollPacks = 1, UnrollPeers = 1;
constexpr int BytePerChunk = MinWarpPerBlock*UnrollPacks*WARP_SIZE*BytePerPack;
uint32_t chunks = (nBytes-cursor)/BytePerChunk;
chunks -= imodFast32(chunks, prim.nBlocks, prim.nBlocks_rcp32);
if (chunks != 0) {
uintptr_t cursorAfter = cursor + uintptr_t(chunks)*BytePerChunk;
bcastDeep<BytePerPack, UnrollPacks, UnrollPeers>(
prim, tn, t, waitNeeded,
(char*)input + cursor, (char*)output + cursor, inPlace,
chunks*MinWarpPerBlock
);
cursor = cursorAfter;
waitNeeded = false;
}
}
if (sizeof(T) == 4 || (sizeof(T) < 4 && (inputUptr-outputUptr)%4 == 0)) {
constexpr int BytePerPack = 4, UnrollPacks = 1, UnrollPeers = 1;
constexpr int BytePerChunk = MinWarpPerBlock*UnrollPacks*WARP_SIZE*BytePerPack;
uint32_t chunks = (nBytes-cursor)/BytePerChunk;
chunks -= imodFast32(chunks, prim.nBlocks, prim.nBlocks_rcp32);
if (chunks != 0) {
uintptr_t cursorAfter = cursor + uintptr_t(chunks)*BytePerChunk;
bcastDeep<(sizeof(T) <= BytePerPack ? BytePerPack : 0), UnrollPacks, UnrollPeers>(
prim, tn, t, waitNeeded,
(char*)input + cursor, (char*)output + cursor, inPlace,
chunks*MinWarpPerBlock
);
cursor = cursorAfter;
waitNeeded = false;
}
}
if (waitNeeded) prim.barrierWait(ncclCoopCta(), /*acquire=*/false);
constexpr int UnrollPeers = 8;
size_t nSufElts = (nBytes-cursor)/sizeof(T);
bcastEnds<UnrollPeers>(prim, tn, t, input, output, inPlace, nElts, nPreBytes/sizeof(T), nSufElts);
}
__device__ __forceinline__ void ncclSymRun_AllGather_ST(ncclSymDevArgs const* args) {
ncclSymPrims prim(args->comm, ncclSymPrims_UseBarrier);
int const& rank = prim.rank;
// Threads numbered over rank.
int bt = flattenIx(threadIdx.x%WARP_SIZE, WARP_SIZE,
prim.block, prim.nBlocks,
threadIdx.x/WARP_SIZE, blockDim.x/WARP_SIZE);
int btn = prim.nBlocks*blockDim.x;
prim.barrierArrive(ncclCoopCta(), /*release=*/false);
//prim.barrierWait(ncclCoopCta(), /*acquire=*/false);
bcast(prim, btn, bt, /*waitNeeded=*/true, (char*)args->input, (char*)args->output + rank*args->nElts, args->nElts);
prim.barrierArrive(ncclCoopCta(), /*release=*/true);
prim.barrierWait(ncclCoopCta(), /*acquire=*/false);
}
template<typename T>
static __device__ void bcastMultimem(
ncclSymPrims& prim, int tn, int t, T* input, T* output, size_t nElts
) {
// Move output to multimem
output = prim.multimemPtr(output);
uintptr_t inputUptr = reinterpret_cast<uintptr_t>(input);
uintptr_t outputUptr = reinterpret_cast<uintptr_t>(output);
size_t nBytes = nElts*sizeof(T);
uint32_t nPreBytes = (16-inputUptr)%16;
nPreBytes = min((size_t)nPreBytes, nBytes);
uintptr_t nSufBytes;
if ((inputUptr-outputUptr)%16 == 0) {
constexpr int BytePerPack = 16, UnrollPacks = 8;
constexpr int BytePerChunk = UnrollPacks*WARP_SIZE*BytePerPack;
uintptr_t cursor = nPreBytes;
uint32_t nChunks = (nBytes-cursor)/BytePerChunk;
uintptr_t cursorAfter = cursor + uintptr_t(nChunks)*BytePerChunk;
nSufBytes = nBytes - cursorAfter;
cursor += (t/WARP_SIZE)*UnrollPacks*WARP_SIZE*BytePerPack;
cursor += (t%WARP_SIZE)*BytePerPack;
int nIters = nChunks - t/WARP_SIZE;
#pragma unroll 1
while (0 < nIters) {
BytePack<BytePerPack> tmp[UnrollPacks];
#pragma unroll
for (int u=0; u < UnrollPacks; u++) {
tmp[u] = *reinterpret_cast<BytePack<BytePerPack>*>(inputUptr + cursor + u*WARP_SIZE*BytePerPack);
}
#pragma unroll
for (int u=0; u < UnrollPacks; u++) {
multimem_st_global(outputUptr + cursor + u*WARP_SIZE*BytePerPack, tmp[u]);
}
cursor += tn*UnrollPacks*BytePerPack;
nIters -= tn/WARP_SIZE;
}
} else {
nPreBytes = 0;
nSufBytes = nBytes;
}
// Get the prefix+suffix element one at a time.
#pragma unroll 4
for (uintptr_t i = t*sizeof(T); i < nPreBytes + nSufBytes; i += tn*sizeof(T)) {
uintptr_t cursor = i < nPreBytes ? i : nBytes-nSufBytes+(i-nPreBytes);
BytePack<sizeof(T)> val = *reinterpret_cast<BytePack<sizeof(T)>*>(inputUptr + cursor);
multimem_st_global(outputUptr + cursor, val);
cursor += tn*sizeof(T);
}
}
__device__ __forceinline__ void ncclSymRun_AllGather_STMC(ncclSymDevArgs const* args) {
ncclSymPrims prim(args->comm, ncclSymPrims_UseBarrier|ncclSymPrims_UseMultimem);
int const& rank = prim.rank;
char* input = args->input;
char* output = args->output;
size_t bytes = args->nElts;
// Round robin memory to blocks.
int t = flattenIx(threadIdx.x%WARP_SIZE, WARP_SIZE,
prim.block, prim.nBlocks,
threadIdx.x/WARP_SIZE, blockDim.x/WARP_SIZE);
int tn = prim.nBlocks*blockDim.x;
prim.barrierArrive(ncclCoopCta(), /*release=*/false);
prim.barrierWait(ncclCoopCta(), /*acquire=*/false);
bcastMultimem(prim, tn, t, input, output + rank*bytes, bytes);
prim.barrierArrive(ncclCoopCta(), /*release=*/true);
prim.barrierWait(ncclCoopCta(), /*acquire=*/false);
}
template<typename EltType>
static __device__ void allgather_LL_body(
ncclSymPrims &prim, EltType* input, EltType* output, int nElts, int nPacks, int nStrideElts
) {
using Pack = BytePack<8>;
constexpr int EltPerPack = 8/sizeof(EltType);
ncclCoopCta cta;
int rank = prim.rank;
int nRanks = prim.nRanks;
constexpr int tn = ncclSymMaxThreads;
int t = threadIdx.x;
#pragma unroll 1
while (0 < nElts) {
int nIterPacks = min(nPacks, tn);
if (t < nIterPacks) {
Pack x = loadPack<Pack>(input, t*EltPerPack, nElts);
prim.bcastLL(/*slot=*/nIterPacks*rank + t, x);
}
int tn_div_nPacks = tn/nIterPacks;
int tn_mod_nPacks = tn%nIterPacks;
int peer = t/nIterPacks;
int pack = t%nIterPacks;
#if 1
// NOTE: Unrolling speedup on eos nranks=8 size=64K: 5.7us vs 6.7us
constexpr int Unroll = 1;
#pragma unroll 1
for (int i = t; i < (nRanks*nIterPacks & -(Unroll*tn)); i += Unroll*tn) {
Pack got[Unroll];
prim.template recvLL<Unroll, Unroll>(i, Unroll, tn, /*&*/got);
#pragma unroll
for (int u=0; u < Unroll; u++) {
storePack<Pack>(output + peer*nStrideElts, pack*EltPerPack, nElts, got[u]);
peer += tn_div_nPacks;
pack += tn_mod_nPacks;
if (nIterPacks <= pack) { peer += 1; pack -= nIterPacks; }
}
}
int i = (nRanks*nIterPacks & -(Unroll*tn)) + t;
int n = (nRanks*nIterPacks)/tn % Unroll;
if (i + n*tn < nRanks*nIterPacks) n += 1;
if (n != 0) {
Pack got[Unroll];
prim.template recvLL<1, Unroll>(i, n, tn, /*&*/got);
#pragma unroll
for (int u=0; u < Unroll; u++) {
if (u != 0 && u == n) break;
storePack(output + peer*nStrideElts, pack*EltPerPack, nElts, got[u]);
peer += tn_div_nPacks;
pack += tn_mod_nPacks;
if (nIterPacks <= pack) { peer += 1; pack -= nIterPacks; }
}
}
#else
// The non-unrolled but "obviously correct" implementation for reference.
#pragma unroll 1
for (int i = t; i < nRanks*nIterPacks; i += tn) {
Pack got = prim.template recvLL<Pack>(i);
storePack(output + peer*nStrideElts, pack*EltPerPack, nElts, got);
peer += tn_div_nPacks;
pack += tn_mod_nPacks;
if (nIterPacks <= pack) { peer += 1; pack -= nIterPacks; }
}
#endif
prim.endLL(cta);
input += tn*EltPerPack;
output += tn*EltPerPack;
nElts -= tn*EltPerPack;
nPacks -= tn;
}
}
static __device__ void ncclSymRun_AllGather_LL_impl(ncclSymDevArgs const* args, bool multimem) {
ncclSymPrims prim(args->comm, ncclSymPrims_UseLL | multimem*ncclSymPrims_UseMultimem);
using Pack = BytePack<8>;
constexpr int BytePerPack = 8;
int nElts = args->nElts;
int nPacks = divUp(nElts, BytePerPack);
uint32_t nPackPerBlock, nPackModBlock;
idivmodFast32(&nPackPerBlock, &nPackModBlock, nPacks, prim.nBlocks, prim.nBlocks_rcp32);
int blockPackBegin = prim.block*nPackPerBlock + minval<int>(prim.block, nPackModBlock);
int blockPackEnd = blockPackBegin + nPackPerBlock + (prim.block < nPackModBlock ? 1 : 0);
int nBlockPacks = blockPackEnd - blockPackBegin;
int nBlockElts = nElts - blockPackBegin*BytePerPack;
nBlockElts = min(nBlockElts, nBlockPacks*BytePerPack);
char* blockInput = args->input + blockPackBegin*BytePerPack;
char* blockOutput = args->output + blockPackBegin*BytePerPack;
uint32_t lowBits = args->nElts;
lowBits |= (uint32_t)reinterpret_cast<uintptr_t>(args->input);
lowBits |= (uint32_t)reinterpret_cast<uintptr_t>(args->output);
if (__builtin_expect(lowBits%8 == 0, true)) {
// NOTE: Specializing for 8-byte alignment in one case help at size=65K: 8.9us vs 5.6us
allgather_LL_body(prim, (BytePack<8>*)blockInput, (BytePack<8>*)blockOutput, nBlockElts/8, nBlockPacks, nElts/8);
} else {
allgather_LL_body(prim, blockInput, blockOutput, nBlockElts, nBlockPacks, nElts);
}
}
__device__ __forceinline__ void ncclSymRun_AllGather_LL(ncclSymDevArgs const* args) {
ncclSymRun_AllGather_LL_impl(args, /*multimem=*/false);
}
__device__ __forceinline__ void ncclSymRun_AllGather_LLMC(ncclSymDevArgs const* args) {
ncclSymRun_AllGather_LL_impl(args, /*multimem=*/true);
}
src/device/symmetric/all_reduce.cuh
+432
Просмотреть файл
@@ -0,0 +1,432 @@
#include "symmetric.h"
#include "symmetric/kernel.h"
#include "symmetric/primitives.h"
template<int BytePerPack, int UnrollPacks, int UnrollPeers, typename T, typename Red>
static __device__ __forceinline__ void allreduceDeep(
ncclSymPrims& prim, int tn, int t, bool waitNeeded,
Red red, char* inputRank0, char* outputRank0, int32_t nIters
) {
using Pack = BytePack<BytePerPack>;
using Acc = typename Red::EltType;
using AccPack = BytePack<BytePerPack*sizeof(Acc)/sizeof(T)>;
int wn = tn/WARP_SIZE;
int w = t/WARP_SIZE;
int lane = t%WARP_SIZE;
int const& rank = prim.rank;
int const& nRanks = prim.nRanks;
uint32_t const& stride4G = prim.stride4G;
Pack* inpRank0 = (Pack*)inputRank0 + intptr_t(w)*UnrollPacks*WARP_SIZE + lane;
Pack* outRank0 = (Pack*)outputRank0 + intptr_t(w)*UnrollPacks*WARP_SIZE + lane;
Pack acc0[UnrollPacks];
nIters -= w;
if (0 < nIters) {
#pragma unroll
for (int u=0; u < UnrollPacks; u++) {
acc0[u] = add4G(inpRank0, rank*stride4G)[u*WARP_SIZE];
}
}
if (waitNeeded) prim.barrierWait(ncclCoopCta(), /*acquire=*/false);
if (0 < nIters) {
while (true) {
AccPack acc1[UnrollPacks];
int r = rank;
if (++r == nRanks) r = 0;
{ Pack tmp1[UnrollPacks];
#pragma unroll
for (int u=0; u < UnrollPacks; u++) {
tmp1[u] = add4G(inpRank0, r*stride4G)[u*WARP_SIZE];
}
#pragma unroll
for (int u=0; u < UnrollPacks; u++) {
acc1[u] = applyReduce(red, applyCast<T, Acc>(acc0[u]), applyCast<T, Acc>(tmp1[u]));
}
}
if (++r == nRanks) r = 0;
int dr = 2;
#pragma unroll 2
for (int partial=0; partial <= 1; partial++) {
#pragma unroll 1
for (int i = 0;
partial ? i < 1 : (dr + UnrollPeers <= nRanks);
partial ? i++ : (dr += UnrollPeers)) {
if (partial && dr == nRanks) break;
Pack tmp1[UnrollPeers][UnrollPacks];
#pragma unroll
for (int ur=0; ur < UnrollPeers-partial; ur++) {
if (partial && ur!=0 && dr+ur == nRanks) break;
#pragma unroll UnrollPacks
for (int u=0; u < UnrollPacks; u++) {
tmp1[ur][u] = add4G(inpRank0, r*stride4G)[u*WARP_SIZE];
}
if (++r == nRanks) r = 0;
}
#pragma unroll
for (int ur=0; ur < UnrollPeers-partial; ur++) {
if (partial && ur!=0 && dr+ur == nRanks) break;
#pragma unroll UnrollPacks
for (int u=0; u < UnrollPacks; u++) {
acc1[u] = applyReduce(red, acc1[u], applyCast<T, Acc>(tmp1[ur][u]));
}
}
}
}
#pragma unroll
for (int u=0; u < UnrollPacks; u++) acc0[u] = applyCast<Acc, T>(acc1[u]);
dr = 0;
r = rank;
#pragma unroll 2
for (int partial=0; partial <= 1; partial++) {
#pragma unroll 1
for (int i = 0;
partial ? i < 1 : (dr + UnrollPeers <= nRanks);
partial ? i++ : (dr += UnrollPeers)) {
#pragma unroll
for (int ur=0; ur < UnrollPeers-partial; ur++) {
if (partial && dr == nRanks) break;
#pragma unroll UnrollPacks
for (int u=0; u < UnrollPacks; u++) {
add4G(outRank0, r*stride4G)[u*WARP_SIZE] = acc0[u];
}
if (++r == nRanks) r = 0;
}
}
}
inpRank0 += intptr_t(wn)*UnrollPacks*WARP_SIZE;
outRank0 += intptr_t(wn)*UnrollPacks*WARP_SIZE;
nIters -= wn;
if (nIters <= 0) break;
// Load data for next iteration.
#pragma unroll
for (int u=0; u < UnrollPacks; u++) {
acc0[u] = add4G(inpRank0, rank*stride4G)[u*WARP_SIZE];
}
}
}
}
template<int UnrollPeers, typename Red, typename T>
static __device__ __forceinline__ void allreduceEnds(
ncclSymPrims& prim, int tn, int t, Red red,
T* inputRank0, T* outputRank0, size_t nElts, uint32_t nPreElts, size_t nSufElts
) {
using Acc = typename Red::EltType;
int const& rank = prim.rank;
int const& nRanks = prim.nRanks;
uint32_t const& stride4G = prim.stride4G;
BytePack<sizeof(T)>* inpRank0 = (BytePack<sizeof(T)>*)inputRank0;
BytePack<sizeof(T)>* outRank0 = (BytePack<sizeof(T)>*)outputRank0;
#pragma unroll 1
for (size_t i = t; i < nPreElts+nSufElts; i += tn) {
size_t elt = i < nPreElts ? i : nElts-nSufElts-nPreElts+i;
BytePack<sizeof(T)> acc0 = *add4G(inpRank0+elt, rank*stride4G);
BytePack<sizeof(Acc)> acc1;
BytePack<sizeof(T)> tmp[UnrollPeers];
int dr = 1;
int r = rank+1;
if (nRanks == r) r = 0;
bool first = true;
#pragma unroll 2
for (int partial=0; partial <= 1; partial++) {
#pragma unroll 1
for (int j = 0;
partial ? j < 1 : (dr + UnrollPeers <= nRanks);
partial ? j++ : (dr += UnrollPeers)) {
if (partial && dr == nRanks) break;
#pragma unroll
for (int u=0; u < UnrollPeers-partial; u++) {
if (partial && u!=0 && dr+u == nRanks) break;
tmp[u] = *add4G(inpRank0+elt, r*stride4G);
r += 1;
if (r == nRanks) r = 0;
}
if (first) {
first = false;
acc1 = applyCast<T, Acc>(acc0);
}
#pragma unroll
for (int u=0; u < UnrollPeers-partial; u++) {
if (partial && u!=0 && dr+u == nRanks) break;
acc1 = applyReduce(red, acc1, applyCast<T, Acc>(tmp[u]));
}
}
}
acc0 = applyCast<Acc, T>(acc1);
dr = 0;
r = rank;
#pragma unroll 2
for (int partial=0; partial <= 1; partial++) {
#pragma unroll 1
for (int j=0;
partial ? j < 1 : (dr + UnrollPeers <= nRanks);
partial ? j++ : (dr += UnrollPeers)) {
#pragma unroll
for (int u=0; u < UnrollPeers-partial; u++) {
if (partial && dr+u == nRanks) break;
*add4G(outRank0+elt, r*stride4G) = acc0;
r += 1;
if (r == nRanks) r = 0;
}
}
}
}
}
template<typename Red, typename T>
static __device__ void allreduce(
ncclSymPrims& prim, int tn, int t, bool waitNeeded,
Red red, T* input, T* output, size_t nElts
) {
int nRanks = prim.nRanks;
int nBlocks = prim.nBlocks;
// Mpve to rank=0
input = prim.peerPtr(0, input);
output = prim.peerPtr(0, output);
uintptr_t inputUptr = reinterpret_cast<uintptr_t>(input);
uintptr_t outputUptr = reinterpret_cast<uintptr_t>(output);
size_t nBytes = nElts*sizeof(T);
uint32_t nPreBytes = (16u - inputUptr)%16u;
nPreBytes = min((size_t)nPreBytes, nBytes);
uintptr_t cursor = nPreBytes;
constexpr int MinWarpPerBlock = 4;
if ((inputUptr-outputUptr)%16 == 0) {
constexpr int BytePerPack = 16, UnrollPacks = 4, UnrollPeers = 2;
constexpr int BytePerChunk = MinWarpPerBlock*UnrollPacks*WARP_SIZE*BytePerPack;
uint32_t chunks = (nBytes-cursor)/BytePerChunk;
chunks -= imodFast32(chunks, nRanks*nBlocks, prim.nRanks_nBlocks_rcp32);
if (chunks != 0) {
uintptr_t cursorAfter = cursor + uintptr_t(chunks)*BytePerChunk;
allreduceDeep<BytePerPack, UnrollPacks, UnrollPeers, T>(
prim, tn, t, waitNeeded, red,
(char*)input + cursor, (char*)output + cursor,
chunks*MinWarpPerBlock
);
cursor = cursorAfter;
waitNeeded = false;
}
}
if (sizeof(T) == 4 || (sizeof(T) < 4 && (inputUptr-outputUptr)%4 == 0)) {
constexpr int BytePerPack = 4, UnrollPacks = 4, UnrollPeers = 4;
constexpr int BytePerChunk = MinWarpPerBlock*UnrollPacks*WARP_SIZE*BytePerPack;
uint32_t chunks = (nBytes-cursor)/BytePerChunk;
chunks -= imodFast32(chunks, nRanks*nBlocks, prim.nRanks_nBlocks_rcp32);
if (chunks != 0) {
uintptr_t cursorAfter = cursor + uintptr_t(chunks)*BytePerChunk;
allreduceDeep<(sizeof(T) <= BytePerPack ? BytePerPack : 0), UnrollPacks, UnrollPeers, T>(
prim, tn, t, waitNeeded, red,
(char*)input + cursor, (char*)output + cursor,
chunks*MinWarpPerBlock
);
cursor = cursorAfter;
waitNeeded = false;
}
}
if (waitNeeded) prim.barrierWait(ncclCoopCta(), /*acquire=*/false);
constexpr int UnrollPeers = 8;
size_t nSufElts = (nBytes-cursor)/sizeof(T);
allreduceEnds<UnrollPeers>(prim, tn, t, red, input, output, nElts, nPreBytes/sizeof(T), nSufElts);
}
template<template<typename> typename Red, typename T>
__device__ __forceinline__ void ncclSymRun_AllReduce_RSxLD_AGxST(ncclSymDevArgs const* args) {
ncclSymPrims prim(args->comm, ncclSymPrims_UseBarrier);
int /*const&*/ rank = prim.rank;
int /*const&*/ nRanks = prim.nRanks;
Red<typename ncclSymAccumType<Red, T, /*nvls=*/false>::Type> red(args->redOpArg);
// Threads numbered globally such that we round robin warps by rank then block.
int gt = flattenIx(threadIdx.x%WARP_SIZE, WARP_SIZE,
rank, nRanks,
prim.block, prim.nBlocks,
threadIdx.x/WARP_SIZE, blockDim.x/WARP_SIZE);
int gtn = nRanks*prim.nBlocks*blockDim.x;
prim.barrierArrive(ncclCoopCta(), /*release=*/false);
//prim.barrierWait(ncclCoopCta(), /*acquire=*/false);
allreduce(prim, gtn, gt, /*waitNeeded=*/true, red, (T*)args->input, (T*)args->output, args->nElts);
prim.barrierArrive(ncclCoopCta(), /*release=*/true);
prim.barrierWait(ncclCoopCta(), /*acquire=*/false);
}
template<typename Red, typename T>
static __device__ void allreduceMultimem(
ncclSymPrims& prim, int tn, int t, Red red, T* input, T* output, size_t nElts
) {
// Mpve to multimem
input = prim.multimemPtr(input);
output = prim.multimemPtr(output);
uintptr_t inputUptr = reinterpret_cast<uintptr_t>(input);
uintptr_t outputUptr = reinterpret_cast<uintptr_t>(output);
size_t nBytes = nElts*sizeof(T);
constexpr int BytePerPack = LoadMultimem_BigPackSize<Red>::BigPackSize;
uint32_t nPreBytes = (BytePerPack - inputUptr)%BytePerPack;
nPreBytes = min((size_t)nPreBytes, nBytes);
uintptr_t nSufBytes;
if (alignof(T) == BytePerPack || (inputUptr-outputUptr)%BytePerPack == 0) {
constexpr int UnrollPacks = 16*8/BytePerPack;
constexpr int BytePerChunk = UnrollPacks*WARP_SIZE*BytePerPack;
uintptr_t cursor = nPreBytes;
int nChunks = (nBytes-cursor)/BytePerChunk;
uintptr_t cursorAfter = cursor + uintptr_t(nChunks)*BytePerChunk;
nSufBytes = nBytes - cursorAfter;
cursor += (t/WARP_SIZE)*UnrollPacks*WARP_SIZE*BytePerPack;
cursor += (t%WARP_SIZE)*BytePerPack;
int nIters = nChunks - t/WARP_SIZE;
#pragma unroll 1
while (0 < nIters) {
BytePack<BytePerPack> tmp[UnrollPacks];
#pragma unroll
for (int u=0; u < UnrollPacks; u++) {
tmp[u] = applyLoadMultimem<Red, BytePerPack>(red, inputUptr + cursor + u*WARP_SIZE*BytePerPack);
}
#pragma unroll
for (int u=0; u < UnrollPacks; u++) {
multimem_st_global(outputUptr + cursor + u*WARP_SIZE*BytePerPack, tmp[u]);
}
cursor += tn*UnrollPacks*BytePerPack;
nIters -= tn/WARP_SIZE;
}
} else {
nPreBytes = 0;
nSufBytes = nBytes;
}
// Get the prefix+suffix element one at a time.
#pragma unroll 4
for (uintptr_t i = t*sizeof(T); i < nPreBytes + nSufBytes; i += tn*sizeof(T)) {
uintptr_t cursor = i < nPreBytes ? i : nBytes-nSufBytes+(i-nPreBytes);
BytePack<sizeof(T)> val = applyLoadMultimem<Red, sizeof(T)>(red, inputUptr + cursor);
multimem_st_global(outputUptr + cursor, val);
cursor += tn*sizeof(T);
}
}
template<template<typename> typename Red, typename T>
__device__ __forceinline__ void ncclSymRun_AllReduce_RSxLDMC_AGxSTMC(ncclSymDevArgs const* args) {
ncclSymPrims prim(args->comm, ncclSymPrims_UseBarrier|ncclSymPrims_UseMultimem);
Red<typename ncclSymAccumType<Red, T, /*nvls=*/true>::Type> red(args->redOpArg);
// Threads numbered globally such that we round robin warps by rank then block.
int gt = flattenIx(threadIdx.x%WARP_SIZE, WARP_SIZE,
prim.rank, prim.nRanks,
prim.block, prim.nBlocks,
threadIdx.x/WARP_SIZE, blockDim.x/WARP_SIZE);
int gtn = prim.nRanks*prim.nBlocks*blockDim.x;
prim.barrierArrive(ncclCoopCta(), /*release=*/false);
prim.barrierWait(ncclCoopCta(), /*acquire=*/false);
allreduceMultimem(prim, gtn, gt, red, (T*)args->input, (T*)args->output, args->nElts);
prim.barrierArrive(ncclCoopCta(), /*release=*/true);
prim.barrierWait(ncclCoopCta(), /*acquire=*/false);
}
template<template<typename> typename Red, typename T>
__device__ __forceinline__ void ncclSymRun_AllReduce_AGxLL_R_impl(ncclSymDevArgs const* args, bool multimem) {
ncclSymPrims prim(args->comm, ncclSymPrims_UseLL | multimem*ncclSymPrims_UseMultimem);
int /*const&*/ rank = prim.rank;
using Acc = typename ncclSymAccumType<Red, T, /*nvls=*/false>::Type;
Red<Acc> red(args->redOpArg);
using Pack = BytePack<8>;
using AccPack = BytePack<8*sizeof(Acc)/sizeof(T)>;
constexpr int EltPerPack = 8/sizeof(T);
int nElts = args->nElts;
int nPacks = divUp(nElts, EltPerPack);
bool packAligned = 8 <= alignof(T) || (
args->nElts*sizeof(T) |
(uint32_t)reinterpret_cast<uintptr_t>(args->input) |
(uint32_t)reinterpret_cast<uintptr_t>(args->output)
)%8 == 0;
uint32_t nPackPerBlock, nPackModBlock;
idivmodFast32(&nPackPerBlock, &nPackModBlock, nPacks, prim.nBlocks, prim.nBlocks_rcp32);
int begin = prim.block*nPackPerBlock + minval<int>(prim.block, nPackModBlock);
int end = begin + nPackPerBlock + (prim.block < nPackModBlock ? 1 : 0);
nPacks = end - begin;
nElts -= begin*EltPerPack;
nElts = min(nElts, nPacks*EltPerPack);
T* input = (T*)args->input + begin*EltPerPack;
T* output = (T*)args->output + begin*EltPerPack;
ncclCoopCta cta;
int t = threadIdx.x;
int tn = ncclSymMaxThreads;
if (__builtin_expect(packAligned, true)) {
#pragma unroll 1
while (0 < nPacks) {
if (t < nPacks) {
int nIterPacks = min(nPacks, tn);
Pack inp = loadPack<Pack>((Pack*)input, t, nPacks);
prim.bcastLL(/*slot=*/nIterPacks*rank + t, inp);
Pack out = prim.template recvReduceLL<Pack, T>(t, nIterPacks, red);
storePack((Pack*)output, t, nPacks, out);
}
prim.endLL(cta);
input += tn*EltPerPack;
output += tn*EltPerPack;
nPacks -= tn;
}
} else {
#pragma unroll 1
while (0 < nElts) {
if (t*EltPerPack < nElts) {
int nIterPacks = min(nPacks, tn);
Pack inp = loadPack<Pack>(input, t*EltPerPack, nElts);
prim.bcastLL(/*slot=*/nIterPacks*rank + t, inp);
Pack out = prim.template recvReduceLL<Pack, T>(t, nIterPacks, red);
storePack(output, t*EltPerPack, nElts, out);
}
prim.endLL(cta);
input += tn*EltPerPack;
output += tn*EltPerPack;
nElts -= tn*EltPerPack;
nPacks -= tn;
}
}
}
template<template<typename> typename Red, typename T>
__device__ __forceinline__ void ncclSymRun_AllReduce_AGxLL_R(ncclSymDevArgs const* args) {
ncclSymRun_AllReduce_AGxLL_R_impl<Red, T>(args, /*multimem=*/false);
}
template<template<typename> typename Red, typename T>
__device__ __forceinline__ void ncclSymRun_AllReduce_AGxLLMC_R(ncclSymDevArgs const* args) {
ncclSymRun_AllReduce_AGxLL_R_impl<Red, T>(args, /*multimem=*/true);
}
src/device/symmetric/generate.py
Исполняемый файл
+247
Просмотреть файл
@@ -0,0 +1,247 @@
#!/usr/bin/env python3
import os
import sys
################################################################################
# The first command line argument is the path to the directory to generate and
# populate.
gensrc = sys.argv[1]
if os.path.exists(gensrc):
for name in os.listdir(gensrc):
os.remove(os.path.join(gensrc, name))
#os.truncate(os.path.join(gensrc, name), 0)
else:
os.mkdir(gensrc)
def paste(sep, *args):
return sep.join(args)
indents = 0
def emitln(f, lines):
global indents
for ln in ((lines,) if isinstance(lines, str) else lines):
f.write(' '*indents + ln + '\n')
def indent(s):
return '\n'.join(' '+l for l in s.splitlines())
class Rec(object):
def __init__(me, **kw):
me.__dict__.update(kw)
def __eq__(x, y):
if len(x) != len(y): return False
for k in x:
if k not in y: return False
if x[k] != y[k]: return False
return True
def __hash__(me):
h = 0
for k in me.__dict__:
h += hash((k, me.__dict__[k]))
return h
################################################################################
# Edit this region for introducing new algos etc
reductions = ["AllReduce","ReduceScatter"]
all_reds = ["sum"]
all_tys = ["f32","f16","bf16","f8e4m3","f8e5m2"]
nvls_algos_by_coll = {
"AllReduce": ["AGxLLMC_R","RSxLDMC_AGxSTMC"],
"ReduceScatter": ["LDMC"]
}
ldmc_algos = ["RSxLDMC_AGxSTMC", "LDMC"]
coll_to_lower = {
"AllGather": "all_gather",
"AllReduce": "all_reduce",
"ReduceScatter": "reduce_scatter"
}
red_to_ncclDevRedOp = {
"sum": "ncclDevSum"
}
red_to_Func = {
"sum": "FuncSum"
}
ty_to_ncclDataType = {
"f32": "ncclFloat32",
"f16": "ncclFloat16",
"bf16": "ncclBfloat16",
"f8e4m3": "ncclFloat8e4m3",
"f8e5m2": "ncclFloat8e5m2"
}
ty_to_cxxtype = {
"f32": "float",
"f16": "half",
"bf16": "hip_bfloat16",
"f8e4m3": "rccl_float8",
"f8e5m2": "rccl_bfloat8"
}
def enumerate_kernels():
for algo in ["LL","ST"]:
yield Rec(coll="AllGather", algo=algo)
for red in all_reds:
for ty in all_tys:
for algo in ["AGxLL_R","RSxLD_AGxST"]:
yield Rec(coll="AllReduce", algo=algo, red=red, ty=ty)
for algo in ["LL","LD"]:
yield Rec(coll="ReduceScatter", algo=algo, red=red, ty=ty)
def required_cuda(k):
cudart, arch, specific_sms = 0, 0, None
is_nvls = k.algo in nvls_algos_by_coll.get(k.coll, [])
if is_nvls:
cudart = max(cudart, 12010)
arch = 900
if k.coll in reductions:
if k.ty == "bf16":
cudart = max(cudart, 11000)
if k.ty.startswith("f8"):
cudart = max(cudart, 11080)
arch = 900
if k.algo in ldmc_algos:
cudart = 12070
arch = None
specific_sms = [100, 120]
return (cudart, arch, specific_sms)
################################################################################
def kernel_fdep(k):
return coll_to_lower[k.coll] + '.cpp'
def kernel_fname(k):
if k.coll in reductions:
if k.algo in ldmc_algos and k.ty.startswith('f8'):
return paste('_', coll_to_lower[k.coll], k.red, k.ty, k.algo) + '.cpp'
else:
return paste('_', coll_to_lower[k.coll], k.red, k.ty) + '.cpp'
else:
return coll_to_lower[k.coll] + '.cpp'
def kernel_gencode(k):
if k.coll in reductions and k.algo in ldmc_algos and k.ty.startswith('f8'):
return "$(NVCC_GENCODE_LDMC_FP8)"
else:
return "$(NVCC_GENCODE)"
def kernel_cname(k):
if k.coll in reductions:
return paste("_", "ncclSymDevKernel", k.coll, k.algo, k.red, k.ty)
else:
return paste("_", "ncclSymDevKernel", k.coll, k.algo)
def kernel_conds(k):
cudart, arch, specific_sms = required_cuda(k)
if cudart == 0: return (None, None)
cudart_cond = "CUDART_VERSION >= %d"%cudart
if not specific_sms:
arch_cond = "__CUDA_ARCH__ >= %d"%arch
else:
arch_cond = " || ".join(["0"] + ["NCCL_CUDA_ARCH_SPECIFIC==%d"%(10*sm) for sm in specific_sms])
return cudart_cond, arch_cond
def instantiate(k):
form_red_ty = (
"__global__ void {cname}(ncclSymDevArgs NCCL_GRID_CONSTANT const *args) {{\n"
" ncclSymRun_{id}<{red}, {ty}>(args);\n"
"}}"
)
form = (
"__global__ void {cname}(ncclSymDevArgs NCCL_GRID_CONSTANT const *args) {{\n"
" ncclSymRun_{id}(args);\n"
"}}"
)
id = k.coll+'_'+k.algo
cname = kernel_cname(k)
if k.coll in reductions:
inst = form_red_ty.format(cname=cname, id=id, red=red_to_Func[k.red], ty=ty_to_cxxtype[k.ty])
else:
inst = form.format(cname=cname, id=id)
return inst
def prototype(k):
return "__global__ void {cname}(ncclSymDevArgs const *args);".format(cname=kernel_cname(k))
################################################################################
def partition(vals, keyfn):
ans = {}
for x in vals:
k = keyfn(x)
if k not in ans:
ans[k] = []
ans[k].append(x)
return ans
kernels_by_file = partition(enumerate_kernels(), lambda k: (kernel_fname(k), k.coll))
# Add dependency only files (e.g. allreduce.cpp)
for coll in set(k.coll for k in enumerate_kernels()):
fname = coll_to_lower[coll]+'.cpp'
if (fname, coll) not in kernels_by_file:
kernels_by_file[fname, coll] = []
# Generate each kernel instantiation file
for (fname, coll), ks in kernels_by_file.items():
with open(os.path.join(gensrc, fname), "w") as f:
print("-- Generating %s" % os.path.join(gensrc, fname))
emitln(f, '#include "symmetric.h"')
emitln(f, '#include "symmetric/kernel.h"')
emitln(f, '#include "symmetric/{coll}.h"'.format(coll=coll_to_lower[coll]))
for k in ks:
emitln(f, instantiate(k))
# Generate <gensrc>/symmetric_host.cc
with open(os.path.join(gensrc, "symmetric_kernels.cc"), "w") as f:
print("-- Generating %s" % os.path.join(gensrc, "symmetric_kernels.cc"))
emitln(f, '#include "symmetric.h"')
emitln(f, '#include "device.h"')
emitln(f, '')
for k in enumerate_kernels():
emitln(f, prototype(k))
emitln(f, '')
emitln(f, 'extern int const ncclSymKernelCount = %d;' % len(list(enumerate_kernels())))
emitln(f, 'extern void* const ncclSymKernelList[] = {')
for k in enumerate_kernels():
emitln(f, '(void*){cname},'.format(cname=kernel_cname(k)))
emitln(f, 'nullptr};')
emitln(f, '')
emitln(f, 'void* ncclSymGetKernelPtr(ncclSymKernelId id, int red, ncclDataType_t ty) {')
indents += 1
emitln(f, 'switch (id) {')
emitln(f, 'default: return nullptr;')
for (coll, algo), coll_algo_ks in partition(enumerate_kernels(), lambda k: (k.coll, k.algo)).items():
emitln(f, 'case ncclSymKernelId_'+coll+'_'+algo+':')
indents += 1
if len(coll_algo_ks) == 1:
emitln(f, 'return (void*)&'+kernel_cname(coll_algo_ks[0])+';')
else:
emitln(f, 'switch ((ncclDevRedOp_t)red) {')
emitln(f, 'default: return nullptr;')
for red, coll_algo_red_ks in partition(coll_algo_ks, lambda k: k.red).items():
emitln(f, 'case '+red_to_ncclDevRedOp[red]+':')
indents += 1
emitln(f, 'switch (ty) {')
emitln(f, 'default: return nullptr;')
for k in coll_algo_red_ks:
emitln(f, 'case '+ty_to_ncclDataType[k.ty]+': return (void*)'+kernel_cname(k)+';')
emitln(f, '}')
indents -= 1
emitln(f, '}')
indents -=1
emitln(f, '}')
indents -= 1
emitln(f, '}')
src/device/symmetric/kernel.cuh
+27
Просмотреть файл
@@ -0,0 +1,27 @@
#ifndef NCCL_DEVICE_SYMMETRIC_KERNEL_H_
#define NCCL_DEVICE_SYMMETRIC_KERNEL_H_
#include "symmetric.h"
template<template<typename> typename Red, typename T>
__device__ __forceinline__ void ncclSymRun_AllReduce_AGxLL_R(struct ncclSymDevArgs const* args);
template<template<typename> typename Red, typename T>
__device__ __forceinline__ void ncclSymRun_AllReduce_AGxLLMC_R(struct ncclSymDevArgs const* args);
template<template<typename> typename Red, typename T>
__device__ __forceinline__ void ncclSymRun_AllReduce_RSxLD_AGxST(struct ncclSymDevArgs const* args);
template<template<typename> typename Red, typename T>
__device__ __forceinline__ void ncclSymRun_AllReduce_RSxLDMC_AGxSTMC(struct ncclSymDevArgs const* args);
__device__ __forceinline__ void ncclSymRun_AllGather_LL(struct ncclSymDevArgs const* args);
__device__ __forceinline__ void ncclSymRun_AllGather_LLMC(struct ncclSymDevArgs const* args);
__device__ __forceinline__ void ncclSymRun_AllGather_ST(struct ncclSymDevArgs const* args);
__device__ __forceinline__ void ncclSymRun_AllGather_STMC(struct ncclSymDevArgs const* args);
template<template<typename> typename Red, typename T>
__device__ __forceinline__ void ncclSymRun_ReduceScatter_LL(struct ncclSymDevArgs const* args);
template<template<typename> typename Red, typename T>
__device__ __forceinline__ void ncclSymRun_ReduceScatter_LD(struct ncclSymDevArgs const* args);
template<template<typename> typename Red, typename T>
__device__ __forceinline__ void ncclSymRun_ReduceScatter_LDMC(struct ncclSymDevArgs const* args);
#endif
src/device/symmetric/primitives.cuh
+477
Просмотреть файл
@@ -0,0 +1,477 @@
#ifndef NCCL_DEVICE_SYMMETRIC_PRIMITIVES_H_
#define NCCL_DEVICE_SYMMETRIC_PRIMITIVES_H_
#include "symmetric.h"
#include "bitops.h"
#include "collectives.h"
#include "op128.h"
#include "reduce_kernel.h"
#include "common.h"
#if __CUDA_ARCH__ >= 700
// __grid_constant__ appears to break cuda-gdb
#define NCCL_GRID_CONSTANT __grid_constant__
#else
#define NCCL_GRID_CONSTANT
#endif
// flattenIx(pos0, dim0, pos1, dim1, pos2, dim2, ...)
// Given a position vector `pos` in a rectangular index space with lengths in the `dim`
// vector, flatten that down to a linear index. The fastest moving dimension is given first.
__device__ __forceinline__ int flattenIx() { return 0; }
template<typename Int0, typename Int1, typename ...Ints>
static __device__ Int0 flattenIx(Int0 pos, Int1 size, Ints ...more) {
return pos + size*flattenIx(more...);
}
// Precomputed integer reciprocoals for denominator values 1..64 inclusive.
// Pass these to idivFast64() for fast division on the GPU.
static __device__ uint64_t idivRcp64_upto64(int x) {
static constexpr uint64_t table[65] = {
idivRcp64(0x01), idivRcp64(0x01), idivRcp64(0x02), idivRcp64(0x03),
idivRcp64(0x04), idivRcp64(0x05), idivRcp64(0x06), idivRcp64(0x07),
idivRcp64(0x08), idivRcp64(0x09), idivRcp64(0x0a), idivRcp64(0x0b),
idivRcp64(0x0c), idivRcp64(0x0d), idivRcp64(0x0e), idivRcp64(0x0f),
idivRcp64(0x10), idivRcp64(0x11), idivRcp64(0x12), idivRcp64(0x13),
idivRcp64(0x14), idivRcp64(0x15), idivRcp64(0x16), idivRcp64(0x17),
idivRcp64(0x18), idivRcp64(0x19), idivRcp64(0x1a), idivRcp64(0x1b),
idivRcp64(0x1c), idivRcp64(0x1d), idivRcp64(0x1e), idivRcp64(0x1f),
idivRcp64(0x20), idivRcp64(0x21), idivRcp64(0x22), idivRcp64(0x23),
idivRcp64(0x24), idivRcp64(0x25), idivRcp64(0x26), idivRcp64(0x27),
idivRcp64(0x28), idivRcp64(0x29), idivRcp64(0x2a), idivRcp64(0x2b),
idivRcp64(0x2c), idivRcp64(0x2d), idivRcp64(0x2e), idivRcp64(0x2f),
idivRcp64(0x30), idivRcp64(0x31), idivRcp64(0x32), idivRcp64(0x33),
idivRcp64(0x34), idivRcp64(0x35), idivRcp64(0x36), idivRcp64(0x37),
idivRcp64(0x38), idivRcp64(0x39), idivRcp64(0x3a), idivRcp64(0x3b),
idivRcp64(0x3c), idivRcp64(0x3d), idivRcp64(0x3e), idivRcp64(0x3f),
idivRcp64(0x40)
};
return table[x];
}
static __device__ uint32_t idivRcp32_upto64(int x) {
return idivRcp64_upto64(x)>>32;
}
namespace {
struct ncclCoopCta {
__device__ void sync() { __syncthreads(); }
__device__ int self() { return threadIdx.x; }
__device__ int count() { return blockDim.x; }
};
struct ncclCoopWarps {
int log2_nWarps;
__device__ void sync() {
asm volatile("barrier.sync %0, %1;" :: "r"(1 + (threadIdx.x>>(5+log2_nWarps))), "r"(32<<log2_nWarps) : "memory");
}
__device__ int self() { return threadIdx.x & ((32<<log2_nWarps)-1); }
__device__ int count() { return 32<<log2_nWarps; }
};
struct ncclCoopWarp {
__device__ void sync() { __syncwarp(); }
__device__ int self() { return threadIdx.x%32; }
__device__ int count() { return 32; }
};
}
namespace {
static constexpr int ncclSymPrims_UseBarrier = 1;
static constexpr int ncclSymPrims_UseLL = 2;
static constexpr int ncclSymPrims_UseMultimem = 4;
struct ncclSymPrims {
int flags;
int const &rank;
int const &nRanks;
uint32_t const &nRanks_rcp32;
int block, nBlocks;
uint32_t nBlocks_rcp32;
uint32_t nBlocks_nWarps_rcp32;
uint32_t nRanks_nBlocks_rcp32;
uint32_t nWarpPerRank, nWarpPerRank_rcp32;
struct ncclSymDevBase* const &base;
uintptr_t offsetMc;
uint32_t const &stride4G;
uint32_t barEpoch;
uint32_t llEpoch;
__device__ ncclSymPrims(ncclSymDevComm const &comm, int flags):
flags(flags),
rank(comm.rank),
nRanks(comm.nRanks),
nRanks_rcp32(comm.nRanks_rcp32),
block(blockIdx.x),
nBlocks(gridDim.x),
nBlocks_rcp32(idivRcp32_upto64(nBlocks)),
nBlocks_nWarps_rcp32(imulRcp32(nBlocks, nBlocks_rcp32, blockDim.x/32, idivRcp32_upto64(blockDim.x/32))),
nRanks_nBlocks_rcp32(imulRcp32(nRanks, nRanks_rcp32, gridDim.x, nBlocks_rcp32)),
nWarpPerRank(idivFast32(nBlocks*blockDim.x/32, nRanks, nRanks_rcp32)),
nWarpPerRank_rcp32(idivRcp32_upto64(nWarpPerRank)),
base(comm.base),
offsetMc((flags & ncclSymPrims_UseMultimem) ? (char*)comm.baseMc - (char*)base : 0x0),
stride4G(comm.stride4G) {
#if CUDART_VERSION >= 12030 && __CUDA_ARCH__ >= 900
cudaGridDependencySynchronize();
#endif
if ((flags & ncclSymPrims_UseBarrier) && threadIdx.x < nRanks) {
barEpoch = (flags & ncclSymPrims_UseMultimem) ? base->barEpochMc[block] : base->barEpochUc[block];
}
if (flags & ncclSymPrims_UseLL) llEpoch = base->llEpoch[block] + 2;
}
__device__ ~ncclSymPrims() {
if (threadIdx.x == 0) {
if (flags & ncclSymPrims_UseBarrier) {
((flags & ncclSymPrims_UseMultimem) ? base->barEpochMc : base->barEpochUc)[block] = barEpoch;
}
if (flags & ncclSymPrims_UseLL) base->llEpoch[block] = llEpoch - 2;
}
}
template<typename T>
__device__ T* peerPtr(int peer, T* selfPtr) {
return add4G(selfPtr, (peer-rank)*stride4G);
}
template<typename T>
__device__ T* multimemPtr(T* selfPtr) {
return reinterpret_cast<T*>(reinterpret_cast<uintptr_t>(selfPtr) + offsetMc);
}
__device__ void barrierArrive(ncclCoopCta cta, bool release) {
cta.sync();
#if __CUDA_ARCH__ < 700
if (release) {
if (cta.self() == 0) __threadfence_system();
cta.sync();
}
#endif
if (flags & ncclSymPrims_UseMultimem) {
#if __CUDA_ARCH__ >= 900 && CUDART_VERSION >= 12010
if (cta.self() == 0) {
uint32_t* inbox = &multimemPtr(base)->barInboxMc[block];
if (release) {
asm volatile("multimem.red.release.sys.add.u32 [%0],1;" :: "l"(inbox));
} else {
asm volatile("multimem.red.relaxed.sys.add.u32 [%0],1;" :: "l"(inbox));
}
}
#endif
} else {
int r = cta.self();
if (r != rank && r < nRanks) {
uint32_t* inbox = &peerPtr(r, base)->barInboxPerPeer[block*nRanks + rank];
#if __CUDA_ARCH__ >= 700
if (release) {
asm volatile("st.release.sys.u32 [%0],%1;" :: "l"(inbox), "r"(barEpoch+1));
} else {
asm volatile("st.relaxed.sys.u32 [%0],%1;" :: "l"(inbox), "r"(barEpoch+1));
}
#else
if (release) {
__atomic_store_n(inbox, barEpoch + 1, __ATOMIC_RELEASE);
} else {
__atomic_store_n(inbox, barEpoch + 1, __ATOMIC_RELAXED);
}
// asm volatile("st.volatile.u32 [%0],%1;" :: "l"(inbox), "r"(barEpoch+1));
#endif
}
}
}
__device__ void barrierWait(ncclCoopCta cta, bool acquire) {
if (flags & ncclSymPrims_UseMultimem) {
#if __CUDA_ARCH__ >= 900
if (cta.self() == 0) {
uint32_t* inbox = &base->barInboxMc[block];
while (true) {
uint32_t got;
if (acquire) {
asm volatile("ld.acquire.sys.u32 %0,[%1];" : "=r"(got) : "l"(inbox));
} else {
asm volatile("ld.relaxed.sys.u32 %0,[%1];" : "=r"(got) : "l"(inbox));
}
if (got-(barEpoch+nRanks) <= uint32_t(-1)>>1) break;
}
barEpoch += nRanks;
}
#endif
} else {
int r = cta.self();
if (r != rank && r < nRanks) {
uint32_t* inbox = &base->barInboxPerPeer[block*nRanks + r];
while (true) {
uint32_t got;
#if __CUDA_ARCH__ >= 700
if (acquire) {
asm volatile("ld.acquire.sys.u32 %0,[%1];" : "=r"(got) : "l"(inbox));
} else {
asm volatile("ld.relaxed.sys.u32 %0,[%1];" : "=r"(got) : "l"(inbox));
}
#else
if (acquire) {
got = __atomic_load_n(inbox, __ATOMIC_ACQUIRE);
} else {
got = __atomic_load_n(inbox, __ATOMIC_RELAXED);
}
// asm volatile("ld.volatile.u32 %0,[%1];" : "=r"(got) : "l"(inbox));
#endif
if (got-(barEpoch+1) <= uint32_t(-1)>>1) break;
}
}
#if __CUDA_ARCH__ < 700
if (acquire) {
cta.sync();
if (cta.self() == 0) __threadfence();
}
#endif
barEpoch += 1;
}
cta.sync();
}
__device__ void endLL(ncclCoopCta cta) {
if (__builtin_expect(llEpoch >= -2u, false)) {
cta.sync();
uint4* buf = ncclSymDevBase_getLLBuf(base, nRanks, block, llEpoch);
int epochSize = ncclSymLLEpochSize(nRanks);
#pragma unroll 4
for (int i=cta.self(); i*16 < epochSize; i += cta.count()) {
buf[i] = uint4{0, 0, 0, 0};
}
}
cta.sync();
llEpoch += (llEpoch == -1u) ? 3 : 1;
}
template<typename T>
__device__ void sendLL(int peer, int slot, T val) {
union { T tmp; uint32_t u32[divUp(sizeof(T),8)][2]; };
tmp = val;
uint4* buf = ncclSymDevBase_getLLBuf(peerPtr(peer, base), nRanks, block, llEpoch) + slot;
#pragma unroll
for (int u=0; u < divUp(sizeof(T),8); u++) {
using Vec = uint32_t __attribute__((ext_vector_type(4)));
Vec i4;
i4[0] = u32[u][0];
i4[1] = llEpoch;
i4[2] = u32[u][1];
i4[3] = llEpoch;
#if defined(__gfx950__)
asm volatile ("flat_store_dwordx4 %0, %1 sc0 sc1 nt" :: "v"(buf + ncclSymLLMaxSlots(sizeof(T))*u), "v"(i4));
#else
__builtin_nontemporal_store(i4, (Vec*)(buf + ncclSymLLMaxSlots(sizeof(T))*u));
#endif
// asm volatile("st.volatile.v4.u32 [%0],{%1,%3,%2,%3};" :: "l"(buf + ncclSymLLMaxSlots(sizeof(T))*u), "r"(u32[u][0]), "r"(u32[u][1]), "r"(llEpoch));
}
}
template<typename T>
__device__ void bcastLL(int slot, T val) {
if (flags & ncclSymPrims_UseMultimem) {
union { T tmp; uint32_t u32[divUp(sizeof(T),8)][2]; };
tmp = val;
uint4* bufmc = ncclSymDevBase_getLLBuf(multimemPtr(base), nRanks, block, llEpoch) + slot;
#pragma unroll
for (int u=0; u < divUp(sizeof(T),8); u++) {
using Vec = uint32_t __attribute__((ext_vector_type(4)));
Vec i4;
i4[0] = u32[u][0];
i4[1] = llEpoch;
i4[2] = u32[u][1];
i4[3] = llEpoch;
#if defined(__gfx950__)
asm volatile ("flat_store_dwordx4 %0, %1 sc0 sc1 nt" :: "v"(bufmc + ncclSymLLMaxSlots(sizeof(T))*u), "v"(i4));
#else
__builtin_nontemporal_store(i4, (Vec*)(bufmc + ncclSymLLMaxSlots(sizeof(T))*u));
#endif
// asm volatile("st.volatile.v4.u32 [%0],{%1,%3,%2,%3};" :: "l"(bufmc + ncclSymLLMaxSlots(sizeof(T))*u), "r"(u32[u][0]), "r"(u32[u][1]), "r"(llEpoch));
}
} else {
union { T tmp; uint32_t u32[divUp(sizeof(T),8)][2]; };
tmp = val;
uint4* buf0 = ncclSymDevBase_getLLBuf(peerPtr(0, base), nRanks, block, llEpoch) + slot;
int dr = 0;
int r = rank;
#pragma unroll 1
for (; dr+8 <= nRanks; dr += 8) {
#pragma unroll
for (int ur=0; ur < 8; ur++) {
uint4* buf = add4G(buf0, r*stride4G);
#pragma unroll
for (int u=0; u < divUp(sizeof(T),8); u++) {
using Vec = uint32_t __attribute__((ext_vector_type(4)));
Vec i4;
i4[0] = u32[u][0];
i4[1] = llEpoch;
i4[2] = u32[u][1];
i4[3] = llEpoch;
#if defined(__gfx950__)
asm volatile ("flat_store_dwordx4 %0, %1 sc0 sc1 nt" :: "v"(buf + ncclSymLLMaxSlots(sizeof(T))*u), "v"(i4));
#else
__builtin_nontemporal_store(i4, (Vec*)((buf + ncclSymLLMaxSlots(sizeof(T))*u)));
#endif
// asm volatile("st.volatile.v4.u32 [%0],{%1,%3,%2,%3};" :: "l"(buf + ncclSymLLMaxSlots(sizeof(T))*u), "r"(u32[u][0]), "r"(u32[u][1]), "r"(llEpoch));
}
r += 1;
if (r == nRanks) r = 0;
}
}
#pragma unroll
for (int ur=0; ur < 8; ur++, dr++) {
if (dr == nRanks) break;
uint4* buf = add4G(buf0, r*stride4G);
#pragma unroll
for (int u=0; u < divUp(sizeof(T),8); u++) {
using Vec = uint32_t __attribute__((ext_vector_type(4)));
Vec i4;
i4[0] = u32[u][0];
i4[1] = llEpoch;
i4[2] = u32[u][1];
i4[3] = llEpoch;
#if defined(__gfx950__)
asm volatile ("flat_store_dwordx4 %0, %1 sc0 sc1 nt" :: "v"(buf + ncclSymLLMaxSlots(sizeof(T))*u), "v"(i4));
#else
__builtin_nontemporal_store(i4, (Vec*)(buf + ncclSymLLMaxSlots(sizeof(T))*u));
#endif
// asm volatile("st.volatile.v4.u32 [%0],{%1,%3,%2,%3};" :: "l"(buf + ncclSymLLMaxSlots(sizeof(T))*u), "r"(u32[u][0]), "r"(u32[u][1]), "r"(llEpoch));
}
r += 1;
if (r == nRanks) r = 0;
}
}
}
template<int nSlotsMin, int nSlotsMax, typename T>
__device__ void recvLL(int slot0, int nSlots, int stride, T(&elts)[nSlotsMax]) {
uint4* buf = ncclSymDevBase_getLLBuf(base, nRanks, block, llEpoch) + slot0;
uint4 tmp[nSlotsMax][divUp(sizeof(T),8)];
//int spins=0;
while (true) {
#pragma unroll
for (int u=0; u < nSlotsMax; u++) {
if (u < nSlotsMin || u < nSlots) {
#pragma unroll
for (int v=0; v < divUp(sizeof(T),8); v++) {
tmp[u][v] = *(buf + u * stride + v * ncclSymLLMaxSlots(sizeof(T)));
// asm volatile("ld.volatile.v4.u32 {%0,%1,%2,%3},[%4];" : "=r"(tmp[u][v].x), "=r"(tmp[u][v].y), "=r"(tmp[u][v].z), "=r"(tmp[u][v].w) : "l"(buf + u*stride + v*ncclSymLLMaxSlots(sizeof(T))));
}
}
}
bool okAll = true;
#pragma unroll
for (int u=0; u < nSlotsMax; u++) {
#pragma unroll
for (int v=0; v < divUp(sizeof(T),8); v++) {
if (u < nSlotsMin || u < nSlots) {
bool ok = tmp[u][v].y == llEpoch &&
tmp[u][v].w == llEpoch;
okAll &= ok;
}
}
}
if (__builtin_expect(okAll, true)) break;
//if (spins++ == 10<<20) spins=0;
}
#pragma unroll
for (int u=0; u < nSlotsMax; u++) {
if (nSlotsMin <= u && u == nSlots) break;
union { T val; uint32_t u32[divUp(sizeof(T),8)][2]; };
#pragma unroll
for (int v=0; v < divUp(sizeof(T),8); v++) {
u32[v][0] = tmp[u][v].x;
u32[v][1] = tmp[u][v].z;
}
elts[u] = val;
}
}
template<typename Pack, typename T, typename Red, int Unroll=8>
__device__ Pack recvReduceLL(int slot, int stride, Red red) {
using Acc = typename Red::EltType;
using AccPack = BytePack<sizeof(Pack)*sizeof(Acc)/sizeof(T)>;
AccPack acc;
bool first = true;
int r = 0;
#pragma unroll 1
for (; r+Unroll <= nRanks; r += Unroll) {
Pack got[Unroll];
this->template recvLL</*Min=*/Unroll>(slot + r*stride, Unroll, stride, got);
AccPack acc0 = applyCast<T, Acc>(got[0]);
acc = first ? acc0 : applyReduce(red, acc, acc0);
first = false;
#pragma unroll
for (int i=1; i < Unroll; i++) acc = applyReduce(red, acc, applyCast<T, Acc>(got[i]));
}
if (r < nRanks) {
Pack got[Unroll];
this->template recvLL</*Min=*/1>(slot + r*stride, nRanks-r, stride, got);
AccPack acc0 = applyCast<T, Acc>(got[0]);
acc = first ? acc0 : applyReduce(red, acc, acc0);
#pragma unroll
for (int i=1; i < Unroll-1; i++) {
if (r+i < nRanks) acc = applyReduce(red, acc, applyCast<T, Acc>(got[i]));
}
}
return applyCast<Acc, T>(acc);
}
template<typename T>
__device__ T recvLL(int slot) {
T one[1];
this->template recvLL<1, 1, T>(slot, 1, 0, one);
return one[0];
}
template<typename Coop, typename T>
__device__ void coopRecvLL(Coop coop, int slot0, int nSlots, T* dst) {
int me = coop.self();
if (me < nSlots) {
uint4* buf = ncclSymDevBase_getLLBuf(base, nRanks, block, llEpoch) + slot0 + me;
uint4 got[divUp(sizeof(T), 8)];
//int spins=0;
#pragma unroll 1
while (true) {
#pragma unroll
for (int u=0; u < divUp(sizeof(T), 8); u++) {
got[u] = *((buf + u * ncclSymLLMaxSlots(sizeof(T))));
// asm volatile("ld.volatile.v4.u32 {%0,%1,%2,%3},[%4];" : "=r"(got[u].x), "=r"(got[u].y), "=r"(got[u].z), "=r"(got[u].w) : "l"(buf + u*ncclSymLLMaxSlots(sizeof(T))));
}
bool ok = true;
#pragma unroll
for (int u=0; u < divUp(sizeof(T), 8); u++) {
ok &= got[u].y == llEpoch;
ok &= got[u].w == llEpoch;
}
if (__builtin_expect(ok, true)) break;
//if (++spins == 10<<20) { spins=0; printf("r=%d LL spin @ ix=%d got=%d want=%d\n", rank, slot0+me, got[0].y, llEpoch); }
}
union { T val; uint32_t u32[divUp(sizeof(T), 8)][2]; };
#pragma unroll
for (int u=0; u < divUp(sizeof(T), 8); u++) {
u32[u][0] = got[u].x;
u32[u][1] = got[u].z;
}
dst[slot0 + me] = val;
}
}
};
}
template<template<typename> typename Red, typename T, bool nvls>
struct ncclSymAccumType { using Type = T; };
// Only Red's whose opArg is invariant w.r.t. the datatype can have a different
// accumulator type. At the moment this excludes integer min/max, sumpostdiv,
// and premulsum.
template<> struct ncclSymAccumType<FuncSum, __half, false> { using Type = float; };
#if defined(__CUDA_BF16_TYPES_EXIST__)
template<> struct ncclSymAccumType<FuncSum, __nv_bfloat16, false> { using Type = float; };
#endif
#if defined(__CUDA_FP8_TYPES_EXIST__)
template<> struct ncclSymAccumType<FuncSum, __nv_fp8_e4m3, false> { using Type = float; };
template<> struct ncclSymAccumType<FuncSum, __nv_fp8_e5m2, false> { using Type = float; };
#endif
#endif
src/device/symmetric/reduce_scatter.cuh
+387
Просмотреть файл
@@ -0,0 +1,387 @@
#include "symmetric.h"
#include "symmetric/kernel.h"
#include "symmetric/primitives.h"
template<int BytePerPack, int UnrollPacks, int UnrollPeers, typename T, typename Red>
static __device__ void reduceDeep(
ncclSymPrims& prim, int tn, int t, bool waitNeeded,
Red red, char* inputRank0, char* outputHere, int32_t nIters
) {
using Pack = BytePack<BytePerPack>;
using Acc = typename Red::EltType;
using AccPack = BytePack<BytePerPack*sizeof(Acc)/sizeof(T)>;
int wn = tn/WARP_SIZE;
int w = t/WARP_SIZE;
int lane = t%WARP_SIZE;
int const& rank = prim.rank;
int const& nRanks = prim.nRanks;
uint32_t const& stride4G = prim.stride4G;
Pack* inpRank0 = (Pack*)inputRank0 + intptr_t(w)*UnrollPacks*WARP_SIZE + lane;
Pack* outHere = (Pack*)outputHere + intptr_t(w)*UnrollPacks*WARP_SIZE + lane;
Pack acc0[UnrollPacks];
nIters -= w;
if (0 < nIters) {
#pragma unroll
for (int u=0; u < UnrollPacks; u++) {
acc0[u] = add4G(inpRank0, rank*stride4G)[u*WARP_SIZE];
}
}
if (waitNeeded) prim.barrierWait(ncclCoopCta(), /*acquire=*/false);
if (0 < nIters) {
while (true) {
AccPack acc1[UnrollPacks];
int r = rank+1;
if (r == nRanks) r = 0;
{ Pack tmp1[UnrollPacks];
#pragma unroll
for (int u=0; u < UnrollPacks; u++) {
tmp1[u] = add4G(inpRank0, r*stride4G)[u*WARP_SIZE];
}
#pragma unroll
for (int u=0; u < UnrollPacks; u++) {
acc1[u] = applyReduce(red, applyCast<T, Acc>(acc0[u]), applyCast<T, Acc>(tmp1[u]));
}
}
r += 1;
if (r == nRanks) r = 0;
int dr = 2;
#pragma unroll 2
for (int partial=0; partial <= 1; partial++) {
#pragma unroll 1
for (int i = 0;
partial ? i < 1 : (dr + UnrollPeers <= nRanks);
partial ? i++ : (dr += UnrollPeers)) {
if (partial && dr == nRanks) break;
Pack tmp1[UnrollPeers][UnrollPacks];
#pragma unroll
for (int ur=0; ur < UnrollPeers-partial; ur++) {
if (partial && ur!=0 && dr+ur == nRanks) break;
#pragma unroll UnrollPacks
for (int u=0; u < UnrollPacks; u++) {
tmp1[ur][u] = add4G(inpRank0, r*stride4G)[u*WARP_SIZE];
}
r += 1;
if (r == nRanks) r = 0;
}
#pragma unroll
for (int ur=0; ur < UnrollPeers-partial; ur++) {
if (partial && ur!=0 && dr+ur == nRanks) break;
#pragma unroll UnrollPacks
for (int u=0; u < UnrollPacks; u++) {
acc1[u] = applyReduce(red, acc1[u], applyCast<T, Acc>(tmp1[ur][u]));
}
}
}
}
#pragma unroll
for (int u=0; u < UnrollPacks; u++) acc0[u] = applyCast<Acc, T>(acc1[u]);
#pragma unroll UnrollPacks
for (int u=0; u < UnrollPacks; u++) outHere[u*WARP_SIZE] = acc0[u];
inpRank0 += intptr_t(wn)*UnrollPacks*WARP_SIZE;
outHere += intptr_t(wn)*UnrollPacks*WARP_SIZE;
nIters -= wn;
if (nIters <= 0) break;
// Load data for next iteration.
#pragma unroll
for (int u=0; u < UnrollPacks; u++) {
acc0[u] = add4G(inpRank0, rank*stride4G)[u*WARP_SIZE];
}
}
}
}
template<int UnrollPeers, typename Red, typename T>
static __device__ void reduceEnds(
ncclSymPrims& prim, int tn, int t, Red red,
T* inputRank0, T* outputHere, size_t nElts, uint32_t nPreElts, size_t nSufElts
) {
using Acc = typename Red::EltType;
int const& rank = prim.rank;
int const& nRanks = prim.nRanks;
uint32_t const& stride4G = prim.stride4G;
BytePack<sizeof(T)>* inpRank0 = (BytePack<sizeof(T)>*)inputRank0;
BytePack<sizeof(T)>* outHere = (BytePack<sizeof(T)>*)outputHere;
#pragma unroll 1
for (size_t i = t; i < nPreElts+nSufElts; i += tn) {
size_t elt = i < nPreElts ? i : nElts-nSufElts-nPreElts+i;
BytePack<sizeof(T)> acc0 = *add4G(inpRank0+elt, rank*stride4G);
BytePack<sizeof(Acc)> acc1;
BytePack<sizeof(T)> tmp[UnrollPeers];
int dr = 1;
int r = rank+1;
if (nRanks == r) r = 0;
bool first = true;
#pragma unroll 2
for (int partial=0; partial <= 1; partial++) {
#pragma unroll 1
for (int j = 0;
partial ? j < 1 : (dr + UnrollPeers <= nRanks);
partial ? j++ : (dr += UnrollPeers)) {
if (partial && dr == nRanks) break;
#pragma unroll
for (int u=0; u < UnrollPeers-partial; u++) {
if (partial && u!=0 && dr+u == nRanks) break;
tmp[u] = *add4G(inpRank0+elt, r*stride4G);
r += 1;
if (r == nRanks) r = 0;
}
if (first) {
first = false;
acc1 = applyCast<T, Acc>(acc0);
}
#pragma unroll
for (int u=0; u < UnrollPeers-partial; u++) {
if (partial && u!=0 && dr+u == nRanks) break;
acc1 = applyReduce(red, acc1, applyCast<T, Acc>(tmp[u]));
}
}
}
acc0 = applyCast<Acc, T>(acc1);
outHere[elt] = acc0;
}
}
template<typename Red, typename T>
static __device__ void reduce(
ncclSymPrims& prim, int tn, int t, bool waitNeeded,
Red red, T* input, T* output, size_t nElts
) {
int nRanks = prim.nRanks;
int nBlocks = prim.nBlocks;
// Mpve input to rank=0
input = prim.peerPtr(0, input);
uintptr_t inputUptr = reinterpret_cast<uintptr_t>(input);
uintptr_t outputUptr = reinterpret_cast<uintptr_t>(output);
uint32_t alignment = uint32_t(inputUptr - outputUptr);
size_t nBytes = nElts*sizeof(T);
uint32_t nPreBytes = (16u - inputUptr)%16u;
nPreBytes = min((size_t)nPreBytes, nBytes);
uintptr_t cursor = nPreBytes;
constexpr int MinWarpPerBlock = 4;
if (alignment%16 == 0) {
constexpr int BytePerPack = 16, UnrollPacks = 4, UnrollPeers = 2;
constexpr int BytePerChunk = MinWarpPerBlock*UnrollPacks*WARP_SIZE*BytePerPack;
uint32_t chunks = (nBytes-cursor)/BytePerChunk;
chunks -= imodFast32(chunks, nRanks*nBlocks, prim.nRanks_nBlocks_rcp32);
if (chunks != 0) {
uintptr_t cursorAfter = cursor + uintptr_t(chunks)*BytePerChunk;
reduceDeep<BytePerPack, UnrollPacks, UnrollPeers, T>(
prim, tn, t, waitNeeded, red,
(char*)input + cursor, (char*)output + cursor,
chunks*MinWarpPerBlock
);
cursor = cursorAfter;
waitNeeded = false;
}
}
if (sizeof(T) == 4 || (sizeof(T) < 4 && alignment%4 == 0)) {
constexpr int BytePerPack = 4, UnrollPacks = 4, UnrollPeers = 4;
constexpr int BytePerChunk = MinWarpPerBlock*UnrollPacks*WARP_SIZE*BytePerPack;
uint32_t chunks = (nBytes-cursor)/BytePerChunk;
chunks -= imodFast32(chunks, nRanks*nBlocks, prim.nRanks_nBlocks_rcp32);
if (chunks != 0) {
uintptr_t cursorAfter = cursor + uintptr_t(chunks)*BytePerChunk;
reduceDeep<(sizeof(T) <= BytePerPack ? BytePerPack : 0), UnrollPacks, UnrollPeers, T>(
prim, tn, t, waitNeeded, red,
(char*)input + cursor, (char*)output + cursor,
chunks*MinWarpPerBlock
);
cursor = cursorAfter;
waitNeeded = false;
}
}
if (waitNeeded) prim.barrierWait(ncclCoopCta(), /*acquire=*/false);
constexpr int UnrollPeers = 8;
size_t nSufElts = (nBytes-cursor)/sizeof(T);
reduceEnds<UnrollPeers>(prim, tn, t, red, input, output, nElts, nPreBytes/sizeof(T), nSufElts);
}
template<template<typename> typename Red, typename T>
__device__ __forceinline__ void ncclSymRun_ReduceScatter_LD(ncclSymDevArgs const* args) {
ncclSymPrims prim(args->comm, ncclSymPrims_UseBarrier);
Red<typename ncclSymAccumType<Red, T, /*nvls=*/false>::Type> red(args->redOpArg);
// Round robin warps over blocks.
int t = flattenIx(threadIdx.x%WARP_SIZE, WARP_SIZE,
prim.block, prim.nBlocks,
threadIdx.x/WARP_SIZE, blockDim.x/WARP_SIZE);
int tn = prim.nBlocks*blockDim.x;
prim.barrierArrive(ncclCoopCta(), /*release=*/false);
//prim.barrierWait(ncclCoopCta(), /*acquire=*/false);
reduce(prim, tn, t, /*waitNeeded=*/true, red, (T*)args->input + prim.rank*args->nElts, (T*)args->output, args->nElts);
prim.barrierArrive(ncclCoopCta(), /*release=*/false);
prim.barrierWait(ncclCoopCta(), /*acquire=*/false);
}
template<typename Red, typename T>
static __device__ void reduceMultimem(
ncclSymPrims& prim, int tn, int t, Red red, T* input, T* output, size_t nElts
) {
// Mpve input to multimem
input = prim.multimemPtr(input);
uintptr_t inputUptr = reinterpret_cast<uintptr_t>(input);
uintptr_t outputUptr = reinterpret_cast<uintptr_t>(output);
size_t nBytes = nElts*sizeof(T);
constexpr int BytePerPack = LoadMultimem_BigPackSize<Red>::BigPackSize;
uint32_t nPreBytes = (BytePerPack - inputUptr)%BytePerPack;
nPreBytes = min((size_t)nPreBytes, nBytes);
uintptr_t nSufBytes;
if (sizeof(T) == BytePerPack || (inputUptr-outputUptr)%BytePerPack == 0) {
constexpr int UnrollPacks = 8*(16/BytePerPack);
constexpr int BytePerChunk = UnrollPacks*WARP_SIZE*BytePerPack;
uintptr_t cursor = nPreBytes;
uint32_t nChunks = (nBytes-cursor)/BytePerChunk;
uintptr_t cursorAfter = cursor + uintptr_t(nChunks)*BytePerChunk;
nSufBytes = nBytes - cursorAfter;
cursor += (t/WARP_SIZE)*UnrollPacks*WARP_SIZE*BytePerPack;
cursor += (t%WARP_SIZE)*BytePerPack;
int nIters = nChunks - t/WARP_SIZE;
#pragma unroll 1
while (0 < nIters) {
BytePack<BytePerPack> tmp[UnrollPacks];
#pragma unroll
for (int u=0; u < UnrollPacks; u++) {
tmp[u] = applyLoadMultimem<Red, BytePerPack>(red, inputUptr + cursor + u*WARP_SIZE*BytePerPack);
}
#pragma unroll
for (int u=0; u < UnrollPacks; u++) {
*reinterpret_cast<BytePack<BytePerPack>*>(outputUptr + cursor + u*WARP_SIZE*BytePerPack) = tmp[u];
}
cursor += tn*UnrollPacks*BytePerPack;
nIters -= tn/WARP_SIZE;
}
} else {
nPreBytes = 0;
nSufBytes = nBytes;
}
// Get the prefix+suffix element one at a time.
#pragma unroll 4
for (uintptr_t i = t*sizeof(T); i < nPreBytes + nSufBytes; i += tn*sizeof(T)) {
uintptr_t cursor = i < nPreBytes ? i : nBytes-nSufBytes+(i-nPreBytes);
BytePack<sizeof(T)> val = applyLoadMultimem<Red, sizeof(T)>(red, inputUptr + cursor);
*reinterpret_cast<BytePack<sizeof(T)>*>(outputUptr + cursor) = val;
cursor += tn*sizeof(T);
}
}
template<template<typename> typename Red, typename T>
__device__ __forceinline__ void ncclSymRun_ReduceScatter_LDMC(ncclSymDevArgs const* args) {
ncclSymPrims prim(args->comm, ncclSymPrims_UseBarrier|ncclSymPrims_UseMultimem);
Red<typename ncclSymAccumType<Red, T, /*nvls=*/true>::Type> red(args->redOpArg);
// Round robin warps over blocks.
int t = flattenIx(threadIdx.x%WARP_SIZE, WARP_SIZE,
prim.block, prim.nBlocks,
threadIdx.x/WARP_SIZE, blockDim.x/WARP_SIZE);
int tn = prim.nBlocks*blockDim.x;
prim.barrierArrive(ncclCoopCta(), /*release=*/false);
prim.barrierWait(ncclCoopCta(), /*acquire=*/false);
reduceMultimem(prim, tn, t, red, (T*)args->input + prim.rank*args->nElts, (T*)args->output, args->nElts);
prim.barrierArrive(ncclCoopCta(), /*release=*/false);
prim.barrierWait(ncclCoopCta(), /*acquire=*/false);
}
// T is user type, EltType is the most aligned type
template<typename T, typename Red, typename EltType>
__device__ __forceinline__ void ncclSymRun_ReduceScatter_LL_body(
ncclSymPrims &prim, Red red, EltType* input, EltType* output, int nElts, int nPacks, int nStrideElts) {
using Pack = BytePack<8>;
constexpr int EltPerPack = 8/sizeof(EltType);
int nRanks = prim.nRanks;
int rank = prim.rank;
int t = threadIdx.x;
int tn = ncclSymMaxThreads;
ncclCoopCta cta;
#pragma unroll 1
while (0 < nElts) {
int nIterPacks = min(nPacks, tn);
int tn_div_nPacks = tn/nIterPacks;
int tn_mod_nPacks = tn%nIterPacks;
int peer = t/nIterPacks;
int pack = t%nIterPacks;
#pragma unroll 1
for (int i = t; i < nRanks*nIterPacks; i += tn) {
Pack got = loadPack<Pack>(input + peer*nStrideElts, pack*EltPerPack, nElts);
prim.sendLL(peer, rank*nIterPacks + pack, got);
peer += tn_div_nPacks;
pack += tn_mod_nPacks;
if (nIterPacks <= pack) { peer += 1; pack -= nIterPacks; }
}
if (t < nIterPacks) {
Pack got = prim.template recvReduceLL<Pack, T>(t, nIterPacks, red);
storePack(output, t*EltPerPack, nElts, got);
}
prim.endLL(cta);
input += tn*EltPerPack;
output += tn*EltPerPack;
nElts -= tn*EltPerPack;
nPacks -= tn;
}
}
template<template<typename> typename Red, typename T>
__device__ __forceinline__ void ncclSymRun_ReduceScatter_LL(ncclSymDevArgs const* args) {
ncclSymPrims prim(args->comm, ncclSymPrims_UseLL);
Red<typename ncclSymAccumType<Red, T, /*nvls=*/false>::Type> red(args->redOpArg);
using Pack = BytePack<8>;
constexpr int EltPerPack = 8/sizeof(T);
int nAllElts = args->nElts;
int nAllPacks = divUp(nAllElts, EltPerPack);
uint32_t nPackPerBlock, nPackModBlock;
idivmodFast32(&nPackPerBlock, &nPackModBlock, nAllPacks, prim.nBlocks, prim.nBlocks_rcp32);
int blockPackBegin = prim.block*nPackPerBlock + minval<int>(prim.block, nPackModBlock);
int blockPackEnd = blockPackBegin + nPackPerBlock + (prim.block < nPackModBlock ? 1 : 0);
int nPacks = blockPackEnd - blockPackBegin;
int nElts = nAllElts - blockPackBegin*EltPerPack;
nElts = min(nElts, nPacks*EltPerPack);
T* input = (T*)args->input + blockPackBegin*EltPerPack;
T* output = (T*)args->output + blockPackBegin*EltPerPack;
uint32_t lowBits = args->nElts*sizeof(T);
lowBits |= (uint32_t)reinterpret_cast<uintptr_t>(args->input);
lowBits |= (uint32_t)reinterpret_cast<uintptr_t>(args->output);
if (__builtin_expect(lowBits%8 == 0, true)) {
ncclSymRun_ReduceScatter_LL_body<T>(prim, red, (Pack*)input, (Pack*)output, nPacks, nPacks, nAllElts/EltPerPack);
} else {
ncclSymRun_ReduceScatter_LL_body<T>(prim, red, input, output, nElts, nPacks, nAllElts);
}
}