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

Template generic kernel for unroll factor (#1419)

Browse Source 
* Template generic kernel for unroll factor
This commit is contained in:
Bertan Dogancay
2024-11-12 18:27:29 -05:00
committed by GitHub
parent 2d07f18696
commit cb175fb0b3
8 changed files with 116 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/include/comm.h
+3 -3
View File
@@ -224,9 +224,6 @@ 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.
@@ -434,6 +431,9 @@ struct ncclComm {
// buffer registration cache
struct ncclRegCache regCache;
uint64_t endMagic;
// Unroll factor for plan [RCCL]
int unroll;
};
enum ncclLaunchMode {
src/include/device.h
+17 -23
View File
@@ -552,64 +552,58 @@ inline bool ncclNvlsSupported(int devRedOp, int type) {
// Map the rowIdx to funcIdx
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, int unroll) {
// `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 / <all_protos> / Sum / int8_t
if (coll == ncclFuncAllGather) {
row += proto * NCCL_NUM_UNROLLS + unroll;
row += proto;
break;
}
row += NCCL_NUM_UNROLLS * NCCL_NUM_PROTOCOLS;
row += 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_UNROLLS + unroll) - NCCL_NUM_FLOATS * (algo * NCCL_NUM_PROTOCOLS + proto) * NCCL_NUM_UNROLLS;
row += (((algo * NCCL_NUM_PROTOCOLS + proto) * ncclNumDevRedOps + devRedOp) * ncclNumTypes + type) - NCCL_NUM_FLOATS * (algo * NCCL_NUM_PROTOCOLS + proto);
break;
}
row += NCCL_NUM_UNROLLS * (NCCL_NUM_ALGORITHMS - 4) * NCCL_NUM_PROTOCOLS * (ncclNumDevRedOps * ncclNumTypes - NCCL_NUM_FLOATS);
row += (NCCL_NUM_ALGORITHMS - 4) * NCCL_NUM_PROTOCOLS * (ncclNumDevRedOps * ncclNumTypes - NCCL_NUM_FLOATS);
// RING / SIMPLE / Sum / int8_t
if (coll == ncclFuncAllToAllPivot) {
row += unroll;
break;
}
row += NCCL_NUM_UNROLLS;
if (coll == ncclFuncAllToAllPivot) break;
row += 1;
// RING / <all_protos> / Sum / int8_t
if (coll == ncclFuncBroadcast) {
row += proto * NCCL_NUM_UNROLLS + unroll;
row += proto;
break;
}
row += NCCL_NUM_UNROLLS * NCCL_NUM_PROTOCOLS;
row += NCCL_NUM_PROTOCOLS;
// RING / <all_protos> / <all_redops> / <all_types>
if (coll == ncclFuncReduce) {
row += (((proto * ncclNumDevRedOps + devRedOp) * ncclNumTypes + type) * NCCL_NUM_UNROLLS + unroll) - NCCL_NUM_FLOATS * proto * NCCL_NUM_UNROLLS;
row += ((proto * ncclNumDevRedOps + devRedOp) * ncclNumTypes + type) - NCCL_NUM_FLOATS * proto;
break;
}
row += NCCL_NUM_UNROLLS * NCCL_NUM_PROTOCOLS * (ncclNumDevRedOps * ncclNumTypes - NCCL_NUM_FLOATS);
row += 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_UNROLLS + unroll) - NCCL_NUM_FLOATS * proto * NCCL_NUM_UNROLLS;
row += ((proto * ncclNumDevRedOps + devRedOp) * ncclNumTypes + type) - NCCL_NUM_FLOATS * proto;
break;
}
row += NCCL_NUM_UNROLLS * NCCL_NUM_PROTOCOLS * (ncclNumDevRedOps * ncclNumTypes - NCCL_NUM_FLOATS);
row += NCCL_NUM_PROTOCOLS * (ncclNumDevRedOps * ncclNumTypes - NCCL_NUM_FLOATS);
// RING / SIMPLE / Sum / int8_t
if (coll == ncclFuncSendRecv) {
row += unroll;
break;
}
row += NCCL_NUM_UNROLLS;
if (coll == ncclFuncSendRecv) break;
row += 1;
} while (false);
return ncclDevFuncRowToId[row];
}
inline int ncclDevFuncId_P2p(int unroll) { return ncclDevFuncRowToId[FUNC_INDEX_TOTAL - NCCL_NUM_ONERANK - (unroll > 0 ? 0 : 1) - 1]; }
inline int ncclDevFuncId_P2p() { return ncclDevFuncRowToId[FUNC_INDEX_TOTAL - NCCL_NUM_ONERANK - 1]; }
#endif
src/include/nccl_common.h
+1 -1
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 1312 + NCCL_NUM_ONERANK
#define FUNC_INDEX_TOTAL 656 + NCCL_NUM_ONERANK
#define NCCL_NUM_FUNCTIONS 5 // Send/Recv not included for now
typedef enum {