def p2p_broadcast(t, src):
    for dst in range(g_env.world_size):
        if ((src == dst) or (g_env.rank not in (src, dst))):
            continue
        dst_adj_nz_col = g_data.nz_col_dict[(dst, src)]
        needed_rows_idx = dst_adj_nz_col
        if (g_env.rank == src):
            p2p_buf = t[needed_rows_idx]
        elif (g_env.rank == dst):
            p2p_buf = torch.zeros((needed_rows_idx.size(0), t.size(1)), device=g_env.device)
        dist.broadcast(p2p_buf, src, group=g_env.p2p_group_dict[(src, dst)])
        if (g_env.rank == dst):
            t[needed_rows_idx] = p2p_buf
            return
def p2p_broadcast(t, src):
    for dst in range(g_env.world_size):
        if ((src == dst) or (g_env.rank not in (src, dst))):
            continue
        dst_adj_nz_col = g_data.nz_col_dict[(dst, src)]
        needed_rows_idx = dst_adj_nz_col
        if (g_env.rank == src):
            p2p_buf = t[needed_rows_idx]
        elif (g_env.rank == dst):
            p2p_buf = torch.zeros((needed_rows_idx.size(0), t.size(1)), device=g_env.device)
        dist.broadcast(p2p_buf, src, group=g_env.p2p_group_dict[(src, dst)])
        if (g_env.rank == dst):
            t[needed_rows_idx] = p2p_buf
            return