From 373f113524ca422d6a356d211dce9a9e9ae16ca4 Mon Sep 17 00:00:00 2001 From: Bertan Dogancay <111835151+BertanDogancay@users.noreply.github.com> Date: Mon, 21 Oct 2024 10:53:11 -0400 Subject: [PATCH] Dynamically select unroll factor to build for when targeting local arch (#1371) * Dynamically select unroll factor to build for when targeting local arch only --- CMakeLists.txt | 2 +- src/device/common.cu | 10 +--- src/device/common.h | 28 +++------ src/device/generate.py | 123 +++++++++++++++++++++----------------- src/enqueue.cc | 35 +++++------ src/include/comm.h | 3 + src/include/device.h | 38 +++++++----- src/include/nccl_common.h | 8 ++- src/init.cc | 2 +- 9 files changed, 124 insertions(+), 125 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 3fac74cedc..e4b3603543 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -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 diff --git a/src/device/common.cu b/src/device/common.cu index ae15e06598..0022b6c233 100644 --- a/src/device/common.cu +++ b/src/device/common.cu @@ -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 diff --git a/src/device/common.h b/src/device/common.h index cb3a7aa3cc..e9f0bed8b3 100644 --- a/src/device/common.h +++ b/src/device/common.h @@ -227,7 +227,7 @@ static __forceinline__ __device__ void ncclRedopPtrDeref(struct ncclWorkElem* we } } -template +template __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, algo, proto, 2>().run(&ncclShmem.work); \ - } \ - __device__ void ncclDevFunc_##suffix##_4() { \ - RunWork, algo, proto, 4>().run(&ncclShmem.work); \ + RunWork, 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, algo, proto, 2>().run(&ncclShmem.work); \ - } \ - __device__ __attribute__((noinline)) void ncclDevFunc_##suffix##_4() { \ - RunWork, algo, proto, 4>().run(&ncclShmem.work); \ + RunWork, algo, proto, unroll>().run(&ncclShmem.work); \ } #endif diff --git a/src/device/generate.py b/src/device/generate.py index 294a449535..892e4b2a7f 100755 --- a/src/device/generate.py +++ b/src/device/generate.py @@ -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 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\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\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::call4(funcIndex) : Caller4::call4(funcIndex);\n" - " }\n" - "};\n" - "\n" - "template\n" - "struct Caller4{\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 /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") diff --git a/src/enqueue.cc b/src/enqueue.cc index 196653c94b..430aa23018 100644 --- a/src/enqueue.cc +++ b/src/enqueue.cc @@ -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(&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; diff --git a/src/include/comm.h b/src/include/comm.h index 160439c048..ccc5572a6f 100644 --- a/src/include/comm.h +++ b/src/include/comm.h @@ -224,6 +224,9 @@ struct ncclKernelPlan { struct ncclIntruQueue proxyOpQueue; } channels[MAXCHANNELS]; size_t maxBytesPerChannel; + + // Unroll factor for plan [RCCL] + int unroll; }; #define NCCL_MAGIC 0x0280028002800280 // Nickel atomic number is 28. diff --git a/src/include/device.h b/src/include/device.h index 54bb7027a2..f808146b5a 100644 --- a/src/include/device.h +++ b/src/include/device.h @@ -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 / / 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; // / / / 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 / / 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 / / / 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 / / / 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 diff --git a/src/include/nccl_common.h b/src/include/nccl_common.h index 68e8d31e02..67ff0959ee 100644 --- a/src/include/nccl_common.h +++ b/src/include/nccl_common.h @@ -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 \ No newline at end of file diff --git a/src/init.cc b/src/init.cc index e6f2c05d54..9c93d8ac51 100644 --- a/src/init.cc +++ b/src/init.cc @@ -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);