AI/rocm-systems
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Dynamically select unroll factor to build for when targeting local arch (#1371)

Browse Source 
* Dynamically select unroll factor to build for when targeting local arch only
This commit is contained in:
Bertan Dogancay
2024-10-21 10:53:11 -04:00
committed by GitHub
parent 7c077db307
commit 373f113524
9 changed files with 124 additions and 125 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/include/comm.h
+3
View File
@@ -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.
src/include/device.h
+22 -16
View File
@@ -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
src/include/nccl_common.h
+6 -2
View File
@@ -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