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.
This commit is contained in:
Mustafa Abduljabbar
committad av
GitHub
GitHub
förälder
6ade5065b4
incheckning
c1b3cd8911
+50
-21
@@ -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 <coll> <algo> <proto> <redop> <ty>
|
||||
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<unsigned short f, unsigned short l>\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 <unordered_map>\n")
|
||||
out("extern std::unordered_map<uint64_t, int> 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> <algo> <proto> <redop> <ty>"
|
||||
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
|
||||
|
||||
+42
-64
@@ -20,6 +20,9 @@
|
||||
#include <algorithm>
|
||||
#include <stdint.h>
|
||||
#include <sys/types.h>
|
||||
#include <unordered_map>
|
||||
#include <string>
|
||||
#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<uint64_t, int> 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 | <all_protos> | 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 | <all_protos> | <all_redops> | <all_types>
|
||||
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 | <all_protos> | Sum | int8_t
|
||||
nAlgos = 1;
|
||||
if (coll == ncclFuncBroadcast) {
|
||||
row += proto;
|
||||
break;
|
||||
}
|
||||
row += nAlgos * NCCL_NUM_PROTOCOLS;
|
||||
|
||||
// RING | <all_protos> | <all_redops> | <all_types>
|
||||
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 | <all_protos> | <all_redops> | <all_types>
|
||||
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
|
||||
|
||||
Referens i nytt ärende
Block a user