def test_group_broadcast():
row_procs = []
for i in range(0,8,2):
row_procs.append(list(range(i,i+2)))
col_procs = []
for i in range(2):
col_procs.append(list(range(i,8,2)))
row_groups = []
for i in range(len(row_procs)):
row_groups.append(ad.new_group_comm(row_procs[i]))
col_groups = []
for i in range(len(col_procs)):
col_groups.append(ad.new_group_comm(col_procs[i]))
rank_row = rank//2
rank_col = rank%2
group_row = row_procs[rank_row]
group_col = col_procs[rank_col]
comm_row = row_groups[rank_row]
comm_col = col_groups[rank_col]
a = ndarray.array(np.array([rank,rank,rank,rank,rank]),ctx=ctx)
comm_row.dlarrayBroadcast(a, a, ncclDataType_t.ncclFloat32, root = group_row[1])
print("Broadcast device=%d, a:"%device_id,a.asnumpy())
b = ndarray.array(np.array([rank,rank,rank,rank,rank]),ctx=ctx)
comm_col.dlarrayBroadcast(b, b, ncclDataType_t.ncclFloat32, root = group_col[1])
print("Broadcast device=%d, b:"%device_id,b.asnumpy())
def test_allgather(group):
comm1 = ad.new_group_comm(group)
a = ndarray.array(np.array([rank,rank]),ctx=ctx)
b = ndarray.array(np.zeros(2*len(group)),ctx=ctx)
if rank in group:
comm1.dlarrayAllGather(a, b, ncclDataType_t.ncclFloat32)
print("Allgather device=%d"%comm1.device_id.value,b.asnumpy())
def test_broadcast(group, root):
comm1 = ad.new_group_comm(group)
a = ndarray.array(np.array([-1,-1,-1,-1,-1]),ctx=ctx)
if rank == root:
a = ndarray.array(np.array([2,3,4,5,6]),ctx=ctx)
if rank in group:
comm1.dlarrayBroadcast(a, a, ncclDataType_t.ncclFloat32, root = root)
print("Broadcast device=%d"%comm1.device_id.value,a.asnumpy())
def get_proc_groups(size, replication):
if replication == 1:
return None, None, None, None
row_procs = []
for i in range(0, size, replication):
row_procs.append(list(range(i, i + replication)))
col_procs = []
for i in range(replication):
col_procs.append(list(range(i, size, replication)))
row_groups = []
for i in range(len(row_procs)):
row_groups.append(ad.new_group_comm(row_procs[i]))
col_groups = []
for i in range(len(col_procs)):
col_groups.append(ad.new_group_comm(col_procs[i]))
return row_procs, col_procs, row_groups, col_groups
def test_default():
comm1 = ad.new_group_comm()
a = ndarray.array(np.array([1,2,3,4,5]),ctx=ctx)
comm1.dlarrayNcclAllReduce(a, a, ncclDataType_t.ncclFloat32, reduceop=ncclRedOp_t.ncclSum)
print("Default Allreduce device=%d"%comm1.device_id.value,a.asnumpy())