def broad_func(node_count, am_partitions, inputs, btype=None):
    global cache_features
    device = g_env.device
    n_per_proc = math.ceil((float(node_count) / g_env.world_size))
    z_loc = torch.zeros((am_partitions[0].size(0), inputs.size(1)), device=device)
    inputs_recv = torch.zeros((n_per_proc, inputs.size(1)), device=device)
    for i in range(g_env.world_size):
        layer1_use_cache = (cur_epoch >= 1)
        layer2_use_cache = ((cur_epoch >= 50) and ((cur_epoch % 5) != 0))
        backward_layer2_use_cache = ((cur_epoch >= 50) and ((cur_epoch % 5) != 0))
        backward_layer2_use_cache = False
        backward_layer1_use_cache = ((cur_epoch >= 50) and ((cur_epoch % 5) != 0))
        backward_layer1_use_cache = False
        if (i == g_env.rank):
            inputs_recv = inputs.clone()
        elif (i == (g_env.world_size - 1)):
            inputs_recv = torch.zeros((am_partitions[i].size(1), inputs.size(1)), device=device)
        g_timer.barrier_all()
        torch.cuda.synchronize()
        g_timer.start(f'gcn_broadcast_{btype}_ep{cur_epoch}')
        if (not cache_enabled):
            dist.broadcast(inputs_recv, src=i, group=g_env.world_group)
        elif (btype == 'layer1'):
            if layer1_use_cache:
                if (g_env.rank != i):
                    inputs_recv = cache_features[i]
                else:
                    pass
            else:
                dist.broadcast(inputs_recv, src=i, group=g_env.world_group)
                cache_features[i] = inputs_recv.clone()
        elif (btype == 'layer2'):
            if layer2_use_cache:
                if (g_env.rank != i):
                    inputs_recv = cache_layer2[i]
            else:
                dist.broadcast(inputs_recv, src=i, group=g_env.world_group)
                cache_layer2[i] = inputs_recv.clone()
        elif (btype == 'backward_layer2'):
            if backward_layer2_use_cache:
                if (g_env.rank != i):
                    inputs_recv = cache_backward_layer2[i]
            else:
                dist.broadcast(inputs_recv, src=i, group=g_env.world_group)
                cache_backward_layer2[i] = inputs_recv.clone()
        elif (btype == 'backward_layer1'):
            if backward_layer1_use_cache:
                if (g_env.rank != i):
                    inputs_recv = cache_backward_layer1[i]
            else:
                dist.broadcast(inputs_recv, src=i, group=g_env.world_group)
                cache_backward_layer1[i] = inputs_recv.clone()
        else:
            dist.broadcast(inputs_recv, src=i, group=g_env.world_group)
        torch.cuda.synchronize()
        g_timer.stop(f'gcn_broadcast_{btype}_ep{cur_epoch}')
        g_logger.log(f'[gcn_broadcast_{btype}_ep{cur_epoch}] size: {utils.mem_report(inputs_recv)} MBytes')
        g_timer.barrier_all()
        torch.cuda.synchronize()
        g_timer.start(f'gcn_spmm_ep{cur_epoch}')
        spmm_gpu(am_partitions[i].indices()[0].int(), am_partitions[i].indices()[1].int(), am_partitions[i].values(), am_partitions[i].size(0), am_partitions[i].size(1), inputs_recv, z_loc)
        torch.cuda.synchronize()
        g_timer.stop(f'gcn_spmm_ep{cur_epoch}')
        g_timer.barrier_all()
    return z_loc
def broad_func(node_count, am_partitions, inputs, btype=None):
    global cache_features
    device = g_env.device
    n_per_proc = math.ceil((float(node_count) / g_env.world_size))
    z_loc = torch.zeros((am_partitions[0].size(0), inputs.size(1)), device=device)
    inputs_recv = torch.zeros((n_per_proc, inputs.size(1)), device=device)
    for i in range(g_env.world_size):
        layer1_use_cache = (cur_epoch >= 1)
        layer2_use_cache = ((cur_epoch >= 50) and ((cur_epoch % 5) != 0))
        backward_layer2_use_cache = ((cur_epoch >= 50) and ((cur_epoch % 5) != 0))
        backward_layer2_use_cache = False
        backward_layer1_use_cache = ((cur_epoch >= 50) and ((cur_epoch % 5) != 0))
        backward_layer1_use_cache = False
        if (i == g_env.rank):
            inputs_recv = inputs.clone()
        elif (i == (g_env.world_size - 1)):
            inputs_recv = torch.zeros((am_partitions[i].size(1), inputs.size(1)), device=device)
        g_timer.barrier_all()
        torch.cuda.synchronize()
        g_timer.start(f'gcn_broadcast_{btype}_ep{cur_epoch}')
        if (not cache_enabled):
            dist.broadcast(inputs_recv, src=i, group=g_env.world_group)
        elif (btype == 'layer1'):
            if layer1_use_cache:
                if (g_env.rank != i):
                    inputs_recv = cache_features[i]
                else:
                    pass
            else:
                dist.broadcast(inputs_recv, src=i, group=g_env.world_group)
                cache_features[i] = inputs_recv.clone()
        elif (btype == 'layer2'):
            if layer2_use_cache:
                if (g_env.rank != i):
                    inputs_recv = cache_layer2[i]
            else:
                dist.broadcast(inputs_recv, src=i, group=g_env.world_group)
                cache_layer2[i] = inputs_recv.clone()
        elif (btype == 'backward_layer2'):
            if backward_layer2_use_cache:
                if (g_env.rank != i):
                    inputs_recv = cache_backward_layer2[i]
            else:
                dist.broadcast(inputs_recv, src=i, group=g_env.world_group)
                cache_backward_layer2[i] = inputs_recv.clone()
        elif (btype == 'backward_layer1'):
            if backward_layer1_use_cache:
                if (g_env.rank != i):
                    inputs_recv = cache_backward_layer1[i]
            else:
                dist.broadcast(inputs_recv, src=i, group=g_env.world_group)
                cache_backward_layer1[i] = inputs_recv.clone()
        else:
            dist.broadcast(inputs_recv, src=i, group=g_env.world_group)
        torch.cuda.synchronize()
        g_timer.stop(f'gcn_broadcast_{btype}_ep{cur_epoch}')
        g_logger.log(f'[gcn_broadcast_{btype}_ep{cur_epoch}] size: {utils.mem_report(inputs_recv)} MBytes')
        g_timer.barrier_all()
        torch.cuda.synchronize()
        g_timer.start(f'gcn_spmm_ep{cur_epoch}')
        spmm_gpu(am_partitions[i].indices()[0].int(), am_partitions[i].indices()[1].int(), am_partitions[i].values(), am_partitions[i].size(0), am_partitions[i].size(1), inputs_recv, z_loc)
        torch.cuda.synchronize()
        g_timer.stop(f'gcn_spmm_ep{cur_epoch}')
        g_timer.barrier_all()
    return z_loc