void NcclCommThread()

in src/nccl_op.cc [145:206]
48 lines of code
12 McCabe index (conditional complexity)

    void NcclCommThread()
    {
        // Activate appropriate context for this thread.
        ScopedActivateExecutorContext scoped_context(executor);

        while (true)
        {
            CommEntry* entry = NULL;
            {
                mutex_lock l(mu_comm);
                while (pending_comm.empty() || shutdown_requested)
                {
                    if (shutdown_requested)
                        return;
                    cv_comm.wait(l);
                }
                entry = pending_comm.back();
                pending_comm.pop_back();
            }
            // printf("NcclCommThread: %03d %d %d\n", entry->op_num, mpi_rank, (int)entry->input->NumElements());

            // Ensure the nccl op doesn't start executing before it would on the main compute stream.
            // This is executed on device and does not block.
            cuStreamWaitEvent(nccl_stream, entry->op_event, 0);

            ncclResult_t nccl_result = ncclSuccess;
            if (mpi_rank != -1)
            {
                if (prereduce)
                {
                    // reduce among each machine's gpu
                    nccl_result = ncclReduce(entry->buffer_in, entry->buffer_out, entry->count, entry->dtype, ncclSum, 0, comm_local, nccl_stream);
                    if (nccl_result == ncclSuccess)
                    {
                        // reduce among all machines with the root gpu on each
                        if (local_rank == 0 && global_size > 1)
                            nccl_result = ncclAllReduce(entry->buffer_in, entry->buffer_out, entry->count, entry->dtype, ncclSum, comm_global, nccl_stream);
                        // broadcast root gpu results out to the rest of the gpus on each machine
                        if (nccl_result == ncclSuccess)
                            nccl_result = ncclBroadcast(entry->buffer_in, entry->buffer_out, entry->count, entry->dtype, 0, comm_local, nccl_stream);
                    }
                }
                else
                {
                    // Simple all-reduce among all mpi ranks
                    nccl_result = ncclAllReduce(entry->buffer_in, entry->buffer_out, entry->count, entry->dtype, ncclSum, comm_global, nccl_stream);
                }
            }

            if (nccl_result != ncclSuccess)
            {
                entry->status = errors::Internal("Error invoking ncclAllReduce (", mpi_rank, "/", mpi_size, "): ", ncclGetErrorString(nccl_result));
                shutdown_requested = true;
            }

            // Wait for kernels to finish on another thread so this thread can continue to queue up work
            cuEventRecord(entry->op_event, nccl_stream);
            mutex_lock l(mu_done);
            pending_done.push_front(entry);
            cv_done.notify_one();
        }
    }