void CreateCustomComm()

in src/nccl_op.cc [332:383]


    void CreateCustomComm(int op_num, const std::set<int>& mpi_ranks, int comm_id, StreamExecutor* executor)
    {
        auto ranks_and_id = std::make_pair(mpi_ranks, comm_id);

        if (op_to_comm.count(op_num))
        {
            return;
        }
        else if (ranks_to_comm.count(ranks_and_id))
        {
            op_to_comm[op_num] = ranks_to_comm[ranks_and_id];
            return;
        }

        ncclUniqueId id;

        if (mpi_rank == 0)
        {
            NCCL_CHECK(ncclGetUniqueId(&id), "ncclGetUniqueId");
        }
        MPI_CHECK(MPI_Bcast((void*)(&id), sizeof(ncclUniqueId), MPI_BYTE, 0, mpi_comm), "MPI_Bcast");

        int custom_mpi_size = static_cast<int>(mpi_ranks.size());
        int custom_mpi_rank = -1;

        if (mpi_ranks.count(mpi_rank))
        {
            custom_mpi_rank = 0;

            for (auto r : mpi_ranks)
            {
                if (mpi_rank == r)
                {
                    break;
                }
                custom_mpi_rank++;
            }
        }

        if (custom_mpi_rank != -1)
        {
            custom_comms.push_back(new NcclComm(custom_mpi_size, custom_mpi_rank, -1, -1, id, id, 0, executor));
        }
        else
        {
            custom_comms.push_back(new NcclComm(-1, -1, -1, -1, id, id, 0, executor));
        }

        int custom_comm_idx = static_cast<int>(custom_comms.size()) - 1;
        ranks_to_comm[ranks_and_id] = custom_comm_idx;
        op_to_comm[op_num] = custom_comm_idx;
    }