Dynamically select unroll factor to build for when targeting local arch (#1371)
* Dynamically select unroll factor to build for when targeting local arch only
This commit is contained in:
Bertan Dogancay
GitHub
+1
-1
@@ -600,7 +600,7 @@ set(GEN_DIR "${HIPIFY_DIR}/gensrc")
|
||||
|
||||
# Execute the python script to generate required files
|
||||
execute_process(
|
||||
COMMAND ${PYTHON_EXECUTABLE} ${CMAKE_SOURCE_DIR}/src/device/generate.py ${GEN_DIR} ${IFC_ENABLED} ${COLLTRACE} ${ENABLE_MSCCL_KERNEL} ${ONLY_FUNCS}
|
||||
COMMAND ${PYTHON_EXECUTABLE} ${CMAKE_SOURCE_DIR}/src/device/generate.py ${GEN_DIR} ${IFC_ENABLED} ${COLLTRACE} ${ENABLE_MSCCL_KERNEL} ${BUILD_LOCAL_GPU_TARGET_ONLY} ${ONLY_FUNCS}
|
||||
WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}
|
||||
RESULT_VARIABLE gen_py_result
|
||||
ERROR_VARIABLE gen_py_error
|
||||
|
||||
@@ -18,17 +18,11 @@ struct RunWorkNop {
|
||||
};
|
||||
|
||||
__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernel_Generic(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead) {
|
||||
ncclKernelMain<-1, RunWorkNop, false, 2>(comm, channelMask, workHead);
|
||||
}
|
||||
__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernel_Generic_4(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead) {
|
||||
ncclKernelMain<-1, RunWorkNop, false, 4>(comm, channelMask, workHead);
|
||||
ncclKernelMain<-1, RunWorkNop, false>(comm, channelMask, workHead);
|
||||
}
|
||||
#ifdef ENABLE_COLLTRACE
|
||||
__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernelDebug_Generic(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead) {
|
||||
ncclKernelMain<-1, RunWorkNop, true, 2>(comm, channelMask, workHead);
|
||||
}
|
||||
__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernelDebug_Generic_4(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead) {
|
||||
ncclKernelMain<-1, RunWorkNop, true, 4>(comm, channelMask, workHead);
|
||||
ncclKernelMain<-1, RunWorkNop, true>(comm, channelMask, workHead);
|
||||
}
|
||||
#endif
|
||||
|
||||
|
||||
+7
-21
@@ -227,7 +227,7 @@ static __forceinline__ __device__ void ncclRedopPtrDeref(struct ncclWorkElem* we
|
||||
}
|
||||
}
|
||||
|
||||
template<int SpecializedFnId, typename SpecializedRunWork, bool COLLTRACE, int COLL_UNROLL>
|
||||
template<int SpecializedFnId, typename SpecializedRunWork, bool COLLTRACE>
|
||||
__forceinline__ __device__ void ncclKernelMain(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead) {
|
||||
const int tid = threadIdx.x;
|
||||
int x = tid;
|
||||
@@ -346,15 +346,9 @@ __forceinline__ __device__ void ncclKernelMain(struct ncclDevComm* comm, struct
|
||||
SpecializedRunWork().run(&ncclShmem.work);
|
||||
} else {
|
||||
#ifdef USE_INDIRECT_FUNCTION_CALL
|
||||
if (COLL_UNROLL == 4)
|
||||
ncclDevFuncTable_4[ncclShmem.work.header.funcIndex]();
|
||||
else
|
||||
ncclDevFuncTable[ncclShmem.work.header.funcIndex]();
|
||||
ncclDevFuncTable[ncclShmem.work.header.funcIndex]();
|
||||
#else
|
||||
if (COLL_UNROLL == 4)
|
||||
NCCL_CALL_FUNCTIONS_4(ncclShmem.work.header.funcIndex);
|
||||
else
|
||||
NCCL_CALL_FUNCTIONS(ncclShmem.work.header.funcIndex);
|
||||
NCCL_CALL_FUNCTIONS(ncclShmem.work.header.funcIndex);
|
||||
#endif
|
||||
}
|
||||
|
||||
@@ -385,27 +379,19 @@ __forceinline__ __device__ void ncclKernelMain(struct ncclDevComm* comm, struct
|
||||
}
|
||||
|
||||
__global__ void ncclDevKernel_Generic(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead);
|
||||
__global__ void ncclDevKernel_Generic_4(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead);
|
||||
#ifdef ENABLE_COLLTRACE
|
||||
__global__ void ncclDevKernelDebug_Generic(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead);
|
||||
__global__ void ncclDevKernelDebug_Generic_4(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead);
|
||||
#endif
|
||||
|
||||
#ifdef USE_INDIRECT_FUNCTION_CALL
|
||||
#define DEFINE_ncclDevFunc(suffix, coll, redop, ty, algo, proto) \
|
||||
#define DEFINE_ncclDevFunc(suffix, coll, redop, ty, algo, proto, unroll) \
|
||||
__device__ void ncclDevFunc_##suffix() { \
|
||||
RunWork<coll, ty, redop<ty>, algo, proto, 2>().run(&ncclShmem.work); \
|
||||
} \
|
||||
__device__ void ncclDevFunc_##suffix##_4() { \
|
||||
RunWork<coll, ty, redop<ty>, algo, proto, 4>().run(&ncclShmem.work); \
|
||||
RunWork<coll, ty, redop<ty>, algo, proto, unroll>().run(&ncclShmem.work); \
|
||||
}
|
||||
#else
|
||||
#define DEFINE_ncclDevFunc(suffix, coll, redop, ty, algo, proto) \
|
||||
#define DEFINE_ncclDevFunc(suffix, coll, redop, ty, algo, proto, unroll) \
|
||||
__device__ __attribute__((noinline)) void ncclDevFunc_##suffix() { \
|
||||
RunWork<coll, ty, redop<ty>, algo, proto, 2>().run(&ncclShmem.work); \
|
||||
} \
|
||||
__device__ __attribute__((noinline)) void ncclDevFunc_##suffix##_4() { \
|
||||
RunWork<coll, ty, redop<ty>, algo, proto, 4>().run(&ncclShmem.work); \
|
||||
RunWork<coll, ty, redop<ty>, algo, proto, unroll>().run(&ncclShmem.work); \
|
||||
}
|
||||
#endif
|
||||
|
||||
|
||||
+67
-56
@@ -1,6 +1,7 @@
|
||||
#!/usr/bin/env python3
|
||||
import os
|
||||
import sys
|
||||
import subprocess
|
||||
|
||||
# Order of redops, tys, protos, algos must match src/include/device.h
|
||||
all_colls = ["AllGather","AllReduce","AllToAllPivot","Broadcast","Reduce","ReduceScatter","SendRecv"]
|
||||
@@ -8,8 +9,9 @@ all_redops = ["Sum","Prod","MinMax","PreMulSum","SumPostDiv"]
|
||||
all_tys = ["i8","u8","i32","u32","i64","u64","f16","f32","f64","bf16", "f8", "bf8"]
|
||||
all_protos = ["LL","LL128","SIMPLE"]
|
||||
all_algos = ["TREE","RING"]
|
||||
all_unroll = ["2", "4"]
|
||||
|
||||
all_params = [all_colls, all_algos, all_protos, all_redops, all_tys]
|
||||
all_params = [all_colls, all_algos, all_protos, all_redops, all_tys, all_unroll]
|
||||
|
||||
################################################################################
|
||||
# The first command line argument is the path to the directory to generate and
|
||||
@@ -65,11 +67,12 @@ else:
|
||||
def paste(sep, *args):
|
||||
return sep.join(x for x in args if x is not None)
|
||||
|
||||
is_ifc = 1 if sys.argv[2] == "ON" else 0
|
||||
is_colltrace = 1 if sys.argv[3] == "ON" else 0
|
||||
is_msccl_kernels = 1 if sys.argv[4] == "ON" else 0
|
||||
is_ifc = 1 if sys.argv[2] == "ON" else 0
|
||||
is_colltrace = 1 if sys.argv[3] == "ON" else 0
|
||||
is_msccl_kernels = 1 if sys.argv[4] == "ON" else 0
|
||||
is_local_arch_only = 1 if sys.argv[5] == "ON" else 0
|
||||
|
||||
func_pattern = sys.argv[5:6]
|
||||
func_pattern = sys.argv[6:7]
|
||||
if func_pattern and func_pattern[0]:
|
||||
func_pattern = func_pattern[0]
|
||||
else:
|
||||
@@ -130,6 +133,40 @@ coll_lower_to_camel = {coll_camel_to_lower[x]: x for x in coll_camel_to_lower}
|
||||
|
||||
################################################################################
|
||||
|
||||
def calc_unroll_for_local_arch():
|
||||
if not is_local_arch_only:
|
||||
return
|
||||
|
||||
res = subprocess.run(['rocminfo'], capture_output=True, text=True)
|
||||
rocminfo_output = res.stdout
|
||||
|
||||
# Parse rocminfo binary output
|
||||
gfx_targets = {}
|
||||
curr_name = None
|
||||
for line in rocminfo_output.splitlines():
|
||||
line = line.strip()
|
||||
|
||||
if line.startswith("Name:"):
|
||||
name = line.split(':')[-1].strip()
|
||||
if "gfx" in name:
|
||||
curr_name = name
|
||||
if line.startswith("Compute Unit:") and curr_name:
|
||||
cu_count = int(line.split(':')[-1].strip())
|
||||
gfx_targets[(curr_name, cu_count)] = None
|
||||
curr_name = None
|
||||
|
||||
# We want to remove duplicates but cannot use a dictionary since same gfx name can have different cu counts
|
||||
# Use (gfx_name, cu_count) as key for dictionary and convert it to list here
|
||||
gfx_targets = list(gfx_targets.keys())
|
||||
|
||||
# Homogeneous system is required to build for only 1 varient of unroll factor
|
||||
if len(gfx_targets) == 1:
|
||||
gfx_name, cu_count = gfx_targets[0]
|
||||
if ("gfx908" == gfx_name or "gfx94" in gfx_name) and cu_count > 80:
|
||||
return 2
|
||||
else:
|
||||
return 4
|
||||
|
||||
# Helper function to check if the conditions for the collective is being met
|
||||
def func_validate(coll, algo, proto, redop, ty):
|
||||
if redop == "SumPostDiv" and ty[0] not in ("i","u"):
|
||||
@@ -182,10 +219,10 @@ def func_filter(function_params, current_idx, item_list=None):
|
||||
# For each loop layer remove the last element in item_list
|
||||
item_list.pop()
|
||||
else:
|
||||
coll, algo, proto, redop, ty = item_list
|
||||
coll, algo, proto, redop, ty, unroll = item_list
|
||||
|
||||
if func_validate(*item_list):
|
||||
yield(coll, algo, proto, redop, ty)
|
||||
if func_validate(coll, algo, proto, redop, ty):
|
||||
yield(coll, algo, proto, redop, ty, unroll)
|
||||
|
||||
# Parse ONLY_FUNCS input and feed it to func_filter
|
||||
def parse_input(func_pattern):
|
||||
@@ -205,7 +242,11 @@ def parse_input(func_pattern):
|
||||
|
||||
# Maps functions to the chosen representative for the equivalence class it
|
||||
# belongs to. For instance (sum, signed int) maps to (sum, unsigned int).
|
||||
def equivalent_primary(coll, algo, proto, redop, ty):
|
||||
def equivalent_primary(coll, algo, proto, redop, ty, unroll):
|
||||
# if local arch only, we only need to build for 1 varient of coll_unroll.
|
||||
# map the other varient of coll_unroll to this one.
|
||||
if coll_unroll:
|
||||
unroll = str(coll_unroll)
|
||||
if coll in ("AllReduce", "Reduce", "ReduceScatter"):
|
||||
# map signed integer sum/prod to unsigned
|
||||
if redop in ("Sum","Prod","PreMulSum") and ty[0]=="i":
|
||||
@@ -213,7 +254,7 @@ def equivalent_primary(coll, algo, proto, redop, ty):
|
||||
# map signed integer min/max to unsigned for non-NVLS
|
||||
elif redop=="MinMax" and ty[0]=="i" and ("NVLS" not in algo):
|
||||
ty = "u"+ty[1:]
|
||||
return (coll, algo, proto, redop, ty)
|
||||
return (coll, algo, proto, redop, ty, unroll)
|
||||
|
||||
# Order rows are enumerated must match formula of `ncclDevFuncId()`:
|
||||
def enumerate_func_rows():
|
||||
@@ -222,23 +263,27 @@ def enumerate_func_rows():
|
||||
for proto in all_protos:
|
||||
for redop in all_redops:
|
||||
for ty in all_tys:
|
||||
if func_validate(coll, algo, proto, redop, ty):
|
||||
yield (coll, algo, proto, redop, ty)
|
||||
for unroll in all_unroll:
|
||||
if func_validate(coll, algo, proto, redop, ty):
|
||||
yield (coll, algo, proto, redop, ty, unroll)
|
||||
|
||||
# Sort the hashmap based on custom key <coll> <algo> <proto> <redop> <ty>
|
||||
def custom_sort_key(fn):
|
||||
coll, algo, proto, redop, ty = fn
|
||||
coll, algo, proto, redop, ty, unroll = fn
|
||||
|
||||
return (
|
||||
all_colls.index(coll),
|
||||
all_algos.index(algo),
|
||||
all_protos.index(proto),
|
||||
all_redops.index(redop),
|
||||
all_tys.index(ty)
|
||||
all_tys.index(ty),
|
||||
all_unroll.index(unroll)
|
||||
)
|
||||
|
||||
################################################################################
|
||||
|
||||
coll_unroll = calc_unroll_for_local_arch()
|
||||
|
||||
# Corresponds to ncclDevFuncRowToId[]
|
||||
func_rows = [fn for fn in enumerate_func_rows()]
|
||||
|
||||
@@ -262,12 +307,12 @@ with open(os.path.join(gensrc, "device_table.h"), "w") as f:
|
||||
sym = paste("_", "ncclDevFunc", *fn)
|
||||
if fn[2] == "LL128":
|
||||
out("#if defined(__gfx90a__) && defined(ENABLE_LL128)\n")
|
||||
out("%s %s();\n%s %s_4();\n#else\n" % (func_declaration, sym, func_declaration, sym))
|
||||
out("%s %s();\n#else\n" % (func_declaration, sym))
|
||||
fn_ll = fn[:2] + ("LL",) + fn[3:]
|
||||
sym_ll = paste("_", "ncclDevFunc", *fn_ll)
|
||||
out("%s %s();\n%s %s_4();\n#endif\n" % (func_declaration, sym_ll, func_declaration, sym_ll))
|
||||
out("%s %s();\n#endif\n" % (func_declaration, sym_ll))
|
||||
else:
|
||||
out("%s %s();\n%s %s_4();\n" % (func_declaration, sym, func_declaration, sym))
|
||||
out("%s %s();\n" % (func_declaration, sym))
|
||||
out("\n")
|
||||
|
||||
out("typedef void(*ncclDevFuncPtr_t)();\n\n")
|
||||
@@ -286,22 +331,6 @@ with open(os.path.join(gensrc, "device_table.h"), "w") as f:
|
||||
index += 1
|
||||
out("nullptr};\n")
|
||||
out("\n")
|
||||
|
||||
out("__device__ ncclDevFuncPtr_t const ncclDevFuncTable_4[] = {\n")
|
||||
index = 0
|
||||
for fn in primary_funcs:
|
||||
sym = paste("_", "ncclDevFunc", *fn)
|
||||
if fn[2] == "LL128":
|
||||
out("#if defined(__gfx90a__) && defined(ENABLE_LL128)\n")
|
||||
out("/*%4d*/ %s_4,\n#else\n" % (index, sym))
|
||||
fn_ll = fn[:2] + ("LL",) + fn[3:]
|
||||
sym_ll = paste("_", "ncclDevFunc", *fn_ll)
|
||||
out("/*%4d*/ %s_4,\n#endif\n" % (index, sym_ll))
|
||||
else:
|
||||
out("/*%4d*/ %s_4,\n" % (index, sym))
|
||||
index += 1
|
||||
out("nullptr};\n")
|
||||
out("\n")
|
||||
|
||||
if not is_ifc:
|
||||
out("template<unsigned short f, unsigned short l>\n"
|
||||
@@ -322,24 +351,6 @@ with open(os.path.join(gensrc, "device_table.h"), "w") as f:
|
||||
out("__forceinline__ __device__ void NCCL_CALL_FUNCTIONS(unsigned short funcIndex) noexcept {\n")
|
||||
out(f" Caller<0, {index}>::call(funcIndex);\n")
|
||||
out("}\n\n")
|
||||
out("template<unsigned short f, unsigned short l>\n"
|
||||
"struct Caller4 {\n"
|
||||
" static __forceinline__ __device__ __host__\n"
|
||||
" void call4(unsigned short funcIndex) noexcept\n"
|
||||
" {\n"
|
||||
" constexpr unsigned short m = f + (l - f) / 2;\n"
|
||||
" return (funcIndex < m) ? Caller4<f, m>::call4(funcIndex) : Caller4<m, l>::call4(funcIndex);\n"
|
||||
" }\n"
|
||||
"};\n"
|
||||
"\n"
|
||||
"template<unsigned short f>\n"
|
||||
"struct Caller4<f, f + 1>{\n"
|
||||
" static __forceinline__ __device__ __host__\n"
|
||||
" void call4(unsigned short funcIndex) noexcept { ncclDevFuncTable_4[f](); }\n"
|
||||
"};\n")
|
||||
out("__forceinline__ __device__ void NCCL_CALL_FUNCTIONS_4(unsigned short funcIndex) noexcept {\n")
|
||||
out(f" Caller4<0, {index}>::call4(funcIndex);\n")
|
||||
out("}\n\n")
|
||||
|
||||
# Generate <gensrc>/device_table.cpp
|
||||
if is_colltrace:
|
||||
@@ -381,7 +392,7 @@ with open(os.path.join(gensrc, "host_table.cpp"), "w") as f:
|
||||
# Maps to .cu filename which implements this func. The only constraint is that
|
||||
# "coll" is reflected in the name: formally that no two funcs having different
|
||||
# coll's map to the same filename.
|
||||
def impl_filename(coll, algo, proto, redop, ty):
|
||||
def impl_filename(coll, algo, proto, redop, ty, unroll):
|
||||
return "%s.cpp" % paste("_", coll_camel_to_lower[coll], redop and redop.lower(), ty)
|
||||
|
||||
# Partition the functions and kernels to the .cu filenames. The partition is
|
||||
@@ -437,14 +448,14 @@ for name in name_to_funcs.keys():
|
||||
)
|
||||
|
||||
for fn in fns:
|
||||
(coll, algo, proto, redop, ty) = fn
|
||||
sym = paste("_", coll, algo, proto, redop, ty)
|
||||
(coll, algo, proto, redop, ty, unroll) = fn
|
||||
sym = paste("_", coll, algo, proto, redop, ty, unroll)
|
||||
if proto == "LL128":
|
||||
out("#if defined(__gfx90a__) && defined(ENABLE_LL128)\n")
|
||||
out(
|
||||
"DEFINE_ncclDevFunc({sym}, ncclFunc{coll}, {redop_cxx}, {ty_cxx}, NCCL_ALGO_{algo}, NCCL_PROTO_{proto})\n"
|
||||
"DEFINE_ncclDevFunc({sym}, ncclFunc{coll}, {redop_cxx}, {ty_cxx}, NCCL_ALGO_{algo}, NCCL_PROTO_{proto}, {unroll})\n"
|
||||
.format(sym=sym, coll=coll, redop_cxx=redop_to_cxx[redop], ty_cxx=ty_to_cxx[ty],
|
||||
algo=(algo or "RING"), proto=(proto or "SIMPLE"))
|
||||
algo=(algo or "RING"), proto=(proto or "SIMPLE"), unroll=unroll)
|
||||
)
|
||||
if proto == "LL128":
|
||||
out("#endif\n")
|
||||
|
||||
+15
-20
@@ -31,16 +31,15 @@ struct ncclKernelMatch {
|
||||
};
|
||||
|
||||
#ifdef ENABLE_COLLTRACE
|
||||
static ncclKernelMatch const ncclKerns[4] = {
|
||||
#define ncclGetKernelIndex(p_comm) ((p_comm)->collTraceThread ? 1 : 0)
|
||||
static ncclKernelMatch const ncclKerns[2] = {
|
||||
{(void *)ncclDevKernel_Generic, true},
|
||||
{(void *)ncclDevKernel_Generic_4, true},
|
||||
{(void *)ncclDevKernelDebug_Generic, true},
|
||||
{(void *)ncclDevKernelDebug_Generic_4, true},
|
||||
};
|
||||
#else
|
||||
static ncclKernelMatch const ncclKerns[2] = {
|
||||
{(void*)ncclDevKernel_Generic, true},
|
||||
{(void*)ncclDevKernel_Generic_4, true},
|
||||
#define ncclGetKernelIndex(p_comm) (0)
|
||||
static ncclKernelMatch const ncclKerns[1] = {
|
||||
{(void*)ncclDevKernel_Generic, true}
|
||||
};
|
||||
#endif
|
||||
|
||||
@@ -56,24 +55,19 @@ static ncclResult_t initCollProxyOp(struct ncclInfo* collInfo, int channelId, ui
|
||||
static ncclResult_t getTunerInfo(struct ncclInfo* collInfo, int collNetSupport, int nvlsSupport, int numPipeOps);
|
||||
static ncclResult_t topoGetAlgoInfo(struct ncclInfo* collInfo, int collNetSupport, int nvlsSupport, int numPipeOps);
|
||||
static ncclResult_t getChannnelThreadInfo(struct ncclInfo* collInfo);
|
||||
static ncclResult_t computeCollWorkFunc(struct ncclInfo* collInfo);
|
||||
static ncclResult_t computeCollWorkFunc(struct ncclInfo* collInfo, int unroll);
|
||||
static ncclResult_t getPatternInfo(struct ncclInfo* collInfo);
|
||||
static ncclResult_t getLoopInfo(struct ncclInfo* collInfo);
|
||||
static ncclResult_t getCollNetSupport(struct ncclInfo* info, int* collNetSupport);
|
||||
|
||||
int ncclGetKernelIndex(struct ncclComm* comm) {
|
||||
#if ENABLE_COLLTRACE
|
||||
int start_idx = comm->collTraceThread ? 2 : 0;
|
||||
#else
|
||||
int start_idx = 0;
|
||||
#endif
|
||||
int getUnrollFactor(struct ncclComm* comm) {
|
||||
hipDeviceProp_t devProp;
|
||||
CUDACHECK(hipGetDeviceProperties(&devProp, comm->cudaDev));
|
||||
if(IsArchMatch(devProp.gcnArchName, "gfx908") || (IsArchMatch(devProp.gcnArchName, "gfx94")
|
||||
&& devProp.multiProcessorCount > 80))
|
||||
return start_idx;
|
||||
return NCCL_UNROLL_2;
|
||||
else
|
||||
return start_idx + 1;
|
||||
return NCCL_UNROLL_4;
|
||||
}
|
||||
|
||||
// Returns maximum kernel stack size of all CUDA kernels
|
||||
@@ -194,7 +188,7 @@ static ncclResult_t appendWorkElemP2p(
|
||||
struct ncclComm* comm, struct ncclKernelPlan* plan, int channelId,
|
||||
struct ncclWorkElemP2p const *elem, bool fuseOk
|
||||
) {
|
||||
int funcIndex = ncclDevFuncId_P2p();
|
||||
int funcIndex = ncclDevFuncId_P2p(plan->unroll);
|
||||
if (funcIndex < 0) {
|
||||
WARN("%s: unsupported collective. Please ensure the collective has been enabled in build.", __func__);
|
||||
return ncclInvalidUsage;
|
||||
@@ -220,7 +214,7 @@ static ncclResult_t appendWorkElemP2p(
|
||||
}
|
||||
q = ncclMemoryStackAlloc<struct ncclWorkList>(&comm->memScoped);
|
||||
q->work.header.type = ncclWorkTypeP2p;
|
||||
q->work.header.funcIndex = ncclDevFuncId_P2p();
|
||||
q->work.header.funcIndex = funcIndex;
|
||||
chan->p2pTailElem[ncclWorkP2pTypeRecv-1] = 0;
|
||||
chan->p2pTailElem[ncclWorkP2pTypeSend-1] = 1;
|
||||
q->work.p2pElems[chan->p2pTailElem[elem->p2pType-1]] = *elem; // C++ struct assignment
|
||||
@@ -851,7 +845,7 @@ static ncclResult_t scheduleCollTasksToPlan(
|
||||
NCCLCHECK(getTunerInfo(aggInfo, collNetSupport, nvlsSupport, 1));
|
||||
NCCLCHECK(topoGetAlgoInfo(aggInfo, collNetSupport, nvlsSupport, 1));
|
||||
NCCLCHECK(getChannnelThreadInfo(aggInfo));
|
||||
NCCLCHECK(computeCollWorkFunc(aggInfo));
|
||||
NCCLCHECK(computeCollWorkFunc(aggInfo, plan->unroll));
|
||||
NCCLCHECK(getPatternInfo(aggInfo));
|
||||
|
||||
// Try to assign algo and proto to all possible collectives
|
||||
@@ -1330,6 +1324,7 @@ ncclResult_t ncclLaunchPrepare(struct ncclComm* comm) {
|
||||
plan->comm = comm;
|
||||
plan->reclaimer.fn = reclaimPlan;
|
||||
plan->persistent = persistent;
|
||||
plan->unroll = getUnrollFactor(comm);
|
||||
|
||||
// Non-persistent kernels fill up at most half of our fifo per kernel.
|
||||
int nWorkBudget = plan->persistent ? INT_MAX : comm->workFifoDepth/2;
|
||||
@@ -1762,8 +1757,8 @@ static ncclResult_t getPatternInfo(struct ncclInfo* collInfo) {
|
||||
|
||||
RCCL_PARAM(IntraNetThreshold, "INTRANET_THRESHOLD", 8388608);
|
||||
|
||||
static ncclResult_t computeCollWorkFunc(struct ncclInfo* collInfo) {
|
||||
collInfo->workFuncIndex = ncclDevFuncId(collInfo->coll, collInfo->opFull.op, collInfo->datatype, collInfo->algorithm, collInfo->protocol);
|
||||
static ncclResult_t computeCollWorkFunc(struct ncclInfo* collInfo, int unroll) {
|
||||
collInfo->workFuncIndex = ncclDevFuncId(collInfo->coll, collInfo->opFull.op, collInfo->datatype, collInfo->algorithm, collInfo->protocol, unroll);
|
||||
if (collInfo->workFuncIndex < 0) {
|
||||
WARN("%s: unsupported collective. Please ensure the collective has been enabled in build.", __func__);
|
||||
return ncclInvalidUsage;
|
||||
|
||||
@@ -224,6 +224,9 @@ struct ncclKernelPlan {
|
||||
struct ncclIntruQueue<struct ncclProxyOp, &ncclProxyOp::enqNext> proxyOpQueue;
|
||||
} channels[MAXCHANNELS];
|
||||
size_t maxBytesPerChannel;
|
||||
|
||||
// Unroll factor for plan [RCCL]
|
||||
int unroll;
|
||||
};
|
||||
|
||||
#define NCCL_MAGIC 0x0280028002800280 // Nickel atomic number is 28.
|
||||
|
||||
+22
-16
@@ -553,57 +553,63 @@ inline bool ncclNvlsSupported(int devRedOp, int type) {
|
||||
extern int const ncclDevFuncRowToId[];
|
||||
|
||||
// `ncclDevFuncId()` needs to be in sync with 'ALL_COLLS' in generate.py
|
||||
inline int ncclDevFuncId(int coll, int devRedOp, int type, int algo, int proto) {
|
||||
inline int ncclDevFuncId(int coll, int devRedOp, int type, int algo, int proto, int unroll) {
|
||||
int row = 0;
|
||||
do {
|
||||
// RING / <all_protos> / Sum / int8_t
|
||||
if (coll == ncclFuncAllGather) {
|
||||
row += proto;
|
||||
row += proto * NCCL_NUM_UNROLLS + unroll;
|
||||
break;
|
||||
}
|
||||
row += NCCL_NUM_PROTOCOLS;
|
||||
row += NCCL_NUM_UNROLLS * NCCL_NUM_PROTOCOLS;
|
||||
|
||||
// <all_algos> / <all_protos> / <all_redops> / <all_types>
|
||||
if (coll == ncclFuncAllReduce) {
|
||||
row += (((algo * NCCL_NUM_PROTOCOLS + proto) * ncclNumDevRedOps + devRedOp) * ncclNumTypes + type) - NCCL_NUM_FLOATS * (algo * NCCL_NUM_PROTOCOLS + proto);
|
||||
row += ((((algo * NCCL_NUM_PROTOCOLS + proto) * ncclNumDevRedOps + devRedOp) * ncclNumTypes + type) * NCCL_NUM_UNROLLS + unroll) - NCCL_NUM_FLOATS * (algo * NCCL_NUM_PROTOCOLS + proto) * NCCL_NUM_UNROLLS;
|
||||
break;
|
||||
}
|
||||
row += (NCCL_NUM_ALGORITHMS - 4) * NCCL_NUM_PROTOCOLS * (ncclNumDevRedOps * ncclNumTypes - NCCL_NUM_FLOATS);
|
||||
row += NCCL_NUM_UNROLLS * (NCCL_NUM_ALGORITHMS - 4) * NCCL_NUM_PROTOCOLS * (ncclNumDevRedOps * ncclNumTypes - NCCL_NUM_FLOATS);
|
||||
|
||||
// RING / SIMPLE / Sum / int8_t
|
||||
if (coll == ncclFuncAllToAllPivot) break;
|
||||
row += 1;
|
||||
if (coll == ncclFuncAllToAllPivot) {
|
||||
row += unroll;
|
||||
break;
|
||||
}
|
||||
row += NCCL_NUM_UNROLLS;
|
||||
|
||||
// RING / <all_protos> / Sum / int8_t
|
||||
if (coll == ncclFuncBroadcast) {
|
||||
row += proto;
|
||||
row += proto * NCCL_NUM_UNROLLS + unroll;
|
||||
break;
|
||||
}
|
||||
row += NCCL_NUM_PROTOCOLS;
|
||||
row += NCCL_NUM_UNROLLS * NCCL_NUM_PROTOCOLS;
|
||||
|
||||
// RING / <all_protos> / <all_redops> / <all_types>
|
||||
if (coll == ncclFuncReduce) {
|
||||
row += ((proto * ncclNumDevRedOps + devRedOp) * ncclNumTypes + type) - NCCL_NUM_FLOATS * proto;
|
||||
row += (((proto * ncclNumDevRedOps + devRedOp) * ncclNumTypes + type) * NCCL_NUM_UNROLLS + unroll) - NCCL_NUM_FLOATS * proto * NCCL_NUM_UNROLLS;
|
||||
break;
|
||||
}
|
||||
row += NCCL_NUM_PROTOCOLS * (ncclNumDevRedOps * ncclNumTypes - NCCL_NUM_FLOATS);
|
||||
row += NCCL_NUM_UNROLLS * NCCL_NUM_PROTOCOLS * (ncclNumDevRedOps * ncclNumTypes - NCCL_NUM_FLOATS);
|
||||
|
||||
// RING / <all_protos> / <all_redops> / <all_types>
|
||||
if (coll == ncclFuncReduceScatter) {
|
||||
row += ((proto * ncclNumDevRedOps + devRedOp) * ncclNumTypes + type) - NCCL_NUM_FLOATS * proto;
|
||||
row += (((proto * ncclNumDevRedOps + devRedOp) * ncclNumTypes + type) * NCCL_NUM_UNROLLS + unroll) - NCCL_NUM_FLOATS * proto * NCCL_NUM_UNROLLS;
|
||||
break;
|
||||
}
|
||||
row += NCCL_NUM_PROTOCOLS * (ncclNumDevRedOps * ncclNumTypes - NCCL_NUM_FLOATS);
|
||||
row += NCCL_NUM_UNROLLS * NCCL_NUM_PROTOCOLS * (ncclNumDevRedOps * ncclNumTypes - NCCL_NUM_FLOATS);
|
||||
|
||||
// RING / SIMPLE / Sum / int8_t
|
||||
if (coll == ncclFuncSendRecv) break;
|
||||
row += 1;
|
||||
if (coll == ncclFuncSendRecv) {
|
||||
row += unroll;
|
||||
break;
|
||||
}
|
||||
row += NCCL_NUM_UNROLLS;
|
||||
|
||||
} while (false);
|
||||
|
||||
return ncclDevFuncRowToId[row];
|
||||
}
|
||||
|
||||
inline int ncclDevFuncId_P2p() { return ncclDevFuncRowToId[FUNC_INDEX_TOTAL - NCCL_NUM_ONERANK - 1]; }
|
||||
inline int ncclDevFuncId_P2p(int unroll) { return ncclDevFuncRowToId[FUNC_INDEX_TOTAL - NCCL_NUM_ONERANK - (unroll > 0 ? 0 : 1) - 1]; }
|
||||
|
||||
#endif
|
||||
|
||||
@@ -13,7 +13,7 @@ typedef enum {NCCL_INIT=1, NCCL_COLL=2, NCCL_P2P=4, NCCL_SHM=8, NCCL_NET=16, NCC
|
||||
typedef void (*ncclDebugLogger_t)(ncclDebugLogLevel level, unsigned long flags, const char *file, int line, const char *fmt, ...);
|
||||
|
||||
#define NCCL_NUM_ONERANK 12
|
||||
#define FUNC_INDEX_TOTAL 656 + NCCL_NUM_ONERANK
|
||||
#define FUNC_INDEX_TOTAL 1312 + NCCL_NUM_ONERANK
|
||||
|
||||
#define NCCL_NUM_FUNCTIONS 5 // Send/Recv not included for now
|
||||
typedef enum {
|
||||
@@ -44,6 +44,10 @@ typedef enum {
|
||||
#define NCCL_PROTO_LL128 1
|
||||
#define NCCL_PROTO_SIMPLE 2
|
||||
|
||||
#define NCCL_NUM_FLOATS 6 // half/float/double/rccl_bfloat16
|
||||
#define NCCL_NUM_UNROLLS 2 // 2/4
|
||||
#define NCCL_UNROLL_2 0
|
||||
#define NCCL_UNROLL_4 1
|
||||
|
||||
#define NCCL_NUM_FLOATS 6 // half/float/double/rccl_bfloat16
|
||||
|
||||
#endif
|
||||
+1
-1
@@ -243,7 +243,7 @@ void *ncclCommThreadMain(void *arg) {
|
||||
(double)(td->timeStamp)/vega_gpu_rtc_freq, comm->rank, td->bid,
|
||||
fIdx, td->data_0, td->opCount, td->data_1);
|
||||
} else {
|
||||
if (fIdx == ncclDevFuncId_P2p() || type == ncclCollTraceP2pElemType)
|
||||
if (type == ncclCollTraceP2pElemType)
|
||||
sprintf(line, "## [%012.6f] [%02d:%02d] %06x-%06x", (double)(td->timeStamp)/vega_gpu_rtc_freq, comm->rank, td->bid, td->p2pOpCount[0], td->p2pOpCount[1]);
|
||||
else
|
||||
sprintf(line, "## [%012.6f] [%02d:%02d] %06lx", (double)(td->timeStamp)/vega_gpu_rtc_freq, comm->rank, td->bid, td->opCount);
|
||||
|
||||
Reference in New Issue
Block a user