From c1b3cd89117ba710890dbda94b3c3079944ab848 Mon Sep 17 00:00:00 2001 From: Mustafa Abduljabbar Date: Tue, 19 Aug 2025 16:41:19 -0400 Subject: [PATCH] Have ncclDevFuncId use 64-Bit keyed map with field packing (#1857) - Updated ncclDevFuncId to use a hash-based lookup with std::unordered_map. - Keys are now 64-bit integers, which pack coll, algo, proto, devRedOp, and type fields. - Improved flexibility and maintainability by moving away from row-based indexing. - Added error handling for missing keys in the hash map. - Aligned key generation logic with generate.py and updated generate.py. --- src/device/generate.py | 71 +++++++++++++++++++-------- src/include/device.h | 106 ++++++++++++++++------------------------- 2 files changed, 92 insertions(+), 85 deletions(-) diff --git a/src/device/generate.py b/src/device/generate.py index a129fc65ca..4f5ca038c8 100755 --- a/src/device/generate.py +++ b/src/device/generate.py @@ -3,15 +3,14 @@ import os import sys import subprocess -# Order of redops, tys, protos, algos must match src/include/device.h -all_colls = ["AllGather","AllReduce","AllReduceWithBias","AllToAllPivot","Broadcast","Reduce","ReduceScatter","SendRecv"] +# Order of colls, redops, tys, protos, algos must match src/include/device.h +all_colls = ["Broadcast", "Reduce", "AllGather", "ReduceScatter", "AllReduce", "AllReduceWithBias", "SendRecv", "", "", "AllToAllPivot"] all_redops = ["Sum","Prod","MinMax","PreMulSum","SumPostDiv"] all_tys = ["i8","u8","i32","u32","i64","u64","f16","f32","f64","bf16","f8e4m3","f8e5m2"] all_protos = ["LL","LL128","SIMPLE"] -all_algos = ["TREE","RING", "PAT"] +all_algos = ["TREE","RING", "", "", "", "", "PAT"] all_unroll = ["1", "2", "4"] use_acc = ["0", "1"] - all_params = [all_colls, all_algos, all_protos, all_redops, all_tys, use_acc, all_unroll] ################################################################################ @@ -187,7 +186,9 @@ def func_validate(coll, algo, proto, redop, ty, acc, unroll): return False if redop == "SumPostDiv" and ty[0] not in ("i","u"): return False - if algo not in algos_of_coll[coll] or proto not in protos_of_coll[coll] or redop not in redops_of_coll[coll] or ty not in tys_of_coll[coll] or unroll not in all_unroll: + if coll == "" or algo == "": + return False + if algo not in algos_of_coll[coll] or proto not in protos_of_coll[coll] or redop not in redops_of_coll[coll] or ty not in tys_of_coll[coll] or acc not in use_acc or unroll not in all_unroll: return False return True @@ -219,12 +220,12 @@ def func_filter(function_params, current_idx, item_list=None): # Check if the current element is recognized elements = current_element.split("/") current_param = all_params[current_idx] - + # Iterate over the elements in the elements list for item in elements: if item not in current_param: raise ValueError(f"Error: {item} is unrecognized or does not belong to this category {current_param}.") - + for item in elements: item_list.append(item) yield from func_filter(function_params, current_idx+1, item_list) @@ -267,8 +268,8 @@ def equivalent_primary(coll, algo, proto, redop, ty, acc, unroll): # Order rows are enumerated must match formula of `ncclDevFuncId()`: def enumerate_func_rows(): - for acc in use_acc: - for unroll in all_unroll: + for unroll in all_unroll: + for acc in use_acc: for coll in all_colls: for algo in all_algos: for proto in all_protos: @@ -280,7 +281,6 @@ def enumerate_func_rows(): # Sort the hashmap based on custom key def custom_sort_key(fn): coll, algo, proto, redop, ty, acc, unroll = fn - return ( all_unroll.index(unroll), use_acc.index(acc), @@ -380,7 +380,7 @@ with open(os.path.join(gensrc, "device_table.h"), "w") as f: index4 += 1 out("nullptr};\n") out("\n") - + if not is_ifc: out("template\n" "struct Caller1 {\n" @@ -445,7 +445,7 @@ if is_colltrace: out = f.write out('#include "nccl_common.h"\n#include "device.h"\n') out("\n") - + seen_fns = set() out("const char* funcNames[FUNC_INDEX_TOTAL] = {\n") for fn in primary_funcs: @@ -464,18 +464,47 @@ with open(os.path.join(gensrc, "host_table.cpp"), "w") as f: out = f.write out('#include "device.h"\n') out("\n") - - # The mapping from function rows to valid primary function ids. - out("extern int const ncclDevFuncRowToId[] = {\n") - index = 0 - for fn in func_rows[:len(func_rows)//len(all_unroll)]: - fn_id, comment = -1, "" + out("// The key for the ncclDevFuncNameToId map is a 64-bit unsigned integer.\n") + out("// Each field (coll, algo, proto, redop, ty) is packed into 4 bits,\n") + out("// This allows up to 16 unique values per field. The layout is:\n") + out("// bits 0-3: coll index\n") + out("// bits 4-7: algo index\n") + out("// bits 8-11: proto index\n") + out("// bits 12-15: redop index\n") + out("// bits 16-19: ty index\n") + out("#include \n") + out("extern std::unordered_map ncclDevFuncNameToId = {\n") + for fn in func_rows: + fn_id = -1 if fn is not None: fn_id = primary_to_index[equivalent_primary(*fn)] comment = " // " + paste(" ", *fn[:-1]) - out("/*%4d*/ %d,%s\n" % (index, fn_id, comment)) - index += 1 - out(f"{index}") + # Build the function signature string: " " + coll_idx = all_colls.index(fn[0]) + algo_idx = all_algos.index(fn[1]) + proto_idx = all_protos.index(fn[2]) + redop_idx = all_redops.index(fn[3]) + ty_idx = all_tys.index(fn[4]) + # Assert that 4 bits (16 values) is enough to map all_colls, all_algos, etc. + assert len(all_colls) <= 16, "Error: all_colls has more than 16 values, which exceeds 4-bit capacity." + assert len(all_algos) <= 16, "Error: all_algos has more than 16 values, which exceeds 4-bit capacity." + assert len(all_protos) <= 16, "Error: all_protos has more than 16 values, which exceeds 4-bit capacity." + assert len(all_redops) <= 16, "Error: all_redops has more than 16 values, which exceeds 4-bit capacity." + assert len(all_tys) <= 16, "Error: all_tys has more than 16 values, which exceeds 4-bit capacity." + # Create a 64-bit unsigned integer key and pack the indices into 4 bits each + key = ( + (coll_idx & 0xF) + | ((algo_idx & 0xF) << 4) + | ((proto_idx & 0xF) << 8) + | ((redop_idx & 0xF) << 12) + | ((ty_idx & 0xF) << 16) + ) + fn_str = f"{coll_idx} {algo_idx} {proto_idx} {redop_idx} {ty_idx}" + if fn[0] == "Broadcast": + key = ((coll_idx & 0x3F) | ((proto_idx & 0x3F) << 8)) + if fn[0] in ["SendRecv", "AllToAllPivot"]: + key = ((coll_idx & 0x3F)) + out(f' {{{key}, {fn_id}}}, {comment}\n') out("};\n") # Maps to .cu filename which implements this func. The only constraint is that diff --git a/src/include/device.h b/src/include/device.h index a1e92a413a..1789203ef1 100644 --- a/src/include/device.h +++ b/src/include/device.h @@ -20,6 +20,9 @@ #include #include #include +#include +#include +#include "debug.h" extern const char* ncclFuncStr[NCCL_NUM_FUNCTIONS+2]; @@ -125,6 +128,13 @@ static_assert(NCCL_LL_CLEAN_MASK % NCCL_STEPS == 0, "Invalid NCCL_LL_CLEAN_MASK #define NCCL_NVLS_REG_BUFFER 0x02 #define NCCL_NET_REG_BUFFER 0x04 +#define RCCL_FUNC_ID_MASK 0xF +#define RCCL_COLL_SHIFT 0 +#define RCCL_ALGO_SHIFT 4 +#define RCCL_PROTO_SHIFT 8 +#define RCCL_REDOP_SHIFT 12 +#define RCCL_DTYPE_SHIFT 16 + struct ncclConnInfo { // Regular comm mechanism char *buffs[NCCL_NUM_PROTOCOLS]; // Local for recv, remote for send @@ -687,74 +697,42 @@ inline bool ncclNvlsSupported(int devRedOp, int type) { } } -// Map the rowIdx to funcIdx -extern int const ncclDevFuncRowToId[]; +// Map the uint64_t key to funcIdx +extern std::unordered_map ncclDevFuncNameToId; // `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) { - int row = 0; - do { - // RING/PAT | | Sum | int8_t - int nAlgos = 2; - if (coll == ncclFuncAllGather) { - int algo1 = algo == NCCL_ALGO_RING ? 0 : - /*algo == NCCL_ALGO_PAT*/ 1; - row += algo1 * NCCL_NUM_PROTOCOLS + proto; - break; - } - row += nAlgos * NCCL_NUM_PROTOCOLS; - - // RING/TREE | | | - nAlgos = 2; - if (coll == ncclFuncAllReduce) { - int algo1 = algo == NCCL_ALGO_TREE ? 0 : - /*algo == NCCL_ALGO_RING*/ 1; - row += (((algo1 * NCCL_NUM_PROTOCOLS + proto) * ncclNumDevRedOps + devRedOp) * ncclNumTypes + type) - NCCL_NUM_FLOATS * (algo1 * NCCL_NUM_PROTOCOLS + proto); - break; - } - row += nAlgos * NCCL_NUM_PROTOCOLS * (ncclNumDevRedOps * ncclNumTypes - NCCL_NUM_FLOATS); - - // RING | SIMPLE | Sum | int8_t - nAlgos = 1; - if (coll == ncclFuncAllToAllPivot) break; - row += nAlgos * 1; - - // RING | | Sum | int8_t - nAlgos = 1; - if (coll == ncclFuncBroadcast) { - row += proto; - break; - } - row += nAlgos * NCCL_NUM_PROTOCOLS; - - // RING | | | - nAlgos = 1; - if (coll == ncclFuncReduce) { - row += ((proto * ncclNumDevRedOps + devRedOp) * ncclNumTypes + type) - NCCL_NUM_FLOATS * proto; - break; - } - row += nAlgos * NCCL_NUM_PROTOCOLS * (ncclNumDevRedOps * ncclNumTypes - NCCL_NUM_FLOATS); - - // RING/PAT | | | - nAlgos = 2; - if (coll == ncclFuncReduceScatter) { - int algo1 = algo == NCCL_ALGO_RING ? 0 : - /*algo == NCCL_ALGO_PAT*/ 1; - row += (((algo1 * NCCL_NUM_PROTOCOLS + proto) * ncclNumDevRedOps + devRedOp) * ncclNumTypes + type) - NCCL_NUM_FLOATS * (algo1 * NCCL_NUM_PROTOCOLS + proto); - break; - } - row += NCCL_NUM_PROTOCOLS * (ncclNumDevRedOps * ncclNumTypes - NCCL_NUM_FLOATS); - - // RING | SIMPLE | Sum | int8_t - nAlgos = 1; - if (coll == ncclFuncSendRecv) break; - row += nAlgos * 1; - - } while (false); - - return ncclDevFuncRowToId[row]; + int row = -1; + uint64_t key; + // Pack 4-bit fields from right (LSB) to left in order: + // coll, algo, proto, devRedOp, type + // This logic must be in sync with the key generation logic in generate.py + if (coll == ncclFuncBroadcast) { + key = ((uint64_t)(coll & RCCL_FUNC_ID_MASK) << RCCL_COLL_SHIFT ) | + ((uint64_t)(proto & RCCL_FUNC_ID_MASK) << RCCL_PROTO_SHIFT); + } else if (coll == ncclFuncSendRecv || coll == ncclFuncAllToAllPivot) { + key = ((uint64_t)(coll & RCCL_FUNC_ID_MASK) << RCCL_COLL_SHIFT ); + } else { + key = ((uint64_t)(coll & RCCL_FUNC_ID_MASK) << RCCL_COLL_SHIFT ) | + ((uint64_t)(algo & RCCL_FUNC_ID_MASK) << RCCL_ALGO_SHIFT ) | + ((uint64_t)(proto & RCCL_FUNC_ID_MASK) << RCCL_PROTO_SHIFT) | + ((uint64_t)(devRedOp & RCCL_FUNC_ID_MASK) << RCCL_REDOP_SHIFT) | + ((uint64_t)(type & RCCL_FUNC_ID_MASK) << RCCL_DTYPE_SHIFT); + } + auto it = ncclDevFuncNameToId.find(key); + if (it != ncclDevFuncNameToId.end()) { + row = it->second; + } + if(row < 0) { + WARN("Fatal error: ncclDevFuncId: %llu not found for coll: %d, algo: %d, proto: %d, devRedOp: %d, type: %d", key, coll, algo, proto, devRedOp, type); + return -1; + } + return row; } -inline int ncclDevFuncId_P2p() { return ncclDevFuncRowToId[FUNC_INDEX_TOTAL - AR_WITH_BIAS_FUNC_COUNTS - NCCL_NUM_ONERANK - 1]; } +inline int ncclDevFuncId_P2p() { + static int ncclDevFuncIdP2p = ncclDevFuncId(ncclFuncSendRecv, -1 , -1 , NCCL_ALGO_UNDEF, NCCL_PROTO_UNDEF); + return ncclDevFuncIdP2p; +} #endif