NCCL Tree源码
https://github.com/NVIDIA/nccl/blob/master/src/misc/trees.c
https://github.com/NVIDIA/nccl/pull/172
为了让代码运行起来，进行了小改
#include <iostream>
using namespace std;
int GetBtree(int nranks, int rank, int* u, int* d0, int* d1);
int GetDtree(int nranks, int rank, int* s0, int* d0_0, int* d0_1, int* s1, int* d1_0, int* d1_1);
main()
{
int nranks=12;
int i=0, u=0, d0=0, d1=0, du=0, dd0=0, dd1=0;
for ( i=0; i<nranks; i++ ) {
GetDtree (nranks, i, &u, &d0, &d1, &du, &dd0, &dd1);
cout << "i=" << i << ";\tu=" << u << ",\td0=" << d0 << ",\td1=" << d1 << endl;
cout << "\tdu=" << du << ",\tdd0=" << dd0 << ",\tdd1=" << dd1 << endl;
}
}
//得到单向tree
int GetBtree(int nranks, int rank, int* u, int* d0, int* d1) {
int up, down0, down1;
int bit;
//和IP掩码类似，从低位往高位，第一个不为0的位，标志了这个节点以下的子树边界
for (bit=1; bit<nranks; bit<<=1) {
if (bit & rank) break;
}
if (rank == 0) {
*u = -1;
*d0 = nranks > 1 ? bit >> 1 : -1;
*d1 = -1;
return 0;
}
up = (rank ^ bit) | (bit << 1);
if (up >= nranks) up = (rank ^ bit);
*u = up;
int lowbit = bit >> 1;
down0 = lowbit == 0 ? -1 : rank-lowbit;
down1 = lowbit == 0 ? -1 : rank+lowbit;
while (down1 >= nranks) {
down1 = lowbit == 0 ? -1 : rank+lowbit;
lowbit >>= 1;
}
*d0 = down0; *d1 = down1;
return 0;
}
int GetDtree(int nranks, int rank, int* s0, int* d0_0, int* d0_1, int* s1, int* d1_0, int* d1_1) {
GetBtree(nranks, rank, s0, d0_0, d0_1);
if (nranks % 2 == 0) {
int shiftrank = (rank-1+nranks) % nranks;
int u, d0, d1;
GetBtree(nranks, shiftrank, &u, &d0, &d1);
*s1 = u == -1 ? -1 : (u+1) % nranks;
*d1_0 = d0 == -1 ? -1 : (d0+1) % nranks;
*d1_1 = d1 == -1 ? -1 : (d1+1) % nranks;
} else {
int shiftrank=nranks-rank;
int u, d0, d1;
GetBtree(nranks, shiftrank, &u, &d0, &d1);
*s1 = u == -1 ? -1 : nranks-u;
*d1_0 = d0 == -1 ? -1 : nranks-d0;
*d1_1 = d1 == -1 ? -1 : nranks-d1;
}
return 0;
}
NCCL-tree生成的ranking方案
代码需要修改一下才能实现注释中提到的ranking方案
奇数节点Ranking
对于奇数个rank，比如下图的13，NV建议对第二个tree采用mirror ranking，也就是结构上镜像
文档中原图
* 8---------0---------5
* ______/ \______ _____/ \______
* 4 12 1 9
* / \ / \ / \
* 2 6 10 3 7 10
* / \ / \ / \ / \ / \ / \
* 1 3 5 7 9 11 2 4 6 8 11 12
*
这里文档的右下角写错了，因为如果按照这个路子画，程序跑出来的树应该是
修正右下角
* 8---------0---------5
* ______/ \______ _____/ \______
* 4 12 1 9
* / \ / \ / \
* 2 6 10 3 7 11
* / \ / \ / \ / \ / \ / \
* 1 3 5 7 9 11 2 4 6 8 12 10
如果不按照NV的做法，直接shift ranking，那么就是下图所示的结构，看上去也不错
在奇数节点中应用shift ranking
* 8---------0-----------------9
* ______/ \______ ______/ \______
* 4 12 5 1
* / \ / / \ /
* 2 6 10 3 7 11
* / \ / \ / \ / \ / \ / \
* 1 3 5 7 9 11 2 4 6 8 10 12
偶数节点ranking
对于偶数个rank，比如下图的12，NV建议采用shift ranking
文档中原图
* or shift it by one rank (if nranks is even)
*
* 8---------0--------------9
* ______/ \ ______/ \
* 4 \ 5 \
* / \ \ / \ \
* 2 6 10 3 7 11
* / \ / \ / \ / \ / \ / \
* 1 3 5 7 9 11 2 4 6 8 10 1
*/
这里文档又写错了，真正运行出来的结果应该是1变成了新root
修正根节点
* 8---------0--1-----------9
* ______/ \ ______/ \
* 4 \ 5 \
* / \ \ / \ \
* 2 6 10 3 7 11
* / \ / \ / \ / \ / \ / \
* 1 3 5 7 9 11 2 4 6 8 10 0
我试试这个结构能不能做mirror ranking
结果发现leaf节点还都是奇数编号，这就是为啥偶数ranks时候不能mirror的原因
在偶数节点中应用mirror ranking
* 8---------0--------4
* ______/ \ / \______
* 4 \ / 8
* / \ \ / / \
* 2 6 10 2 6 10
* / \ / \ / \ / \ / \ / \
* 1 3 5 7 9 11 1 3 5 7 9 11
*/
测试程序输出
nrank=13的时候：
i=0; u=-1, d0=8, d1=-1
du=-1, dd0=5, dd1=-1
i=1; u=2, d0=-1, d1=-1
du=5, dd0=3, dd1=-1
i=2; u=4, d0=1, d1=3
du=3, dd0=-1, dd1=-1
i=3; u=2, d0=-1, d1=-1
du=1, dd0=4, dd1=2
i=4; u=8, d0=2, d1=6
du=3, dd0=-1, dd1=-1
i=5; u=6, d0=-1, d1=-1
du=13, dd0=9, dd1=1
i=6; u=4, d0=5, d1=7
du=7, dd0=-1, dd1=-1
i=7; u=6, d0=-1, d1=-1
du=9, dd0=8, dd1=6
i=8; u=0, d0=4, d1=12
du=7, dd0=-1, dd1=-1
i=9; u=10, d0=-1, d1=-1
du=5, dd0=11, dd1=7
i=10; u=12, d0=9, d1=11
du=11, dd0=-1, dd1=-1
i=11; u=10, d0=-1, d1=-1
du=9, dd0=12, dd1=10
i=12; u=8, d0=10, d1=-1
du=11, dd0=-1, dd1=-1
nrank=12的时候
i=0; u=-1, d0=8, d1=-1
du=11, dd0=-1, dd1=-1
i=1; u=2, d0=-1, d1=-1
du=-1, dd0=9, dd1=-1
i=2; u=4, d0=1, d1=3
du=3, dd0=-1, dd1=-1
i=3; u=2, d0=-1, d1=-1
du=5, dd0=2, dd1=4
i=4; u=8, d0=2, d1=6
du=3, dd0=-1, dd1=-1
i=5; u=6, d0=-1, d1=-1
du=9, dd0=3, dd1=7
i=6; u=4, d0=5, d1=7
du=7, dd0=-1, dd1=-1
i=7; u=6, d0=-1, d1=-1
du=5, dd0=6, dd1=8
i=8; u=0, d0=4, d1=10
du=7, dd0=-1, dd1=-1
i=9; u=10, d0=-1, d1=-1
du=1, dd0=5, dd1=11
i=10; u=8, d0=9, d1=11
du=11, dd0=-1, dd1=-1
i=11; u=10, d0=-1, d1=-1
du=9, dd0=10, dd1=0
