Enable tree kernels in build

Need to tune and specify NCCL_TREE_THRESHOLD to allow usage [ROCm/rccl commit: 76976c9e2e]
2019-10-04 17:45:40 +00:00
parent de13a48f7b
commit f86ee41415
3 changed files with 22 additions and 14 deletions
@@ -17,7 +17,8 @@
  NCCL_KERN_NAME(coll##LL, op, dtype)

 #define NCCL_FUNC4(coll, op, dtype) \
-  NCCL_FUNC5(coll##Ring, op, dtype)
+  NCCL_FUNC5(coll##Ring, op, dtype), \
+  NCCL_FUNC5(coll##Tree, op, dtype)

 // Must be consistent with ncclDataType_t
 #define NCCL_FUNCS3A(coll, op) \
@@ -55,7 +56,7 @@

 typedef void(*ncclKern_t)(struct ncclColl);
 // Must be consistent with the ncclFuncSet enum
-static ncclKern_t const ncclKerns[ncclCollCount*ncclNumOps*ncclNumTypes*2] = {
+static ncclKern_t const ncclKerns[ncclCollCount*ncclNumOps*ncclNumTypes*2*2] = {
  NCCL_FUNCS2B(ncclBroadcast),
  NCCL_FUNCS2A(ncclReduce),
  NCCL_FUNCS2B(ncclAllGather),